mochi model/pipeline review

Commit tested: 0f1abc4ae8b0eb2a3b40e82a310507281144c423

Review performed against the repository review rules. I read the available referenced rule files; AGENTS.md is referenced by .ai/review-rules.md but is not present in this checkout.

Duplicate search status: searched GitHub Issues and PRs for mochi, pipeline_mochi.py, AutoencoderKLMochi, MochiPipelineOutput, linear_quadratic_schedule, force_zeros_for_empty_prompt, MochiAttention, and the failure modes below. I did not find likely duplicates. Related but not duplicates: closed issue #11291 covers a different Mochi VAE tiling-width bug, and open PR #13348 refactors Mochi transformer tests only.

Issue 1: Custom timesteps are unusable and one-step schedules crash

Affected code:

Problem:
__call__ always builds Mochi’s internal sigmas schedule and passes it to retrieve_timesteps, even when the user also passes timesteps. retrieve_timesteps rejects receiving both. Separately, linear_quadratic_schedule(num_steps=1, ...) divides by zero.

Impact:
The public timesteps argument documented on MochiPipeline.__call__ cannot be used. num_inference_steps=1 also crashes before denoising.

Reproduction:

from diffusers import FlowMatchEulerDiscreteScheduler
from diffusers.pipelines.mochi.pipeline_mochi import linear_quadratic_schedule, retrieve_timesteps

for name, fn in {
    "one step": lambda: linear_quadratic_schedule(1, 0.025),
    "custom timesteps": lambda: retrieve_timesteps(
        FlowMatchEulerDiscreteScheduler(),
        num_inference_steps=2,
        device="cpu",
        timesteps=[999, 500],
        sigmas=linear_quadratic_schedule(2, 0.025),
    ),
}.items():
    try:
        fn()
    except Exception as e:
        print(name, type(e).__name__, e)

Relevant precedent:
retrieve_timesteps already supports either custom timesteps or sigmas; the pipeline should not pass both.

Suggested fix:

sigmas = None
if timesteps is None:
    sigmas = np.array(linear_quadratic_schedule(num_inference_steps, threshold_noise), dtype=np.float32)

timesteps, num_inference_steps = retrieve_timesteps(
    self.scheduler, num_inference_steps, timestep_device, timesteps, sigmas
)

Also either support num_inference_steps=1 in linear_quadratic_schedule or reject it with a clear check_inputs error.

Issue 2: Spatial and temporal input validation accepts unsupported requests

Affected code:

Problem:
The pipeline only checks height/width divisibility by 8, but the transformer also patches the VAE latents by patch_size=2, so image dimensions need to be divisible by 16 for the default model. num_frames is also accepted even when (num_frames - 1) is not divisible by the VAE temporal scale; those requests silently decode fewer frames than requested.

Impact:
Users get late transformer reshape errors for sizes like 24x24, or silently receive fewer frames than requested for values like num_frames=8.

Reproduction:

import torch
from diffusers import AutoencoderKLMochi, MochiPipeline, MochiTransformer3DModel

pipe = MochiPipeline(scheduler=None, vae=None, text_encoder=None, tokenizer=None, transformer=None)
pipe.check_inputs(prompt="x", height=24, width=24)
print("24x24 accepted")

model = MochiTransformer3DModel(
    patch_size=2, num_attention_heads=2, attention_head_dim=8, num_layers=1,
    pooled_projection_dim=16, in_channels=4, text_embed_dim=16, time_embed_dim=4, max_sequence_length=8,
)
try:
    model(torch.randn(1, 4, 1, 3, 3), torch.randn(1, 8, 16), torch.tensor([1.0]), torch.ones(1, 8).bool())
except RuntimeError as e:
    print(type(e).__name__, e)

vae = AutoencoderKLMochi(
    latent_channels=4, encoder_block_out_channels=(32, 32, 32, 32),
    decoder_block_out_channels=(32, 32, 32, 32), layers_per_block=(1, 1, 1, 1, 1),
)
pipe = MochiPipeline(scheduler=None, vae=vae, text_encoder=None, tokenizer=None, transformer=None)
latents = pipe.prepare_latents(1, 4, 16, 16, 8, torch.float32, "cpu", torch.Generator().manual_seed(0))
print("requested 8, decoded", vae.decode(latents, return_dict=False)[0].shape[2])

Relevant precedent:
Wan derives and enforces a 16 spatial factor and rounds invalid frame counts:

Suggested fix:

self.vae_spatial_scale_factor = self.vae.spatial_compression_ratio if getattr(self, "vae", None) else 8
self.vae_temporal_scale_factor = self.vae.temporal_compression_ratio if getattr(self, "vae", None) else 6
self.patch_size = self.transformer.config.patch_size if getattr(self, "transformer", None) else 2

spatial_factor = self.vae_spatial_scale_factor * self.patch_size
if height % spatial_factor != 0 or width % spatial_factor != 0:
    raise ValueError(f"`height` and `width` have to be divisible by {spatial_factor}.")

if num_frames % self.vae_temporal_scale_factor != 1:
    logger.warning(...)
    num_frames = num_frames // self.vae_temporal_scale_factor * self.vae_temporal_scale_factor + 1

Issue 3: Empty-prompt zeroing is applied to the wrong batch rows

Affected code:

Problem:
_get_t5_prompt_embeds checks prompt == "" or prompt[-1] == "". For batched prompts, this only looks at the last prompt and then zeroes the entire batch. If the first prompt is empty and the last is not, the empty row is not zeroed.

Impact:
Batched negative prompts can condition the wrong samples, especially with force_zeros_for_empty_prompt=True.

Reproduction:

import torch
from transformers import AutoConfig, AutoTokenizer, T5EncoderModel
from diffusers import MochiPipeline

config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5")
pipe = MochiPipeline(
    scheduler=None, vae=None, transformer=None,
    text_encoder=T5EncoderModel(config),
    tokenizer=AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5"),
    force_zeros_for_empty_prompt=True,
)

_, mask = pipe._get_t5_prompt_embeds(["not empty", ""], max_sequence_length=8, device=torch.device("cpu"))
print(mask.sum(dim=1).tolist())  # [0, 0], but only row 2 should be zeroed

_, mask = pipe._get_t5_prompt_embeds(["", "not empty"], max_sequence_length=8, device=torch.device("cpu"))
print(mask.sum(dim=1).tolist())  # first row is not zeroed

Relevant precedent:
The current behavior was made configurable in PR #10284, but I found no duplicate for the batched-row bug.

Suggested fix:

if self.config.force_zeros_for_empty_prompt:
    empty_prompt_mask = torch.tensor([p == "" for p in prompt], device=device)
    text_input_ids = text_input_ids.to(device)
    text_input_ids[empty_prompt_mask] = 0
    prompt_attention_mask[empty_prompt_mask] = False

Issue 4: AutoencoderKLMochi.forward(return_dict=False) returns a DecoderOutput

Affected code:

Problem:
forward calls self.decode(z) without taking .sample. With return_dict=False, it returns (DecoderOutput(...),) instead of (tensor,).

Impact:
Generic model callers and tests expecting tuple outputs get the wrong type.

Reproduction:

import torch
from diffusers import AutoencoderKLMochi

vae = AutoencoderKLMochi(
    latent_channels=4, encoder_block_out_channels=(32, 32, 32, 32),
    decoder_block_out_channels=(32, 32, 32, 32), layers_per_block=(1, 1, 1, 1, 1),
)
out = vae(torch.randn(1, 3, 7, 16, 16), return_dict=False)[0]
print(type(out), hasattr(out, "sample"))

Relevant precedent:

Suggested fix:

dec = self.decode(z).sample
if not return_dict:
    return (dec,)
return DecoderOutput(sample=dec)

Issue 5: Mochi attention ignores set_attention_backend

Affected code:

Problem:
MochiAttnProcessor2_0 has no _attention_backend attribute and calls F.scaled_dot_product_attention directly. ModelMixin.set_attention_backend() therefore skips the processor and the forward path ignores the selected backend.

Impact:
Users can call model.set_attention_backend("...") and receive no Mochi attention behavior change. This also leaves Mochi outside the current dispatcher/context-parallel attention conventions.

Reproduction:

from diffusers import MochiTransformer3DModel

model = MochiTransformer3DModel(
    patch_size=2, num_attention_heads=2, attention_head_dim=8, num_layers=1,
    pooled_projection_dim=16, in_channels=4, text_embed_dim=16, time_embed_dim=4, max_sequence_length=8,
)
processor = model.transformer_blocks[0].attn1.processor
print(hasattr(processor, "_attention_backend"), getattr(processor, "_attention_backend", None))
model.set_attention_backend("_native_math")
processor = model.transformer_blocks[0].attn1.processor
print(hasattr(processor, "_attention_backend"), getattr(processor, "_attention_backend", None))

Relevant precedent:
QwenImage routes attention through dispatch_attention_fn:

Suggested fix:
This is larger than an inline patch: move Mochi attention to the current AttentionModuleMixin pattern, give the processor _attention_backend / _parallel_config, and route the joint valid-token attention through dispatch_attention_fn.

Issue 6: MochiPipelineOutput is not exported from the Mochi package

Affected code:

Problem:
pipeline_output.py defines MochiPipelineOutput, but src/diffusers/pipelines/mochi/__init__.py only exports MochiPipeline.

Impact:
from diffusers.pipelines.mochi import MochiPipelineOutput fails, unlike many newer pipeline families that expose their output dataclasses through the package lazy loader.

Reproduction:

from diffusers.pipelines.mochi import MochiPipelineOutput

Relevant precedent:

Suggested fix:

_import_structure = {"pipeline_output": ["MochiPipelineOutput"]}

# TYPE_CHECKING / slow import branch
from .pipeline_output import MochiPipelineOutput

Issue 7: Slow coverage is missing and existing expected-output checks are placeholders

Affected code:

Problem:
There are fast tests and a @nightly pipeline integration test, but no @slow Mochi tests in the pipeline/model/LoRA Mochi test files. The fast test compares to torch.randn(...) with max_diff <= 1e10, and the nightly test compares the full model output to a random tensor instead of a fixed expected slice.

Impact:
Shape-level smoke tests exist, but numerical regressions in Mochi pipeline/model behavior are unlikely to be caught by normal slow test runs.

Reproduction:

from pathlib import Path

paths = [
    Path("tests/pipelines/mochi/test_mochi.py"),
    Path("tests/models/autoencoders/test_models_autoencoder_mochi.py"),
    Path("tests/models/transformers/test_models_transformer_mochi.py"),
    Path("tests/lora/test_lora_layers_mochi.py"),
]
assert any("@slow" in path.read_text() for path in paths), "No @slow Mochi tests found"

Relevant precedent:
Other mature pipeline tests use fixed expected slices rather than random placeholders for integration assertions.

Suggested fix:
Add at least one @slow Mochi pipeline test using genmo/mochi-1-preview with a small deterministic run and fixed expected output slice, and replace the random placeholder assertions in the fast/nightly tests with meaningful deterministic assertions.