From 5bc4e31a8001c6af7ecdac7976b4acf05c5fb388 Mon Sep 17 00:00:00 2001 From: lavinal712 Date: Sat, 12 Sep 2026 21:08:09 +0800 Subject: [PATCH] init: add MAGI-1 --- docs/source/en/_toctree.yml | 8 + .../en/api/models/autoencoderkl_magi.md | 97 ++ .../en/api/models/magi_transformer3d.md | 114 +++ docs/source/en/api/modular_diffusers/magi.md | 240 +++++ docs/source/en/api/schedulers/magi_euler.md | 66 ++ examples/magi/README-prefix.md | 64 ++ examples/magi/README.md | 69 ++ examples/magi/inference_magi.py | 210 +++++ examples/magi/requirements-dev.txt | 3 + examples/magi/requirements.txt | 14 + scripts/convert_magi_to_diffusers.py | 264 ++++++ src/diffusers/__init__.py | 26 + src/diffusers/guiders/__init__.py | 1 + .../guiders/magi_classifier_free_guidance.py | 104 +++ src/diffusers/models/__init__.py | 6 + src/diffusers/models/autoencoders/__init__.py | 1 + .../autoencoders/autoencoder_kl_magi.py | 570 +++++++++++ src/diffusers/models/magi_conditioning.py | 91 ++ src/diffusers/models/transformers/__init__.py | 1 + .../models/transformers/transformer_magi.py | 608 ++++++++++++ src/diffusers/modular_pipelines/__init__.py | 20 + .../modular_pipelines/magi/__init__.py | 67 ++ .../modular_pipelines/magi/before_denoise.py | 254 +++++ .../modular_pipelines/magi/decoders.py | 143 +++ .../modular_pipelines/magi/denoise.py | 884 ++++++++++++++++++ .../modular_pipelines/magi/encoders.py | 303 ++++++ .../magi/modular_blocks_magi.py | 410 ++++++++ .../magi/modular_pipeline.py | 21 + .../modular_pipelines/modular_pipeline.py | 1 + src/diffusers/schedulers/__init__.py | 2 + .../schedulers/scheduling_magi_euler.py | 194 ++++ src/diffusers/utils/dummy_pt_objects.py | 75 ++ .../dummy_torch_and_transformers_objects.py | 120 +++ .../test_models_autoencoder_kl_magi.py | 279 ++++++ tests/models/test_models_magi_conditioning.py | 125 +++ .../test_models_transformer_magi.py | 446 +++++++++ tests/modular_pipelines/magi/__init__.py | 13 + .../magi/test_magi_denoise.py | 240 +++++ .../magi/test_modular_pipeline_magi.py | 321 +++++++ .../magi/test_modular_pipeline_magi_prefix.py | 242 +++++ tests/modular_pipelines/magi/testing_utils.py | 57 ++ tests/schedulers/test_scheduler_magi_euler.py | 179 ++++ 42 files changed, 6953 insertions(+) create mode 100644 docs/source/en/api/models/autoencoderkl_magi.md create mode 100644 docs/source/en/api/models/magi_transformer3d.md create mode 100644 docs/source/en/api/modular_diffusers/magi.md create mode 100644 docs/source/en/api/schedulers/magi_euler.md create mode 100644 examples/magi/README-prefix.md create mode 100644 examples/magi/README.md create mode 100644 examples/magi/inference_magi.py create mode 100644 examples/magi/requirements-dev.txt create mode 100644 examples/magi/requirements.txt create mode 100644 scripts/convert_magi_to_diffusers.py create mode 100644 src/diffusers/guiders/magi_classifier_free_guidance.py create mode 100644 src/diffusers/models/autoencoders/autoencoder_kl_magi.py create mode 100644 src/diffusers/models/magi_conditioning.py create mode 100644 src/diffusers/models/transformers/transformer_magi.py create mode 100644 src/diffusers/modular_pipelines/magi/__init__.py create mode 100644 src/diffusers/modular_pipelines/magi/before_denoise.py create mode 100644 src/diffusers/modular_pipelines/magi/decoders.py create mode 100644 src/diffusers/modular_pipelines/magi/denoise.py create mode 100644 src/diffusers/modular_pipelines/magi/encoders.py create mode 100644 src/diffusers/modular_pipelines/magi/modular_blocks_magi.py create mode 100644 src/diffusers/modular_pipelines/magi/modular_pipeline.py create mode 100644 src/diffusers/schedulers/scheduling_magi_euler.py create mode 100644 tests/models/autoencoders/test_models_autoencoder_kl_magi.py create mode 100644 tests/models/test_models_magi_conditioning.py create mode 100644 tests/models/transformers/test_models_transformer_magi.py create mode 100644 tests/modular_pipelines/magi/__init__.py create mode 100644 tests/modular_pipelines/magi/test_magi_denoise.py create mode 100644 tests/modular_pipelines/magi/test_modular_pipeline_magi.py create mode 100644 tests/modular_pipelines/magi/test_modular_pipeline_magi_prefix.py create mode 100644 tests/modular_pipelines/magi/testing_utils.py create mode 100644 tests/schedulers/test_scheduler_magi_euler.py diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index f05667986f11..6458ba9b73ad 100644 --- a/docs/source/en/_toctree.yml +++ b/docs/source/en/_toctree.yml @@ -260,6 +260,8 @@ title: Components and configs - local: api/modular_diffusers/guiders title: Guiders + - local: api/modular_diffusers/magi + title: MAGI chunk denoising title: Modular - sections: - local: api/loaders/ip_adapter @@ -373,6 +375,8 @@ title: Lumina2Transformer2DModel - local: api/models/lumina_nextdit2d title: LuminaNextDiT2DModel + - local: api/models/magi_transformer3d + title: MagiTransformer3DModel - local: api/models/minimax_h3_transformer3d title: MiniMaxH3Transformer3DModel - local: api/models/minimax_music3_transformer @@ -463,6 +467,8 @@ title: AutoencoderKLLTX2Video - local: api/models/autoencoderkl_ltx_video title: AutoencoderKLLTXVideo + - local: api/models/autoencoderkl_magi + title: AutoencoderKLMagi - local: api/models/autoencoderkl_magvit title: AutoencoderKLMagvit - local: api/models/autoencoderkl_minimax_h3 @@ -790,6 +796,8 @@ title: LCMScheduler - local: api/schedulers/lms_discrete title: LMSDiscreteScheduler + - local: api/schedulers/magi_euler + title: MagiEulerScheduler - local: api/schedulers/minimax_h3 title: MiniMaxH3Scheduler - local: api/schedulers/pndm diff --git a/docs/source/en/api/models/autoencoderkl_magi.md b/docs/source/en/api/models/autoencoderkl_magi.md new file mode 100644 index 000000000000..4c2ce939f440 --- /dev/null +++ b/docs/source/en/api/models/autoencoderkl_magi.md @@ -0,0 +1,97 @@ +# AutoencoderKLMagi + +MAGI-1 uses a transformer VAE with 16 latent channels, 8x spatial compression, and 4x temporal compression. The encoder +and decoder each contain 24 transformer blocks. Both use full attention within the input tile, learned position +embeddings, and the reference normalization formula `(x - mean) / (std + eps)` on Q, K, and V. + +The implementation retains the official fused `qkv`, `proj`, `mlp.fc1`, `mlp.fc2`, and `patch_embed.proj` layers. CLS and +position parameters are grouped in a small module so device placement and offloading hooks move them together with +their inputs. The conversion script maps these four parameter names and translates the configuration; it does not +split or reshape the checkpoint tensors. The encoder preserves the reference's channel-last storage layout. After +tiling, posterior means are packed separately from log-variance storage, as in the reference's mean-only concatenation; +this avoids changing the reduced-precision decoder's matrix multiplication path. + +## Convert and load + +Convert the official `ckpt/vae` directory before loading it: + +```bash +python scripts/convert_magi_to_diffusers.py vae --checkpoint_path /path/to/ckpt/vae --output_path /path/to/magi-vae-diffusers +``` + +The script loads all weights strictly and verifies that saving and reloading preserves every tensor. + +```python +import torch +from diffusers import AutoencoderKLMagi + +vae = AutoencoderKLMagi.from_pretrained("/path/to/magi-vae-diffusers", torch_dtype=torch.bfloat16).to("cuda") +vae.set_attention_backend("flash") + +# video has shape (batch, 3, frames, height, width) and values normalized to [-1, 1]. +video = video.to(device="cuda", dtype=torch.bfloat16) +with torch.no_grad(): + latents = vae.encode(video).latent_dist.mode() + reconstruction = vae.decode(latents, num_frames=video.shape[2]).sample +``` + +QKV is always fused, as in the reference. The `flash` backend requires FlashAttention and matches the official attention +kernel. The default backend works without this optional dependency; reduced-precision results can differ between +kernels. Latents here are raw VAE latents. The diffusion pipeline's latent scaling is applied outside this model. + +## Frame counts and sampling + +Spatial dimensions must be divisible by the spatial patch size. Without tiling, frame counts must be divisible by the +temporal patch length, except that a single input frame is repeated to fill one patch. With temporal tiling, the final +tile may also contain exactly one frame. Other incomplete temporal patches are rejected instead of silently dropping +frames. + +`encode` returns a posterior distribution. Use `.mode()` to match the official inference pipeline, which patches the +original VAE to encode deterministically. `.sample(generator=...)` follows Diffusers' generator and dtype conventions. +The original standalone VAE samples CPU float32 noise, so its stochastic outputs are not guaranteed to match for the +same seed. The deterministic mode is the reference inference path. + +By default, `decode` follows the original VAE convention: a tile containing one latent time position returns only the +first decoded frame. This also applies to the last tile of a longer video. For example, a four-frame video becomes one +latent time position and decodes to one frame by default. Pass `num_frames=4` to retain all four frames. The supplied +length must fit the latent patch count. `vae(video).sample` automatically supplies the original frame count. + +## Temporal and spatial tiling + +Enable single-device tiling to limit the attention sequence length: + +```python +vae.enable_tiling( + tile_sample_min_length=12, + tile_sample_min_height=256, + tile_sample_min_width=256, + temporal_tile_overlap_factor=0.0, + spatial_tile_overlap_factor=0.25, +) +with torch.no_grad(): + latents = vae.encode(video).latent_dist.mode() + reconstruction = vae.decode(latents, num_frames=video.shape[2]).sample +``` + +All tile dimensions are expressed in input video pixels or frames. The official video helper uses half the configured +FPS as its temporal tile length and enables spatial tiling; 12 corresponds to a configured FPS of 24. Match these +settings when comparing against the reference. Set `allow_spatial_tiling=False` for temporal tiling only. + +Tile dimensions and overlaps must align with the latent grid. Tiles at the boundary may be shorter. The tiler follows +the official frame/height/width iteration and blending order: encoding blends against preceding already-blended tiles, +while decoding blends against the original decoded neighbors. Decoder blending uses FP32 intermediates for low-precision +tensors, matching the accumulation precision of the compiled reference blend operations. In the reference, compiler +fallbacks can instead use low-precision eager arithmetic; numerical comparisons should isolate compiler state between +different tiling configurations. + +Tiling changes the attention context, so tiled outputs need not match whole-input outputs. The tiled posterior blends +mean and log-variance channels independently; sampling this distribution is not equivalent to blending independently +sampled tiles. Use `.mode()` for reference inference parity. + +Use `vae.enable_slicing()` to process batch items individually, `vae.disable_slicing()` to restore full-batch processing, +and `vae.disable_tiling()` to restore whole-input processing. Distributed tile processing is not implemented. + +## AutoencoderKLMagi + +[[autodoc]] AutoencoderKLMagi + - all diff --git a/docs/source/en/api/models/magi_transformer3d.md b/docs/source/en/api/models/magi_transformer3d.md new file mode 100644 index 000000000000..b7ecef8050c1 --- /dev/null +++ b/docs/source/en/api/models/magi_transformer3d.md @@ -0,0 +1,114 @@ +# MagiTransformer3DModel + +MAGI-1 uses parallel video self-attention and text cross-attention, grouped-query attention, learned 3D rotary +embeddings, and chunk-level timestep conditioning. The 4.5B model has 34 blocks and GELU feed-forward layers; the +24B model has 48 blocks, SwiGLU, and duplicated latent input channels. + +This implementation covers the non-quantized base and distilled architectures. It does not include MAGI's transport +scheduler, three-way guidance, distributed context/pipeline parallelism, or the original FP8 execution engine. + +## Convert and load + +Pass the official inference weight directory and its matching example configuration: + +```bash +python scripts/convert_magi_to_diffusers.py transformer \ + --checkpoint_path /path/to/MAGI-1/ckpt/magi/4.5B_base/inference_weight \ + --config_path /path/to/MAGI-1/example/4.5B/4.5B_base_config.json \ + --output_path /path/to/magi-transformer-diffusers +``` + +The converter loads weights strictly and checks every tensor after saving and reloading. It retains the reference's +self/cross-attention output interleave. The text key/value projection is split into the eight linear calls used by the +reference; the other projection tensors are not reordered. + +```python +import torch +from diffusers import MagiTransformer3DModel + +transformer = MagiTransformer3DModel.from_pretrained( + "/path/to/magi-transformer-diffusers", torch_dtype=torch.bfloat16 +).to("cuda") +with torch.no_grad(): + prediction = transformer( + hidden_states=latents, + encoder_hidden_states=prompt_embeds, + timestep=timesteps, + encoder_attention_mask=prompt_mask, + ).sample +``` + +Latents use `(batch, 16, frames, height, width)` for both model sizes. The 24B channel duplication, output truncation, +and internal `x_rescale_factor` are handled by the model. The VAE-to-diffusion latent scaling still belongs outside it. +Embedding and output projections, rotary frequencies, self-attention Q/K normalization, and residual post-normalization +retain the reference's high-precision behavior. The attention output projection also computes in FP32 while retaining +its BF16 checkpoint weights, and timestep frequencies are rounded to the model dtype before their FP32 MLP. Load with `torch_dtype` instead of casting all weights with `.bfloat16()`. + +The output `sample` is flow velocity, not a clean-latent prediction. Apply guidance to this velocity and pass it +directly to `MagiEulerScheduler.step`; do not divide a prediction residual by `1 - timestep`. + +## Chunks and conditioning + +`timestep` has shape `(batch,)` for one chunk or `(batch, chunks)` for equally sized temporal chunks. Values follow +the reference's [0, 1] convention; the model applies the factor of 1000 in its sinusoidal embedding. Text features can +be shared across chunks with shape `(batch, length, channels)` or supplied per chunk as `(batch, chunks, length, channels)`. +Boolean text masks have the corresponding shape without the channel dimension; `True` keeps a token. + +Self-attention is chunk-causal by default: tokens attend to their entire current chunk and all preceding chunks. +`kv_ranges` overrides this with one exclusive `(start, end)` token range per current chunk. Ranges are shared across +batch items and index the concatenation of cached and current video tokens. The generation scheduler is responsible +for choosing the reference's timestep-dependent sliding windows. + +`caption_dropout_mask` selects the learned conditional or unconditional *adaptive* embedding. As in official inference, +it does not replace text features in cross-attention. Pass the appropriate text features separately for guidance. + +Distilled checkpoints also require `timestep_delta`. This is the extra timestep passed to the same embedding MLP, +not a difference between consecutive diffusion timesteps. In the official helper it is `num_steps / 2`, except when +`num_steps == 12`, where it is `8 / distill_interval`. These are the official distilled sampler's conventions. +The caller must supply this value; `MagiEulerScheduler` and the current base pipeline do not implement distilled sampling. + +## Prefix cache + +Use `use_cache=True` to return a tuple of `(key, value)` tensors, one pair per layer. Each tensor has shape +`(batch, tokens, key_value_heads, head_dim)`. Passing the tuple as `kv_cache` prepends those entries to current keys +and values and offsets temporal rotary positions accordingly. Input caches are not modified in place. + +Only retain and reuse entries for clean, finalized prefix chunks. By default the output cache also includes current +chunks; `cache_token_count` can retain only a leading clean prefix and `cache_device="cpu"` can offload it per layer. +Caches computed with different conditioning are not interchangeable. The MAGI base pipeline intentionally shares one +null-caption clean-prefix cache between its two prefix-conditioned branches, matching the official sampler. Its +independent branch uses no cache. Do not reuse a cache across different spatial resolutions or batches. + +The default attention backend works without MAGI-specific CUDA extensions. Explicitly selecting `flash` or +`flash_varlen` with `set_attention_backend` also uses FlashAttention's rotary kernel, matching the reference's fused +rounding. The native PyTorch rotary path can produce different low-precision outputs. Text padding masks must be honored; +backends that reject masks, such as `flash`, require inputs without padding and `encoder_attention_mask=None`. +Use a mask-capable backend when padding is present. Gradient checkpointing, standard device mapping, and group +offloading are supported; MAGI's distributed cache engine is not included. + +## Compilation with text masks + +Packing valid text tokens before projection preserves the reference GEMM shapes, but its output length depends on +the mask values. Enable Dynamo's dynamic-output-shape capture when compiling masked inputs with `fullgraph=True`: + +```python +with torch._dynamo.config.patch(capture_dynamic_output_shape_ops=True): + compiled_transformer = torch.compile(transformer, fullgraph=True) + with torch.no_grad(): + prediction = compiled_transformer( + hidden_states=latents, + encoder_hidden_states=prompt_embeds, + timestep=timesteps, + encoder_attention_mask=prompt_mask, + ).sample +``` + +Keep this context active during execution, including calls that may trigger recompilation. The same requirement applies +to `compile_repeated_blocks(fullgraph=True)`. This is an explicit caller setting; importing MAGI does not change global +Dynamo configuration. Use the native attention backend for this path. Fused FlashAttention compilation is not validated. +Compilation can change floating-point rounding; eager FlashAttention remains the official numerical-parity path. + +## MagiTransformer3DModel + +[[autodoc]] MagiTransformer3DModel + - all diff --git a/docs/source/en/api/modular_diffusers/magi.md b/docs/source/en/api/modular_diffusers/magi.md new file mode 100644 index 000000000000..ab488a789395 --- /dev/null +++ b/docs/source/en/api/modular_diffusers/magi.md @@ -0,0 +1,240 @@ +# MAGI-1 base generation + +`MagiTextToVideoBlocks` connects T5 text encoding, official HQ/duration conditioning, chunk denoising, and chunk-wise +VAE decoding. Separate `MagiImageToVideoBlocks` and `MagiVideoToVideoBlocks` support image/video prefixes with the +same base checkpoint. Distilled sampling is not supported by these workflows. + +## Convert and load + +Convert the official base checkpoint directly into a self-contained pipeline with one command: + +```bash +python scripts/convert_magi_to_diffusers.py pipeline \ + --transformer_path /path/to/ckpt/magi/4.5B_base/inference_weight \ + --config_path /path/to/MAGI-1/example/4.5B/4.5B_base_config.json \ + --vae_path /path/to/ckpt/vae \ + --t5_path /path/to/t5-v1_1-xxl \ + --special_tokens_path /path/to/MAGI-1/example/assets/special_tokens.npz \ + --output_path /path/to/magi-base-pipeline +``` + +The converter saves T5, the tokenizer, learned null-caption features, and the official HQ/duration vectors with the +models. Runtime inference does not import the official repository or read its special-token archive. +The script converts sharded T5 `.bin` checkpoints with `torch.load(weights_only=True)` into temporary safetensors +before loading them. Only convert trusted checkpoints. The destination must be empty. After saving, the script +reloads all model components and verifies their tensor values exactly. Use the `transformer` or `vae` subcommand +to convert those components separately. + +```python +import torch +from diffusers import MagiModularPipeline +from diffusers.utils import export_to_video + +pipeline = MagiModularPipeline.from_pretrained("/path/to/magi-base-pipeline") +pipeline.load_components(dtype={ + "default": torch.float32, + "transformer": torch.bfloat16, + "vae": torch.bfloat16, +}) +pipeline.to("cuda") +pipeline.vae.enable_tiling(tile_sample_min_length=12) + +videos = pipeline( + prompt="A cat walks through a sunlit garden.", + height=512, + width=512, + num_frames=96, + num_inference_steps=64, + generator=torch.Generator("cpu").manual_seed(42), + output_type="np", + output="videos", +) +export_to_video(videos[0], "magi.mp4", fps=24) +``` + +This example requires enough memory for all components. For smaller devices, arrange component offloading before +inference. Keep T5, text-conditioning features, initial noise, and Euler state in FP32. Load the Transformer with +the desired dtype instead of casting the entire pipeline: some Transformer weights must remain FP32. +The decoder uses autocast when the VAE runs in FP16 or BF16, matching the reference decode precision context. +For official CUDA BF16 comparisons, select `pipeline.transformer.set_attention_backend("flash_varlen")` and +`pipeline.vae.set_attention_backend("flash")`; use the same seed and generator device as the reference. + +Prompt cleaning requires `ftfy` and `beautifulsoup4`. The default follows the official two-pass cleaning, including +its removal of CJK characters. Set `clean_caption=False` to use only lowercasing and whitespace trimming. +T5 processes each prompt separately and pads to 800 tokens. Each chunk receives +`[duration, HQ, T5 tokens...]`, truncated back to 800 positions. Duration counts remaining chunks and saturates at +eight. Unconditional features come from the learned null caption, with the first 50 tokens unmasked; no negative +prompt is encoded. + +Height and width must be compatible with both the VAE compression ratio and Transformer patch size. +`num_frames` must be divisible by the temporal compression ratio; generation rounds up to complete latent chunks. +The defaults are six latent frames per chunk and four active chunks per window. `num_images_per_prompt` +controls how many videos are generated per prompt. + +The top-level blocks are `text_encoder`, `prepare_latents`, `denoise`, and `decode`. They can be initialized +independently with `block.init_pipeline(checkpoint_path)` and `load_components()`. +The decoder divides latents by 0.18215 and decodes each generated chunk independently. Enable VAE tiling before +running large videos. Outputs follow Diffusers conventions: `pt` is normalized float video with shape +`(batch, frames, channels, height, width)`, `np` uses channels last, `pil` is a list of frame lists, and +`latent` bypasses decoding. Standard PIL conversion is not the original MAGI uint8 truncation path; compressed +video bytes are not a parity target. + +## Prepared-latent denoising + +`MagiDenoiseStep` generates latents for MAGI-1 base models. It accepts prepared text features and initial noise; +it does not encode prompts, encode prefix videos, or decode generated latents. Distilled models and prefixes that +do not contain a whole number of latent chunks are not supported by this block. + +### Usage + +Load a converted `MagiTransformer3DModel` checkpoint and provide it alongside the scheduler and guider: + +```python +import torch +from diffusers import MagiClassifierFreeGuidance, MagiDenoiseStep, MagiEulerScheduler, MagiTransformer3DModel + +transformer = MagiTransformer3DModel.from_pretrained(transformer_path, torch_dtype=torch.bfloat16).to("cuda") +pipeline = MagiDenoiseStep().init_pipeline() +pipeline.update_components( + transformer=transformer, + scheduler=MagiEulerScheduler(), + guider=MagiClassifierFreeGuidance(), +) +result = pipeline( + latents=initial_latents, + prompt_embeds=prompt_embeds, + prompt_attention_mask=prompt_attention_mask, + negative_prompt_embeds=null_embeds, + negative_prompt_attention_mask=null_mask, + num_inference_steps=64, + chunk_width=6, + window_size=4, + output=["latents", "clean_kv_cache", "completed_chunks"], +) +``` + +The example assumes that the input tensors are already on the Transformer's execution device. Keep initial noise +and sampling state in FP32. The Transformer handles its internal mixed precision; do not cast the whole model +with `.bfloat16()`, which would also cast weights that need FP32. + +The official mixed-precision parity checks use the `flash_varlen` attention backend. If Flash Attention is +installed, select it with `transformer.set_attention_backend("flash_varlen")` before running the pipeline. +The default native backend remains available, but its BF16 arithmetic is not guaranteed to match the official +Flash Attention path bit for bit. + +`initial_latents` has shape `(batch, channels, chunks * chunk_width, height, width)`. Conditional text features have +shape `(batch, length, caption_channels)` or `(batch, chunks, length, caption_channels)`, with a matching boolean +keep-mask. The latter form supports chunk-dependent text features prepared by the caller. + +The negative inputs are the model's **learned null-caption features**, not an empty prompt passed through a text +encoder. They have shape `(batch, length, caption_channels)` and `(batch, length)` and are shared across chunks. +Use the same padded text length for conditional and null features. The official text preparation uses 800 padded +tokens and keeps the first 50 tokens for the null caption. Text preprocessing and special tokens must be prepared +separately; this block does not reconstruct them. + +### Window and guidance + +`num_inference_steps` counts Euler updates per generated chunk. It must be divisible by both `window_size` and the +length of `noise2clean_kvrange`, whose default is `(5, 4, 3, 2)`. Each entry is a positive attention-window size in +chunks, including the current chunk. The active generation window first expands, then moves forward, then shrinks. +For `K` total chunks, `P` prefix chunks, `W` window size, and `N` updates per chunk, there are +`(N // W) * (K + W - 1 - P)` window iterations. + +The guider combines three velocity predictions at each active chunk's own timestep: + +```python +velocity = (1 - prefix_scale) * independent_velocity +velocity += (prefix_scale - text_scale) * prefix_velocity + text_scale * text_and_prefix_velocity +``` + +The independent branch treats each active chunk as a separate batch item, with no cached prefix and temporal +positions starting at zero. Configure thresholds and scales through `MagiClassifierFreeGuidance`, not through +pipeline inputs. Guidance is applied before the FP32 Euler update. + +### Clean-prefix cache + +Optionally pass `prefix_latents` with shape `(batch, channels, prefix_chunks * chunk_width, height, width)`. +These replace the corresponding leading slots of `initial_latents` and remain unchanged. At least one chunk must +remain to generate. Caller-owned latent tensors are not modified. + +Initial prefix KV is computed with null-caption conditioning at `clean_t=0.9999`. When a generated chunk leaves +the active window, the next iteration recomputes it at `clean_t`. Both prefix-conditioned branches read the same +previous cache; only the null-text branch supplies newly reusable KV. Each layer copies only the requested clean +prefix before processing the next layer, so full noisy-window caches do not accumulate across the model. +`clean_chunk_kvrange` defaults to 1 and must be positive. + +Set `cache_device="cpu"` on a pipeline call to keep clean KV on CPU. Each Transformer layer transfers only its own +input cache to the compute device and returns its updated clean slice to CPU. This changes storage, not attention +ranges or sampling math. The example maps the official engine config's `kv_offload` setting to this argument. +Component CPU offload and clean-KV offload are separate controls; long videos may need both. + +`clean_kv_cache` contains per-layer `(key, value)` tensors with shape `(batch, tokens, kv_heads, head_dim)`. It keeps +the complete clean prefix to preserve absolute temporal positions; there is no cache eviction. The final generated +chunk is not refreshed because no later chunk needs it, so the returned cache covers all but the last chunk (or is +`None` for a one-chunk video). It is an inspection output, not a resumable generation checkpoint. + +Every fresh pipeline call resets the schedule and prefix cache. `completed_chunks` includes supplied prefix chunks +and all finalized generated chunks. Only `result["latents"]` should be passed to later VAE decoding, with the +appropriate MAGI latent scaling applied separately. + +## MagiDenoiseStep + +[[autodoc]] MagiDenoiseStep + +## MagiClassifierFreeGuidance + +[[autodoc]] MagiClassifierFreeGuidance + + +## Image and video prefixes + +The same base checkpoint can run image- or video-conditioned continuation with `MagiImageToVideoBlocks` or +`MagiVideoToVideoBlocks`. Select the blockset explicitly; loading the existing checkpoint normally still selects +text-to-video. + +```python +from diffusers import MagiImageToVideoBlocks + +pipe = MagiImageToVideoBlocks().init_pipeline("/path/to/MAGI-1-diffusers") +pipe.load_components() +pipe.vae.enable_tiling(tile_sample_min_length=12) +# Configure precision, attention backends, and device/offload as in the example script. +videos = pipe( + prompt="Good Boy", + image=image_tensor, + height=256, + width=256, + num_frames=24, + output_type="pt", + output="videos", +) +``` + +The image input is pre-resized RGB `torch.uint8` with shape `(batch, 3, height, width)`. The video blockset instead +accepts `video` with shape `(batch, 3, frames, height, width)`. File decoding and fps resampling are outside the +core blocks. The example script `examples/magi/inference_magi.py` handles these with FFmpeg; `--image` selects I2V +and `--video` selects V2V. The video loader uses the first 32 frames after fps resampling, as in the official loader. + +The VAE encoder emits deterministic, scaled `conditioning_latents`. Complete prefix chunks initialize the clean +cache. Partial-prefix values are injected before every evaluation, including clean-cache refresh, but only active +chunks are updated by Euler integration. Finalized output is not overwritten by a later cache refresh. + +For these workflows, `num_frames` requests new frames. Prefix plus new frames rounds up to full latent chunks. +Prefix latents are removed before each output chunk is decoded, except for the first chunk of a one-latent-frame +image prefix. The official decoder returns one frame for a tile containing only one latent position; the prefix +decoder preserves that convention. Actual decoded length therefore depends on both chunk rounding and VAE tiling. +At chunk width 6 with 12-frame VAE tiles, an image plus 24 requested new frames produces 48 frames; a 32-frame video +prefix plus 24 requested new frames produces 37 frames. + +The `latents` intermediate retains prefix slots. `output_type="latent", output="videos"` returns the cropped latent +suffix, with the image-prefix exception above. The `vae_encoder`, `prepare_latents`, `denoise`, and `decode` blocks +are independently reusable through `init_pipeline()`. The standalone `MagiDenoiseStep` described above still accepts +only complete prefix chunks; use the new workflow's `denoise` block for partial prefixes. + +## MagiImageToVideoBlocks + +[[autodoc]] MagiImageToVideoBlocks + +## MagiVideoToVideoBlocks + +[[autodoc]] MagiVideoToVideoBlocks diff --git a/docs/source/en/api/schedulers/magi_euler.md b/docs/source/en/api/schedulers/magi_euler.md new file mode 100644 index 000000000000..9ca6291106df --- /dev/null +++ b/docs/source/en/api/schedulers/magi_euler.md @@ -0,0 +1,66 @@ +# MagiEulerScheduler + +MAGI-1 predicts flow velocity directly. Its sampler advances from noise at time 0 toward clean data at time 1: + +```python +next_sample = sample + velocity * (next_timestep - timestep) +``` + +Do not convert the Transformer output from a clean-sample prediction to velocity. Apply guidance to velocity before +calling `step`. Keep the sampling state in FP32; the Transformer manages its internal mixed-precision computation. + +## Time schedule + +The default schedule squares a uniform grid and then applies the official inverse shift of 3. The scheduler also +supports the reference's square, piecewise and linear schedules, and both 12-step shortcut grid orderings. +`num_inference_steps` is the number of updates **per chunk**, not the total number of sliding-window model calls. + +`timesteps` contains normalized model evaluation times without a factor of 1000. `timestep_schedule` includes the +last integration endpoint. Build the grid on the execution device to match the reference's arithmetic. The FP32 +default endpoint can be slightly above 1; it is intentionally not clamped or reconstructed from `1 - sigma`. + +## Sequential and chunk-wise updates + +For a single chunk, use sequential steps: + +```python +from diffusers import MagiEulerScheduler + +scheduler = MagiEulerScheduler() +scheduler.set_timesteps(64, device=latents.device) +for timestep in scheduler.timesteps: + velocity = transformer( + latents, prompt_embeds, timestep.expand(latents.shape[0]) + ).sample + latents = scheduler.step(velocity, timestep, latents).prev_sample +``` + +This illustrates the scheduler interface, not the complete MAGI generation loop: text preparation, guidance and +chunk-window management must be supplied separately. + +For an active window of equally sized video chunks, pass both endpoint tensors: + +```python +latents = scheduler.step( + guided_velocity, + timestep=current_chunk_times, + sample=latents, + next_timestep=next_chunk_times, +).prev_sample +``` + +Endpoints can be scalars, `(chunks,)` shared across the batch, or `(batch, chunks)`. A one-dimensional tensor indexes +chunks, not batch items. Explicit endpoint updates leave `step_index` unchanged so chunks can follow different +parts of the same schedule. Sequential calls advance it; `set_timesteps` resets it. Outputs always remain FP32. + +The scheduler does not select active chunks, construct attention ranges, manage prefix caches or apply three-way +guidance. The 12-step time grid alone is not a complete distilled workflow; distillation conditioning and the +near-clean-chunk branch belong to the generation loop. + +## MagiEulerScheduler + +[[autodoc]] MagiEulerScheduler + +## MagiEulerSchedulerOutput + +[[autodoc]] schedulers.scheduling_magi_euler.MagiEulerSchedulerOutput diff --git a/examples/magi/README-prefix.md b/examples/magi/README-prefix.md new file mode 100644 index 000000000000..6c7a9a08712d --- /dev/null +++ b/examples/magi/README-prefix.md @@ -0,0 +1,64 @@ +# MAGI-1 image and video prefixes + +The same converted base checkpoint supports T2V, I2V, and V2V. The default saved workflow remains T2V. +Select `MagiImageToVideoBlocks` or `MagiVideoToVideoBlocks` explicitly to use a prefix. + +## Official example inputs + +Run from the repository root using the environment described in `README.md`: + +```bash +PYTHONPATH=src /path/to/magi-env/bin/python examples/magi/inference_magi.py \ + --model /path/to/MAGI-1-diffusers \ + --config /path/to/MAGI-1/example/4.5B/4.5B_base_config.json \ + --image /path/to/MAGI-1/example/assets/image.jpeg \ + --height 256 --width 256 --num-frames 24 \ + --output /path/to/new-i2v-output --save-latents +``` + +For V2V, replace `--image ...` with `--video /path/to/input.mp4`. The example uses FFmpeg to stretch frames to +the requested resolution. For video it resamples to the config's fps and uses at most the first 32 frames, +matching the official prefix loader. FFmpeg must be installed. Image and video inputs are mutually exclusive. +These small dimensions are smoke-test settings, not a visual-quality recommendation. + +## Python interface + +```python +from diffusers import MagiImageToVideoBlocks + +pipe = MagiImageToVideoBlocks().init_pipeline("/path/to/MAGI-1-diffusers") +pipe.load_components() +# Configure device, per-component precision, offload, and attention backends as in inference_magi.py. +pipe.vae.enable_tiling(tile_sample_min_length=12) +videos = pipe( + prompt="Good Boy", image=image_tensor, height=256, width=256, + num_frames=24, output_type="pt", output="videos", +) +``` + +The core image encoder takes pre-resized RGB `torch.uint8` pixels shaped `(batch, 3, height, width)`. +The video encoder takes `(batch, 3, frames, height, width)`. It encodes exactly the supplied frames; file +decoding, fps resampling, and the 32-frame limit belong to the example loader. One prefix may be broadcast +over several prompts. Otherwise its batch must match the prompt batch. Spatial dimensions must match +`height` and `width`; video length must satisfy the VAE's temporal patch requirements. + +The `vae_encoder` block can run independently and returns `conditioning_latents`. The `prepare_latents` +block expands these per-prompt latents to the generated video batch. The `denoise` block accepts the full +scaled prefix, including a partial final chunk. Complete prefix chunks initialize the clean KV cache; +partial-prefix values are reinjected at every model evaluation. The partial prefix still participates in the +Euler update, as in the official sampler: it is not permanently frozen in the output state. + +`num_frames` requests **new** frames. Prefix plus requested frames rounds up to a full latent chunk. V2V +omits prefix latents before decoding each output chunk. A one-latent-frame prefix retains the first four +decoded frames, matching official I2V behavior. Thus, at the default chunk width of 6: + +- I2V with 24 requested new frames returns 48 frames. +- V2V with a 32-frame prefix and 24 requested new frames returns 37 frames with 12-frame VAE tiles. + +The official VAE emits only one image frame for a final tile containing one latent position. For V2V's +first cropped chunk, this can make the decoded output shorter than four times its latent length. The prefix +decoder preserves this behavior rather than forcing a frame count. + +The `latents` intermediate retains prefix slots. `output_type="latent", output="videos"` returns the cropped +output suffix, with the I2V exception above. `--save-latents` saves full latents, conditioning latents, and the +exact input pixels for reference comparisons. diff --git a/examples/magi/README.md b/examples/magi/README.md new file mode 100644 index 000000000000..6dc1ab13be3c --- /dev/null +++ b/examples/magi/README.md @@ -0,0 +1,69 @@ +# MAGI-1 base inference + +This example loads a converted Diffusers checkpoint and uses the sampling parameters from an official MAGI +base config. It does not import the official implementation. The default prompt, `Good Boy`, is the prompt in +the official 4.5B `example/4.5B/run.sh`. + +## Environment + +The tested base environment is Python 3.10, PyTorch 2.5.1+cu121, torchvision 0.20.1+cu121, and Flash Attention +2.7.4.post1 on CUDA. Flash Attention must be built for the installed PyTorch/CUDA combination. +The requirements file pins the additional Python dependencies. It does not install or replace PyTorch or +Flash Attention. To preserve an existing compatible environment, create a separate overlay from its Python: + +```bash +/path/to/diffusion/bin/python -m venv --system-site-packages /path/to/magi-env +/path/to/magi-env/bin/python -m pip install -r examples/magi/requirements.txt +``` + +Run commands below from the Diffusers repository root with `PYTHONPATH=src`. This ensures the local MAGI +implementation is used without changing the installed Diffusers package in the base environment. +The overlay still depends on its base environment; it is not a standalone container or a complete system lock. + +For repository checks, install `requirements-dev.txt` instead and activate the overlay so subprocesses find +the pinned Ruff and documentation builder: + +```bash +/path/to/magi-env/bin/python -m pip install -r examples/magi/requirements-dev.txt +source /path/to/magi-env/bin/activate +make quality +``` + +## Official text-to-video example + +First convert the original checkpoint with `scripts/convert_magi_to_diffusers.py pipeline` as described in the +MAGI modular pipeline documentation. Then run: + +```bash +PYTHONPATH=src /path/to/magi-env/bin/python examples/magi/inference_magi.py \ + --model /path/to/MAGI-1-diffusers \ + --config /path/to/MAGI-1/example/4.5B/4.5B_base_config.json \ + --output /path/to/new-output-directory +``` + +The official 4.5B base config selects 720 × 720, 96 frames, 64 steps, seed 1234, and 24 fps. The example maps +the config's guidance thresholds/scales and clean-prefix attention settings to the Diffusers components. +Official checkpoint paths and distributed-engine settings in the config are not used: `--model` selects the +converted weights, and this example runs on one GPU with component CPU offload. + +Use `--device cuda:1` to select another GPU. `--height`, `--width`, `--num-frames`, `--seed`, and `--prompt` +explicitly override the official sampling inputs. Overrides are recorded in the output metadata. +For example, `--height 512 --width 512 --num-frames 192 --seed 1235` tests a longer clip. + +The output directory must be new or empty. It receives `settings.json`, `output_t2v.mp4`, and `metrics.json`. +Use `--save-latents` to also retain the final latent tensor. Generation time includes text encoding, denoising, +and VAE decoding, but excludes loading and MP4 encoding. MP4 frame count and finite/range checks run before +the final metrics are written. A successful run is not, by itself, a visual-quality or official numerical-parity +claim. + +## Precision and scope + +T5 and sampling state remain FP32; the Transformer and VAE load in BF16 with the model's FP32 exceptions. +The Transformer uses `flash_varlen`, the VAE uses `flash`, and TF32 is disabled. VAE temporal tiles use +`fps // 2` input frames (12 at 24 fps), matching the official pipeline entry point rather than the VAE helper's +16-frame default. The temporal tile size is recorded in the output metadata. +Changing backend, precision, or generator device can change the output. + +This entry point supports non-quantized base T2V, I2V, and V2V. See [README-prefix.md](README-prefix.md) for +image/video prefix inputs and output-length conventions. Distilled sampling, FP8, and distributed execution +are not supported by this example. diff --git a/examples/magi/inference_magi.py b/examples/magi/inference_magi.py new file mode 100644 index 000000000000..0fb25cbe3d2c --- /dev/null +++ b/examples/magi/inference_magi.py @@ -0,0 +1,210 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import argparse +import importlib.metadata +import json +import subprocess +import time +from pathlib import Path + +import imageio.v2 as imageio +import numpy as np +import torch + +from diffusers import ( + ComponentsManager, + MagiClassifierFreeGuidance, + MagiImageToVideoBlocks, + MagiVideoToVideoBlocks, + ModularPipeline, +) +from diffusers.utils import export_to_video + + +def main(): + parser = argparse.ArgumentParser(description="Run MAGI base generation using an official example config.") + parser.add_argument("--model", required=True) + parser.add_argument("--config", type=Path, required=True) + parser.add_argument("--prompt", default="Good Boy") + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--device", default="cuda:0") + parser.add_argument("--height", type=int) + parser.add_argument("--width", type=int) + parser.add_argument("--num-frames", type=int) + parser.add_argument("--seed", type=int) + parser.add_argument("--save-latents", action="store_true") + prefix = parser.add_mutually_exclusive_group() + prefix.add_argument("--image", type=Path) + prefix.add_argument("--video", type=Path) + args = parser.parse_args() + if args.output.exists() and (not args.output.is_dir() or any(args.output.iterdir())): + parser.error("Choose a new or empty output directory.") + config = json.loads(args.config.read_text()) + if config["engine_config"]["distill"] or config["engine_config"].get("fp8_quant", False): + parser.error("This example supports non-quantized base checkpoints only.") + runtime = config["runtime_config"] + if runtime["cfg_number"] != 3: + parser.error("MAGI base requires three-way guidance.") + height = runtime["video_size_h"] if args.height is None else args.height + width = runtime["video_size_w"] if args.width is None else args.width + num_frames = runtime["num_frames"] if args.num_frames is None else args.num_frames + seed = runtime["seed"] if args.seed is None else args.seed + torch.set_num_threads(4) + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + manager = ComponentsManager() + manager.enable_auto_cpu_offload(device=args.device) + workflow = "i2v" if args.image else "v2v" if args.video else "t2v" + pixels = {} + if args.image or args.video: + filters = f"scale={width}:{height}" + if args.video: + filters = f"fps={runtime['fps']}," + filters + decoded = subprocess.run( + [ + "ffmpeg", + "-v", + "error", + "-i", + str(args.image or args.video), + "-vf", + filters, + "-frames:v", + "1" if args.image else "32", + "-f", + "rawvideo", + "-pix_fmt", + "rgb24", + "pipe:1", + ], + check=True, + stdout=subprocess.PIPE, + ).stdout + frames = np.frombuffer(decoded, dtype=np.uint8).reshape(-1, height, width, 3).copy() + if len(frames) == 0: + raise ValueError("The prefix contains no decodable frames.") + tensor = torch.from_numpy(frames).permute(3, 0, 1, 2).unsqueeze(0) + pixels = {"image": tensor[:, :, 0]} if args.image else {"video": tensor} + blocks = MagiImageToVideoBlocks() if args.image else MagiVideoToVideoBlocks() + pipe = blocks.init_pipeline(args.model, components_manager=manager) + else: + pipe = ModularPipeline.from_pretrained(args.model, components_manager=manager) + pipe.load_components(dtype={"default": torch.float32, "transformer": torch.bfloat16, "vae": torch.bfloat16}) + pipe.update_components( + guider=MagiClassifierFreeGuidance( + timestep_thresholds=runtime["cfg_t_range"], + prefix_scales=runtime["prev_chunk_scales"], + text_scales=runtime["text_scales"], + ) + ) + if pipe.config.latent_scaling_factor != runtime["scale_factor"]: + raise ValueError("The checkpoint and official config have different latent scaling factors.") + pipe.transformer.set_attention_backend("flash_varlen") + pipe.vae.set_attention_backend("flash") + pipe.vae.enable_tiling(tile_sample_min_length=runtime["fps"] // 2) + args.output.mkdir(parents=True, exist_ok=True) + call = { + "prompt": args.prompt, + "height": height, + "width": width, + "num_frames": num_frames, + "num_inference_steps": runtime["num_steps"], + "chunk_width": runtime["chunk_width"], + "window_size": runtime["window_size"], + "noise2clean_kvrange": runtime["noise2clean_kvrange"], + "clean_chunk_kvrange": runtime["clean_chunk_kvrange"], + "clean_t": runtime["clean_t"], + "cache_device": "cpu" if config["engine_config"].get("kv_offload", False) else None, + "output_type": "pt", + } + metadata = { + "model": args.model, + "workflow": workflow, + "prefix_path": str(args.image or args.video) if pixels else None, + "prefix_frames": int(tensor.shape[2]) if pixels else 0, + "official_config": str(args.config), + "call": call, + "seed": seed, + "fps": runtime["fps"], + "vae_tile_sample_min_length": runtime["fps"] // 2, + "device": torch.cuda.get_device_name(args.device), + "versions": { + name: importlib.metadata.version(name) + for name in ("torch", "transformers", "huggingface-hub", "accelerate", "flash-attn") + }, + } + (args.output / "settings.json").write_text(json.dumps(metadata, indent=2)) + print(json.dumps(metadata, indent=2), flush=True) + started = time.perf_counter() + count = 0 + + def progress(module, inputs): + nonlocal count + count += 1 + if count % 24 == 1: + print(f"Transformer call {count}, elapsed {time.perf_counter() - started:.1f}s", flush=True) + + handle = pipe.transformer.register_forward_pre_hook(progress) + torch.cuda.reset_peak_memory_stats(args.device) + try: + result = pipe( + **call, + **pixels, + generator=torch.Generator(args.device).manual_seed(seed), + output=["videos", "latents", "completed_chunks"] + (["conditioning_latents"] if pixels else []), + ) + except Exception as error: + metadata.update( + status="failed", + error_type=type(error).__name__, + error=str(error), + elapsed_seconds=time.perf_counter() - started, + ) + (args.output / "failure.json").write_text(json.dumps(metadata, indent=2)) + raise + finally: + handle.remove() + torch.cuda.synchronize(args.device) + elapsed = time.perf_counter() - started + video = result["videos"].cpu() + assert video.isfinite().all() and result["latents"].isfinite().all() + assert video.min() >= 0 and video.max() <= 1 + assert video.shape[0] == 1 and video.shape[2:] == (3, height, width) + if args.save_latents: + torch.save(result["latents"].cpu(), args.output / "latents.pt") + if pixels: + torch.save(result["conditioning_latents"].cpu(), args.output / "conditioning_latents.pt") + torch.save(pixels, args.output / "prefix_pixels.pt") + video_path = args.output / f"output_{workflow}.mp4" + export_to_video(video[0].permute(0, 2, 3, 1).numpy(), str(video_path), fps=runtime["fps"]) + reader = imageio.get_reader(video_path) + assert reader.count_frames() == video.shape[1] + reader.close() + metadata.update( + elapsed_seconds=elapsed, + peak_allocated_gib=torch.cuda.max_memory_allocated(args.device) / 2**30, + transformer_calls=count, + completed_chunks=result["completed_chunks"], + shape=list(video.shape), + video_std=video.std().item(), + mean_frame_difference=(video[:, 1:] - video[:, :-1]).abs().mean().item(), + finite=True, + ) + (args.output / "metrics.json").write_text(json.dumps(metadata, indent=2)) + print(json.dumps(metadata, indent=2), flush=True) + + +if __name__ == "__main__": + main() diff --git a/examples/magi/requirements-dev.txt b/examples/magi/requirements-dev.txt new file mode 100644 index 000000000000..261a5fed2614 --- /dev/null +++ b/examples/magi/requirements-dev.txt @@ -0,0 +1,3 @@ +-r requirements.txt +hf-doc-builder==0.5.0 +ruff==0.9.10 diff --git a/examples/magi/requirements.txt b/examples/magi/requirements.txt new file mode 100644 index 000000000000..21ccf001623a --- /dev/null +++ b/examples/magi/requirements.txt @@ -0,0 +1,14 @@ +accelerate==1.3.0 +beautifulsoup4==4.15.0 +ffmpeg-python==0.2.0 +ftfy==6.3.1 +huggingface-hub==1.31.0 +imageio==2.37.4 +imageio-ffmpeg==0.6.0 +numpy==1.26.4 +peft==0.20.0 +safetensors==0.5.2 +sentencepiece==0.2.2 +socksio==1.0.0 +tokenizers==0.22.2 +transformers==5.3.0 diff --git a/scripts/convert_magi_to_diffusers.py b/scripts/convert_magi_to_diffusers.py new file mode 100644 index 000000000000..f4102e0b9847 --- /dev/null +++ b/scripts/convert_magi_to_diffusers.py @@ -0,0 +1,264 @@ +# Copyright 2025 SandAI and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import argparse +import json +import shutil +import tempfile +from pathlib import Path + +import numpy as np +import torch +from safetensors.torch import load_file, save_file +from transformers import AutoTokenizer, T5EncoderModel + +from diffusers import ( + AutoencoderKLMagi, + MagiClassifierFreeGuidance, + MagiEulerScheduler, + MagiTextConditioningModel, + MagiTextToVideoBlocks, + MagiTransformer3DModel, + ModularPipeline, +) + + +def convert_magi_transformer_config(config): + if config["engine_config"].get("fp8_quant", False): + raise ValueError("Convert the non-quantized MAGI checkpoint; official FP8 weights are not supported.") + model = config["model_config"] + if model["xattn_cond_hidden_ratio"] != 1 or model["cond_gating_ratio"] != 1: + raise ValueError("Only the published MAGI attention and gating ratios are supported.") + if model["hidden_size"] != model["num_attention_heads"] * model["kv_channels"]: + raise ValueError("hidden_size must equal num_attention_heads * kv_channels.") + duplicate = model["half_channel_vae"] + return { + "in_channels": model["in_channels"] // (2 if duplicate else 1), + "out_channels": model["out_channels"] // (2 if duplicate else 1), + "num_layers": model["num_layers"], + "num_attention_heads": model["num_attention_heads"], + "num_key_value_heads": model["num_query_groups"], + "attention_head_dim": model["kv_channels"], + "ffn_dim": model["ffn_hidden_size"], + "condition_dim": int(model["hidden_size"] * model["cond_hidden_ratio"]), + "caption_channels": model["caption_channels"], + "caption_max_length": model["caption_max_length"], + "patch_size": (model["t_patch_size"], model["patch_size"], model["patch_size"]), + "gated_linear_unit": model["gated_linear_unit"], + "norm_eps": model["layernorm_epsilon"], + "zero_centered_gamma": model["apply_layernorm_1p"], + "x_rescale_factor": model["x_rescale_factor"], + "duplicate_channels": duplicate, + "distilled": config["engine_config"]["distill"], + } + + +def convert_magi_transformer_state_dict(state_dict): + converted = {} + for name, tensor in state_dict.items(): + if name.startswith(("t_embedder.", "y_embedder.")): + name = "condition_embedder." + name + name = name.replace("videodit_blocks.layers.", "transformer_blocks.") + name = name.replace("videodit_blocks.final_layernorm.", "final_layernorm.") + name = name.replace("final_linear.linear.", "final_linear.") + if name.endswith("linear_kv_xattn.weight"): + for index, part in enumerate(tensor.chunk(8, dim=0)): + converted[name.replace(".weight", f".projections.{index}.weight")] = part.contiguous() + else: + converted[name] = tensor + return converted + + +def convert_magi_transformer(checkpoint_path, config_path): + checkpoint_path = Path(checkpoint_path) + with Path(config_path).open() as handle: + config = convert_magi_transformer_config(json.load(handle)) + index_path = checkpoint_path / "model.safetensors.index.json" + if index_path.exists(): + with index_path.open() as handle: + shards = sorted(set(json.load(handle)["weight_map"].values())) + else: + shards = ["model.safetensors"] + state_dict = {} + for shard in shards: + state_dict.update(convert_magi_transformer_state_dict(load_file(checkpoint_path / shard))) + with torch.device("meta"): + model = MagiTransformer3DModel(**config) + model.load_state_dict(state_dict, strict=True, assign=True) + return model.eval() + + +def convert_magi_vae_config(config): + if config.get("model_type", "vit") != "vit": + raise ValueError("Only the official vit VAE is supported.") + config = config["ddconfig"] + required = {"double_z": True, "ln_in_attn": True, "qkv_bias": True, "conv_last_layer": True} + for name, value in required.items(): + if config.get(name) != value: + raise ValueError(f"Unsupported MAGI VAE configuration: {name} must be {value}.") + for name in ("norm_code", "use_rope", "use_final_proj"): + if config.get(name, False): + raise ValueError(f"Unsupported MAGI VAE configuration: {name} must be False.") + if not config.get("with_cls_token", True): + raise ValueError("The MAGI VAE requires a CLS token.") + return { + "in_channels": config["in_chans"], + "out_channels": 3, + "latent_channels": config["z_chans"], + "embed_dim": config["embed_dim"], + "num_layers": config["depth"], + "num_attention_heads": config["num_heads"], + "mlp_ratio": config["mlp_ratio"], + "patch_size": config["patch_size"], + "patch_length": config["patch_length"], + "sample_size": config["video_size"], + "sample_frames": config["video_length"], + } + + +def convert_magi_vae_state_dict(state_dict): + converted = {} + for name, tensor in state_dict.items(): + if name.endswith((".cls_token", ".pos_embed")): + component, parameter = name.split(".") + name = f"{component}.position_embedding.{parameter}" + converted[name] = tensor + return converted + + +def convert_magi_vae(checkpoint_path): + checkpoint_path = Path(checkpoint_path) + with (checkpoint_path / "config.json").open() as config_file: + config = convert_magi_vae_config(json.load(config_file)) + state_dict = convert_magi_vae_state_dict(load_file(checkpoint_path / "diffusion_pytorch_model.safetensors")) + with torch.device("meta"): + model = AutoencoderKLMagi(**config) + model.load_state_dict(state_dict, strict=True, assign=True) + return model.eval() + + +def load_t5(t5_path): + t5_path = Path(t5_path) + index_path = t5_path / "pytorch_model.bin.index.json" + if not index_path.exists() or (t5_path / "model.safetensors.index.json").exists(): + return T5EncoderModel.from_pretrained(t5_path, torch_dtype=torch.float32) + with tempfile.TemporaryDirectory(prefix="magi-t5-") as temporary: + target = Path(temporary) + index = json.loads(index_path.read_text()) + mapping = {} + for shard in sorted(set(index["weight_map"].values())): + print(f"Converting T5 shard: {shard}", flush=True) + state = torch.load(t5_path / shard, map_location="cpu", weights_only=True, mmap=True) + state = {key: value.clone().contiguous() for key, value in state.items()} + name = Path(shard).with_suffix(".safetensors").name + save_file(state, target / name, metadata={"format": "pt"}) + del state + mapping.update({key: name for key, value in index["weight_map"].items() if value == shard}) + (target / "model.safetensors.index.json").write_text( + json.dumps({"metadata": index.get("metadata", {}), "weight_map": mapping}) + ) + shutil.copy2(t5_path / "config.json", target / "config.json") + return T5EncoderModel.from_pretrained(target, torch_dtype=torch.float32) + + +def convert_text_conditioning(transformer, special_tokens_path): + if transformer.config.distilled: + raise ValueError("This pipeline supports base checkpoints only.") + model = MagiTextConditioningModel( + caption_channels=transformer.config.caption_channels, + caption_max_length=transformer.config.caption_max_length, + ) + with np.load(special_tokens_path, allow_pickle=False) as features: + other = torch.from_numpy(features["other_tokens"].astype(np.float16)).float() + special = torch.cat([other[1:2], other[7:15]], dim=0) + with torch.no_grad(): + model.null_embedding.weight.copy_(transformer.condition_embedder.y_embedder.null_caption_embedding) + model.special_embedding.weight.copy_(special) + return model.eval() + + +def convert_pipeline(transformer, vae, t5_path, special_tokens_path): + text_conditioning = convert_text_conditioning(transformer, special_tokens_path) + pipe = MagiTextToVideoBlocks().init_pipeline() + pipe.update_components( + transformer=transformer, + vae=vae, + text_encoder=load_t5(t5_path), + tokenizer=AutoTokenizer.from_pretrained(t5_path), + text_conditioning=text_conditioning, + scheduler=MagiEulerScheduler(), + guider=MagiClassifierFreeGuidance(), + ) + pipe.load_components() + return pipe + + +def main(): + parser = argparse.ArgumentParser(description="Convert official MAGI weights and verify the saved components.") + subparsers = parser.add_subparsers(dest="component", required=True) + for component in ("pipeline", "transformer", "vae"): + command = subparsers.add_parser(component) + command.add_argument("--output_path", type=Path, required=True) + if component == "pipeline": + command.add_argument("--transformer_path", type=Path, required=True) + command.add_argument("--vae_path", type=Path, required=True) + command.add_argument("--t5_path", type=Path, required=True) + command.add_argument("--special_tokens_path", type=Path, required=True) + else: + command.add_argument("--checkpoint_path", type=Path, required=True) + if component != "vae": + command.add_argument("--config_path", type=Path, required=True) + args = parser.parse_args() + if args.output_path.exists() and (not args.output_path.is_dir() or any(args.output_path.iterdir())): + parser.error("output_path must be an empty directory to avoid overwriting existing files.") + if args.component == "pipeline": + with args.config_path.open() as handle: + if json.load(handle)["engine_config"]["distill"]: + parser.error("The complete pipeline supports base checkpoints only.") + model = convert_pipeline( + convert_magi_transformer(args.transformer_path, args.config_path), + convert_magi_vae(args.vae_path), + args.t5_path, + args.special_tokens_path, + ) + model.save_pretrained(str(args.output_path), max_shard_size="5GB", overwrite_modular_index=True) + restored = ModularPipeline.from_pretrained(str(args.output_path)) + restored.load_components(dtype=torch.float32) + pairs = [ + (name, component, restored.components[name]) + for name, component in model.components.items() + if isinstance(component, torch.nn.Module) + ] + assert model.tokenizer.get_vocab() == restored.tokenizer.get_vocab() + else: + model = ( + convert_magi_vae(args.checkpoint_path) + if args.component == "vae" + else convert_magi_transformer(args.checkpoint_path, args.config_path) + ) + model.save_pretrained(str(args.output_path), max_shard_size="5GB") + restored = type(model).from_pretrained(str(args.output_path), torch_dtype=torch.float32) + pairs = [(args.component, model, restored)] + for component_name, original, reloaded in pairs: + expected = original.state_dict() + actual = reloaded.state_dict() + assert expected.keys() == actual.keys() + for name, tensor in expected.items(): + torch.testing.assert_close(tensor.float(), actual[name].float(), rtol=0, atol=0) + print(f"Verified {component_name}: {len(expected)} tensors match exactly.", flush=True) + print(f"Saved and verified MAGI {args.component}: {args.output_path}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index 3ed956c75e49..45a9c412a982 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -190,6 +190,7 @@ "ClassifierFreeZeroStarGuidance", "FrequencyDecoupledGuidance", "LTX2Guidance", + "MagiClassifierFreeGuidance", "PerturbedAttentionGuidance", "SkipLayerGuidance", "SmoothedEnergyGuidance", @@ -248,6 +249,7 @@ "AutoencoderKLLTX2Audio", "AutoencoderKLLTX2Video", "AutoencoderKLLTXVideo", + "AutoencoderKLMagi", "AutoencoderKLMagvit", "AutoencoderKLMiniMaxH3", "AutoencoderKLMiniMaxH3Audio", @@ -314,6 +316,8 @@ "LTXVideoTransformer3DModel", "Lumina2Transformer2DModel", "LuminaNextDiT2DModel", + "MagiTextConditioningModel", + "MagiTransformer3DModel", "MiniMaxH3Transformer3DModel", "MiniMaxMusic3ConditionEncoder", "MiniMaxMusic3RVQDepthDecoder", @@ -456,6 +460,7 @@ "KDPM2DiscreteScheduler", "LCMScheduler", "LTXEulerAncestralRFScheduler", + "MagiEulerScheduler", "MiniMaxH3Scheduler", "PNDMScheduler", "RePaintScheduler", @@ -549,6 +554,14 @@ "LTX2ModularPipeline", "LTXAutoBlocks", "LTXModularPipeline", + "MagiDenoiseStep", + "MagiImageToVideoBlocks", + "MagiModularPipeline", + "MagiPrepareLatentsStep", + "MagiTextEncoderStep", + "MagiTextToVideoBlocks", + "MagiVaeDecoderStep", + "MagiVideoToVideoBlocks", "MiniMaxH3Blocks", "MiniMaxH3ModularPipeline", "MiniMaxMusic3Blocks", @@ -1071,6 +1084,7 @@ ClassifierFreeZeroStarGuidance, FrequencyDecoupledGuidance, LTX2Guidance, + MagiClassifierFreeGuidance, PerturbedAttentionGuidance, SkipLayerGuidance, SmoothedEnergyGuidance, @@ -1125,6 +1139,7 @@ AutoencoderKLLTX2Audio, AutoencoderKLLTX2Video, AutoencoderKLLTXVideo, + AutoencoderKLMagi, AutoencoderKLMagvit, AutoencoderKLMiniMaxH3, AutoencoderKLMiniMaxH3Audio, @@ -1191,6 +1206,8 @@ LTXVideoTransformer3DModel, Lumina2Transformer2DModel, LuminaNextDiT2DModel, + MagiTextConditioningModel, + MagiTransformer3DModel, MiniMaxH3Transformer3DModel, MiniMaxMusic3ConditionEncoder, MiniMaxMusic3RVQDepthDecoder, @@ -1329,6 +1346,7 @@ KDPM2DiscreteScheduler, LCMScheduler, LTXEulerAncestralRFScheduler, + MagiEulerScheduler, MiniMaxH3Scheduler, PNDMScheduler, RePaintScheduler, @@ -1405,6 +1423,14 @@ LTX25ModularPipeline, LTXAutoBlocks, LTXModularPipeline, + MagiDenoiseStep, + MagiImageToVideoBlocks, + MagiModularPipeline, + MagiPrepareLatentsStep, + MagiTextEncoderStep, + MagiTextToVideoBlocks, + MagiVaeDecoderStep, + MagiVideoToVideoBlocks, MiniMaxH3Blocks, MiniMaxH3ModularPipeline, MiniMaxMusic3Blocks, diff --git a/src/diffusers/guiders/__init__.py b/src/diffusers/guiders/__init__.py index fe2a07858e71..ae9cdc97b023 100644 --- a/src/diffusers/guiders/__init__.py +++ b/src/diffusers/guiders/__init__.py @@ -25,6 +25,7 @@ from .frequency_decoupled_guidance import FrequencyDecoupledGuidance from .guider_utils import BaseGuidance from .ltx2_guidance import LTX2Guidance + from .magi_classifier_free_guidance import MagiClassifierFreeGuidance from .magnitude_aware_guidance import MagnitudeAwareGuidance from .perturbed_attention_guidance import PerturbedAttentionGuidance from .skip_layer_guidance import SkipLayerGuidance diff --git a/src/diffusers/guiders/magi_classifier_free_guidance.py b/src/diffusers/guiders/magi_classifier_free_guidance.py new file mode 100644 index 000000000000..afb85b04145b --- /dev/null +++ b/src/diffusers/guiders/magi_classifier_free_guidance.py @@ -0,0 +1,104 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import math + +import torch + +from ..configuration_utils import register_to_config +from .guider_utils import BaseGuidance, GuiderOutput + + +class MagiClassifierFreeGuidance(BaseGuidance): + """ + Combine text-and-prefix, prefix-only, and independent-chunk velocities for MAGI base models. + + Args: + timestep_thresholds (`tuple[float]`, defaults to `(0.0, 0.0217, 0.1, 0.3, 0.999)`): + Increasing interval boundaries in normalized noise-to-clean time, starting at zero. + prefix_scales (`tuple[float]`, defaults to `(1.5, 1.5, 1.5, 1.0, 1.0)`): + Prefix guidance scales, one per timestep interval. + text_scales (`tuple[float]`, defaults to `(7.5, 7.5, 7.5, 0.0, 0.0)`): + Text guidance scales, one per timestep interval. + + Call `set_state` with times shaped `(batch, chunks)` before applying guidance. All three branches remain available + for clean-cache management, even when a guidance coefficient is zero. + """ + + _input_predictions = ["pred_cond", "pred_prefix", "pred_uncond"] + + @register_to_config + def __init__( + self, + timestep_thresholds=(0.0, 0.0217, 0.1, 0.3, 0.999), + prefix_scales=(1.5, 1.5, 1.5, 1.0, 1.0), + text_scales=(7.5, 7.5, 7.5, 0.0, 0.0), + ): + super().__init__() + if ( + not timestep_thresholds + or len(prefix_scales) != len(timestep_thresholds) + or len(text_scales) != len(timestep_thresholds) + ): + raise ValueError("Thresholds and guidance scales must have the same nonzero length.") + if not all( + math.isfinite(value) for values in (timestep_thresholds, prefix_scales, text_scales) for value in values + ): + raise ValueError("Thresholds and guidance scales must be finite.") + if timestep_thresholds[0] != 0 or any(a >= b for a, b in zip(timestep_thresholds, timestep_thresholds[1:])): + raise ValueError("Timestep thresholds must start at zero and increase strictly.") + + @property + def num_conditions(self): + return 3 + + @property + def is_conditional(self): + return self._count_prepared == 1 + + def prepare_inputs(self, data): + return [self._prepare_batch(data, i, name) for i, name in enumerate(self._input_predictions)] + + def prepare_inputs_from_block_state(self, data, input_fields): + return [ + self._prepare_batch_from_block_state(input_fields, data, i, name) + for i, name in enumerate(self._input_predictions) + ] + + def forward(self, pred_cond, pred_prefix, pred_uncond): + if self._timestep is None: + raise ValueError("Set the current chunk timesteps with set_state before applying guidance.") + times = self._timestep.to(device=pred_cond.device, dtype=torch.float32) + if times.ndim == 1: + times = times[None].expand(pred_cond.shape[0], -1) + if ( + times.ndim != 2 + or times.shape[0] != pred_cond.shape[0] + or times.shape[1] == 0 + or pred_cond.shape[2] % times.shape[1] + ): + raise ValueError("Guidance timesteps must match the batch and divide the latent frames into chunks.") + thresholds = torch.tensor(self.config.timestep_thresholds, device=times.device, dtype=torch.float32) + indices = torch.searchsorted(thresholds - 1e-7, times.contiguous()) - 1 + indices = indices.clamp(0, len(thresholds) - 1) + frames_per_chunk = pred_cond.shape[2] // times.shape[1] + prefix_scale = torch.tensor(self.config.prefix_scales, device=times.device, dtype=torch.float32)[indices] + text_scale = torch.tensor(self.config.text_scales, device=times.device, dtype=torch.float32)[indices] + prefix_scale = prefix_scale.repeat_interleave(frames_per_chunk, dim=1)[:, None, :, None, None] + text_scale = text_scale.repeat_interleave(frames_per_chunk, dim=1)[:, None, :, None, None] + pred = (1 - prefix_scale) * pred_uncond.float() + pred = pred + (prefix_scale - text_scale) * pred_prefix.float() + text_scale * pred_cond.float() + if not self._enabled: + pred = pred_cond + return GuiderOutput(pred=pred, pred_cond=pred_cond, pred_uncond=pred_uncond) diff --git a/src/diffusers/models/__init__.py b/src/diffusers/models/__init__.py index 8ba17d896434..d44f8e5aa08a 100755 --- a/src/diffusers/models/__init__.py +++ b/src/diffusers/models/__init__.py @@ -45,6 +45,7 @@ _import_structure["autoencoders.autoencoder_kl_ltx"] = ["AutoencoderKLLTXVideo"] _import_structure["autoencoders.autoencoder_kl_ltx2"] = ["AutoencoderKLLTX2Video"] _import_structure["autoencoders.autoencoder_kl_ltx2_audio"] = ["AutoencoderKLLTX2Audio"] + _import_structure["autoencoders.autoencoder_kl_magi"] = ["AutoencoderKLMagi"] _import_structure["autoencoders.autoencoder_kl_magvit"] = ["AutoencoderKLMagvit"] _import_structure["autoencoders.autoencoder_kl_minimax_h3"] = ["AutoencoderKLMiniMaxH3"] _import_structure["autoencoders.autoencoder_kl_minimax_h3_audio"] = ["AutoencoderKLMiniMaxH3Audio"] @@ -85,6 +86,7 @@ _import_structure["controlnets.multicontrolnet"] = ["MultiControlNetModel"] _import_structure["controlnets.multicontrolnet_union"] = ["MultiControlNetUnionModel"] _import_structure["embeddings"] = ["ImageProjection"] + _import_structure["magi_conditioning"] = ["MagiTextConditioningModel"] _import_structure["modeling_utils"] = ["ModelMixin"] _import_structure["transformers.ace_step_transformer"] = ["AceStepTransformer1DModel"] _import_structure["transformers.auraflow_transformer_2d"] = ["AuraFlowTransformer2DModel"] @@ -135,6 +137,7 @@ _import_structure["transformers.transformer_ltx"] = ["LTXVideoTransformer3DModel"] _import_structure["transformers.transformer_ltx2"] = ["LTX2VideoTransformer3DModel"] _import_structure["transformers.transformer_lumina2"] = ["Lumina2Transformer2DModel"] + _import_structure["transformers.transformer_magi"] = ["MagiTransformer3DModel"] _import_structure["transformers.transformer_minimax_h3"] = ["MiniMaxH3Transformer3DModel"] _import_structure["transformers.transformer_minimax_music3"] = ["MiniMaxMusic3Transformer1DModel"] _import_structure["transformers.transformer_mochi"] = ["MochiTransformer3DModel"] @@ -190,6 +193,7 @@ AutoencoderKLLTX2Audio, AutoencoderKLLTX2Video, AutoencoderKLLTXVideo, + AutoencoderKLMagi, AutoencoderKLMagvit, AutoencoderKLMiniMaxH3, AutoencoderKLMiniMaxH3Audio, @@ -232,6 +236,7 @@ ZImageControlNetModel, ) from .embeddings import ImageProjection + from .magi_conditioning import MagiTextConditioningModel from .modeling_utils import ModelMixin from .transformers import ( AceStepTransformer1DModel, @@ -276,6 +281,7 @@ LTXVideoTransformer3DModel, Lumina2Transformer2DModel, LuminaNextDiT2DModel, + MagiTransformer3DModel, MiniMaxH3Transformer3DModel, MiniMaxMusic3RVQDepthDecoder, MiniMaxMusic3Transformer1DModel, diff --git a/src/diffusers/models/autoencoders/__init__.py b/src/diffusers/models/autoencoders/__init__.py index 607704343743..ab4900d0928d 100644 --- a/src/diffusers/models/autoencoders/__init__.py +++ b/src/diffusers/models/autoencoders/__init__.py @@ -15,6 +15,7 @@ from .autoencoder_kl_ltx import AutoencoderKLLTXVideo from .autoencoder_kl_ltx2 import AutoencoderKLLTX2Video from .autoencoder_kl_ltx2_audio import AutoencoderKLLTX2Audio +from .autoencoder_kl_magi import AutoencoderKLMagi from .autoencoder_kl_magvit import AutoencoderKLMagvit from .autoencoder_kl_minimax_h3 import AutoencoderKLMiniMaxH3 from .autoencoder_kl_minimax_h3_audio import AutoencoderKLMiniMaxH3Audio diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_magi.py b/src/diffusers/models/autoencoders/autoencoder_kl_magi.py new file mode 100644 index 000000000000..074bd7b03778 --- /dev/null +++ b/src/diffusers/models/autoencoders/autoencoder_kl_magi.py @@ -0,0 +1,570 @@ +# Copyright 2025 SandAI and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from itertools import product + +import torch +from torch import nn +from torch.nn import functional as F + +from ...configuration_utils import ConfigMixin, register_to_config +from ...utils.accelerate_utils import apply_forward_hook +from ..attention import AttentionMixin, AttentionModuleMixin +from ..attention_dispatch import dispatch_attention_fn +from ..modeling_outputs import AutoencoderKLOutput +from ..modeling_utils import ModelMixin +from .vae import DecoderOutput, DiagonalGaussianDistribution + + +class MagiVAEManualLayerNorm(nn.Module): + def __init__(self, eps=1e-5): + super().__init__() + self.eps = eps + + def forward(self, x): + mean = x.mean(dim=-1, keepdim=True) + std = x.std(dim=-1, keepdim=True, unbiased=False) + x_normalized = (x - mean) / (std + self.eps) + return x_normalized + + +class MagiVAEAttnProcessor: + _attention_backend = None + _parallel_config = None + + def __call__(self, attn, hidden_states): + batch_size, sequence_length, dim = hidden_states.shape + qkv = attn.qkv(hidden_states).reshape(batch_size, sequence_length, 3, attn.num_heads, dim // attn.num_heads) + qkv = attn.qkv_norm(qkv) + query, key, value = qkv.unbind(dim=2) + hidden_states = dispatch_attention_fn( + query, key, value, backend=self._attention_backend, parallel_config=self._parallel_config + ) + return attn.proj(hidden_states.flatten(2, 3)) + + +class MagiVAEAttention(nn.Module, AttentionModuleMixin): + _supports_qkv_fusion = False + _default_processor_cls = MagiVAEAttnProcessor + _available_processors = [MagiVAEAttnProcessor] + + def __init__(self, dim, num_heads, eps): + super().__init__() + self.num_heads = num_heads + self.qkv_norm = MagiVAEManualLayerNorm(eps) + self.qkv = nn.Linear(dim, dim * 3) + self.proj = nn.Linear(dim, dim) + self.set_processor(MagiVAEAttnProcessor()) + + def forward(self, hidden_states): + return self.processor(self, hidden_states) + + +class MagiVAEMlp(nn.Module): + def __init__(self, in_features, hidden_features, out_features): + super().__init__() + self.fc1 = nn.Linear(in_features, hidden_features) + self.act = nn.GELU() + self.fc2 = nn.Linear(hidden_features, out_features) + + def forward(self, x): + x = self.fc1(x) + x = self.act(x) + x = self.fc2(x) + return x + + +class MagiVAEBlock(nn.Module): + def __init__(self, dim, num_heads, mlp_ratio, eps): + super().__init__() + self.attn = MagiVAEAttention(dim, num_heads, eps) + self.norm2 = nn.LayerNorm(dim, eps=eps) + self.mlp = MagiVAEMlp(in_features=dim, hidden_features=int(dim * mlp_ratio), out_features=dim) + + def forward(self, hidden_states): + hidden_states = hidden_states + self.attn(hidden_states) + hidden_states = hidden_states + self.mlp(self.norm2(hidden_states)) + return hidden_states + + +class MagiVAEPatchEmbed(nn.Module): + def __init__(self, in_chans, embed_dim, patch_size, patch_length): + super().__init__() + self.proj = nn.Conv3d( + in_chans, + embed_dim, + kernel_size=(patch_length, patch_size, patch_size), + stride=(patch_length, patch_size, patch_size), + ) + + def forward(self, x): + """Project a video onto a five-dimensional patch feature grid.""" + x = self.proj(x) + return x + + +def resize_pos_embed(posemb, src_shape, target_shape): + posemb = posemb.reshape(1, src_shape[0], src_shape[1], src_shape[2], -1) + posemb = posemb.permute(0, 4, 1, 2, 3) + posemb = F.interpolate(posemb, size=target_shape, mode="trilinear", align_corners=False) + posemb = posemb.permute(0, 2, 3, 4, 1) + posemb = posemb.reshape(1, target_shape[0] * target_shape[1] * target_shape[2], -1) + return posemb + + +class MagiVAEPositionEmbedding(nn.Module): + def __init__(self, dim, latent_shape): + super().__init__() + self.latent_shape = latent_shape + self.cls_token = nn.Parameter(torch.zeros(1, 1, dim)) + self.pos_embed = nn.Parameter(torch.zeros(1, latent_shape[0] * latent_shape[1] * latent_shape[2] + 1, dim)) + + def forward(self, hidden_states, latent_shape): + cls_tokens = self.cls_token.expand(hidden_states.shape[0], -1, -1) + hidden_states = torch.cat((cls_tokens, hidden_states), dim=1) + if latent_shape != self.latent_shape: + pos_embed = resize_pos_embed(self.pos_embed[:, 1:], self.latent_shape, latent_shape) + pos_embed = torch.cat((self.pos_embed[:, :1], pos_embed), dim=1) + else: + pos_embed = self.pos_embed + return hidden_states + pos_embed + + +class MagiVAEEncoder(nn.Module): + def __init__(self, in_channels, latent_channels, dim, depth, num_heads, mlp_ratio, patch_shape, latent_shape, eps): + super().__init__() + self.patch_embed = MagiVAEPatchEmbed(in_channels, dim, patch_shape[1], patch_shape[0]) + self.position_embedding = MagiVAEPositionEmbedding(dim, latent_shape) + self.blocks = nn.ModuleList([MagiVAEBlock(dim, num_heads, mlp_ratio, eps) for _ in range(depth)]) + self.norm = nn.LayerNorm(dim, eps=eps) + self.last_layer = nn.Linear(dim, latent_channels * 2) + self.gradient_checkpointing = False + + def forward(self, sample): + hidden_states = self.patch_embed(sample) + batch_size, _, num_frames, height, width = hidden_states.shape + hidden_states = hidden_states.flatten(2).transpose(1, 2) + hidden_states = self.position_embedding(hidden_states, (num_frames, height, width)) + for block in self.blocks: + if torch.is_grad_enabled() and self.gradient_checkpointing: + hidden_states = self._gradient_checkpointing_func(block, hidden_states) + else: + hidden_states = block(hidden_states) + hidden_states = self.last_layer(self.norm(hidden_states))[:, 1:] + return hidden_states.reshape(batch_size, num_frames, height, width, -1).permute(0, 4, 1, 2, 3) + + +class MagiVAEDecoder(nn.Module): + def __init__( + self, out_channels, latent_channels, dim, depth, num_heads, mlp_ratio, patch_shape, latent_shape, eps + ): + super().__init__() + self.patch_shape = patch_shape + self.unpatch_channels = dim // (patch_shape[0] * patch_shape[1] * patch_shape[2]) + self.proj_in = nn.Linear(latent_channels, dim) + self.position_embedding = MagiVAEPositionEmbedding(dim, latent_shape) + self.blocks = nn.ModuleList([MagiVAEBlock(dim, num_heads, mlp_ratio, eps) for _ in range(depth)]) + self.norm = nn.LayerNorm(dim, eps=eps) + self.last_layer = nn.Conv3d(self.unpatch_channels, out_channels, kernel_size=3, padding=1) + self.gradient_checkpointing = False + + def forward(self, latent): + batch_size, _, num_frames, height, width = latent.shape + hidden_states = latent.permute(0, 2, 3, 4, 1).flatten(1, 3) + hidden_states = self.proj_in(hidden_states) + hidden_states = self.position_embedding(hidden_states, (num_frames, height, width)) + for block in self.blocks: + if torch.is_grad_enabled() and self.gradient_checkpointing: + hidden_states = self._gradient_checkpointing_func(block, hidden_states) + else: + hidden_states = block(hidden_states) + hidden_states = self.norm(hidden_states)[:, 1:] + hidden_states = hidden_states.reshape( + batch_size, num_frames, height, width, *self.patch_shape, self.unpatch_channels + ) + hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6).reshape( + batch_size, + self.unpatch_channels, + num_frames * self.patch_shape[0], + height * self.patch_shape[1], + width * self.patch_shape[2], + ) + return self.last_layer(hidden_states) + + +class AutoencoderKLMagi(ModelMixin, AttentionMixin, ConfigMixin): + r""" + Transformer VAE used by MAGI-1. + + The encoder returns a diagonal Gaussian posterior over video latents. Both encoder and decoder use full attention + over the input patch sequence. The learned position embeddings are interpolated for different video sizes. + + Parameters: + in_channels (`int`, defaults to 3): + Number of input video channels. + out_channels (`int`, defaults to 3): + Number of reconstructed video channels. + latent_channels (`int`, defaults to 16): + Number of latent channels. + embed_dim (`int`, defaults to 1024): + Transformer hidden size. Must be divisible by the attention head count and the patch volume. + num_layers (`int`, defaults to 24): + Number of transformer blocks in each of the encoder and decoder. + num_attention_heads (`int`, defaults to 16): + Number of attention heads. + mlp_ratio (`float`, defaults to 4.0): + Feed-forward hidden size relative to the transformer hidden size. + patch_size (`int`, defaults to 8): + Spatial patch size and compression ratio. + patch_length (`int`, defaults to 4): + Temporal patch size and compression ratio. + sample_size (`int`, defaults to 256): + Spatial size used to define the learned position embedding grid. + sample_frames (`int`, defaults to 16): + Frame count used to define the learned position embedding grid. + norm_eps (`float`, defaults to 1e-5): + Epsilon for normalization. Attention divides by the standard deviation plus epsilon. + """ + + _supports_gradient_checkpointing = True + _no_split_modules = ["MagiVAEBlock", "MagiVAEPositionEmbedding"] + _repeated_blocks = ["MagiVAEBlock"] + _skip_layerwise_casting_patterns = ["patch_embed", "pos_embed", "cls_token", "norm"] + + @register_to_config + def __init__( + self, + in_channels: int = 3, + out_channels: int = 3, + latent_channels: int = 16, + embed_dim: int = 1024, + num_layers: int = 24, + num_attention_heads: int = 16, + mlp_ratio: float = 4.0, + patch_size: int = 8, + patch_length: int = 4, + sample_size: int = 256, + sample_frames: int = 16, + norm_eps: float = 1e-5, + ): + super().__init__() + if embed_dim % num_attention_heads: + raise ValueError("embed_dim must be divisible by num_attention_heads.") + if embed_dim % (patch_length * patch_size * patch_size): + raise ValueError("embed_dim must be divisible by the spatiotemporal patch volume.") + if sample_frames % patch_length or sample_size % patch_size: + raise ValueError("The reference sample dimensions must be divisible by the patch dimensions.") + patch_shape = (patch_length, patch_size, patch_size) + latent_shape = (sample_frames // patch_length, sample_size // patch_size, sample_size // patch_size) + args = ( + latent_channels, + embed_dim, + num_layers, + num_attention_heads, + mlp_ratio, + patch_shape, + latent_shape, + norm_eps, + ) + self.encoder = MagiVAEEncoder(in_channels, *args) + self.decoder = MagiVAEDecoder(out_channels, *args) + self.spatial_compression_ratio = patch_size + self.temporal_compression_ratio = patch_length + self.use_tiling = False + self.use_slicing = False + self.tile_sample_min_length = sample_frames + self.tile_sample_min_height = sample_size + self.tile_sample_min_width = sample_size + self.temporal_tile_overlap_factor = 0.0 + self.spatial_tile_overlap_factor = 0.25 + self.allow_spatial_tiling = True + + def enable_slicing(self): + """Encode and decode one batch item at a time.""" + self.use_slicing = True + + def disable_slicing(self): + """Encode and decode the full batch together.""" + self.use_slicing = False + + def enable_tiling( + self, + tile_sample_min_length: int | None = None, + tile_sample_min_height: int | None = None, + tile_sample_min_width: int | None = None, + temporal_tile_overlap_factor: float = 0.0, + spatial_tile_overlap_factor: float = 0.25, + allow_spatial_tiling: bool = True, + ): + """ + Enable single-device temporal and optional spatial tiling. + + Tile sizes are measured in input video pixels and frames. Overlaps must give integral latent strides. The + default spatial overlap is 25%; temporal tiles do not overlap by default. Tiling changes the attention context + and is not numerically equivalent to whole-video encoding or decoding. + """ + lengths = ( + self.config.sample_frames if tile_sample_min_length is None else tile_sample_min_length, + self.config.sample_size if tile_sample_min_height is None else tile_sample_min_height, + self.config.sample_size if tile_sample_min_width is None else tile_sample_min_width, + ) + overlaps = (temporal_tile_overlap_factor, spatial_tile_overlap_factor, spatial_tile_overlap_factor) + factors = (self.config.patch_length, self.config.patch_size, self.config.patch_size) + for axis, (length, overlap, factor) in enumerate(zip(lengths, overlaps, factors)): + if not isinstance(length, int) or isinstance(length, bool) or length <= 0 or length % factor: + raise ValueError("Tile dimensions must be positive integer multiples of the patch dimensions.") + if not 0 <= overlap < 1: + raise ValueError("Tile overlap factors must be in [0, 1).") + if axis > 0 and not allow_spatial_tiling: + continue + latent_overlap = length // factor * overlap + if not float(latent_overlap).is_integer(): + raise ValueError("Tile overlaps must align with the latent grid.") + self.tile_sample_min_length, self.tile_sample_min_height, self.tile_sample_min_width = lengths + self.temporal_tile_overlap_factor = temporal_tile_overlap_factor + self.spatial_tile_overlap_factor = spatial_tile_overlap_factor + self.allow_spatial_tiling = allow_spatial_tiling + self.use_tiling = True + + def disable_tiling(self): + """Restore whole-input encoding and decoding.""" + self.use_tiling = False + + def _tile_parameters(self, shape): + sample_sizes = (self.tile_sample_min_length, self.tile_sample_min_height, self.tile_sample_min_width) + factors = (self.config.patch_length, self.config.patch_size, self.config.patch_size) + latent_sizes = tuple(size // factor for size, factor in zip(sample_sizes, factors)) + overlaps = ( + self.temporal_tile_overlap_factor, + self.spatial_tile_overlap_factor, + self.spatial_tile_overlap_factor, + ) + if not self.allow_spatial_tiling: + latent_sizes = (latent_sizes[0], shape[1], shape[2]) + overlaps = (overlaps[0], 0.0, 0.0) + blend_extents = tuple(int(size * overlap) for size, overlap in zip(latent_sizes, overlaps)) + strides = tuple(size - blend for size, blend in zip(latent_sizes, blend_extents)) + return latent_sizes, strides, blend_extents + + @staticmethod + def _blend(previous, current, extent, dim, upcast=False): + extent = min(previous.shape[dim], current.shape[dim], extent) + for index in range(extent): + before = previous.select(dim, previous.shape[dim] - extent + index) + after = current.select(dim, index) + if upcast: + # The compiled reference decoder blends low-precision tiles with FP32 intermediates. + blend_dtype = torch.promote_types(current.dtype, torch.float32) + before = before.to(blend_dtype) + after = after.to(blend_dtype) + current.select(dim, index).copy_(before * (1 - index / extent) + after * (index / extent)) + return current + + def _encode(self, x): + if x.shape[2] == 1: + x = x.expand(-1, -1, self.config.patch_length, -1, -1) + if x.shape[2] % self.config.patch_length: + raise ValueError("Each temporal tile must contain complete patches or a single frame.") + return self.encoder(x) + + def _tiled_encode(self, x): + factors = (self.config.patch_length, self.config.patch_size, self.config.patch_size) + latent_shape = tuple((size + factor - 1) // factor for size, factor in zip(x.shape[2:], factors)) + sizes, strides, blends = self._tile_parameters(latent_shape) + positions = [range(0, size, stride) for size, stride in zip(latent_shape, strides)] + tiles = {} + for frame, height, width in product(*positions): + tile = x[ + :, + :, + frame * factors[0] : (frame + sizes[0]) * factors[0], + height * factors[1] : (height + sizes[1]) * factors[1], + width * factors[2] : (width + sizes[2]) * factors[2], + ] + tile = self._encode(tile) + for axis, position in enumerate((frame, height, width)): + if position > 0 and blends[axis] > 0: + previous = [frame, height, width] + previous[axis] -= strides[axis] + tile = self._blend(tiles[tuple(previous)], tile.clone(), blends[axis], axis + 2) + tiles[(frame, height, width)] = tile + return torch.cat( + [ + torch.cat( + [ + torch.cat( + [ + tiles[(frame, height, width)][:, :, : strides[0], : strides[1], : strides[2]] + for width in positions[2] + ], + dim=4, + ) + for height in positions[1] + ], + dim=3, + ) + for frame in positions[0] + ], + dim=2, + ) + + def _tiled_decode(self, z, preserve_frames): + sizes, strides, blends = self._tile_parameters(z.shape[2:]) + factors = (self.config.patch_length, self.config.patch_size, self.config.patch_size) + limits = tuple(stride * factor for stride, factor in zip(strides, factors)) + positions = [range(0, size, stride) for size, stride in zip(z.shape[2:], strides)] + tiles = {} + results = {} + for frame, height, width in product(*positions): + latent = z[:, :, frame : frame + sizes[0], height : height + sizes[1], width : width + sizes[2]] + decoded = self.decoder(latent) + if latent.shape[2] == 1 and not preserve_frames: + decoded = decoded[:, :, :1] + tiles[(frame, height, width)] = decoded + tile = decoded.clone() + for axis, position in enumerate((frame, height, width)): + if position > 0 and blends[axis] > 0: + previous = [frame, height, width] + previous[axis] -= strides[axis] + tile = self._blend( + tiles[tuple(previous)], tile, blends[axis] * factors[axis], axis + 2, upcast=True + ) + results[(frame, height, width)] = tile[:, :, : limits[0], : limits[1], : limits[2]] + return torch.cat( + [ + torch.cat( + [ + torch.cat([results[(frame, height, width)] for width in positions[2]], dim=4) + for height in positions[1] + ], + dim=3, + ) + for frame in positions[0] + ], + dim=2, + ) + + @apply_forward_hook + def encode(self, x: torch.Tensor, return_dict: bool = True) -> AutoencoderKLOutput | tuple: + """ + Encode an image or video into a posterior without sampling it. + + Args: + x (`torch.Tensor`): + Input of shape `(batch, channels, frames, height, width)`, normalized to [-1, 1]. Spatial dimensions + must be divisible by the spatial patch size. A single frame is repeated to fill one temporal patch. + Tiling additionally supports a final one-frame temporal tile. + return_dict (`bool`, defaults to `True`): + Return an `AutoencoderKLOutput` instead of a tuple. + + Returns: + `AutoencoderKLOutput` or `tuple`: + The posterior. With tiling, its mean and log-variance are blended independently. Use `.mode()` to + reproduce the official deterministic inference encoder. + """ + if x.ndim != 5 or x.shape[1] != self.config.in_channels or any(size <= 0 for size in x.shape): + raise ValueError( + "Expected a nonempty video tensor with shape (batch, in_channels, frames, height, width)." + ) + if any(size % self.config.patch_size for size in x.shape[3:]): + raise ValueError("Spatial dimensions must be divisible by the patch size.") + remainder = x.shape[2] % self.config.patch_length + if x.shape[2] != 1 and remainder and not (self.use_tiling and remainder == 1): + raise ValueError("Frame count must be divisible by patch_length, or end in a single frame when tiling.") + encode = self._tiled_encode if self.use_tiling else self._encode + if self.use_slicing and x.shape[0] > 1: + moments = torch.cat([encode(sample) for sample in x.split(1)], dim=0) + else: + moments = encode(x) + posterior = DiagonalGaussianDistribution(moments) + if self.use_tiling: + # The reference concatenates means, without interleaved log-variance storage. + posterior.mean = posterior.mean.clone() + if not return_dict: + return (posterior,) + return AutoencoderKLOutput(latent_dist=posterior) + + @apply_forward_hook + def decode( + self, z: torch.Tensor, return_dict: bool = True, num_frames: int | None = None + ) -> DecoderOutput | tuple: + """ + Decode latents using the official MAGI image/video convention. + + Args: + z (`torch.Tensor`): + Latents of shape `(batch, latent_channels, frames, height, width)`. + return_dict (`bool`, defaults to `True`): + Return a `DecoderOutput` instead of a tuple. + num_frames (`int`, optional): + Requested output length. If omitted, each one-position latent tile returns only its first decoded + frame, matching the reference. Set this to the original video length to retain all frames. + + Returns: + `DecoderOutput` or `tuple`: + The reconstructed image or video. + """ + if z.ndim != 5 or z.shape[1] != self.config.latent_channels or any(size <= 0 for size in z.shape): + raise ValueError("Expected nonempty latents with shape (batch, latent_channels, frames, height, width).") + if num_frames is not None and ( + not isinstance(num_frames, int) + or isinstance(num_frames, bool) + or not (z.shape[2] - 1) * self.config.patch_length < num_frames <= z.shape[2] * self.config.patch_length + ): + raise ValueError("num_frames must fit the number of latent temporal patches.") + batches = z.split(1) if self.use_slicing else (z,) + outputs = [] + for latent in batches: + if self.use_tiling: + decoded = self._tiled_decode(latent, preserve_frames=num_frames is not None) + else: + decoded = self.decoder(latent) + if latent.shape[2] == 1 and num_frames is None: + decoded = decoded[:, :, :1] + outputs.append(decoded) + decoded = torch.cat(outputs, dim=0) if len(outputs) > 1 else outputs[0] + if num_frames is not None: + decoded = decoded[:, :, :num_frames] + if not return_dict: + return (decoded,) + return DecoderOutput(sample=decoded) + + def forward( + self, + sample: torch.Tensor, + sample_posterior: bool = False, + return_dict: bool = True, + generator: torch.Generator | None = None, + ) -> DecoderOutput | tuple: + """ + Reconstruct an image or video, preserving its frame count. + + Args: + sample (`torch.Tensor`): + Input of shape `(batch, channels, frames, height, width)`, normalized to [-1, 1]. + sample_posterior (`bool`, defaults to `False`): + Sample the posterior instead of using its mode. Sampling follows Diffusers generator and dtype + conventions; the official inference pipeline uses the mode. + return_dict (`bool`, defaults to `True`): + Return a `DecoderOutput` instead of a tuple. + generator (`torch.Generator`, optional): + Random generator used when sampling the posterior. + + Returns: + `DecoderOutput` or `tuple`: + The reconstruction with the original frame count. + """ + posterior = self.encode(sample).latent_dist + latent = posterior.sample(generator=generator) if sample_posterior else posterior.mode() + return self.decode(latent, return_dict=return_dict, num_frames=sample.shape[2]) diff --git a/src/diffusers/models/magi_conditioning.py b/src/diffusers/models/magi_conditioning.py new file mode 100644 index 000000000000..00808e95f3fd --- /dev/null +++ b/src/diffusers/models/magi_conditioning.py @@ -0,0 +1,91 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dataclasses import dataclass + +import torch +from torch import nn + +from ..configuration_utils import ConfigMixin, register_to_config +from ..utils import BaseOutput +from .modeling_utils import ModelMixin + + +@dataclass +class MagiTextConditioningOutput(BaseOutput): + sample: torch.Tensor + attention_mask: torch.Tensor + negative_prompt_embeds: torch.Tensor + negative_prompt_attention_mask: torch.Tensor + + +class MagiTextConditioningModel(ModelMixin, ConfigMixin): + """Store MAGI's learned null caption and the official HQ/duration feature vectors.""" + + _no_split_modules = ["MagiTextConditioningModel"] + _keep_in_fp32_modules = ["null_embedding", "special_embedding"] + + @register_to_config + def __init__(self, caption_channels=4096, caption_max_length=800, null_token_length=50): + super().__init__() + if not 0 < null_token_length <= caption_max_length: + raise ValueError("null_token_length must be between one and caption_max_length.") + self.null_embedding = nn.Embedding(caption_max_length, caption_channels) + self.special_embedding = nn.Embedding(9, caption_channels) + nn.init.zeros_(self.null_embedding.weight) + nn.init.zeros_(self.special_embedding.weight) + + def forward(self, hidden_states, attention_mask, num_chunks=1, return_dict=True): + """ + Prepare chunk-dependent conditional and stationary null text features. + + Args: + hidden_states (`torch.Tensor`): T5 features shaped `(batch, length, caption_channels)`. + attention_mask (`torch.Tensor`): Boolean keep-mask shaped `(batch, length)`. + num_chunks (`int`, defaults to `1`): Number of video chunks to condition. + return_dict (`bool`, defaults to `True`): Return structured conditioning outputs. + + Returns: + `MagiTextConditioningOutput` or `tuple`: Conditional features and masks, followed by null features and + masks. + """ + batch, length, channels = hidden_states.shape + if (length, channels) != (self.config.caption_max_length, self.config.caption_channels): + raise ValueError("T5 features must match the conditioning model's caption length and channels.") + if attention_mask.shape != (batch, length) or num_chunks < 1: + raise ValueError("Expected a matching text mask and at least one chunk.") + indices = torch.arange(length, device=hidden_states.device) + null = self.null_embedding(indices).to(hidden_states.dtype)[None].expand(batch, -1, -1) + hq = self.special_embedding(torch.zeros(1, device=hidden_states.device, dtype=torch.long)) + duration_indices = torch.arange(num_chunks, 0, -1, device=hidden_states.device).clamp(max=8) + duration = self.special_embedding(duration_indices).to(hidden_states.dtype) + conditional = torch.cat( + [ + duration[None, :, None].expand(batch, -1, -1, -1), + hq.to(hidden_states.dtype)[None, None].expand(batch, num_chunks, -1, -1), + hidden_states[:, None].expand(-1, num_chunks, -1, -1), + ], + dim=2, + )[:, :, :length] + mask = torch.cat( + [ + torch.ones(batch, num_chunks, 2, device=hidden_states.device, dtype=torch.bool), + attention_mask.bool()[:, None].expand(-1, num_chunks, -1), + ], + dim=2, + )[:, :, :length] + null_mask = (indices < self.config.null_token_length)[None].expand(batch, -1) + if not return_dict: + return conditional, mask, null, null_mask + return MagiTextConditioningOutput(conditional, mask, null, null_mask) diff --git a/src/diffusers/models/transformers/__init__.py b/src/diffusers/models/transformers/__init__.py index 0e167812ad88..db532707a0c0 100755 --- a/src/diffusers/models/transformers/__init__.py +++ b/src/diffusers/models/transformers/__init__.py @@ -51,6 +51,7 @@ from .transformer_ltx import LTXVideoTransformer3DModel from .transformer_ltx2 import LTX2VideoTransformer3DModel from .transformer_lumina2 import Lumina2Transformer2DModel + from .transformer_magi import MagiTransformer3DModel from .transformer_minimax_h3 import MiniMaxH3Transformer3DModel from .transformer_minimax_music3 import MiniMaxMusic3Transformer1DModel from .transformer_mochi import MochiTransformer3DModel diff --git a/src/diffusers/models/transformers/transformer_magi.py b/src/diffusers/models/transformers/transformer_magi.py new file mode 100644 index 000000000000..85f0881b6e5a --- /dev/null +++ b/src/diffusers/models/transformers/transformer_magi.py @@ -0,0 +1,608 @@ +# Copyright 2025 SandAI and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import math +from dataclasses import dataclass + +import torch +from torch import nn +from torch.nn import functional as F + +from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import PeftAdapterMixin +from ...utils import BaseOutput, apply_lora_scale, is_flash_attn_available +from ..attention import AttentionMixin, AttentionModuleMixin +from ..attention_dispatch import dispatch_attention_fn +from ..modeling_utils import ModelMixin + + +if is_flash_attn_available(): + from flash_attn.layers.rotary import apply_rotary_emb as flash_apply_rotary_emb + + +@dataclass +class MagiTransformer3DModelOutput(BaseOutput): + """Predicted flow velocities and optional per-layer self-attention keys and values.""" + + sample: torch.Tensor + kv_cache: tuple[tuple[torch.Tensor, torch.Tensor], ...] | None = None + + +class MagiFP32Linear(nn.Linear): + def forward(self, hidden_states): + bias = self.bias.float() if self.bias is not None else None + with torch.autocast(device_type=hidden_states.device.type, enabled=False): + return F.linear(hidden_states.float(), self.weight.float(), bias) + + +class MagiPatchEmbedding(nn.Conv3d): + def forward(self, hidden_states): + with torch.autocast(device_type=hidden_states.device.type, enabled=False): + return F.conv3d(hidden_states.float(), self.weight.float(), stride=self.stride) + + +class MagiLayerNorm(nn.Module): + def __init__(self, dim, eps, zero_centered_gamma=True, upcast=False): + super().__init__() + self.weight = nn.Parameter(torch.zeros(dim) if zero_centered_gamma else torch.ones(dim)) + self.bias = nn.Parameter(torch.zeros(dim)) + self.eps = eps + self.zero_centered_gamma = zero_centered_gamma + self.upcast = upcast + + def forward(self, hidden_states): + if self.upcast: + hidden_states = hidden_states.float() + weight = self.weight.to(hidden_states.dtype) + weight = weight + 1 if self.zero_centered_gamma else weight + return F.layer_norm(hidden_states, self.weight.shape, weight, self.bias.to(hidden_states.dtype), self.eps) + + +class MagiTimestepEmbedding(nn.Module): + def __init__(self, dim, frequency_embedding_size): + super().__init__() + self.frequency_embedding_size = frequency_embedding_size + self.mlp = nn.Sequential(MagiFP32Linear(frequency_embedding_size, dim), nn.SiLU(), MagiFP32Linear(dim, dim)) + + def forward(self, timestep, hidden_dtype): + half = self.frequency_embedding_size // 2 + frequencies = torch.exp(-math.log(10000) * torch.arange(half, device="cpu").float() / half).to(timestep.device) + angles = timestep.flatten()[:, None].float() * frequencies[None] * 1000 + embedding = torch.cat([angles.cos(), angles.sin()], dim=-1) + if self.frequency_embedding_size % 2: + embedding = F.pad(embedding, (0, 1)) + return self.mlp(embedding.to(hidden_dtype)) + + +class MagiCaptionEmbedding(nn.Module): + def __init__(self, caption_channels, caption_max_length, dim, condition_dim): + super().__init__() + self.null_caption_embedding = nn.Parameter(torch.zeros(caption_max_length, caption_channels)) + self.y_proj_xattn = nn.Sequential(MagiFP32Linear(caption_channels, dim), nn.SiLU()) + self.y_proj_adaln = nn.Sequential(MagiFP32Linear(caption_channels, condition_dim)) + + def forward(self, encoder_hidden_states, caption_dropout_mask): + encoder_hidden_states = self.y_proj_xattn(encoder_hidden_states.contiguous()) + caption = torch.where( + caption_dropout_mask[:, None], self.null_caption_embedding[-1], self.null_caption_embedding[-2] + ) + condition = self.y_proj_adaln(caption) + return encoder_hidden_states, condition + + +class MagiConditionEmbedding(nn.Module): + def __init__(self, caption_channels, caption_max_length, dim, condition_dim, frequency_embedding_size): + super().__init__() + self.t_embedder = MagiTimestepEmbedding(condition_dim, frequency_embedding_size) + self.y_embedder = MagiCaptionEmbedding(caption_channels, caption_max_length, dim, condition_dim) + + def forward(self, timestep, encoder_hidden_states, caption_dropout_mask, timestep_delta, hidden_dtype): + condition = self.t_embedder(timestep, hidden_dtype) + if timestep_delta is not None: + condition = condition + self.t_embedder(timestep_delta, hidden_dtype) + condition = condition.reshape(*timestep.shape, -1) + encoder_hidden_states, caption_condition = self.y_embedder(encoder_hidden_states, caption_dropout_mask) + return condition + caption_condition[:, None], encoder_hidden_states + + +class MagiRotaryEmbedding(nn.Module): + def __init__(self, head_dim): + super().__init__() + num_bands = head_dim // 8 + self.bands = nn.Parameter(1.0 / (10000 ** (torch.arange(num_bands).float() / num_bands))) + + def forward(self, num_frames, height, width, frame_offset=0): + total_frames = num_frames + frame_offset + rescale_factor = math.sqrt(height * width / 256) + shapes = (total_frames, height, width) + reference_shapes = (total_frames, height / rescale_factor, width / rescale_factor) + coordinates = [] + for axis, (size, reference_size) in enumerate(zip(shapes, reference_shapes)): + coordinate = torch.arange(size, device=self.bands.device).float() + if axis > 0: + coordinate = coordinate - (size - 1) / 2 + if size > 1: + coordinate = coordinate / (size - 1) * (reference_size - 1) + coordinates.append(coordinate) + grid = torch.stack(torch.meshgrid(*coordinates, indexing="ij"), dim=-1) + angles = grid.unsqueeze(-1) * self.bands.float() + angles = angles[frame_offset:].reshape(num_frames * height * width, -1) + return angles.cos(), angles.sin() + + +class MagiQueryKeyValueProjection(nn.Module): + def __init__(self, dim, kv_dim, eps): + super().__init__() + self.layer_norm = nn.LayerNorm(dim, eps=eps) + self.q = nn.Linear(dim, dim, bias=False) + self.qx = nn.Linear(dim, dim, bias=False) + self.k = nn.Linear(dim, kv_dim, bias=False) + self.v = nn.Linear(dim, kv_dim, bias=False) + + def forward(self, hidden_states): + hidden_states = self.layer_norm(hidden_states) + return self.q(hidden_states), self.qx(hidden_states), self.k(hidden_states), self.v(hidden_states) + + +class MagiTextKeyValueProjection(nn.Module): + def __init__(self, dim, kv_dim): + super().__init__() + self.projections = nn.ModuleList([nn.Linear(dim, 2 * kv_dim // 8, bias=False) for _ in range(8)]) + + def forward(self, hidden_states): + return torch.cat([projection(hidden_states) for projection in self.projections], dim=-1) + + +class MagiAttnProcessor: + _attention_backend = None + _parallel_config = None + + def __call__( + self, attn, hidden_states, encoder_hidden_states, rotary_emb, encoder_attention_mask, kv_ranges, kv_cache + ): + batch_size, sequence_length, _ = hidden_states.shape + num_chunks = encoder_hidden_states.shape[1] + chunk_length = sequence_length // num_chunks + query, cross_query, key, value = attn.linear_qkv(hidden_states) + query = query.unflatten(-1, (attn.num_heads, attn.head_dim)) + key = key.unflatten(-1, (attn.num_kv_heads, attn.head_dim)) + value = value.unflatten(-1, (attn.num_kv_heads, attn.head_dim)) + query = attn.q_layernorm(query) + key = attn.k_layernorm(key) + if self._attention_backend in ("flash", "flash_varlen"): + cosine, sine = rotary_emb + query = flash_apply_rotary_emb(query.contiguous(), cosine, sine).to(hidden_states.dtype) + key = flash_apply_rotary_emb(key.contiguous(), cosine, sine).to(hidden_states.dtype) + else: + cosine, sine = (embedding[None, :, None] for embedding in rotary_emb) + rotary_dim = cosine.shape[-1] * 2 + rotated = [] + for tensor in (query, key): + first, second = tensor[..., :rotary_dim].chunk(2, dim=-1) + tensor = torch.cat( + [first * cosine - second * sine, first * sine + second * cosine, tensor[..., rotary_dim:]], + dim=-1, + ) + rotated.append(tensor.to(hidden_states.dtype)) + query, key = rotated + if kv_cache is not None: + key = torch.cat([kv_cache[0].to(key), key], dim=1) + value = torch.cat([kv_cache[1].to(value), value], dim=1) + new_cache = (key, value) + groups = attn.num_heads // attn.num_kv_heads + key = key.repeat_interleave(groups, dim=2) + value = value.repeat_interleave(groups, dim=2) + + cross_query = cross_query.unflatten(-1, (attn.num_heads, attn.head_dim)) + cross_query = attn.q_layernorm_xattn(cross_query) + if encoder_attention_mask is None: + text_states = encoder_hidden_states.flatten(0, 2) + else: + # Match the reference GEMM shape by packing valid tokens before the KV projections. + text_states = encoder_hidden_states[encoder_attention_mask] + torch._check(text_states.shape[0] > 0) + cross_kv = attn.linear_kv_xattn(text_states) + cross_key, cross_value = cross_kv.unflatten(-1, (attn.num_kv_heads, 2 * attn.head_dim)).chunk(2, dim=-1) + cross_key = attn.k_layernorm_xattn(cross_key) + padded_shape = (*encoder_hidden_states.shape[:-1], attn.num_kv_heads, attn.head_dim) + if encoder_attention_mask is None: + cross_key = cross_key.reshape(padded_shape) + cross_value = cross_value.reshape(padded_shape) + else: + padded_key = cross_key.new_zeros(padded_shape) + padded_value = cross_value.new_zeros(padded_shape) + padded_key[encoder_attention_mask] = cross_key + padded_value[encoder_attention_mask] = cross_value + cross_key, cross_value = padded_key, padded_value + cross_key = cross_key.repeat_interleave(groups, dim=3) + cross_value = cross_value.repeat_interleave(groups, dim=3) + if self._attention_backend == "flash_varlen" and encoder_attention_mask is not None: + # This backend accepts right padding; preserve the order of valid tokens when packing other masks. + indices = (~encoder_attention_mask).to(torch.int32).argsort(dim=-1, stable=True) + gather_indices = indices[..., None, None].expand_as(cross_key) + cross_key = cross_key.gather(2, gather_indices) + cross_value = cross_value.gather(2, gather_indices) + encoder_attention_mask = encoder_attention_mask.gather(2, indices) + self_outputs, cross_outputs = [], [] + for chunk, (start, end) in enumerate(kv_ranges): + query_slice = slice(chunk * chunk_length, (chunk + 1) * chunk_length) + self_outputs.append( + dispatch_attention_fn( + query[:, query_slice], + key[:, start:end], + value[:, start:end], + backend=self._attention_backend, + parallel_config=self._parallel_config, + ) + ) + mask = None if encoder_attention_mask is None else encoder_attention_mask[:, chunk, None, None, :] + cross_outputs.append( + dispatch_attention_fn( + cross_query[:, query_slice], + cross_key[:, chunk], + cross_value[:, chunk], + attn_mask=mask, + backend=self._attention_backend, + parallel_config=self._parallel_config, + ) + ) + self_output = torch.cat(self_outputs, dim=1).flatten(2) + cross_output = torch.cat(cross_outputs, dim=1).flatten(2) + hidden_states = torch.cat([self_output, cross_output], dim=-1) + # Preserve the TP8 self/cross-attention interleave used by the checkpoint. + hidden_states = hidden_states.unflatten(-1, (2, 8, attn.dim // 8)).transpose(-3, -2).flatten(-3) + return attn.linear_proj(hidden_states), new_cache + + +class MagiAttention(nn.Module, AttentionModuleMixin): + _default_processor_cls = MagiAttnProcessor + _available_processors = [MagiAttnProcessor] + _supports_qkv_fusion = False + + def __init__(self, dim, num_heads, num_kv_heads, eps, zero_centered_gamma): + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.num_kv_heads = num_kv_heads + self.head_dim = dim // num_heads + self.linear_qkv = MagiQueryKeyValueProjection(dim, num_kv_heads * self.head_dim, eps) + self.linear_kv_xattn = MagiTextKeyValueProjection(dim, num_kv_heads * self.head_dim) + self.linear_proj = MagiFP32Linear(2 * dim, dim, bias=False) + self.q_layernorm = MagiLayerNorm(self.head_dim, eps, zero_centered_gamma, upcast=True) + self.k_layernorm = MagiLayerNorm(self.head_dim, eps, zero_centered_gamma, upcast=True) + self.q_layernorm_xattn = MagiLayerNorm(self.head_dim, eps, zero_centered_gamma) + self.k_layernorm_xattn = MagiLayerNorm(self.head_dim, eps, zero_centered_gamma) + self.set_processor(MagiAttnProcessor()) + + def forward(self, hidden_states, encoder_hidden_states, rotary_emb, encoder_attention_mask, kv_ranges, kv_cache): + return self.processor( + self, hidden_states, encoder_hidden_states, rotary_emb, encoder_attention_mask, kv_ranges, kv_cache + ) + + +class MagiAdaModulateLayer(nn.Module): + def __init__(self, condition_dim, dim): + super().__init__() + self.act = nn.SiLU() + self.proj = nn.Sequential(nn.Linear(condition_dim, 2 * dim)) + + def forward(self, condition): + return self.proj(self.act(condition)) + + +class MagiMLP(nn.Module): + def __init__(self, dim, ffn_dim, gated_linear_unit, eps): + super().__init__() + self.gated_linear_unit = gated_linear_unit + self.layer_norm = nn.LayerNorm(dim, eps=eps) + self.linear_fc1 = nn.Linear(dim, ffn_dim * (2 if gated_linear_unit else 1), bias=False) + self.linear_fc2 = nn.Linear(ffn_dim, dim, bias=False) + + def forward(self, hidden_states): + hidden_states = self.linear_fc1(self.layer_norm(hidden_states)) + if self.gated_linear_unit: + gate, value = hidden_states.chunk(2, dim=-1) + hidden_states = F.silu(gate) * value + else: + hidden_states = F.gelu(hidden_states) + return self.linear_fc2(hidden_states) + + +class MagiTransformerBlock(nn.Module): + def __init__( + self, dim, condition_dim, ffn_dim, num_heads, num_kv_heads, gated_linear_unit, eps, zero_centered_gamma + ): + super().__init__() + self.ada_modulate_layer = MagiAdaModulateLayer(condition_dim, dim) + self.self_attention = MagiAttention(dim, num_heads, num_kv_heads, eps, zero_centered_gamma) + self.self_attn_post_norm = MagiLayerNorm(dim, eps, zero_centered_gamma, upcast=True) + self.mlp = MagiMLP(dim, ffn_dim, gated_linear_unit, eps) + self.mlp_post_norm = MagiLayerNorm(dim, eps, zero_centered_gamma, upcast=True) + + def forward( + self, hidden_states, condition, encoder_hidden_states, rotary_emb, encoder_attention_mask, kv_ranges, kv_cache + ): + residual = hidden_states + hidden_states, new_cache = self.self_attention( + hidden_states, encoder_hidden_states, rotary_emb, encoder_attention_mask, kv_ranges, kv_cache + ) + gates = self.ada_modulate_layer(condition) + gates = gates.float().tanh().to(gates.dtype) + gates = gates.repeat_interleave(hidden_states.shape[1] // condition.shape[1], dim=1) + gate_msa, gate_mlp = gates.chunk(2, dim=-1) + hidden_states = self.self_attn_post_norm(hidden_states.float() * gate_msa.float()) + hidden_states = (hidden_states + residual.float()).to(residual.dtype) + residual = hidden_states + hidden_states = self.mlp(hidden_states) + hidden_states = self.mlp_post_norm(hidden_states.float() * gate_mlp.float()) + return (hidden_states + residual.float()).to(residual.dtype), new_cache + + +class MagiTransformer3DModel(ModelMixin, ConfigMixin, AttentionMixin, PeftAdapterMixin): + """ + MAGI-1 video Transformer with parallel self/cross-attention and chunk-level conditioning. + + Parameters: + in_channels (`int`, defaults to 16): Number of input latent channels. + out_channels (`int`, defaults to 16): Number of output latent channels. + num_layers (`int`, defaults to 34): Number of Transformer blocks. + num_attention_heads (`int`, defaults to 24): Number of query heads. + num_key_value_heads (`int`, defaults to 8): Number of key/value heads. + attention_head_dim (`int`, defaults to 128): Channels per head. + ffn_dim (`int`, defaults to 12288): Feed-forward intermediate dimension. + condition_dim (`int`, defaults to 768): Timestep and adaptive gating embedding dimension. + caption_channels (`int`, defaults to 4096): Text encoder output dimension. + caption_max_length (`int`, defaults to 800): Length of the learned null caption. + patch_size (`tuple[int, int, int]`, defaults to `(1, 2, 2)`): Temporal and spatial patch dimensions. + frequency_embedding_size (`int`, defaults to 256): Sinusoidal timestep embedding dimension. + gated_linear_unit (`bool`, defaults to `False`): Use SwiGLU instead of GELU, as in the 24B model. + norm_eps (`float`, defaults to 1e-6): Layer normalization epsilon. + zero_centered_gamma (`bool`, defaults to `True`): Add one to the custom normalization weights. + x_rescale_factor (`float`, defaults to 1.0): Internal input scaling, inverted at the output. + duplicate_channels (`bool`, defaults to `False`): + Duplicate input channels and retain half the output, as in 24B. + distilled (`bool`, defaults to `False`): Require the additional distillation timestep embedding. + """ + + _supports_gradient_checkpointing = True + _no_split_modules = ["MagiTransformerBlock", "MagiConditionEmbedding"] + _repeated_blocks = ["MagiTransformerBlock"] + _skip_layerwise_casting_patterns = ["x_embedder", "condition_embedder", "rope", "norm", "final_linear"] + _keep_in_fp32_modules = [ + "x_embedder", + "condition_embedder", + "rope", + "self_attn_post_norm", + "q_layernorm", + "k_layernorm", + "mlp_post_norm", + "final_layernorm", + "final_linear", + ] + _skip_keys = ["kv_cache"] + + @register_to_config + def __init__( + self, + in_channels=16, + out_channels=16, + num_layers=34, + num_attention_heads=24, + num_key_value_heads=8, + attention_head_dim=128, + ffn_dim=12288, + condition_dim=768, + caption_channels=4096, + caption_max_length=800, + patch_size=(1, 2, 2), + frequency_embedding_size=256, + gated_linear_unit=False, + norm_eps=1e-6, + zero_centered_gamma=True, + x_rescale_factor=1.0, + duplicate_channels=False, + distilled=False, + ): + super().__init__() + dim = num_attention_heads * attention_head_dim + if dim % 8 or attention_head_dim % 8 or num_attention_heads % num_key_value_heads: + raise ValueError("Hidden/head dimensions must be divisible by 8, and query heads by key/value heads.") + self.x_embedder = MagiPatchEmbedding( + in_channels * (2 if duplicate_channels else 1), dim, kernel_size=patch_size, stride=patch_size, bias=False + ) + self.condition_embedder = MagiConditionEmbedding( + caption_channels, caption_max_length, dim, condition_dim, frequency_embedding_size + ) + self.rope = MagiRotaryEmbedding(attention_head_dim) + self.transformer_blocks = nn.ModuleList( + [ + MagiTransformerBlock( + dim, + condition_dim, + ffn_dim, + num_attention_heads, + num_key_value_heads, + gated_linear_unit, + norm_eps, + zero_centered_gamma, + ) + for _ in range(num_layers) + ] + ) + self.final_layernorm = MagiLayerNorm(dim, norm_eps, zero_centered_gamma, upcast=True) + self.final_linear = MagiFP32Linear( + dim, math.prod(patch_size) * out_channels * (2 if duplicate_channels else 1), bias=False + ) + self.gradient_checkpointing = False + + @apply_lora_scale("attention_kwargs") + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + timestep: torch.Tensor, + encoder_attention_mask: torch.Tensor | None = None, + caption_dropout_mask: torch.Tensor | None = None, + timestep_delta: torch.Tensor | None = None, + kv_ranges: tuple[tuple[int, int], ...] | None = None, + kv_cache: tuple[tuple[torch.Tensor, torch.Tensor], ...] | None = None, + use_cache: bool = False, + cache_token_count: int | None = None, + cache_device: str | torch.device | None = None, + attention_kwargs: dict | None = None, + return_dict: bool = True, + ): + """ + Predict flow velocities for equally sized video chunks. + + Args: + hidden_states (`torch.Tensor`): Latents of shape `(batch, channels, frames, height, width)`. + encoder_hidden_states (`torch.Tensor`): + Text features `(batch, length, channels)` or `(batch, chunks, length, channels)`. + timestep (`torch.Tensor`): Timesteps in [0, 1], shaped `(batch,)` or `(batch, chunks)`. + encoder_attention_mask (`torch.Tensor`, optional): + Boolean text keep-mask `(batch, length)` or `(batch, chunks, length)`. + caption_dropout_mask (`torch.Tensor`, optional): + Select the unconditional adaptive embedding per batch item. Shape `(1,)` broadcasts a single projected + condition across the batch. This does not replace cross-attention text features. + timestep_delta (`torch.Tensor`, optional): Extra distillation timesteps, broadcastable to `timestep`. + kv_ranges (`tuple`, optional): + Exclusive `(start, end)` token ranges, one per current chunk, indexing cached plus current tokens. + Defaults to chunk-causal attention over all preceding chunks. + kv_cache (`tuple`, optional): + Per-layer `(key, value)` tensors shaped `(batch, cached_tokens, kv_heads, head_dim)`. Contains a + complete clean prefix and is never modified in place. + use_cache (`bool`, defaults to `False`): + Return cached prefix plus current keys and values. Only reuse entries belonging to clean, finalized + chunks. + cache_token_count (`int`, optional): + Retain only this many leading tokens in each returned cache, copying them before the next layer. + Requires `use_cache=True` and complete temporal patches. Does not change attention or predictions. + cache_device (`str` or `torch.device`, optional): + Device for returned cache tensors. Set to `"cpu"` to offload each layer's cache as it is produced. + Requires `use_cache=True`; input caches are moved to the compute device one layer at a time. + attention_kwargs (`dict`, optional): Keyword arguments for LoRA scaling. + return_dict (`bool`, defaults to `True`): Return a structured output instead of a tuple. + + Returns: + `MagiTransformer3DModelOutput` or `tuple`: Predicted velocities and, when requested, keys and values. + """ + if ( + hidden_states.ndim != 5 + or hidden_states.shape[1] != self.config.in_channels + or any(size <= 0 for size in hidden_states.shape) + ): + raise ValueError("Expected nonempty latents shaped (batch, in_channels, frames, height, width).") + batch_size, _, frames, height, width = hidden_states.shape + patch_t, patch_h, patch_w = self.config.patch_size + if frames % patch_t or height % patch_h or width % patch_w: + raise ValueError("Video dimensions must be divisible by the patch dimensions.") + timestep = timestep.reshape(batch_size, -1) + num_chunks = timestep.shape[1] + if frames // patch_t % num_chunks: + raise ValueError("Temporal patches must divide evenly into timestep chunks.") + if self.config.distilled and timestep_delta is None: + raise ValueError("Distilled models require timestep_delta from the distillation schedule.") + if timestep_delta is not None: + timestep_delta = torch.broadcast_to(timestep_delta, timestep.shape) + if encoder_hidden_states.ndim == 3: + encoder_hidden_states = encoder_hidden_states[:, None].expand(-1, num_chunks, -1, -1) + if encoder_hidden_states.shape[:2] != (batch_size, num_chunks): + raise ValueError("Text features must match the batch and timestep chunk dimensions.") + if encoder_attention_mask is not None: + if encoder_attention_mask.ndim == 2: + encoder_attention_mask = encoder_attention_mask[:, None].expand(-1, num_chunks, -1) + encoder_attention_mask = encoder_attention_mask.bool() + if caption_dropout_mask is None: + caption_dropout_mask = torch.zeros(batch_size, device=hidden_states.device, dtype=torch.bool) + spatial_tokens = (height // patch_h) * (width // patch_w) + if kv_cache is not None and len(kv_cache) != len(self.transformer_blocks): + raise ValueError("kv_cache must contain one key/value pair per Transformer block.") + cached_tokens = 0 if kv_cache is None else kv_cache[0][0].shape[1] + if cached_tokens % spatial_tokens: + raise ValueError("The cached prefix must contain complete temporal patches at the current resolution.") + sequence_length = frames // patch_t * spatial_tokens + if (cache_token_count is not None or cache_device is not None) and not use_cache: + raise ValueError("Cache retention options require use_cache=True.") + if cache_token_count is not None and ( + not isinstance(cache_token_count, int) + or isinstance(cache_token_count, bool) + or not 0 < cache_token_count <= cached_tokens + sequence_length + or cache_token_count % spatial_tokens + ): + raise ValueError("cache_token_count must retain complete temporal patches within the available tokens.") + chunk_length = sequence_length // num_chunks + if kv_ranges is None: + kv_ranges = tuple((0, cached_tokens + (chunk + 1) * chunk_length) for chunk in range(num_chunks)) + if len(kv_ranges) != num_chunks or any( + not 0 <= start < end <= cached_tokens + sequence_length for start, end in kv_ranges + ): + raise ValueError("Each chunk must have a nonempty key/value range within the available tokens.") + hidden_states = hidden_states * self.config.x_rescale_factor + if self.config.duplicate_channels: + hidden_states = torch.cat([hidden_states, hidden_states], dim=1) + hidden_states = self.x_embedder(hidden_states) + hidden_states = hidden_states.flatten(2).transpose(1, 2).contiguous().to(self.dtype) + rotary_emb = self.rope(frames // patch_t, height // patch_h, width // patch_w, cached_tokens // spatial_tokens) + condition, encoder_hidden_states = self.condition_embedder( + timestep, encoder_hidden_states, caption_dropout_mask, timestep_delta, hidden_states.dtype + ) + condition = condition.to(hidden_states.dtype) + encoder_hidden_states = encoder_hidden_states.to(hidden_states.dtype) + new_cache = [] + for index, block in enumerate(self.transformer_blocks): + layer_cache = None if kv_cache is None else kv_cache[index] + if torch.is_grad_enabled() and self.gradient_checkpointing: + hidden_states, cached = self._gradient_checkpointing_func( + block, + hidden_states, + condition, + encoder_hidden_states, + rotary_emb, + encoder_attention_mask, + kv_ranges, + layer_cache, + ) + else: + hidden_states, cached = block( + hidden_states, + condition, + encoder_hidden_states, + rotary_emb, + encoder_attention_mask, + kv_ranges, + layer_cache, + ) + if use_cache: + if cache_token_count is not None or cache_device is not None: + cached = tuple( + tensor[:, :cache_token_count].to( + device=cache_device if cache_device is not None else tensor.device, copy=True + ) + for tensor in cached + ) + new_cache.append(cached) + del cached + hidden_states = self.final_layernorm(hidden_states) + hidden_states = self.final_linear(hidden_states) + hidden_states = hidden_states.reshape( + batch_size, frames // patch_t, height // patch_h, width // patch_w, patch_t, patch_h, patch_w, -1 + ) + hidden_states = ( + hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6).reshape(batch_size, -1, frames, height, width).contiguous() + ) + hidden_states = hidden_states[:, : self.config.out_channels] / self.config.x_rescale_factor + if not return_dict: + return (hidden_states, tuple(new_cache)) if use_cache else (hidden_states,) + return MagiTransformer3DModelOutput(sample=hidden_states, kv_cache=tuple(new_cache) if use_cache else None) diff --git a/src/diffusers/modular_pipelines/__init__.py b/src/diffusers/modular_pipelines/__init__.py index 8c3f9ccc62c8..fbb0614206c1 100644 --- a/src/diffusers/modular_pipelines/__init__.py +++ b/src/diffusers/modular_pipelines/__init__.py @@ -47,6 +47,16 @@ "WanAnimate2DistilledModularPipeline", "WanAnimate2ModularPipeline", ] + _import_structure["magi"] = [ + "MagiDenoiseStep", + "MagiModularPipeline", + "MagiTextToVideoBlocks", + "MagiImageToVideoBlocks", + "MagiVideoToVideoBlocks", + "MagiTextEncoderStep", + "MagiPrepareLatentsStep", + "MagiVaeDecoderStep", + ] _import_structure["wan"] = [ "WanBlocks", "Wan22Blocks", @@ -192,6 +202,16 @@ ) from .ltx import LTXAutoBlocks, LTXModularPipeline from .ltx2 import LTX2AutoBlocks, LTX2ModularPipeline, LTX25AutoBlocks, LTX25ModularPipeline + from .magi import ( + MagiDenoiseStep, + MagiImageToVideoBlocks, + MagiModularPipeline, + MagiPrepareLatentsStep, + MagiTextEncoderStep, + MagiTextToVideoBlocks, + MagiVaeDecoderStep, + MagiVideoToVideoBlocks, + ) from .minimax_h3 import ( MiniMaxH3Blocks, MiniMaxH3ModularPipeline, diff --git a/src/diffusers/modular_pipelines/magi/__init__.py b/src/diffusers/modular_pipelines/magi/__init__.py new file mode 100644 index 000000000000..b10a7b8c2d49 --- /dev/null +++ b/src/diffusers/modular_pipelines/magi/__init__.py @@ -0,0 +1,67 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import TYPE_CHECKING + +from ...utils import ( + DIFFUSERS_SLOW_IMPORT, + OptionalDependencyNotAvailable, + _LazyModule, + get_objects_from_module, + is_torch_available, + is_transformers_available, +) + + +_dummy_objects = {} +_import_structure = {} + +try: + if not (is_transformers_available() and is_torch_available()): + raise OptionalDependencyNotAvailable() +except OptionalDependencyNotAvailable: + from ...utils import dummy_torch_and_transformers_objects + + _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects)) +else: + _import_structure["before_denoise"] = ["MagiPrepareLatentsStep"] + _import_structure["decoders"] = ["MagiVaeDecoderStep"] + _import_structure["denoise"] = ["MagiDenoiseLoop", "MagiDenoiseStep", "MagiPrepareDenoiseStep"] + _import_structure["encoders"] = ["MagiTextEncoderStep"] + _import_structure["modular_blocks_magi"] = [ + "MagiImageToVideoBlocks", + "MagiTextToVideoBlocks", + "MagiVideoToVideoBlocks", + ] + _import_structure["modular_pipeline"] = ["MagiModularPipeline"] + +if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: + try: + if not (is_transformers_available() and is_torch_available()): + raise OptionalDependencyNotAvailable() + except OptionalDependencyNotAvailable: + from ...utils.dummy_torch_and_transformers_objects import * # noqa F403 + else: + from .before_denoise import MagiPrepareLatentsStep + from .decoders import MagiVaeDecoderStep + from .denoise import MagiDenoiseLoop, MagiDenoiseStep, MagiPrepareDenoiseStep + from .encoders import MagiTextEncoderStep + from .modular_blocks_magi import MagiImageToVideoBlocks, MagiTextToVideoBlocks, MagiVideoToVideoBlocks + from .modular_pipeline import MagiModularPipeline +else: + import sys + + sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__) + for name, value in _dummy_objects.items(): + setattr(sys.modules[__name__], name, value) diff --git a/src/diffusers/modular_pipelines/magi/before_denoise.py b/src/diffusers/modular_pipelines/magi/before_denoise.py new file mode 100644 index 000000000000..3aa027de74c7 --- /dev/null +++ b/src/diffusers/modular_pipelines/magi/before_denoise.py @@ -0,0 +1,254 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import math + +import torch + +from ...models import AutoencoderKLMagi, MagiTransformer3DModel +from ...models.magi_conditioning import MagiTextConditioningModel +from ...utils.torch_utils import randn_tensor +from ..modular_pipeline import ModularPipelineBlocks +from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam + + +class MagiPrepareLatentsStep(ModularPipelineBlocks): + model_name = "magi" + + @property + def description(self): + return "Prepare FP32 text-to-video noise and expand text features for the requested video batch." + + @property + def expected_components(self): + return [ + ComponentSpec("transformer", MagiTransformer3DModel), + ComponentSpec("vae", AutoencoderKLMagi), + ComponentSpec("text_conditioning", MagiTextConditioningModel), + ] + + @property + def inputs(self): + return [ + InputParam( + "text_embeds", required=True, type_hint=torch.Tensor, description="Per-prompt T5 hidden states." + ), + InputParam( + "text_attention_mask", required=True, type_hint=torch.Tensor, description="Per-prompt T5 keep-mask." + ), + InputParam("height", default=720, type_hint=int, description="Video height in pixels."), + InputParam("width", default=720, type_hint=int, description="Video width in pixels."), + InputParam( + "num_frames", + default=96, + type_hint=int, + description="Requested video frames; generation rounds up to full chunks.", + ), + InputParam("chunk_width", default=6, type_hint=int, description="Latent frames per chunk."), + InputParam.template("num_images_per_prompt"), + InputParam.template("generator"), + InputParam( + "latents", + default=None, + type_hint=torch.Tensor, + description="Optional initial FP32 noise for all generated chunks.", + ), + InputParam( + "prefix_latents", + default=None, + type_hint=torch.Tensor, + description="Not supported by this text-to-video preparation block.", + ), + ] + + @property + def intermediate_outputs(self): + return [ + OutputParam(name, type_hint=torch.Tensor, description=description) + for name, description in [ + ("latents", "Initial FP32 noise."), + ("prompt_embeds", "HQ/duration conditioned chunk text features."), + ("prompt_attention_mask", "Conditional text keep-mask."), + ("negative_prompt_embeds", "Learned null text features."), + ("negative_prompt_attention_mask", "Null text keep-mask."), + ] + ] + + @torch.no_grad() + def __call__(self, components, state): + block_state = self.get_block_state(state) + if block_state.prefix_latents is not None: + raise ValueError( + "This workflow is text-to-video; use MagiDenoiseStep for prepared full-chunk prefix latents." + ) + for name in ("height", "width", "num_frames", "chunk_width", "num_images_per_prompt"): + value = getattr(block_state, name) + if not isinstance(value, int) or isinstance(value, bool) or value <= 0: + raise ValueError(f"{name} must be a positive integer.") + spatial, temporal = components.vae.spatial_compression_ratio, components.vae.temporal_compression_ratio + if block_state.height % spatial or block_state.width % spatial or block_state.num_frames % temporal: + raise ValueError("Video dimensions must be divisible by the VAE compression ratios.") + if components.vae.config.latent_channels != components.transformer.config.in_channels: + raise ValueError("VAE latent channels must match the Transformer input channels.") + num_chunks = math.ceil(block_state.num_frames // temporal / block_state.chunk_width) + count = block_state.num_images_per_prompt + device = components._execution_device + text = block_state.text_embeds.repeat_interleave(count, dim=0).to(device=device, dtype=torch.float32) + mask = block_state.text_attention_mask.repeat_interleave(count, dim=0).to(device) + conditioning = components.text_conditioning(text, mask, num_chunks=num_chunks) + block_state.prompt_embeds = conditioning.sample + block_state.prompt_attention_mask = conditioning.attention_mask + block_state.negative_prompt_embeds = conditioning.negative_prompt_embeds + block_state.negative_prompt_attention_mask = conditioning.negative_prompt_attention_mask + shape = ( + text.shape[0], + components.transformer.config.in_channels, + num_chunks * block_state.chunk_width, + block_state.height // spatial, + block_state.width // spatial, + ) + if block_state.latents is None: + block_state.latents = randn_tensor( + shape, generator=block_state.generator, device=device, dtype=torch.float32 + ) + elif tuple(block_state.latents.shape) != shape: + raise ValueError(f"Initial latents must have shape {shape}.") + else: + block_state.latents = block_state.latents.to(device=device, dtype=torch.float32) + self.set_block_state(state, block_state) + return components, state + + +class MagiPrepareConditionedLatentsStep(MagiPrepareLatentsStep): + model_name = "magi" + + @property + def description(self): + return "Prepare prefix-aware FP32 noise and duration conditioning for newly generated chunks." + + @property + def expected_components(self): + return [ + ComponentSpec("transformer", MagiTransformer3DModel), + ComponentSpec("vae", AutoencoderKLMagi), + ComponentSpec("text_conditioning", MagiTextConditioningModel), + ] + + @property + def inputs(self): + return [param for param in super().inputs if param.name not in ("prefix_latents", "num_frames")] + [ + InputParam( + "conditioning_latents", + required=True, + type_hint=torch.Tensor, + description="Per-prompt scaled VAE prefix.", + ), + InputParam( + "num_frames", + default=96, + type_hint=int, + description="Requested new frames; prefix plus new frames rounds up to full latent chunks.", + ), + ] + + @property + def intermediate_outputs(self): + return [ + OutputParam(name, type_hint=torch.Tensor, description=description) + for name, description in [ + ("latents", "Initial FP32 noise."), + ("conditioning_latents", "Scaled prefix expanded to the generated video batch."), + ("prompt_embeds", "HQ/duration conditioned chunk text features."), + ("prompt_attention_mask", "Conditional text keep-mask."), + ("negative_prompt_embeds", "Learned null text features."), + ("negative_prompt_attention_mask", "Null text keep-mask."), + ] + ] + + @torch.no_grad() + def __call__(self, components, state): + block_state = self.get_block_state(state) + for name in ("height", "width", "num_frames", "chunk_width", "num_images_per_prompt"): + value = getattr(block_state, name) + if not isinstance(value, int) or isinstance(value, bool) or value <= 0: + raise ValueError(f"{name} must be a positive integer.") + spatial, temporal = components.vae.spatial_compression_ratio, components.vae.temporal_compression_ratio + if block_state.height % spatial or block_state.width % spatial or block_state.num_frames % temporal: + raise ValueError("Video dimensions must be divisible by the VAE compression ratios.") + if components.vae.config.latent_channels != components.transformer.config.in_channels: + raise ValueError("VAE latent channels must match the Transformer input channels.") + prefix = block_state.conditioning_latents + expected = (block_state.height // spatial, block_state.width // spatial) + if ( + not isinstance(prefix, torch.Tensor) + or prefix.ndim != 5 + or prefix.shape[0] not in (1, block_state.text_embeds.shape[0]) + or prefix.shape[1] != components.transformer.config.in_channels + or prefix.shape[2] < 1 + or tuple(prefix.shape[3:]) != expected + or not prefix.is_floating_point() + or not prefix.isfinite().all() + ): + raise ValueError( + "conditioning_latents must match the prompt batch, latent channels, and requested dimensions." + ) + prefix_chunks = prefix.shape[2] // block_state.chunk_width + num_chunks = math.ceil((block_state.num_frames // temporal + prefix.shape[2]) / block_state.chunk_width) + count = block_state.num_images_per_prompt + device = components._execution_device + text = block_state.text_embeds.repeat_interleave(count, dim=0).to(device=device, dtype=torch.float32) + mask = block_state.text_attention_mask.repeat_interleave(count, dim=0).to(device) + conditioning = components.text_conditioning(text, mask, num_chunks=num_chunks - prefix_chunks) + block_state.conditioning_latents = ( + prefix.expand(block_state.text_embeds.shape[0], -1, -1, -1, -1) + .repeat_interleave(count, dim=0) + .to(device=device, dtype=torch.float32) + ) + block_state.prompt_embeds = conditioning.sample + block_state.prompt_attention_mask = conditioning.attention_mask + block_state.negative_prompt_embeds = conditioning.negative_prompt_embeds + block_state.negative_prompt_attention_mask = conditioning.negative_prompt_attention_mask + if prefix_chunks: + # Clean-prefix conditional slots are never evaluated; use valid null captions for the shared validator. + block_state.prompt_embeds = torch.cat( + [ + conditioning.negative_prompt_embeds[:, None].expand(-1, prefix_chunks, -1, -1), + block_state.prompt_embeds, + ], + dim=1, + ) + block_state.prompt_attention_mask = torch.cat( + [ + conditioning.negative_prompt_attention_mask[:, None].expand(-1, prefix_chunks, -1), + block_state.prompt_attention_mask, + ], + dim=1, + ) + shape = ( + text.shape[0], + components.transformer.config.in_channels, + num_chunks * block_state.chunk_width, + block_state.height // spatial, + block_state.width // spatial, + ) + if block_state.latents is None: + block_state.latents = randn_tensor( + shape, generator=block_state.generator, device=device, dtype=torch.float32 + ) + elif tuple(block_state.latents.shape) != shape: + raise ValueError(f"Initial latents must have shape {shape}.") + else: + block_state.latents = block_state.latents.to(device=device, dtype=torch.float32) + self.set_block_state(state, block_state) + return components, state diff --git a/src/diffusers/modular_pipelines/magi/decoders.py b/src/diffusers/modular_pipelines/magi/decoders.py new file mode 100644 index 000000000000..d0f385403486 --- /dev/null +++ b/src/diffusers/modular_pipelines/magi/decoders.py @@ -0,0 +1,143 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch + +from ...configuration_utils import FrozenDict +from ...models import AutoencoderKLMagi +from ...video_processor import VideoProcessor +from ..modular_pipeline import ModularPipelineBlocks +from ..modular_pipeline_utils import ComponentSpec, ConfigSpec, InputParam, OutputParam + + +class MagiVaeDecoderStep(ModularPipelineBlocks): + model_name = "magi" + + @property + def description(self): + return "Decode each generated chunk independently, then assemble the output video." + + @property + def expected_components(self): + return [ + ComponentSpec("vae", AutoencoderKLMagi), + ComponentSpec( + "video_processor", + VideoProcessor, + config=FrozenDict({"vae_scale_factor": 8}), + default_creation_method="from_config", + ), + ] + + @property + def expected_configs(self): + return [ConfigSpec("latent_scaling_factor", 0.18215)] + + @property + def inputs(self): + return [ + InputParam.template("latents", required=True), + InputParam("chunk_width", default=6, type_hint=int, description="Latent frames per generated chunk."), + InputParam( + "output_type", default="np", type_hint=str, description="Output format: pt, np, pil, or latent." + ), + ] + + @property + def intermediate_outputs(self): + return [OutputParam.template("videos")] + + @torch.no_grad() + def __call__(self, components, state): + block_state = self.get_block_state(state) + if block_state.output_type == "latent": + block_state.videos = block_state.latents + else: + if block_state.output_type not in ("pt", "np", "pil"): + raise ValueError("output_type must be pt, np, pil, or latent.") + if not isinstance(block_state.chunk_width, int) or block_state.chunk_width < 1: + raise ValueError("chunk_width must be a positive integer.") + if components.config.latent_scaling_factor <= 0: + raise ValueError("latent_scaling_factor must be positive.") + chunks = [] + for chunk in block_state.latents.split(block_state.chunk_width, dim=2): + chunk = (chunk.float() / components.config.latent_scaling_factor).to(components.vae.dtype) + num_frames = chunk.shape[2] * components.vae.temporal_compression_ratio + with torch.autocast( + device_type=chunk.device.type, + dtype=chunk.dtype, + enabled=chunk.dtype in (torch.float16, torch.bfloat16), + ): + chunks.append(components.vae.decode(chunk, num_frames=num_frames, return_dict=False)[0]) + video = torch.cat(chunks, dim=2).float() + block_state.videos = components.video_processor.postprocess_video( + video, output_type=block_state.output_type + ) + self.set_block_state(state, block_state) + return components, state + + +class MagiPrefixVaeDecoderStep(MagiVaeDecoderStep): + model_name = "magi" + + @property + def description(self): + return "Trim prefix latents before per-chunk decoding, retaining the first four frames for a one-frame prefix." + + @property + def inputs(self): + return super().inputs + [ + InputParam( + "conditioning_latents", + required=True, + type_hint=torch.Tensor, + description="Scaled input prefix; determines the latent frames omitted from the output.", + ) + ] + + @torch.no_grad() + def __call__(self, components, state): + block_state = self.get_block_state(state) + prefix_length = block_state.conditioning_latents.shape[2] + if not 0 < prefix_length < block_state.latents.shape[2]: + raise ValueError("The prefix must leave at least one generated latent frame.") + output_start = 0 if prefix_length == 1 else prefix_length + if block_state.output_type == "latent": + block_state.videos = block_state.latents[:, :, output_start:] + else: + if block_state.output_type not in ("pt", "np", "pil"): + raise ValueError("output_type must be pt, np, pil, or latent.") + if not isinstance(block_state.chunk_width, int) or block_state.chunk_width < 1: + raise ValueError("chunk_width must be a positive integer.") + if components.config.latent_scaling_factor <= 0: + raise ValueError("latent_scaling_factor must be positive.") + chunks = [] + for start in range(0, block_state.latents.shape[2], block_state.chunk_width): + end = start + block_state.chunk_width + if end <= output_start: + continue + chunk = block_state.latents[:, :, max(start, output_start) : end] + chunk = (chunk.float() / components.config.latent_scaling_factor).to(components.vae.dtype) + with torch.autocast( + device_type=chunk.device.type, + dtype=chunk.dtype, + enabled=chunk.dtype in (torch.float16, torch.bfloat16), + ): + chunks.append(components.vae.decode(chunk, return_dict=False)[0]) + video = torch.cat(chunks, dim=2).float() + block_state.videos = components.video_processor.postprocess_video( + video, output_type=block_state.output_type + ) + self.set_block_state(state, block_state) + return components, state diff --git a/src/diffusers/modular_pipelines/magi/denoise.py b/src/diffusers/modular_pipelines/magi/denoise.py new file mode 100644 index 000000000000..fa77f84b39a0 --- /dev/null +++ b/src/diffusers/modular_pipelines/magi/denoise.py @@ -0,0 +1,884 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch + +from ...guiders.magi_classifier_free_guidance import MagiClassifierFreeGuidance +from ...models import MagiTransformer3DModel +from ...schedulers import MagiEulerScheduler +from ..modular_pipeline import LoopSequentialPipelineBlocks, ModularPipelineBlocks, SequentialPipelineBlocks +from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam + + +_STATE_FIELDS = { + "latents": ( + torch.Tensor, + "FP32 latent state shaped (batch, channels, frames, height, width), including prefix slots.", + ), + "prompt_embeds": (torch.Tensor, "Prepared conditional text features, shared across chunks or provided per chunk."), + "prompt_attention_mask": (torch.Tensor, "Boolean keep-mask matching the conditional text features."), + "negative_prompt_embeds": ( + torch.Tensor, + "Learned null-caption features shaped (batch, length, caption_channels), shared across chunks.", + ), + "negative_prompt_attention_mask": (torch.Tensor, "Boolean keep-mask for the learned null-caption features."), + "prefix_latents": ( + torch.Tensor, + "Optional full-chunk clean prefix; replaces the leading latent slots and remains unchanged.", + ), + "num_inference_steps": (int, "Number of Euler updates per generated chunk."), + "chunk_width": (int, "Number of latent frames in each chunk."), + "window_size": (int, "Maximum number of simultaneously denoised chunks."), + "noise2clean_kvrange": ( + tuple, + "Positive attention-window lengths in chunks, from early to late denoising stages.", + ), + "clean_chunk_kvrange": (int, "Positive attention-window length used when recomputing clean chunks."), + "clean_t": (float, "Model evaluation time for clean-prefix cache extraction."), + "cache_device": ( + str, + "Optional device for clean-prefix KV storage; use cpu to offload between layer evaluations.", + ), + "attention_kwargs": (dict, "Optional keyword arguments passed to Transformer attention."), + "num_chunks": (int, "Total number of latent chunks, including the supplied prefix."), + "prefix_chunks": (int, "Number of supplied full-chunk prefix chunks."), + "chunk_tokens": (int, "Number of Transformer tokens per latent chunk."), + "steps_per_stage": (int, "Number of iterations before the active chunk window moves forward."), + "num_window_steps": (int, "Total number of asynchronous window iterations."), + "timestep_schedule": (torch.Tensor, "FP32 schedule including the final Euler integration endpoint."), + "clean_kv_cache": ( + tuple, + "Per-layer clean-prefix key/value tensors, or None before any prefix is cached; excludes the final generated chunk.", + ), + "completed_chunks": (list, "Indices of supplied prefix chunks and finalized generated chunks."), + "chunk_start": (int, "First active chunk index."), + "chunk_end": (int, "Exclusive end index of the active chunk window."), + "refresh_cache": (bool, "Whether this iteration prepends a finalized chunk to refresh its clean KV."), + "chunk_step_indices": (list, "Denoising step indices for active chunks, ordered from oldest to newest."), + "chunk_times": (torch.Tensor, "Current per-chunk model times shaped (batch, active_chunks)."), + "next_chunk_times": (torch.Tensor, "Next Euler endpoints shaped (batch, active_chunks)."), + "model_times": (torch.Tensor, "Model times including an optional clean-refresh chunk."), + "latent_model_input": (torch.Tensor, "Current latent window including an optional clean-refresh chunk."), + "window_prompt_embeds": ( + torch.Tensor, + "Conditional features for the current window, with null features for a clean-refresh chunk.", + ), + "window_prompt_attention_mask": (torch.Tensor, "Boolean keep-mask matching the current window text features."), + "kv_ranges": (tuple, "Exclusive token attention ranges indexing the cached prefix plus the current window."), + "velocity": (torch.Tensor, "Three-way guided FP32 velocities for active chunks only."), +} + + +def _input(name, **kwargs): + type_hint, description = _STATE_FIELDS[name] + return InputParam(name, type_hint=type_hint, description=description, **kwargs) + + +def _output(name): + type_hint, description = _STATE_FIELDS[name] + return OutputParam(name, type_hint=type_hint, description=description) + + +def _ranges(chunk_indices, limits, chunk_tokens): + return tuple( + (max(0, index + 1 - limit) * chunk_tokens, (index + 1) * chunk_tokens) + for index, limit in zip(chunk_indices, limits) + ) + + +class MagiPrepareDenoiseStep(ModularPipelineBlocks): + model_name = "magi" + + @property + def description(self): + return "Prepare the MAGI base-model chunk schedule and optional full-chunk clean prefix." + + @property + def expected_components(self): + return [ComponentSpec("transformer", MagiTransformer3DModel), ComponentSpec("scheduler", MagiEulerScheduler)] + + @property + def inputs(self): + return [ + _input("latents", required=True), + _input("prompt_embeds", required=True), + _input("prompt_attention_mask", required=True), + _input("negative_prompt_embeds", required=True), + _input("negative_prompt_attention_mask", required=True), + _input("prefix_latents", default=None), + _input("num_inference_steps", default=64), + _input("chunk_width", default=6), + _input("window_size", default=4), + _input("noise2clean_kvrange", default=(5, 4, 3, 2)), + _input("clean_chunk_kvrange", default=1), + _input("clean_t", default=0.9999), + _input("attention_kwargs", default=None), + _input("cache_device", default=None), + ] + + @property + def intermediate_outputs(self): + return [ + _output(name) + for name in [ + "num_chunks", + "prefix_chunks", + "chunk_tokens", + "steps_per_stage", + "num_window_steps", + "timestep_schedule", + "clean_kv_cache", + "completed_chunks", + ] + ] + + @torch.no_grad() + def __call__(self, components, state): + s = self.get_block_state(state) + config = components.transformer.config + if config.distilled: + raise ValueError("MagiDenoiseStep supports base models only, not distilled models.") + for name in ("chunk_width", "window_size", "num_inference_steps", "clean_chunk_kvrange"): + value = getattr(s, name) + if not isinstance(value, int) or isinstance(value, bool) or value <= 0: + raise ValueError(f"{name} must be a positive integer.") + if not s.noise2clean_kvrange or any(not isinstance(x, int) or x <= 0 for x in s.noise2clean_kvrange): + raise ValueError("noise2clean_kvrange must contain positive chunk counts.") + if s.num_inference_steps % s.window_size or s.num_inference_steps % len(s.noise2clean_kvrange): + raise ValueError("num_inference_steps must be divisible by window_size and the number of KV ranges.") + if not 0 <= s.clean_t <= 1: + raise ValueError("clean_t must be in [0, 1].") + if s.latents.ndim != 5 or min(s.latents.shape) <= 0 or s.latents.shape[1] != config.in_channels: + raise ValueError("latents must have shape (batch, in_channels, frames, height, width).") + batch, channels, frames, height, width = s.latents.shape + pt, ph, pw = config.patch_size + if frames % s.chunk_width or s.chunk_width % pt or height % ph or width % pw: + raise ValueError("Latents must contain full chunks and be divisible by the Transformer patch size.") + s.num_chunks = frames // s.chunk_width + for name, mask_name, per_chunk in [ + ("prompt_embeds", "prompt_attention_mask", True), + ("negative_prompt_embeds", "negative_prompt_attention_mask", False), + ]: + features, mask = getattr(s, name), getattr(s, mask_name) + valid_shape = (features.ndim == 3 and features.shape[0] == batch) or ( + per_chunk and features.ndim == 4 and features.shape[:2] == (batch, s.num_chunks) + ) + if not valid_shape or features.shape[-1] != config.caption_channels or features.shape[-2] <= 0: + raise ValueError(f"{name} must match the batch, caption channels, and optional chunk count.") + if mask.shape != features.shape[:-1] or mask.dtype != torch.bool or not mask.any(dim=-1).all(): + raise ValueError(f"{mask_name} must be a boolean keep-mask with at least one valid token per caption.") + if features.device != s.latents.device or mask.device != s.latents.device: + raise ValueError("All latents, text features, and masks must be on the same device.") + if s.prompt_embeds.shape[-2] != s.negative_prompt_embeds.shape[-2]: + raise ValueError("Conditional and null text features must have the same padded length.") + s.latents = s.latents.float().clone() + if s.prompt_embeds.ndim == 3: + s.prompt_embeds = s.prompt_embeds[:, None].expand(-1, s.num_chunks, -1, -1) + s.prompt_attention_mask = s.prompt_attention_mask[:, None].expand(-1, s.num_chunks, -1) + s.prefix_chunks = 0 + s.clean_kv_cache = None + s.chunk_tokens = (s.chunk_width // pt) * (height // ph) * (width // pw) + if s.prefix_latents is not None: + prefix = s.prefix_latents + if ( + prefix.ndim != 5 + or prefix.shape[:2] != (batch, channels) + or prefix.shape[3:] != (height, width) + or prefix.device != s.latents.device + ): + raise ValueError( + "prefix_latents must match the latent batch, channels, spatial dimensions, and device." + ) + if prefix.shape[2] <= 0 or prefix.shape[2] % s.chunk_width or prefix.shape[2] >= frames: + raise ValueError("prefix_latents must contain full chunks and leave at least one chunk to generate.") + s.prefix_chunks = prefix.shape[2] // s.chunk_width + s.latents[:, :, : prefix.shape[2]] = prefix.float() + s.clean_kv_cache = components.transformer( + hidden_states=s.latents[:, :, : prefix.shape[2]], + encoder_hidden_states=s.negative_prompt_embeds, + encoder_attention_mask=s.negative_prompt_attention_mask, + caption_dropout_mask=torch.ones(1, device=s.latents.device, dtype=torch.bool), + timestep=s.latents.new_full((batch, s.prefix_chunks), s.clean_t), + kv_ranges=_ranges(range(s.prefix_chunks), [s.clean_chunk_kvrange] * s.prefix_chunks, s.chunk_tokens), + use_cache=True, + cache_device=s.cache_device, + attention_kwargs=s.attention_kwargs, + ).kv_cache + s.steps_per_stage = s.num_inference_steps // s.window_size + s.num_window_steps = s.steps_per_stage * (s.num_chunks + s.window_size - 1 - s.prefix_chunks) + s.completed_chunks = list(range(s.prefix_chunks)) + components.scheduler.set_timesteps(s.num_inference_steps, device=s.latents.device) + s.timestep_schedule = components.scheduler.timestep_schedule + self.set_block_state(state, s) + return components, state + + +class MagiLoopBeforeDenoiser(ModularPipelineBlocks): + model_name = "magi" + + @property + def description(self): + return "Select the moving chunk window, per-chunk times, and attention ranges." + + @property + def inputs(self): + return [ + _input(name, required=True) + for name in [ + "latents", + "prompt_embeds", + "prompt_attention_mask", + "negative_prompt_embeds", + "negative_prompt_attention_mask", + "prefix_chunks", + "num_chunks", + "chunk_width", + "chunk_tokens", + "steps_per_stage", + "window_size", + "num_inference_steps", + "noise2clean_kvrange", + "clean_chunk_kvrange", + "clean_t", + "timestep_schedule", + ] + ] + + @property + def intermediate_outputs(self): + return [ + _output(name) + for name in [ + "chunk_start", + "chunk_end", + "refresh_cache", + "chunk_step_indices", + "chunk_times", + "next_chunk_times", + "model_times", + "latent_model_input", + "window_prompt_embeds", + "window_prompt_attention_mask", + "kv_ranges", + ] + ] + + def __call__(self, components, s, i): + stage, inner = divmod(i, s.steps_per_stage) + position = s.prefix_chunks + stage + s.chunk_start = max(s.prefix_chunks, position - s.window_size + 1) + s.chunk_end = min(s.num_chunks, position + 1) + t_start = max(0, position - s.num_chunks + 1) + t_end = min(s.window_size, stage + 1) + s.chunk_step_indices = [j * s.steps_per_stage + inner for j in reversed(range(t_start, t_end))] + batch = s.latents.shape[0] + s.chunk_times = s.timestep_schedule[s.chunk_step_indices][None].expand(batch, -1) + s.next_chunk_times = s.timestep_schedule[[j + 1 for j in s.chunk_step_indices]][None].expand(batch, -1) + s.refresh_cache = s.chunk_start > s.prefix_chunks and inner == 0 + first = s.chunk_start - int(s.refresh_cache) + s.latent_model_input = s.latents[:, :, first * s.chunk_width : s.chunk_end * s.chunk_width] + s.window_prompt_embeds = s.prompt_embeds[:, s.chunk_start : s.chunk_end] + s.window_prompt_attention_mask = s.prompt_attention_mask[:, s.chunk_start : s.chunk_end] + s.model_times = s.chunk_times + limits = [ + s.noise2clean_kvrange[j // (s.num_inference_steps // len(s.noise2clean_kvrange))] + for j in s.chunk_step_indices + ] + if s.refresh_cache: + s.model_times = torch.cat([s.latents.new_full((batch, 1), s.clean_t), s.chunk_times], dim=1) + s.window_prompt_embeds = torch.cat([s.negative_prompt_embeds[:, None], s.window_prompt_embeds], dim=1) + s.window_prompt_attention_mask = torch.cat( + [s.negative_prompt_attention_mask[:, None], s.window_prompt_attention_mask], dim=1 + ) + limits.insert(0, s.clean_chunk_kvrange) + s.kv_ranges = _ranges(range(first, s.chunk_end), limits, s.chunk_tokens) + return components, s + + +class MagiLoopDenoiser(ModularPipelineBlocks): + model_name = "magi" + + @property + def description(self): + return "Evaluate three CFG branches and retain only refreshed, unconditional clean-prefix KV." + + @property + def expected_components(self): + return [ + ComponentSpec("transformer", MagiTransformer3DModel), + ComponentSpec("guider", MagiClassifierFreeGuidance), + ] + + @property + def inputs(self): + return [ + _input(name, required=True) + for name in [ + "latent_model_input", + "model_times", + "chunk_times", + "window_prompt_embeds", + "window_prompt_attention_mask", + "negative_prompt_embeds", + "negative_prompt_attention_mask", + "kv_ranges", + "refresh_cache", + "chunk_start", + "chunk_end", + "chunk_width", + "chunk_tokens", + "num_inference_steps", + ] + ] + [ + _input("clean_kv_cache", default=None), + _input("attention_kwargs", default=None), + _input("cache_device", default=None), + ] + + @property + def intermediate_outputs(self): + return [_output("velocity"), _output("clean_kv_cache")] + + @torch.no_grad() + def __call__(self, components, s, i): + batch, channels, _, height, width = s.latent_model_input.shape + count = s.chunk_end - s.chunk_start + skip = int(s.refresh_cache) * s.chunk_width + independent = s.latent_model_input[:, :, skip:].reshape(batch, channels, count, s.chunk_width, height, width) + independent = independent.permute(0, 2, 1, 3, 4, 5).reshape( + batch * count, channels, s.chunk_width, height, width + ) + guider = components.guider + guider.set_state(step=i, num_inference_steps=s.num_inference_steps, timestep=s.chunk_times) + branches = guider.prepare_inputs( + { + "hidden_states": (s.latent_model_input, s.latent_model_input, independent), + "encoder_hidden_states": ( + s.window_prompt_embeds, + s.negative_prompt_embeds, + s.negative_prompt_embeds.repeat_interleave(count, dim=0), + ), + "encoder_attention_mask": ( + s.window_prompt_attention_mask, + s.negative_prompt_attention_mask, + s.negative_prompt_attention_mask.repeat_interleave(count, dim=0), + ), + "timestep": (s.model_times, s.model_times, s.chunk_times.reshape(-1, 1)), + } + ) + refreshed_cache = None + for branch_index, branch in enumerate(branches): + guider.prepare_models(components.transformer) + try: + result = components.transformer( + hidden_states=branch.hidden_states, + encoder_hidden_states=branch.encoder_hidden_states, + encoder_attention_mask=branch.encoder_attention_mask, + timestep=branch.timestep, + caption_dropout_mask=torch.full( + (1,), + branch_index != 0, + device=independent.device, + dtype=torch.bool, + ), + kv_ranges=s.kv_ranges if branch_index < 2 else None, + kv_cache=s.clean_kv_cache if branch_index < 2 else None, + use_cache=branch_index == 1 and s.refresh_cache, + cache_token_count=s.chunk_start * s.chunk_tokens + if branch_index == 1 and s.refresh_cache + else None, + cache_device=s.cache_device if branch_index == 1 and s.refresh_cache else None, + attention_kwargs=s.attention_kwargs, + ) + finally: + guider.cleanup_models(components.transformer) + if branch_index < 2: + branch.noise_pred = result.sample[:, :, skip:] + else: + branch.noise_pred = ( + result.sample.reshape(batch, count, channels, s.chunk_width, height, width) + .permute(0, 2, 1, 3, 4, 5) + .reshape(batch, channels, count * s.chunk_width, height, width) + ) + if branch_index == 1 and s.refresh_cache: + refreshed_cache = result.kv_cache + del result + s.velocity = guider(branches).pred + if refreshed_cache is not None: + s.clean_kv_cache = refreshed_cache + return components, s + + +class MagiLoopAfterDenoiser(ModularPipelineBlocks): + model_name = "magi" + + @property + def description(self): + return "Integrate active chunks in FP32 without changing finalized prefix latents." + + @property + def expected_components(self): + return [ComponentSpec("scheduler", MagiEulerScheduler)] + + @property + def inputs(self): + return [ + _input(name, required=True) + for name in [ + "latents", + "velocity", + "chunk_times", + "next_chunk_times", + "chunk_start", + "chunk_end", + "chunk_width", + "chunk_step_indices", + "num_inference_steps", + "completed_chunks", + ] + ] + + @property + def intermediate_outputs(self): + return [_output("latents"), _output("completed_chunks")] + + def __call__(self, components, s, i): + start, end = s.chunk_start * s.chunk_width, s.chunk_end * s.chunk_width + s.latents[:, :, start:end] = components.scheduler.step( + s.velocity, s.chunk_times, s.latents[:, :, start:end], next_timestep=s.next_chunk_times + ).prev_sample + if s.chunk_step_indices[0] == s.num_inference_steps - 1: + s.completed_chunks.append(s.chunk_start) + return components, s + + +# auto_docstring +class MagiDenoiseLoop(LoopSequentialPipelineBlocks): + """ + Run the asynchronous MAGI chunk-denoising window with a clean-prefix cache. + + Components: + transformer (`MagiTransformer3DModel`) guider (`MagiClassifierFreeGuidance`) scheduler (`MagiEulerScheduler`) + + Inputs: + num_window_steps (`int`): + Total number of asynchronous window iterations. + latents (`Tensor`): + FP32 latent state shaped (batch, channels, frames, height, width), including prefix slots. + prompt_embeds (`Tensor`): + Prepared conditional text features, shared across chunks or provided per chunk. + prompt_attention_mask (`Tensor`): + Boolean keep-mask matching the conditional text features. + negative_prompt_embeds (`Tensor`): + Learned null-caption features shaped (batch, length, caption_channels), shared across chunks. + negative_prompt_attention_mask (`Tensor`): + Boolean keep-mask for the learned null-caption features. + prefix_chunks (`int`): + Number of supplied full-chunk prefix chunks. + num_chunks (`int`): + Total number of latent chunks, including the supplied prefix. + chunk_width (`int`): + Number of latent frames in each chunk. + chunk_tokens (`int`): + Number of Transformer tokens per latent chunk. + steps_per_stage (`int`): + Number of iterations before the active chunk window moves forward. + window_size (`int`): + Maximum number of simultaneously denoised chunks. + num_inference_steps (`int`): + Number of Euler updates per generated chunk. + noise2clean_kvrange (`tuple`): + Positive attention-window lengths in chunks, from early to late denoising stages. + clean_chunk_kvrange (`int`): + Positive attention-window length used when recomputing clean chunks. + clean_t (`float`): + Model evaluation time for clean-prefix cache extraction. + timestep_schedule (`Tensor`): + FP32 schedule including the final Euler integration endpoint. + clean_kv_cache (`tuple`, *optional*): + Per-layer clean-prefix key/value tensors, or None before any prefix is cached; excludes the final + generated chunk. + attention_kwargs (`dict`, *optional*): + Optional keyword arguments passed to Transformer attention. + cache_device (`str`, *optional*): + Optional device for clean-prefix KV storage; use cpu to offload between layer evaluations. + completed_chunks (`list`): + Indices of supplied prefix chunks and finalized generated chunks. + + Outputs: + latents (`Tensor`): + FP32 latent state shaped (batch, channels, frames, height, width), including prefix slots. + completed_chunks (`list`): + Indices of supplied prefix chunks and finalized generated chunks. + """ + + model_name = "magi" + block_classes = [MagiLoopBeforeDenoiser, MagiLoopDenoiser, MagiLoopAfterDenoiser] + block_names = ["before_denoiser", "denoiser", "after_denoiser"] + + @property + def description(self): + return "Run the asynchronous MAGI chunk-denoising window with a clean-prefix cache." + + @property + def loop_inputs(self): + return [_input("num_window_steps", required=True)] + + @torch.no_grad() + def __call__(self, components, state): + s = self.get_block_state(state) + with self.progress_bar(total=s.num_window_steps) as progress: + for i in range(s.num_window_steps): + components, s = self.loop_step(components, s, i=i) + progress.update() + self.set_block_state(state, s) + return components, state + + +# auto_docstring +class MagiDenoiseStep(SequentialPipelineBlocks): + """ + MAGI base-model latent generation; text encoding, partial-chunk prefixes, and VAE decoding are not included. + + Components: + transformer (`MagiTransformer3DModel`) scheduler (`MagiEulerScheduler`) guider (`MagiClassifierFreeGuidance`) + + Inputs: + latents (`Tensor`): + FP32 latent state shaped (batch, channels, frames, height, width), including prefix slots. + prompt_embeds (`Tensor`): + Prepared conditional text features, shared across chunks or provided per chunk. + prompt_attention_mask (`Tensor`): + Boolean keep-mask matching the conditional text features. + negative_prompt_embeds (`Tensor`): + Learned null-caption features shaped (batch, length, caption_channels), shared across chunks. + negative_prompt_attention_mask (`Tensor`): + Boolean keep-mask for the learned null-caption features. + prefix_latents (`Tensor`, *optional*): + Optional full-chunk clean prefix; replaces the leading latent slots and remains unchanged. + num_inference_steps (`int`, *optional*, defaults to 64): + Number of Euler updates per generated chunk. + chunk_width (`int`, *optional*, defaults to 6): + Number of latent frames in each chunk. + window_size (`int`, *optional*, defaults to 4): + Maximum number of simultaneously denoised chunks. + noise2clean_kvrange (`tuple`, *optional*, defaults to (5, 4, 3, 2)): + Positive attention-window lengths in chunks, from early to late denoising stages. + clean_chunk_kvrange (`int`, *optional*, defaults to 1): + Positive attention-window length used when recomputing clean chunks. + clean_t (`float`, *optional*, defaults to 0.9999): + Model evaluation time for clean-prefix cache extraction. + attention_kwargs (`dict`, *optional*): + Optional keyword arguments passed to Transformer attention. + cache_device (`str`, *optional*): + Optional device for clean-prefix KV storage; use cpu to offload between layer evaluations. + + Outputs: + num_chunks (`int`): + Total number of latent chunks, including the supplied prefix. + prefix_chunks (`int`): + Number of supplied full-chunk prefix chunks. + chunk_tokens (`int`): + Number of Transformer tokens per latent chunk. + steps_per_stage (`int`): + Number of iterations before the active chunk window moves forward. + num_window_steps (`int`): + Total number of asynchronous window iterations. + timestep_schedule (`Tensor`): + FP32 schedule including the final Euler integration endpoint. + clean_kv_cache (`tuple`): + Per-layer clean-prefix key/value tensors, or None before any prefix is cached; excludes the final + generated chunk. + completed_chunks (`list`): + Indices of supplied prefix chunks and finalized generated chunks. + chunk_start (`int`): + First active chunk index. + chunk_end (`int`): + Exclusive end index of the active chunk window. + refresh_cache (`bool`): + Whether this iteration prepends a finalized chunk to refresh its clean KV. + chunk_step_indices (`list`): + Denoising step indices for active chunks, ordered from oldest to newest. + chunk_times (`Tensor`): + Current per-chunk model times shaped (batch, active_chunks). + next_chunk_times (`Tensor`): + Next Euler endpoints shaped (batch, active_chunks). + model_times (`Tensor`): + Model times including an optional clean-refresh chunk. + latent_model_input (`Tensor`): + Current latent window including an optional clean-refresh chunk. + window_prompt_embeds (`Tensor`): + Conditional features for the current window, with null features for a clean-refresh chunk. + window_prompt_attention_mask (`Tensor`): + Boolean keep-mask matching the current window text features. + kv_ranges (`tuple`): + Exclusive token attention ranges indexing the cached prefix plus the current window. + velocity (`Tensor`): + Three-way guided FP32 velocities for active chunks only. + latents (`Tensor`): + FP32 latent state shaped (batch, channels, frames, height, width), including prefix slots. + """ + + model_name = "magi" + block_classes = [MagiPrepareDenoiseStep, MagiDenoiseLoop] + block_names = ["prepare", "denoise"] + + @property + def description(self): + return "MAGI base-model latent generation; text encoding, partial-chunk prefixes, and VAE decoding are not included." + + +class MagiPreparePrefixStep(ModularPipelineBlocks): + model_name = "magi" + + @property + def description(self): + return "Select complete prefix chunks for clean-cache initialization." + + @property + def inputs(self): + return [ + _input("latents", required=True), + _input("chunk_width", default=6), + InputParam( + "conditioning_latents", + required=True, + type_hint=torch.Tensor, + description="Scaled prefix matching the generated latent batch and device.", + ), + ] + + @property + def intermediate_outputs(self): + return [_output("prefix_latents")] + + def __call__(self, components, state): + s = self.get_block_state(state) + prefix = s.conditioning_latents + if not isinstance(s.chunk_width, int) or isinstance(s.chunk_width, bool) or s.chunk_width <= 0: + raise ValueError("chunk_width must be a positive integer.") + if ( + not isinstance(prefix, torch.Tensor) + or prefix.ndim != 5 + or s.latents.ndim != 5 + or prefix.shape[:2] != s.latents.shape[:2] + or prefix.shape[3:] != s.latents.shape[3:] + or not 0 < prefix.shape[2] < s.latents.shape[2] + or prefix.device != s.latents.device + or not prefix.is_floating_point() + or not prefix.isfinite().all() + ): + raise ValueError( + "conditioning_latents must be a finite prefix matching the latent batch, shape, and device." + ) + length = prefix.shape[2] // s.chunk_width * s.chunk_width + s.prefix_latents = prefix[:, :, :length] if length else None + self.set_block_state(state, s) + return components, state + + +class MagiPrefixLoopBeforeDenoiser(MagiLoopBeforeDenoiser): + @property + def inputs(self): + return super().inputs + [ + InputParam( + "conditioning_latents", + required=True, + type_hint=torch.Tensor, + description="Original scaled prefix, reinjected before every model evaluation.", + ) + ] + + @property + def intermediate_outputs(self): + return super().intermediate_outputs + [_output("latents")] + + def __call__(self, components, s, i): + components, s = super().__call__(components, s, i) + first = (s.chunk_start - int(s.refresh_cache)) * s.chunk_width + end = min(s.conditioning_latents.shape[2], s.chunk_end * s.chunk_width) + if first < end: + s.latent_model_input = s.latent_model_input.clone() + s.latent_model_input[:, :, : end - first] = s.conditioning_latents[:, :, first:end] + start = s.chunk_start * s.chunk_width + if start < end: + # Euler integrates from the injected prefix, but cache refresh must not overwrite finalized output. + s.latents[:, :, start:end] = s.conditioning_latents[:, :, start:end] + return components, s + + +# auto_docstring +class MagiPrefixDenoiseLoop(MagiDenoiseLoop): + """ + Denoise with per-evaluation prefix injection and separate clean-prefix cache refresh. + + Components: + transformer (`MagiTransformer3DModel`) guider (`MagiClassifierFreeGuidance`) scheduler (`MagiEulerScheduler`) + + Inputs: + num_window_steps (`int`): + Total number of asynchronous window iterations. + latents (`Tensor`): + FP32 latent state shaped (batch, channels, frames, height, width), including prefix slots. + prompt_embeds (`Tensor`): + Prepared conditional text features, shared across chunks or provided per chunk. + prompt_attention_mask (`Tensor`): + Boolean keep-mask matching the conditional text features. + negative_prompt_embeds (`Tensor`): + Learned null-caption features shaped (batch, length, caption_channels), shared across chunks. + negative_prompt_attention_mask (`Tensor`): + Boolean keep-mask for the learned null-caption features. + prefix_chunks (`int`): + Number of supplied full-chunk prefix chunks. + num_chunks (`int`): + Total number of latent chunks, including the supplied prefix. + chunk_width (`int`): + Number of latent frames in each chunk. + chunk_tokens (`int`): + Number of Transformer tokens per latent chunk. + steps_per_stage (`int`): + Number of iterations before the active chunk window moves forward. + window_size (`int`): + Maximum number of simultaneously denoised chunks. + num_inference_steps (`int`): + Number of Euler updates per generated chunk. + noise2clean_kvrange (`tuple`): + Positive attention-window lengths in chunks, from early to late denoising stages. + clean_chunk_kvrange (`int`): + Positive attention-window length used when recomputing clean chunks. + clean_t (`float`): + Model evaluation time for clean-prefix cache extraction. + timestep_schedule (`Tensor`): + FP32 schedule including the final Euler integration endpoint. + conditioning_latents (`Tensor`): + Original scaled prefix, reinjected before every model evaluation. + clean_kv_cache (`tuple`, *optional*): + Per-layer clean-prefix key/value tensors, or None before any prefix is cached; excludes the final + generated chunk. + attention_kwargs (`dict`, *optional*): + Optional keyword arguments passed to Transformer attention. + cache_device (`str`, *optional*): + Optional device for clean-prefix KV storage; use cpu to offload between layer evaluations. + completed_chunks (`list`): + Indices of supplied prefix chunks and finalized generated chunks. + + Outputs: + latents (`Tensor`): + FP32 latent state shaped (batch, channels, frames, height, width), including prefix slots. + completed_chunks (`list`): + Indices of supplied prefix chunks and finalized generated chunks. + """ + + model_name = "magi" + block_classes = [MagiPrefixLoopBeforeDenoiser, MagiLoopDenoiser, MagiLoopAfterDenoiser] + block_names = ["before_denoiser", "denoiser", "after_denoiser"] + + @property + def description(self): + return "Denoise with per-evaluation prefix injection and separate clean-prefix cache refresh." + + +# auto_docstring +class MagiPrefixDenoiseStep(SequentialPipelineBlocks): + """ + Generate latent continuations from complete or partial-chunk prefixes. + + Components: + transformer (`MagiTransformer3DModel`) scheduler (`MagiEulerScheduler`) guider (`MagiClassifierFreeGuidance`) + + Inputs: + latents (`Tensor`): + FP32 latent state shaped (batch, channels, frames, height, width), including prefix slots. + chunk_width (`int`, *optional*, defaults to 6): + Number of latent frames in each chunk. + conditioning_latents (`Tensor`): + Scaled prefix matching the generated latent batch and device. + prompt_embeds (`Tensor`): + Prepared conditional text features, shared across chunks or provided per chunk. + prompt_attention_mask (`Tensor`): + Boolean keep-mask matching the conditional text features. + negative_prompt_embeds (`Tensor`): + Learned null-caption features shaped (batch, length, caption_channels), shared across chunks. + negative_prompt_attention_mask (`Tensor`): + Boolean keep-mask for the learned null-caption features. + num_inference_steps (`int`, *optional*, defaults to 64): + Number of Euler updates per generated chunk. + window_size (`int`, *optional*, defaults to 4): + Maximum number of simultaneously denoised chunks. + noise2clean_kvrange (`tuple`, *optional*, defaults to (5, 4, 3, 2)): + Positive attention-window lengths in chunks, from early to late denoising stages. + clean_chunk_kvrange (`int`, *optional*, defaults to 1): + Positive attention-window length used when recomputing clean chunks. + clean_t (`float`, *optional*, defaults to 0.9999): + Model evaluation time for clean-prefix cache extraction. + attention_kwargs (`dict`, *optional*): + Optional keyword arguments passed to Transformer attention. + cache_device (`str`, *optional*): + Optional device for clean-prefix KV storage; use cpu to offload between layer evaluations. + + Outputs: + prefix_latents (`Tensor`): + Optional full-chunk clean prefix; replaces the leading latent slots and remains unchanged. + num_chunks (`int`): + Total number of latent chunks, including the supplied prefix. + prefix_chunks (`int`): + Number of supplied full-chunk prefix chunks. + chunk_tokens (`int`): + Number of Transformer tokens per latent chunk. + steps_per_stage (`int`): + Number of iterations before the active chunk window moves forward. + num_window_steps (`int`): + Total number of asynchronous window iterations. + timestep_schedule (`Tensor`): + FP32 schedule including the final Euler integration endpoint. + clean_kv_cache (`tuple`): + Per-layer clean-prefix key/value tensors, or None before any prefix is cached; excludes the final + generated chunk. + completed_chunks (`list`): + Indices of supplied prefix chunks and finalized generated chunks. + chunk_start (`int`): + First active chunk index. + chunk_end (`int`): + Exclusive end index of the active chunk window. + refresh_cache (`bool`): + Whether this iteration prepends a finalized chunk to refresh its clean KV. + chunk_step_indices (`list`): + Denoising step indices for active chunks, ordered from oldest to newest. + chunk_times (`Tensor`): + Current per-chunk model times shaped (batch, active_chunks). + next_chunk_times (`Tensor`): + Next Euler endpoints shaped (batch, active_chunks). + model_times (`Tensor`): + Model times including an optional clean-refresh chunk. + latent_model_input (`Tensor`): + Current latent window including an optional clean-refresh chunk. + window_prompt_embeds (`Tensor`): + Conditional features for the current window, with null features for a clean-refresh chunk. + window_prompt_attention_mask (`Tensor`): + Boolean keep-mask matching the current window text features. + kv_ranges (`tuple`): + Exclusive token attention ranges indexing the cached prefix plus the current window. + latents (`Tensor`): + FP32 latent state shaped (batch, channels, frames, height, width), including prefix slots. + velocity (`Tensor`): + Three-way guided FP32 velocities for active chunks only. + """ + + model_name = "magi" + block_classes = [MagiPreparePrefixStep, MagiPrepareDenoiseStep, MagiPrefixDenoiseLoop] + block_names = ["prepare_prefix", "prepare", "denoise"] + + @property + def description(self): + return "Generate latent continuations from complete or partial-chunk prefixes." diff --git a/src/diffusers/modular_pipelines/magi/encoders.py b/src/diffusers/modular_pipelines/magi/encoders.py new file mode 100644 index 000000000000..98b3fa045712 --- /dev/null +++ b/src/diffusers/modular_pipelines/magi/encoders.py @@ -0,0 +1,303 @@ +# Copyright (c) 2025 SandAI. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import html +import re +import urllib.parse as ul + +import torch +from transformers import AutoTokenizer, T5EncoderModel + +from ...models import AutoencoderKLMagi +from ...utils import is_bs4_available, is_ftfy_available, requires_backends +from ..modular_pipeline import ModularPipelineBlocks +from ..modular_pipeline_utils import ComponentSpec, ConfigSpec, InputParam, OutputParam + + +if is_bs4_available(): + from bs4 import BeautifulSoup +if is_ftfy_available(): + import ftfy + + +class MagiTextEncoderStep(ModularPipelineBlocks): + model_name = "magi" + bad_punct_regex = re.compile(r"[#®•©™&@·º½¾¿¡§~\)\(\]\[\}\{\|\\\/\*]{1,}") + + @property + def description(self): + return "Encode prompts with T5 after the official two-pass caption cleaning." + + @property + def expected_components(self): + return [ComponentSpec("text_encoder", T5EncoderModel), ComponentSpec("tokenizer", AutoTokenizer)] + + @property + def inputs(self): + return [ + InputParam.template("prompt", required=True), + InputParam("max_sequence_length", default=800, type_hint=int, description="Padded T5 caption length."), + InputParam( + "clean_caption", default=True, type_hint=bool, description="Apply the official two-pass text cleaning." + ), + ] + + @property + def intermediate_outputs(self): + return [ + OutputParam( + "text_embeds", + type_hint=torch.Tensor, + description="Per-prompt FP32 T5 features before special-token insertion.", + ), + OutputParam("text_attention_mask", type_hint=torch.Tensor, description="Per-prompt boolean T5 keep-mask."), + ] + + @torch.no_grad() + def __call__(self, components, state): + block_state = self.get_block_state(state) + prompts = [block_state.prompt] if isinstance(block_state.prompt, str) else block_state.prompt + if not isinstance(prompts, list) or not prompts or not all(isinstance(prompt, str) for prompt in prompts): + raise ValueError("prompt must be a string or a nonempty list of strings.") + if block_state.clean_caption: + requires_backends(self, ["bs4", "ftfy"]) + prompts = [self.clean_caption(self.clean_caption(prompt)) for prompt in prompts] + else: + prompts = [prompt.lower().strip() for prompt in prompts] + if not isinstance(block_state.max_sequence_length, int) or block_state.max_sequence_length < 1: + raise ValueError("max_sequence_length must be a positive integer.") + embeddings, masks = [], [] + for prompt in prompts: + tokens = components.tokenizer( + [prompt], + max_length=block_state.max_sequence_length, + padding="max_length", + truncation=True, + return_attention_mask=True, + add_special_tokens=True, + return_tensors="pt", + ) + tokens = tokens.to(components.text_encoder.device) + embeddings.append( + components.text_encoder( + input_ids=tokens.input_ids, attention_mask=tokens.attention_mask + ).last_hidden_state.float() + ) + masks.append(tokens.attention_mask.bool()) + block_state.text_embeds = torch.cat(embeddings, dim=0) + block_state.text_attention_mask = torch.cat(masks, dim=0) + self.set_block_state(state, block_state) + return components, state + + @staticmethod + def basic_clean(text): + text = ftfy.fix_text(text) + text = html.unescape(html.unescape(text)) + return text.strip() + + def clean_caption(self, caption): + caption = str(caption) + caption = ul.unquote_plus(caption) + caption = caption.strip().lower() + caption = re.sub("", "person", caption) + # urls: + caption = re.sub( + r"\b((?:https?:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))", # noqa + "", + caption, + ) # regex for urls + caption = re.sub( + r"\b((?:www:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))", # noqa + "", + caption, + ) # regex for urls + # html: + caption = BeautifulSoup(caption, features="html.parser").text + + # @ + caption = re.sub(r"@[\w\d]+\b", "", caption) + + # 31C0—31EF CJK Strokes + # 31F0—31FF Katakana Phonetic Extensions + # 3200—32FF Enclosed CJK Letters and Months + # 3300—33FF CJK Compatibility + # 3400—4DBF CJK Unified Ideographs Extension A + # 4DC0—4DFF Yijing Hexagram Symbols + # 4E00—9FFF CJK Unified Ideographs + caption = re.sub(r"[\u31c0-\u31ef]+", "", caption) + caption = re.sub(r"[\u31f0-\u31ff]+", "", caption) + caption = re.sub(r"[\u3200-\u32ff]+", "", caption) + caption = re.sub(r"[\u3300-\u33ff]+", "", caption) + caption = re.sub(r"[\u3400-\u4dbf]+", "", caption) + caption = re.sub(r"[\u4dc0-\u4dff]+", "", caption) + caption = re.sub(r"[\u4e00-\u9fff]+", "", caption) + ####################################################### + + # все виды тире / all types of dash --> "-" + caption = re.sub( + r"[\u002D\u058A\u05BE\u1400\u1806\u2010-\u2015\u2E17\u2E1A\u2E3A\u2E3B\u2E40\u301C\u3030\u30A0\uFE31\uFE32\uFE58\uFE63\uFF0D]+", # noqa + "-", + caption, + ) + + # кавычки к одному стандарту + caption = re.sub(r"[`´«»“”¨]", '"', caption) + caption = re.sub(r"[‘’]", "'", caption) + + # " + caption = re.sub(r""?", "", caption) + # & + caption = re.sub(r"&", "", caption) + + # ip adresses: + caption = re.sub(r"\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}", " ", caption) + + # article ids: + caption = re.sub(r"\d:\d\d\s+$", "", caption) + + # \n + caption = re.sub(r"\\n", " ", caption) + + # "#123" + caption = re.sub(r"#\d{1,3}\b", "", caption) + # "#12345.." + caption = re.sub(r"#\d{5,}\b", "", caption) + # "123456.." + caption = re.sub(r"\b\d{6,}\b", "", caption) + # filenames: + caption = re.sub(r"[\S]+\.(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)", "", caption) + + # + caption = re.sub(r"[\"\']{2,}", r'"', caption) # """AUSVERKAUFT""" + caption = re.sub(r"[\.]{2,}", r" ", caption) # """AUSVERKAUFT""" + + caption = re.sub(self.bad_punct_regex, r" ", caption) # ***AUSVERKAUFT***, #AUSVERKAUFT + caption = re.sub(r"\s+\.\s+", r" ", caption) # " . " + + # this-is-my-cute-cat / this_is_my_cute_cat + regex2 = re.compile(r"(?:\-|\_)") + if len(re.findall(regex2, caption)) > 3: + caption = re.sub(regex2, " ", caption) + + caption = self.basic_clean(caption) + + caption = re.sub(r"\b[a-zA-Z]{1,3}\d{3,15}\b", "", caption) # jc6640 + caption = re.sub(r"\b[a-zA-Z]+\d+[a-zA-Z]+\b", "", caption) # jc6640vc + caption = re.sub(r"\b\d+[a-zA-Z]+\d+\b", "", caption) # 6640vc231 + + caption = re.sub(r"(worldwide\s+)?(free\s+)?shipping", "", caption) + caption = re.sub(r"(free\s)?download(\sfree)?", "", caption) + caption = re.sub(r"\bclick\b\s(?:for|on)\s\w+", "", caption) + caption = re.sub(r"\b(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)(\simage[s]?)?", "", caption) + caption = re.sub(r"\bpage\s+\d+\b", "", caption) + + caption = re.sub(r"\b\d*[a-zA-Z]+\d+[a-zA-Z]+\d+[a-zA-Z\d]*\b", r" ", caption) # j2d1a2a... + + caption = re.sub(r"\b\d+\.?\d*[xх×]\d+\.?\d*\b", "", caption) + + caption = re.sub(r"\b\s+\:\s+", r": ", caption) + caption = re.sub(r"(\D[,\./])\b", r"\1 ", caption) + caption = re.sub(r"\s+", " ", caption) + + caption = re.sub(r"^[\"\']([\w\W]+)[\"\']$", r"\1", caption) + caption = re.sub(r"^[\'\_,\-\:;]", r"", caption) + caption = re.sub(r"[\'\_,\-\:\-\+]$", r"", caption) + caption = re.sub(r"^\.\S+$", "", caption) + + return caption.strip() + + +def encode_magi_prefix(vae, video, scaling_factor, device): + if not isinstance(video, torch.Tensor) or video.ndim != 5 or video.dtype != torch.uint8: + raise ValueError("Prefix pixels must be a uint8 tensor shaped (batch, 3, frames, height, width).") + if min(video.shape) <= 0 or video.shape[1] != 3: + raise ValueError("Prefix pixels must be nonempty RGB frames.") + if scaling_factor <= 0: + raise ValueError("latent_scaling_factor must be positive.") + pixels = (video.to(device=device, dtype=torch.float32) / 127.5 - 1).to(vae.dtype) + return vae.encode(pixels).latent_dist.mode() * scaling_factor + + +class MagiVideoVaeEncoderStep(ModularPipelineBlocks): + model_name = "magi" + + @property + def description(self): + return "Encode pre-resized uint8 RGB video frames as a deterministic, scaled VAE prefix." + + @property + def expected_components(self): + return [ComponentSpec("vae", AutoencoderKLMagi)] + + @property + def expected_configs(self): + return [ConfigSpec("latent_scaling_factor", 0.18215)] + + @property + def inputs(self): + return [ + InputParam( + "video", + required=True, + type_hint=torch.Tensor, + description="Pre-resized uint8 RGB prefix, shaped (batch, 3, frames, height, width).", + ) + ] + + @property + def intermediate_outputs(self): + return [ + OutputParam( + "conditioning_latents", + type_hint=torch.Tensor, + description="Per-prompt scaled VAE prefix, before video-batch expansion.", + ) + ] + + @torch.no_grad() + def __call__(self, components, state): + s = self.get_block_state(state) + s.conditioning_latents = encode_magi_prefix( + components.vae, s.video, components.config.latent_scaling_factor, components._execution_device + ) + self.set_block_state(state, s) + return components, state + + +class MagiImageVaeEncoderStep(MagiVideoVaeEncoderStep): + @property + def description(self): + return "Encode pre-resized uint8 RGB images as a one-latent-frame prefix." + + @property + def inputs(self): + return [ + InputParam( + "image", + required=True, + type_hint=torch.Tensor, + description="Pre-resized uint8 RGB images, shaped (batch, 3, height, width).", + ) + ] + + @torch.no_grad() + def __call__(self, components, state): + s = self.get_block_state(state) + if not isinstance(s.image, torch.Tensor) or s.image.ndim != 4: + raise ValueError("image must have shape (batch, 3, height, width).") + s.conditioning_latents = encode_magi_prefix( + components.vae, s.image.unsqueeze(2), components.config.latent_scaling_factor, components._execution_device + ) + self.set_block_state(state, s) + return components, state diff --git a/src/diffusers/modular_pipelines/magi/modular_blocks_magi.py b/src/diffusers/modular_pipelines/magi/modular_blocks_magi.py new file mode 100644 index 000000000000..6fff51505cb3 --- /dev/null +++ b/src/diffusers/modular_pipelines/magi/modular_blocks_magi.py @@ -0,0 +1,410 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from ..modular_pipeline import SequentialPipelineBlocks +from ..modular_pipeline_utils import OutputParam +from .before_denoise import MagiPrepareConditionedLatentsStep, MagiPrepareLatentsStep +from .decoders import MagiPrefixVaeDecoderStep, MagiVaeDecoderStep +from .denoise import MagiDenoiseStep, MagiPrefixDenoiseStep +from .encoders import MagiImageVaeEncoderStep, MagiTextEncoderStep, MagiVideoVaeEncoderStep + + +# auto_docstring +class MagiTextToVideoBlocks(SequentialPipelineBlocks): + """ + Generate videos with a MAGI base model, using official HQ and duration conditioning. + + Components: + text_encoder (`T5EncoderModel`) tokenizer (`AutoTokenizer`) transformer (`MagiTransformer3DModel`) vae + (`AutoencoderKLMagi`) text_conditioning (`MagiTextConditioningModel`) scheduler (`MagiEulerScheduler`) guider + (`MagiClassifierFreeGuidance`) video_processor (`VideoProcessor`) + + Configs: + latent_scaling_factor (default: 0.18215) + + Inputs: + prompt (`str`): + The prompt or prompts to guide image generation. + max_sequence_length (`int`, *optional*, defaults to 800): + Padded T5 caption length. + clean_caption (`bool`, *optional*, defaults to True): + Apply the official two-pass text cleaning. + height (`int`, *optional*, defaults to 720): + Video height in pixels. + width (`int`, *optional*, defaults to 720): + Video width in pixels. + num_frames (`int`, *optional*, defaults to 96): + Requested video frames; generation rounds up to full chunks. + chunk_width (`int`, *optional*, defaults to 6): + Latent frames per chunk. + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + latents (`Tensor`, *optional*): + Optional initial FP32 noise for all generated chunks. + prefix_latents (`Tensor`, *optional*): + Not supported by this text-to-video preparation block. + num_inference_steps (`int`, *optional*, defaults to 64): + Number of Euler updates per generated chunk. + window_size (`int`, *optional*, defaults to 4): + Maximum number of simultaneously denoised chunks. + noise2clean_kvrange (`tuple`, *optional*, defaults to (5, 4, 3, 2)): + Positive attention-window lengths in chunks, from early to late denoising stages. + clean_chunk_kvrange (`int`, *optional*, defaults to 1): + Positive attention-window length used when recomputing clean chunks. + clean_t (`float`, *optional*, defaults to 0.9999): + Model evaluation time for clean-prefix cache extraction. + attention_kwargs (`dict`, *optional*): + Optional keyword arguments passed to Transformer attention. + cache_device (`str`, *optional*): + Optional device for clean-prefix KV storage; use cpu to offload between layer evaluations. + output_type (`str`, *optional*, defaults to np): + Output format: pt, np, pil, or latent. + + Outputs: + text_embeds (`Tensor`): + Per-prompt FP32 T5 features before special-token insertion. + text_attention_mask (`Tensor`): + Per-prompt boolean T5 keep-mask. + latents (`Tensor`): + Denoised latents. + prompt_embeds (`Tensor`): + HQ/duration conditioned chunk text features. + prompt_attention_mask (`Tensor`): + Conditional text keep-mask. + negative_prompt_embeds (`Tensor`): + Learned null text features. + negative_prompt_attention_mask (`Tensor`): + Null text keep-mask. + num_chunks (`int`): + Total number of latent chunks, including the supplied prefix. + prefix_chunks (`int`): + Number of supplied full-chunk prefix chunks. + chunk_tokens (`int`): + Number of Transformer tokens per latent chunk. + steps_per_stage (`int`): + Number of iterations before the active chunk window moves forward. + num_window_steps (`int`): + Total number of asynchronous window iterations. + timestep_schedule (`Tensor`): + FP32 schedule including the final Euler integration endpoint. + clean_kv_cache (`tuple`): + Per-layer clean-prefix key/value tensors, or None before any prefix is cached; excludes the final + generated chunk. + completed_chunks (`list`): + Indices of supplied prefix chunks and finalized generated chunks. + chunk_start (`int`): + First active chunk index. + chunk_end (`int`): + Exclusive end index of the active chunk window. + refresh_cache (`bool`): + Whether this iteration prepends a finalized chunk to refresh its clean KV. + chunk_step_indices (`list`): + Denoising step indices for active chunks, ordered from oldest to newest. + chunk_times (`Tensor`): + Current per-chunk model times shaped (batch, active_chunks). + next_chunk_times (`Tensor`): + Next Euler endpoints shaped (batch, active_chunks). + model_times (`Tensor`): + Model times including an optional clean-refresh chunk. + latent_model_input (`Tensor`): + Current latent window including an optional clean-refresh chunk. + window_prompt_embeds (`Tensor`): + Conditional features for the current window, with null features for a clean-refresh chunk. + window_prompt_attention_mask (`Tensor`): + Boolean keep-mask matching the current window text features. + kv_ranges (`tuple`): + Exclusive token attention ranges indexing the cached prefix plus the current window. + velocity (`Tensor`): + Three-way guided FP32 velocities for active chunks only. + videos (`list`): + The generated videos. + """ + + model_name = "magi" + block_classes = [MagiTextEncoderStep, MagiPrepareLatentsStep, MagiDenoiseStep, MagiVaeDecoderStep] + block_names = ["text_encoder", "prepare_latents", "denoise", "decode"] + + @property + def outputs(self): + return [OutputParam.template("latents") if param.name == "latents" else param for param in super().outputs] + + @property + def description(self): + return "Generate videos with a MAGI base model, using official HQ and duration conditioning." + + +# auto_docstring +class MagiImageToVideoBlocks(MagiTextToVideoBlocks): + """ + Generate a MAGI base-model video conditioned on a pre-resized uint8 RGB image. + + Components: + text_encoder (`T5EncoderModel`) tokenizer (`AutoTokenizer`) vae (`AutoencoderKLMagi`) transformer + (`MagiTransformer3DModel`) text_conditioning (`MagiTextConditioningModel`) scheduler (`MagiEulerScheduler`) + guider (`MagiClassifierFreeGuidance`) video_processor (`VideoProcessor`) + + Configs: + latent_scaling_factor (default: 0.18215) + + Inputs: + prompt (`str`): + The prompt or prompts to guide image generation. + max_sequence_length (`int`, *optional*, defaults to 800): + Padded T5 caption length. + clean_caption (`bool`, *optional*, defaults to True): + Apply the official two-pass text cleaning. + image (`Tensor`): + Pre-resized uint8 RGB images, shaped (batch, 3, height, width). + height (`int`, *optional*, defaults to 720): + Video height in pixels. + width (`int`, *optional*, defaults to 720): + Video width in pixels. + chunk_width (`int`, *optional*, defaults to 6): + Latent frames per chunk. + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + latents (`Tensor`, *optional*): + Optional initial FP32 noise for all generated chunks. + num_frames (`int`, *optional*, defaults to 96): + Requested new frames; prefix plus new frames rounds up to full latent chunks. + num_inference_steps (`int`, *optional*, defaults to 64): + Number of Euler updates per generated chunk. + window_size (`int`, *optional*, defaults to 4): + Maximum number of simultaneously denoised chunks. + noise2clean_kvrange (`tuple`, *optional*, defaults to (5, 4, 3, 2)): + Positive attention-window lengths in chunks, from early to late denoising stages. + clean_chunk_kvrange (`int`, *optional*, defaults to 1): + Positive attention-window length used when recomputing clean chunks. + clean_t (`float`, *optional*, defaults to 0.9999): + Model evaluation time for clean-prefix cache extraction. + attention_kwargs (`dict`, *optional*): + Optional keyword arguments passed to Transformer attention. + cache_device (`str`, *optional*): + Optional device for clean-prefix KV storage; use cpu to offload between layer evaluations. + output_type (`str`, *optional*, defaults to np): + Output format: pt, np, pil, or latent. + + Outputs: + text_embeds (`Tensor`): + Per-prompt FP32 T5 features before special-token insertion. + text_attention_mask (`Tensor`): + Per-prompt boolean T5 keep-mask. + conditioning_latents (`Tensor`): + Per-prompt scaled VAE prefix, before video-batch expansion. + latents (`Tensor`): + Denoised latents. + prompt_embeds (`Tensor`): + HQ/duration conditioned chunk text features. + prompt_attention_mask (`Tensor`): + Conditional text keep-mask. + negative_prompt_embeds (`Tensor`): + Learned null text features. + negative_prompt_attention_mask (`Tensor`): + Null text keep-mask. + prefix_latents (`Tensor`): + Optional full-chunk clean prefix; replaces the leading latent slots and remains unchanged. + num_chunks (`int`): + Total number of latent chunks, including the supplied prefix. + prefix_chunks (`int`): + Number of supplied full-chunk prefix chunks. + chunk_tokens (`int`): + Number of Transformer tokens per latent chunk. + steps_per_stage (`int`): + Number of iterations before the active chunk window moves forward. + num_window_steps (`int`): + Total number of asynchronous window iterations. + timestep_schedule (`Tensor`): + FP32 schedule including the final Euler integration endpoint. + clean_kv_cache (`tuple`): + Per-layer clean-prefix key/value tensors, or None before any prefix is cached; excludes the final + generated chunk. + completed_chunks (`list`): + Indices of supplied prefix chunks and finalized generated chunks. + chunk_start (`int`): + First active chunk index. + chunk_end (`int`): + Exclusive end index of the active chunk window. + refresh_cache (`bool`): + Whether this iteration prepends a finalized chunk to refresh its clean KV. + chunk_step_indices (`list`): + Denoising step indices for active chunks, ordered from oldest to newest. + chunk_times (`Tensor`): + Current per-chunk model times shaped (batch, active_chunks). + next_chunk_times (`Tensor`): + Next Euler endpoints shaped (batch, active_chunks). + model_times (`Tensor`): + Model times including an optional clean-refresh chunk. + latent_model_input (`Tensor`): + Current latent window including an optional clean-refresh chunk. + window_prompt_embeds (`Tensor`): + Conditional features for the current window, with null features for a clean-refresh chunk. + window_prompt_attention_mask (`Tensor`): + Boolean keep-mask matching the current window text features. + kv_ranges (`tuple`): + Exclusive token attention ranges indexing the cached prefix plus the current window. + velocity (`Tensor`): + Three-way guided FP32 velocities for active chunks only. + videos (`list`): + The generated videos. + """ + + model_name = "magi" + block_classes = [ + MagiTextEncoderStep, + MagiImageVaeEncoderStep, + MagiPrepareConditionedLatentsStep, + MagiPrefixDenoiseStep, + MagiPrefixVaeDecoderStep, + ] + block_names = ["text_encoder", "vae_encoder", "prepare_latents", "denoise", "decode"] + + @property + def description(self): + return "Generate a MAGI base-model video conditioned on a pre-resized uint8 RGB image." + + +# auto_docstring +class MagiVideoToVideoBlocks(MagiImageToVideoBlocks): + """ + Continue pre-resized uint8 RGB video frames with a MAGI base model. + + Components: + text_encoder (`T5EncoderModel`) tokenizer (`AutoTokenizer`) vae (`AutoencoderKLMagi`) transformer + (`MagiTransformer3DModel`) text_conditioning (`MagiTextConditioningModel`) scheduler (`MagiEulerScheduler`) + guider (`MagiClassifierFreeGuidance`) video_processor (`VideoProcessor`) + + Configs: + latent_scaling_factor (default: 0.18215) + + Inputs: + prompt (`str`): + The prompt or prompts to guide image generation. + max_sequence_length (`int`, *optional*, defaults to 800): + Padded T5 caption length. + clean_caption (`bool`, *optional*, defaults to True): + Apply the official two-pass text cleaning. + video (`Tensor`): + Pre-resized uint8 RGB prefix, shaped (batch, 3, frames, height, width). + height (`int`, *optional*, defaults to 720): + Video height in pixels. + width (`int`, *optional*, defaults to 720): + Video width in pixels. + chunk_width (`int`, *optional*, defaults to 6): + Latent frames per chunk. + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + latents (`Tensor`, *optional*): + Optional initial FP32 noise for all generated chunks. + num_frames (`int`, *optional*, defaults to 96): + Requested new frames; prefix plus new frames rounds up to full latent chunks. + num_inference_steps (`int`, *optional*, defaults to 64): + Number of Euler updates per generated chunk. + window_size (`int`, *optional*, defaults to 4): + Maximum number of simultaneously denoised chunks. + noise2clean_kvrange (`tuple`, *optional*, defaults to (5, 4, 3, 2)): + Positive attention-window lengths in chunks, from early to late denoising stages. + clean_chunk_kvrange (`int`, *optional*, defaults to 1): + Positive attention-window length used when recomputing clean chunks. + clean_t (`float`, *optional*, defaults to 0.9999): + Model evaluation time for clean-prefix cache extraction. + attention_kwargs (`dict`, *optional*): + Optional keyword arguments passed to Transformer attention. + cache_device (`str`, *optional*): + Optional device for clean-prefix KV storage; use cpu to offload between layer evaluations. + output_type (`str`, *optional*, defaults to np): + Output format: pt, np, pil, or latent. + + Outputs: + text_embeds (`Tensor`): + Per-prompt FP32 T5 features before special-token insertion. + text_attention_mask (`Tensor`): + Per-prompt boolean T5 keep-mask. + conditioning_latents (`Tensor`): + Per-prompt scaled VAE prefix, before video-batch expansion. + latents (`Tensor`): + Denoised latents. + prompt_embeds (`Tensor`): + HQ/duration conditioned chunk text features. + prompt_attention_mask (`Tensor`): + Conditional text keep-mask. + negative_prompt_embeds (`Tensor`): + Learned null text features. + negative_prompt_attention_mask (`Tensor`): + Null text keep-mask. + prefix_latents (`Tensor`): + Optional full-chunk clean prefix; replaces the leading latent slots and remains unchanged. + num_chunks (`int`): + Total number of latent chunks, including the supplied prefix. + prefix_chunks (`int`): + Number of supplied full-chunk prefix chunks. + chunk_tokens (`int`): + Number of Transformer tokens per latent chunk. + steps_per_stage (`int`): + Number of iterations before the active chunk window moves forward. + num_window_steps (`int`): + Total number of asynchronous window iterations. + timestep_schedule (`Tensor`): + FP32 schedule including the final Euler integration endpoint. + clean_kv_cache (`tuple`): + Per-layer clean-prefix key/value tensors, or None before any prefix is cached; excludes the final + generated chunk. + completed_chunks (`list`): + Indices of supplied prefix chunks and finalized generated chunks. + chunk_start (`int`): + First active chunk index. + chunk_end (`int`): + Exclusive end index of the active chunk window. + refresh_cache (`bool`): + Whether this iteration prepends a finalized chunk to refresh its clean KV. + chunk_step_indices (`list`): + Denoising step indices for active chunks, ordered from oldest to newest. + chunk_times (`Tensor`): + Current per-chunk model times shaped (batch, active_chunks). + next_chunk_times (`Tensor`): + Next Euler endpoints shaped (batch, active_chunks). + model_times (`Tensor`): + Model times including an optional clean-refresh chunk. + latent_model_input (`Tensor`): + Current latent window including an optional clean-refresh chunk. + window_prompt_embeds (`Tensor`): + Conditional features for the current window, with null features for a clean-refresh chunk. + window_prompt_attention_mask (`Tensor`): + Boolean keep-mask matching the current window text features. + kv_ranges (`tuple`): + Exclusive token attention ranges indexing the cached prefix plus the current window. + velocity (`Tensor`): + Three-way guided FP32 velocities for active chunks only. + videos (`list`): + The generated videos. + """ + + model_name = "magi" + block_classes = [ + MagiTextEncoderStep, + MagiVideoVaeEncoderStep, + MagiPrepareConditionedLatentsStep, + MagiPrefixDenoiseStep, + MagiPrefixVaeDecoderStep, + ] + block_names = ["text_encoder", "vae_encoder", "prepare_latents", "denoise", "decode"] + + @property + def description(self): + return "Continue pre-resized uint8 RGB video frames with a MAGI base model." diff --git a/src/diffusers/modular_pipelines/magi/modular_pipeline.py b/src/diffusers/modular_pipelines/magi/modular_pipeline.py new file mode 100644 index 000000000000..86d3e536425c --- /dev/null +++ b/src/diffusers/modular_pipelines/magi/modular_pipeline.py @@ -0,0 +1,21 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from ..modular_pipeline import ModularPipeline + + +class MagiModularPipeline(ModularPipeline): + """MAGI-1 base text-to-video pipeline with separate text, denoising, and decoding blocks.""" + + default_blocks_name = "MagiTextToVideoBlocks" diff --git a/src/diffusers/modular_pipelines/modular_pipeline.py b/src/diffusers/modular_pipelines/modular_pipeline.py index 69a8f730284f..5df705d02ac0 100644 --- a/src/diffusers/modular_pipelines/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/modular_pipeline.py @@ -131,6 +131,7 @@ def _helios_pyramid_map_fn(config_dict=None): [ ("stable-diffusion-xl", _create_default_map_fn("StableDiffusionXLModularPipeline")), ("stable-diffusion-3", _create_default_map_fn("StableDiffusion3ModularPipeline")), + ("magi", _create_default_map_fn("MagiModularPipeline")), ("wan", _wan_map_fn), ("wan-animate-2", _create_default_map_fn("WanAnimate2ModularPipeline")), ("wan-animate-2-distilled", _create_default_map_fn("WanAnimate2DistilledModularPipeline")), diff --git a/src/diffusers/schedulers/__init__.py b/src/diffusers/schedulers/__init__.py index c0e46ef445df..532c6e2f3b3f 100644 --- a/src/diffusers/schedulers/__init__.py +++ b/src/diffusers/schedulers/__init__.py @@ -72,6 +72,7 @@ _import_structure["scheduling_k_dpm_2_discrete"] = ["KDPM2DiscreteScheduler"] _import_structure["scheduling_lcm"] = ["LCMScheduler"] _import_structure["scheduling_ltx_euler_ancestral_rf"] = ["LTXEulerAncestralRFScheduler"] + _import_structure["scheduling_magi_euler"] = ["MagiEulerScheduler"] _import_structure["scheduling_minimax_h3"] = ["MiniMaxH3Scheduler"] _import_structure["scheduling_pndm"] = ["PNDMScheduler"] _import_structure["scheduling_repaint"] = ["RePaintScheduler"] @@ -156,6 +157,7 @@ from .scheduling_k_dpm_2_discrete import KDPM2DiscreteScheduler from .scheduling_lcm import LCMScheduler from .scheduling_ltx_euler_ancestral_rf import LTXEulerAncestralRFScheduler + from .scheduling_magi_euler import MagiEulerScheduler from .scheduling_minimax_h3 import MiniMaxH3Scheduler from .scheduling_pndm import PNDMScheduler from .scheduling_repaint import RePaintScheduler diff --git a/src/diffusers/schedulers/scheduling_magi_euler.py b/src/diffusers/schedulers/scheduling_magi_euler.py new file mode 100644 index 000000000000..e14336908e44 --- /dev/null +++ b/src/diffusers/schedulers/scheduling_magi_euler.py @@ -0,0 +1,194 @@ +# Copyright 2025 SandAI and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import math +from dataclasses import dataclass + +import torch + +from ..configuration_utils import ConfigMixin, register_to_config +from ..utils import BaseOutput +from .scheduling_utils import SchedulerMixin + + +@dataclass +class MagiEulerSchedulerOutput(BaseOutput): + """ + Output of a MAGI Euler update. + + Args: + prev_sample (`torch.Tensor`): The FP32 sample at the next, cleaner timestep. + """ + + prev_sample: torch.Tensor + + +class MagiEulerScheduler(SchedulerMixin, ConfigMixin): + """ + Euler integration of MAGI velocity predictions, with time increasing from noise (0) to clean data (1). + + Args: + shift (`float`, defaults to 3.0): Shift applied after squaring time when `time_schedule="sd3"`. + time_schedule (`str`, defaults to `"sd3"`): One of `"sd3"`, `"square"`, `"piecewise"`, or `"linear"`. + shortcut_mode (`str`, defaults to `"8,16,16"`): The 12-step grid ordering, either `"8,16,16"` or `"16,16,8"`. + """ + + order = 1 + _compatibles = [] + + @register_to_config + def __init__(self, shift: float = 3.0, time_schedule: str = "sd3", shortcut_mode: str = "8,16,16"): + if not math.isfinite(shift) or shift < 1: + raise ValueError("shift must be finite and at least 1.") + if time_schedule not in {"sd3", "square", "piecewise", "linear"}: + raise ValueError("time_schedule must be sd3, square, piecewise, or linear.") + if shortcut_mode not in {"8,16,16", "16,16,8"}: + raise ValueError("shortcut_mode must be 8,16,16 or 16,16,8.") + self.init_noise_sigma = 1.0 + self.num_inference_steps = None + self.timesteps = None + self.timestep_schedule = None + self._step_index = None + + @property + def step_index(self): + """Index of the next sequential update; explicit endpoint updates do not change it.""" + return self._step_index + + def set_timesteps(self, num_inference_steps: int, device: str | torch.device = None): + """ + Build the official FP32 grid on the execution device and reset sequential stepping. + + Args: + num_inference_steps (`int`): Positive number of updates per chunk, not total chunk-window model calls. + device (`str` or `torch.device`, optional): Device on which to compute the time grid. + + `timesteps` contains model evaluation times. `timestep_schedule` also includes the integration endpoint. The + endpoint retains the reference's floating-point rounding instead of being clamped to exactly 1. + """ + if ( + isinstance(num_inference_steps, bool) + or not isinstance(num_inference_steps, int) + or num_inference_steps <= 0 + ): + raise ValueError("num_inference_steps must be a positive integer.") + if num_inference_steps == 12: + base_t = torch.linspace(0, 1, 5, device=device, dtype=torch.float32) / 4 + offsets = torch.linspace(0, 1, 5, device=device, dtype=torch.float32) + if self.config.shortcut_mode == "16,16,8": + base_t = base_t[:3] + else: + base_t = torch.cat([base_t[:1], base_t[2:4]], dim=0) + timesteps = torch.cat([base_t + offset for offset in offsets], dim=0)[:13] + else: + timesteps = torch.linspace(0, 1, num_inference_steps + 1, device=device, dtype=torch.float32) + if self.config.time_schedule == "sd3": + timesteps = timesteps**2 + inverse_shift = 1.0 / self.config.shift + timesteps = inverse_shift * timesteps / (1 + (inverse_shift - 1) * timesteps) + elif self.config.time_schedule == "square": + timesteps = timesteps**2 + elif self.config.time_schedule == "piecewise": + mask = timesteps < 0.875 + timesteps[mask] = timesteps[mask] * (0.5 / 0.875) + timesteps[~mask] = 0.5 + (timesteps[~mask] - 0.875) * (0.5 / (1 - 0.875)) + self.num_inference_steps = num_inference_steps + self.timestep_schedule = timesteps + self.timesteps = timesteps[:-1] + self._step_index = None + + def step( + self, + model_output: torch.Tensor, + timestep: float | torch.Tensor, + sample: torch.Tensor, + next_timestep: float | torch.Tensor | None = None, + return_dict: bool = True, + ) -> MagiEulerSchedulerOutput | tuple: + """ + Advance the sample with `sample + velocity * (next_timestep - timestep)` in FP32. + + Args: + model_output (`torch.Tensor`): Predicted velocity, after guidance, with the same shape as `sample`. + timestep (`float` or `torch.Tensor`): Normalized time, not a loop index or a time multiplied by 1000. + sample (`torch.Tensor`): + Current sample. Chunk-wise updates require `(batch, channels, frames, height, width)`. + next_timestep (`float` or `torch.Tensor`, optional): + Explicit target time. Endpoint tensors can be scalars, `(chunks,)`, or `(batch, chunks)` and must + broadcast together. Chunks split frames equally. Explicit updates do not advance `step_index`. If + omitted, advance sequentially from a scalar schedule timestep. + return_dict (`bool`, defaults to `True`): Return a structured output instead of a tuple. + + Returns: + `MagiEulerSchedulerOutput` or `tuple`: The next FP32 sample. No prediction-to-velocity conversion is + applied. + """ + if self.num_inference_steps is None: + raise ValueError("Call set_timesteps before step.") + if model_output.shape != sample.shape: + raise ValueError("model_output and sample must have the same shape.") + timestep = torch.as_tensor(timestep, dtype=torch.float32, device=sample.device) + sequential = next_timestep is None + step_index = self._step_index + if sequential: + if timestep.numel() != 1: + raise ValueError("Chunk-wise updates require next_timestep.") + timestep = timestep.reshape(()) + if step_index is None: + indices = (self.timesteps == timestep.to(self.timesteps.device)).nonzero().flatten() + if indices.numel() != 1: + raise ValueError("timestep must be a scalar from scheduler.timesteps.") + step_index = indices.item() + if step_index >= self.num_inference_steps: + raise ValueError("The schedule is complete; call set_timesteps to restart.") + if timestep != self.timesteps[step_index].to(sample.device): + raise ValueError("timestep does not match the next sequential step.") + next_timestep = self.timestep_schedule[step_index + 1] + next_timestep = torch.as_tensor(next_timestep, dtype=torch.float32, device=sample.device) + timestep, next_timestep = torch.broadcast_tensors(timestep, next_timestep) + if ( + not torch.isfinite(timestep).all() + or not torch.isfinite(next_timestep).all() + or (timestep < 0).any() + or (next_timestep > 1 + 1e-5).any() + or (next_timestep < timestep).any() + ): + raise ValueError("Timesteps must be finite, increasing from 0 to 1 (allowing endpoint rounding).") + delta_t = next_timestep - timestep + if delta_t.ndim == 0: + prev_sample = sample.float() + model_output.float() * delta_t + else: + if delta_t.ndim == 1: + delta_t = delta_t[None] + if ( + sample.ndim != 5 + or delta_t.ndim != 2 + or delta_t.shape[0] not in (1, sample.shape[0]) + or delta_t.shape[1] == 0 + or sample.shape[2] % delta_t.shape[1] + ): + raise ValueError( + "Chunk times must have shape (chunks,) or (batch, chunks), dividing video frames evenly." + ) + chunk_shape = (*sample.shape[:2], delta_t.shape[1], -1, *sample.shape[3:]) + prev_sample = ( + sample.float().reshape(chunk_shape) + + model_output.float().reshape(chunk_shape) * delta_t[:, None, :, None, None, None] + ) + prev_sample = prev_sample.reshape(sample.shape) + if sequential: + self._step_index = step_index + 1 + if not return_dict: + return (prev_sample,) + return MagiEulerSchedulerOutput(prev_sample=prev_sample) diff --git a/src/diffusers/utils/dummy_pt_objects.py b/src/diffusers/utils/dummy_pt_objects.py index 1598814f835a..6e3dc405a0a3 100644 --- a/src/diffusers/utils/dummy_pt_objects.py +++ b/src/diffusers/utils/dummy_pt_objects.py @@ -122,6 +122,21 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) +class MagiClassifierFreeGuidance(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + class PerturbedAttentionGuidance(metaclass=DummyObject): _backends = ["torch"] @@ -765,6 +780,21 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) +class AutoencoderKLMagi(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + class AutoencoderKLMagvit(metaclass=DummyObject): _backends = ["torch"] @@ -1755,6 +1785,36 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) +class MagiTextConditioningModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + +class MagiTransformer3DModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + class MiniMaxH3Transformer3DModel(metaclass=DummyObject): _backends = ["torch"] @@ -3617,6 +3677,21 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) +class MagiEulerScheduler(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + class MiniMaxH3Scheduler(metaclass=DummyObject): _backends = ["torch"] diff --git a/src/diffusers/utils/dummy_torch_and_transformers_objects.py b/src/diffusers/utils/dummy_torch_and_transformers_objects.py index 17bca9f23414..5dfed892f148 100644 --- a/src/diffusers/utils/dummy_torch_and_transformers_objects.py +++ b/src/diffusers/utils/dummy_torch_and_transformers_objects.py @@ -572,6 +572,126 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch", "transformers"]) +class MagiDenoiseStep(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + +class MagiImageToVideoBlocks(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + +class MagiModularPipeline(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + +class MagiPrepareLatentsStep(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + +class MagiTextEncoderStep(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + +class MagiTextToVideoBlocks(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + +class MagiVaeDecoderStep(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + +class MagiVideoToVideoBlocks(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + class MiniMaxH3Blocks(metaclass=DummyObject): _backends = ["torch", "transformers"] diff --git a/tests/models/autoencoders/test_models_autoencoder_kl_magi.py b/tests/models/autoencoders/test_models_autoencoder_kl_magi.py new file mode 100644 index 000000000000..148d588d16b8 --- /dev/null +++ b/tests/models/autoencoders/test_models_autoencoder_kl_magi.py @@ -0,0 +1,279 @@ +# Copyright 2025 SandAI and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +import torch + +from diffusers import AutoencoderKLMagi +from diffusers.utils.torch_utils import randn_tensor + +from ...testing_utils import enable_full_determinism, require_torch_multi_accelerator, torch_device +from ..testing_utils import ( + AttentionTesterMixin, + BaseModelTesterConfig, + MemoryTesterMixin, + ModelTesterMixin, + TorchCompileTesterMixin, + TrainingTesterMixin, +) + + +enable_full_determinism() + + +class AutoencoderKLMagiTesterConfig(BaseModelTesterConfig): + main_input_name = "sample" + + @property + def model_class(self): + return AutoencoderKLMagi + + @property + def pretrained_model_name_or_path(self): + return None + + @property + def pretrained_model_kwargs(self): + return {} + + @property + def generator(self): + return torch.Generator("cpu").manual_seed(0) + + def get_init_dict(self): + return { + "latent_channels": 4, + "embed_dim": 32, + "num_layers": 2, + "num_attention_heads": 4, + "mlp_ratio": 2, + "patch_size": 2, + "patch_length": 4, + "sample_size": 8, + "sample_frames": 8, + } + + def get_dummy_inputs(self): + return {"sample": randn_tensor((2, 3, 8, 8, 8), generator=self.generator, device=torch_device)} + + @property + def input_shape(self): + return (3, 8, 8, 8) + + @property + def output_shape(self): + return (3, 8, 8, 8) + + +class TestAutoencoderKLMagiModel(AutoencoderKLMagiTesterConfig, ModelTesterMixin): + @require_torch_multi_accelerator + @torch.no_grad() + def test_model_parallelism(self, base_model_output, tmp_path): + torch.manual_seed(0) + model = self.model_class(**self.get_init_dict()).eval() + model.save_pretrained(tmp_path) + model = self.model_class.from_pretrained(tmp_path, device_map={"encoder": 0, "decoder": 1}) + output = model(**self.get_dummy_inputs()).sample + assert next(model.encoder.parameters()).device == torch.device("cuda:0") + assert next(model.decoder.parameters()).device == torch.device("cuda:1") + torch.testing.assert_close(output.cpu(), base_model_output.cpu(), atol=1e-5, rtol=0) + + @torch.no_grad() + def test_single_frame_encode_and_decode(self): + model = self.model_class(**self.get_init_dict()).to(torch_device).eval() + image = self.get_dummy_inputs()["sample"][:, :, :1] + posterior = model.encode(image).latent_dist + repeated = model.encode(image.repeat(1, 1, 4, 1, 1)).latent_dist + torch.testing.assert_close(posterior.parameters, repeated.parameters) + latent = posterior.mode() + assert latent.shape == (2, 4, 1, 4, 4) + decoded = model.decode(latent).sample + assert decoded.shape == image.shape + torch.testing.assert_close(model(image).sample, decoded) + torch.testing.assert_close(decoded, model.decoder(latent)[:, :, :1]) + + @torch.no_grad() + def test_encoder_preserves_channel_last_storage(self): + model = self.model_class(**self.get_init_dict()).to(torch_device).eval() + moments = model.encoder(self.get_dummy_inputs()["sample"]) + assert moments.stride(1) == 1 + assert not moments.is_contiguous() + model.enable_tiling() + posterior = model.encode(self.get_dummy_inputs()["sample"][:1]).latent_dist + assert posterior.mean.is_contiguous(memory_format=torch.channels_last_3d) + torch.testing.assert_close(posterior.mean, posterior.parameters.chunk(2, dim=1)[0], atol=0, rtol=0) + + @torch.no_grad() + def test_four_frame_forward_preserves_frames(self): + model = self.model_class(**self.get_init_dict()).to(torch_device).eval() + video = self.get_dummy_inputs()["sample"][:, :, :4] + assert model(video).sample.shape == video.shape + + @torch.no_grad() + def test_position_interpolation(self): + model = self.model_class(**self.get_init_dict()).to(torch_device).eval() + sample = randn_tensor((1, 3, 12, 12, 8), generator=self.generator, device=torch_device) + posterior = model.encode(sample).latent_dist + assert posterior.mean.shape == (1, 4, 3, 6, 4) + assert model.decode(posterior.mode()).sample.shape == sample.shape + + @pytest.mark.parametrize("shape", [(1, 3, 6, 8, 8), (1, 3, 8, 7, 8)]) + def test_invalid_video_shape(self, shape): + model = self.model_class(**self.get_init_dict()).to(torch_device) + with pytest.raises(ValueError, match="divisible"): + model.encode(torch.zeros(shape, device=torch_device)) + + @torch.no_grad() + def test_posterior_sampling_generator(self): + model = self.model_class(**self.get_init_dict()).to(torch_device).eval() + sample = self.get_dummy_inputs()["sample"] + first = model(sample, sample_posterior=True, generator=self.generator).sample + second = model(sample, sample_posterior=True, generator=self.generator).sample + torch.testing.assert_close(first, second, rtol=0, atol=0) + + @pytest.mark.parametrize("frames", [1, 4, 8, 12, 16, 17]) + @torch.no_grad() + def test_temporal_tiles_and_frame_count(self, frames): + model = self.model_class(**self.get_init_dict()).to(torch_device).eval() + sample = randn_tensor((1, 3, frames, 8, 8), generator=self.generator, device=torch_device) + expected_latents = torch.cat([model.encode(tile).latent_dist.mode() for tile in sample.split(8, dim=2)], dim=2) + expected_video = torch.cat([model.decode(tile).sample for tile in expected_latents.split(2, dim=2)], dim=2) + model.enable_tiling(tile_sample_min_length=8, allow_spatial_tiling=False) + actual_latents = model.encode(sample).latent_dist.mode() + torch.testing.assert_close(actual_latents, expected_latents) + torch.testing.assert_close(model.decode(actual_latents).sample, expected_video) + assert model(sample).sample.shape == sample.shape + assert model.decode(actual_latents, num_frames=frames).sample.shape == sample.shape + + @torch.no_grad() + def test_spatial_tiling_and_slicing(self): + model = self.model_class(**self.get_init_dict()).to(torch_device).eval() + sample = randn_tensor((2, 3, 16, 12, 12), generator=self.generator, device=torch_device) + model.enable_tiling(tile_sample_min_length=8, temporal_tile_overlap_factor=0.5) + expected_latents = model.encode(sample).latent_dist.mode() + expected_video = model.decode(expected_latents, num_frames=16).sample + assert expected_latents.shape == (2, 4, 4, 6, 6) + assert expected_video.shape == sample.shape + assert torch.isfinite(expected_video).all() + model.enable_slicing() + torch.testing.assert_close(model.encode(sample).latent_dist.mode(), expected_latents, atol=1e-5, rtol=1e-5) + torch.testing.assert_close( + model.decode(expected_latents, num_frames=16).sample, expected_video, atol=1e-5, rtol=1e-5 + ) + model.disable_slicing() + model.disable_tiling() + untiled = model.encode(sample).latent_dist.mode() + assert not torch.allclose(untiled, expected_latents) + assert model(sample).sample.shape == sample.shape + + @pytest.mark.parametrize("dim", [2, 3, 4]) + def test_overlap_blending(self, dim): + shape = [1, 1, 1, 1, 1] + shape[dim] = 4 + before = torch.full(shape, 10.0, device=torch_device) + after = torch.full(shape, 30.0, device=torch_device) + output = self.model_class._blend(before, after, 4, dim) + torch.testing.assert_close(output.flatten(), torch.tensor([10.0, 15.0, 20.0, 25.0], device=torch_device)) + + @pytest.mark.parametrize("dim", [2, 3, 4]) + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) + def test_low_precision_decoder_blending(self, dim, dtype): + before = randn_tensor((1, 3, 8, 8, 8), generator=self.generator, device=torch_device).to(dtype) + after = randn_tensor((1, 3, 8, 8, 8), generator=self.generator, device=torch_device).to(dtype) + expected = after.clone() + for index in range(4): + previous = before.select(dim, 4 + index).float() + current = after.select(dim, index).float() + expected.select(dim, index).copy_(previous * (1 - index / 4) + current * (index / 4)) + actual = self.model_class._blend(before, after, 4, dim, upcast=True) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + @pytest.mark.parametrize( + "kwargs", + [ + {"tile_sample_min_length": 0}, + {"tile_sample_min_length": 6}, + {"tile_sample_min_height": 7}, + {"temporal_tile_overlap_factor": 1.0}, + {"spatial_tile_overlap_factor": -0.1}, + {"temporal_tile_overlap_factor": 0.3}, + ], + ) + def test_invalid_tiling_settings(self, kwargs): + model = self.model_class(**self.get_init_dict()) + with pytest.raises(ValueError): + model.enable_tiling(**kwargs) + assert not model.use_tiling + + def test_partial_temporal_patch_is_rejected(self): + model = self.model_class(**self.get_init_dict()).to(torch_device) + model.enable_tiling(tile_sample_min_length=8, temporal_tile_overlap_factor=0.5, allow_spatial_tiling=False) + with pytest.raises(ValueError, match="complete patches"): + model.encode(torch.zeros((1, 3, 9, 8, 8), device=torch_device)) + + @pytest.mark.parametrize("num_frames", [0, 4, 9]) + def test_invalid_decode_frame_count(self, num_frames): + model = self.model_class(**self.get_init_dict()).to(torch_device) + latent = torch.zeros((1, 4, 2, 4, 4), device=torch_device) + with pytest.raises(ValueError, match="num_frames"): + model.decode(latent, num_frames=num_frames) + + @torch.no_grad() + def test_posterior_sampling_dtype_and_seed(self): + model = self.model_class(**self.get_init_dict()).to(torch_device, dtype=torch.bfloat16).eval() + posterior = model.encode(self.get_dummy_inputs()["sample"].bfloat16()).latent_dist + first = posterior.sample(generator=torch.Generator("cpu").manual_seed(12)) + second = posterior.sample(generator=torch.Generator("cpu").manual_seed(12)) + different = posterior.sample(generator=torch.Generator("cpu").manual_seed(13)) + assert first.dtype == torch.bfloat16 and first.device == posterior.mean.device + torch.testing.assert_close(first, second, rtol=0, atol=0) + assert not torch.equal(first, different) + + +class TestAutoencoderKLMagiMemory(AutoencoderKLMagiTesterConfig, MemoryTesterMixin): + @torch.no_grad() + def test_tiled_group_offload(self): + model = self.model_class(**self.get_init_dict()).to(torch_device).eval() + model.enable_tiling(temporal_tile_overlap_factor=0.5) + sample = self.get_dummy_inputs()["sample"] + expected = model(sample).sample + model.enable_group_offload( + onload_device=torch_device, offload_device="cpu", offload_type="block_level", num_blocks_per_group=1 + ) + torch.testing.assert_close(model(sample).sample, expected, atol=1e-5, rtol=0) + + +class TestAutoencoderKLMagiTorchCompile(AutoencoderKLMagiTesterConfig, TorchCompileTesterMixin): + @property + def different_shapes_for_compilation(self): + return [(4, 4), (4, 8), (8, 8)] + + def get_dummy_inputs(self, height=4, width=4): + return {"sample": randn_tensor((2, 3, 8, height, width), generator=self.generator, device=torch_device)} + + +class TestAutoencoderKLMagiTraining(AutoencoderKLMagiTesterConfig, TrainingTesterMixin): + def test_gradient_checkpointing_is_applied(self): + super().test_gradient_checkpointing_is_applied(expected_set={"MagiVAEEncoder", "MagiVAEDecoder"}) + + def test_tiled_backward(self): + model = self.model_class(**self.get_init_dict()).train() + model.enable_tiling(temporal_tile_overlap_factor=0.5) + model(self.get_dummy_inputs()["sample"].cpu()).sample.square().mean().backward() + assert model.encoder.patch_embed.proj.weight.grad.isfinite().all() + assert model.decoder.last_layer.weight.grad.isfinite().all() + + +class TestAutoencoderKLMagiAttention(AutoencoderKLMagiTesterConfig, AttentionTesterMixin): + pass diff --git a/tests/models/test_models_magi_conditioning.py b/tests/models/test_models_magi_conditioning.py new file mode 100644 index 000000000000..4612a518bd07 --- /dev/null +++ b/tests/models/test_models_magi_conditioning.py @@ -0,0 +1,125 @@ +# coding=utf-8 +# Copyright 2026 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +import torch + +from diffusers import MagiTextConditioningModel +from diffusers.utils.torch_utils import randn_tensor + +from ..testing_utils import enable_full_determinism, torch_device +from .testing_utils import ( + BaseModelTesterConfig, + MemoryTesterMixin, + ModelTesterMixin, + TorchCompileTesterMixin, +) + + +enable_full_determinism() + + +class MagiTextConditioningTesterConfig(BaseModelTesterConfig): + main_input_name = "hidden_states" + + @property + def model_class(self): + return MagiTextConditioningModel + + @property + def pretrained_model_name_or_path(self): + return None + + @property + def pretrained_model_kwargs(self): + return {} + + @property + def generator(self): + return torch.Generator("cpu").manual_seed(0) + + def get_init_dict(self) -> dict[str, int | list[int]]: + return {"caption_channels": 16, "caption_max_length": 8, "null_token_length": 4} + + def get_dummy_inputs(self) -> dict[str, torch.Tensor]: + return { + "hidden_states": randn_tensor((2, 8, 16), generator=self.generator, device=torch_device), + "attention_mask": torch.ones(2, 8, device=torch_device, dtype=torch.bool), + "num_chunks": 2, + } + + @property + def input_shape(self) -> tuple[int, ...]: + return (2, 8, 16) + + @property + def output_shape(self) -> tuple[int, ...]: + return (2, 2, 8, 16) + + +class TestMagiTextConditioningModel(MagiTextConditioningTesterConfig, ModelTesterMixin): + @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) + def test_from_save_pretrained_dtype(self, tmp_path, dtype): + self.check_conditioning_dtype(tmp_path, torch_dtype=dtype) + + @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) + def test_from_pretrained_dtype_alias(self, tmp_path, dtype): + self.check_conditioning_dtype(tmp_path, dtype=dtype) + + @pytest.mark.skip(reason="The small conditioning tables are kept together, not sharded across devices.") + def test_model_parallelism(self): + pass + + def check_conditioning_dtype(self, tmp_path, **kwargs): + model = self.model_class(**self.get_init_dict()) + with torch.no_grad(): + for parameter in model.parameters(): + parameter.normal_() + model.save_pretrained(tmp_path) + restored = self.model_class.from_pretrained(tmp_path, **kwargs) + assert all(parameter.dtype == torch.float32 for parameter in restored.parameters()) + for name, value in model.state_dict().items(): + torch.testing.assert_close(value, restored.state_dict()[name], rtol=0, atol=0) + + +class TestMagiTextConditioningMemory(MagiTextConditioningTesterConfig, MemoryTesterMixin): + @pytest.mark.skip(reason="Embedding tables are excluded from layerwise casting.") + def test_layerwise_casting_memory(self): + pass + + @pytest.mark.skip(reason="The small conditioning tables are offloaded as one component, not split.") + def test_cpu_offload(self): + pass + + @pytest.mark.skip(reason="The small conditioning tables are offloaded as one component, not split.") + def test_disk_offload_without_safetensors(self): + pass + + @pytest.mark.skip(reason="The small conditioning tables are offloaded as one component, not split.") + def test_disk_offload_with_safetensors(self): + pass + + +class TestMagiTextConditioningTorchCompile(MagiTextConditioningTesterConfig, TorchCompileTesterMixin): + @property + def different_shapes_for_compilation(self): + return [(1, 2), (2, 2), (3, 2)] + + def get_dummy_inputs(self, height: int = 4, width: int = 4) -> dict[str, torch.Tensor]: + return { + "hidden_states": randn_tensor((height, 8, 16), generator=self.generator, device=torch_device), + "attention_mask": torch.ones(height, 8, device=torch_device, dtype=torch.bool), + "num_chunks": width, + } diff --git a/tests/models/transformers/test_models_transformer_magi.py b/tests/models/transformers/test_models_transformer_magi.py new file mode 100644 index 000000000000..7313104a3fc0 --- /dev/null +++ b/tests/models/transformers/test_models_transformer_magi.py @@ -0,0 +1,446 @@ +# Copyright 2026 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import math + +import pytest +import torch + +from diffusers import AutoencoderKLMagi, MagiTransformer3DModel +from diffusers.models.attention_dispatch import attention_backend +from diffusers.models.transformers.transformer_magi import MagiTimestepEmbedding +from diffusers.utils import is_flash_attn_available +from diffusers.utils.torch_utils import randn_tensor + +from ...testing_utils import enable_full_determinism, torch_device +from ..testing_utils import ( + AttentionTesterMixin, + BaseModelTesterConfig, + MemoryTesterMixin, + ModelTesterMixin, + TorchCompileTesterMixin, + TrainingTesterMixin, +) + + +enable_full_determinism() + + +class MagiTransformerTesterConfig(BaseModelTesterConfig): + main_input_name = "hidden_states" + + @property + def model_class(self): + return MagiTransformer3DModel + + @property + def pretrained_model_name_or_path(self): + return None + + @property + def pretrained_model_kwargs(self): + return {} + + @property + def generator(self): + return torch.Generator("cpu").manual_seed(0) + + def get_init_dict(self): + return { + "in_channels": 4, + "out_channels": 4, + "num_layers": 2, + "num_attention_heads": 2, + "num_key_value_heads": 1, + "attention_head_dim": 32, + "ffn_dim": 96, + "condition_dim": 16, + "caption_channels": 16, + "caption_max_length": 8, + "frequency_embedding_size": 16, + } + + def get_dummy_inputs(self): + return { + "hidden_states": randn_tensor((2, 4, 4, 4, 4), generator=self.generator, device=torch_device), + "encoder_hidden_states": randn_tensor((2, 8, 16), generator=self.generator, device=torch_device), + "timestep": torch.tensor([0.2, 0.5], device=torch_device), + } + + @property + def input_shape(self): + return (4, 4, 4, 4) + + @property + def output_shape(self): + return self.input_shape + + +class TestMagiTransformerModel(MagiTransformerTesterConfig, ModelTesterMixin): + @torch.no_grad() + def test_empty_text_mask(self): + model = self.model_class(**self.get_init_dict()).to(torch_device).eval() + inputs = self.get_dummy_inputs() + inputs["encoder_attention_mask"] = torch.zeros(2, 8, dtype=torch.bool, device=torch_device) + with pytest.raises(RuntimeError): + model(**inputs) + + @pytest.mark.parametrize("regional", [False, True]) + @torch.no_grad() + def test_masked_fullgraph_capture(self, regional): + from copy import deepcopy + + torch.compiler.reset() + reference = self.model_class(**self.get_init_dict()).to(torch_device).eval() + model = deepcopy(reference) + inputs = self.get_dummy_inputs() + masks = [ + [[True, False] * 4, [True, True, False, False] * 2], + [[True] * 8, [True, False, False, False] * 2], + [[True] * 8] * 2, + ] + try: + with torch._dynamo.config.patch(capture_dynamic_output_shape_ops=True): + if regional: + model.compile_repeated_blocks(backend="eager", fullgraph=True) + else: + model = torch.compile(model, backend="eager", fullgraph=True) + for mask in masks: + inputs["encoder_attention_mask"] = torch.tensor(mask, device=torch_device) + expected = reference(**inputs).sample + torch.testing.assert_close(model(**inputs).sample, expected, atol=0, rtol=0) + changed = dict(inputs) + changed["encoder_hidden_states"] = inputs["encoder_hidden_states"].clone() + changed["encoder_hidden_states"][~inputs["encoder_attention_mask"]] += 100 + torch.testing.assert_close(model(**changed).sample, expected, atol=0, rtol=0) + finally: + torch.compiler.reset() + + @pytest.mark.parametrize("frames", [4, 16]) + @pytest.mark.parametrize( + "duplicate_channels,distilled", [(False, False), (False, True), (True, False), (True, True)] + ) + @torch.no_grad() + def test_vae_transformer_dataflow(self, frames, duplicate_channels, distilled, tmp_path): + torch.manual_seed(0) + vae = ( + AutoencoderKLMagi( + latent_channels=4, + embed_dim=32, + num_layers=1, + num_attention_heads=4, + mlp_ratio=2, + patch_size=2, + patch_length=4, + sample_size=8, + sample_frames=8, + ) + .to(torch_device) + .eval() + ) + model = ( + self.model_class( + **self.get_init_dict(), + duplicate_channels=duplicate_channels, + gated_linear_unit=duplicate_channels, + x_rescale_factor=0.1 if duplicate_channels else 1.0, + distilled=distilled, + ) + .to(torch_device) + .eval() + ) + video = randn_tensor((2, 3, frames, 8, 12), generator=self.generator, device=torch_device) + latent = vae.encode(video).latent_dist.mode() + assert latent.shape == (2, 4, frames // 4, 4, 6) + scale_factor = 0.18215 + inputs = self.get_dummy_inputs() + inputs["hidden_states"] = latent * scale_factor + if distilled: + inputs["timestep_delta"] = torch.tensor(8.0, device=torch_device) + predicted = model(**inputs).sample + assert predicted.shape == latent.shape + decoded = vae.decode(predicted / scale_factor, num_frames=frames).sample + assert decoded.shape == video.shape + assert decoded.isfinite().all() + vae.save_pretrained(tmp_path / "vae") + model.save_pretrained(tmp_path / "transformer") + restored_vae = AutoencoderKLMagi.from_pretrained(tmp_path / "vae").to(torch_device) + restored_model = self.model_class.from_pretrained(tmp_path / "transformer").to(torch_device) + inputs["hidden_states"] = restored_vae.encode(video).latent_dist.mode() * scale_factor + restored = restored_vae.decode(restored_model(**inputs).sample / scale_factor, num_frames=frames).sample + torch.testing.assert_close(restored, decoded, atol=0, rtol=0) + + @torch.no_grad() + def test_three_chunk_cache_with_per_chunk_masks(self): + model = self.model_class(**self.get_init_dict()).to(torch_device).eval() + latent = randn_tensor((2, 4, 6, 4, 6), generator=self.generator, device=torch_device) + text = randn_tensor((2, 3, 8, 16), generator=self.generator, device=torch_device) + mask = torch.tensor( + [[[True, False] * 4, [False, True] * 4, [True] * 3 + [False] * 5]] * 2, + device=torch_device, + ) + timestep = torch.tensor([[0.9999, 0.5, 0.1], [0.9999, 0.7, 0.3]], device=torch_device) + expected = model(latent, text, timestep, encoder_attention_mask=mask, use_cache=True) + cache = None + outputs = [] + for index in range(3): + output = model( + latent[:, :, index * 2 : (index + 1) * 2], + text[:, index], + timestep[:, index], + encoder_attention_mask=mask[:, index], + kv_cache=cache, + use_cache=True, + ) + cache = output.kv_cache + assert len(cache) == model.config.num_layers + assert all(key.shape == value.shape == (2, (index + 1) * 12, 1, 32) for key, value in cache) + outputs.append(output.sample) + torch.testing.assert_close(torch.cat(outputs, dim=2), expected.sample, atol=1e-5, rtol=1e-5) + for actual, full in zip(cache, expected.kv_cache): + torch.testing.assert_close(actual, full, atol=1e-5, rtol=1e-5) + + @torch.no_grad() + def test_timestep_frequency_rounding(self): + embedder = MagiTimestepEmbedding(16, 256).to(torch_device).eval() + timesteps = torch.tensor([0.9999, 0.3, 0.7], device=torch_device) + frequencies = torch.exp(-math.log(10000) * torch.arange(128).float() / 128).to(torch_device) + angles = timesteps[:, None] * frequencies[None] * 1000 + expected = embedder.mlp(torch.cat([angles.cos(), angles.sin()], dim=-1).bfloat16()) + actual = embedder(timesteps, torch.bfloat16) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + @torch.no_grad() + def test_attention_output_projection_uses_fp32(self): + model = self.model_class(**self.get_init_dict()).to(torch_device, dtype=torch.bfloat16).eval() + projection = model.transformer_blocks[0].self_attention.linear_proj + hidden_states = randn_tensor((1, 4, 128), generator=self.generator, device=torch_device).bfloat16() + actual = projection(hidden_states) + expected = torch.nn.functional.linear(hidden_states.float(), projection.weight.float()) + assert actual.dtype == torch.float32 + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + @torch.no_grad() + def test_expanded_caption_matches_contiguous_projection(self): + torch.manual_seed(0) + config = self.get_init_dict() + config["caption_channels"] = 128 + model = self.model_class(**config).to(torch_device).eval() + embedder = model.condition_embedder.y_embedder + captions = torch.randn(1, 8, 128, device=torch_device)[:, None].expand(-1, 3, -1, -1) + mask = torch.zeros(1, device=torch_device, dtype=torch.bool) + expected = embedder(captions.contiguous(), mask) + actual = embedder(captions, mask) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + @torch.no_grad() + def test_caption_dropout_only_changes_adaptive_embedding(self): + model = self.model_class(**self.get_init_dict()).to(torch_device).eval() + embedder = model.condition_embedder.y_embedder + embedder.null_caption_embedding[-2].fill_(0.1) + embedder.null_caption_embedding[-1].fill_(-0.1) + captions = self.get_dummy_inputs()["encoder_hidden_states"][:, None] + conditional, condition = embedder(captions, torch.zeros(2, device=torch_device, dtype=torch.bool)) + unconditional, uncondition = embedder(captions, torch.ones(2, device=torch_device, dtype=torch.bool)) + torch.testing.assert_close(conditional, unconditional, atol=0, rtol=0) + assert not torch.allclose(condition, uncondition) + + def test_invalid_cache_and_shapes(self): + model = self.model_class(**self.get_init_dict()).to(torch_device).eval() + inputs = self.get_dummy_inputs() + with pytest.raises(ValueError, match="one key/value pair"): + model(**inputs, kv_cache=()) + inputs["hidden_states"] = inputs["hidden_states"][:, :, :, :3] + with pytest.raises(ValueError, match="divisible"): + model(**inputs) + + @torch.no_grad() + def test_chunk_causality(self): + model = self.model_class(**self.get_init_dict()).to(torch_device).eval() + inputs = self.get_dummy_inputs() + inputs["timestep"] = inputs["timestep"][:, None].expand(-1, 2) + expected = model(**inputs).sample + inputs["hidden_states"][:, :, 2:] += 10 + actual = model(**inputs).sample + torch.testing.assert_close(actual[:, :, :2], expected[:, :, :2], atol=1e-5, rtol=0) + assert not torch.allclose(actual[:, :, 2:], expected[:, :, 2:]) + + @torch.no_grad() + @pytest.mark.parametrize("cache_device", [None, "cpu"]) + def test_prefix_cache_matches_full_forward(self, cache_device): + model = self.model_class(**self.get_init_dict()).to(torch_device).eval() + inputs = self.get_dummy_inputs() + inputs["timestep"] = inputs["timestep"][:, None].expand(-1, 2) + expected = model(**inputs, use_cache=True) + prefix = model( + hidden_states=inputs["hidden_states"][:, :, :2], + timestep=inputs["timestep"][:, :1], + encoder_hidden_states=inputs["encoder_hidden_states"], + use_cache=True, + cache_device=cache_device, + ) + cached = tuple((key.clone(), value.clone()) for key, value in prefix.kv_cache) + suffix = model( + hidden_states=inputs["hidden_states"][:, :, 2:], + timestep=inputs["timestep"][:, 1:], + encoder_hidden_states=inputs["encoder_hidden_states"], + kv_cache=prefix.kv_cache, + use_cache=True, + ) + torch.testing.assert_close(suffix.sample, expected.sample[:, :, 2:], atol=1e-5, rtol=1e-5) + for (key, value), (expected_key, expected_value) in zip(suffix.kv_cache, expected.kv_cache): + torch.testing.assert_close(key.to(expected_key), expected_key, atol=1e-5, rtol=1e-5) + torch.testing.assert_close(value.to(expected_value), expected_value, atol=1e-5, rtol=1e-5) + for actual, previous in zip(prefix.kv_cache, cached): + torch.testing.assert_close(actual, previous, atol=0, rtol=0) + + @torch.no_grad() + @pytest.mark.parametrize("cache_device", [None, "cpu"]) + @pytest.mark.parametrize("retained_tokens", [4, 8]) + def test_compact_cache_matches_full_forward(self, cache_device, retained_tokens): + model = self.model_class(**self.get_init_dict()).to(torch_device).eval() + inputs = self.get_dummy_inputs() + expected = model(**inputs, use_cache=True) + actual = model(**inputs, use_cache=True, cache_token_count=retained_tokens, cache_device=cache_device) + torch.testing.assert_close(actual.sample, expected.sample, atol=0, rtol=0) + for full_pair, compact_pair in zip(expected.kv_cache, actual.kv_cache): + for full, compact in zip(full_pair, compact_pair): + torch.testing.assert_close(compact.to(full), full[:, :retained_tokens], atol=0, rtol=0) + assert compact.untyped_storage().nbytes() == compact.numel() * compact.element_size() + if cache_device == "cpu": + assert compact.device.type == "cpu" + + @pytest.mark.parametrize( + "options", + [ + {"cache_token_count": 4}, + {"cache_device": "cpu"}, + {"cache_token_count": 3, "use_cache": True}, + {"cache_token_count": 100, "use_cache": True}, + {"cache_token_count": True, "use_cache": True}, + ], + ) + def test_invalid_cache_retention(self, options): + model = self.model_class(**self.get_init_dict()).to(torch_device).eval() + with pytest.raises(ValueError): + model(**self.get_dummy_inputs(), **options) + + @torch.no_grad() + def test_text_padding_is_ignored(self): + model = self.model_class(**self.get_init_dict()).to(torch_device).eval() + inputs = self.get_dummy_inputs() + inputs["encoder_attention_mask"] = torch.tensor([[True] * 4 + [False] * 4] * 2, device=torch_device) + expected = model(**inputs).sample + inputs["encoder_hidden_states"][:, 4:] += 100 + torch.testing.assert_close(model(**inputs).sample, expected, atol=0, rtol=0) + + @torch.no_grad() + def test_per_chunk_text_conditioning(self): + model = self.model_class(**self.get_init_dict()).to(torch_device).eval() + inputs = self.get_dummy_inputs() + inputs["timestep"] = inputs["timestep"][:, None].expand(-1, 2) + inputs["encoder_hidden_states"] = inputs["encoder_hidden_states"][:, None].repeat(1, 2, 1, 1) + expected = model(**inputs).sample + inputs["encoder_hidden_states"][:, 1] += 1 + actual = model(**inputs).sample + torch.testing.assert_close(actual[:, :, :2], expected[:, :, :2], atol=0, rtol=0) + assert not torch.allclose(actual[:, :, 2:], expected[:, :, 2:]) + + @torch.no_grad() + def test_distillation_condition(self): + config = self.get_init_dict() + model = self.model_class(**config, distilled=True).to(torch_device).eval() + inputs = self.get_dummy_inputs() + with pytest.raises(ValueError, match="timestep_delta"): + model(**inputs) + first = model(**inputs, timestep_delta=torch.tensor(2.0, device=torch_device)).sample + second = model(**inputs, timestep_delta=torch.tensor(4.0, device=torch_device)).sample + assert first.shape == second.shape == inputs["hidden_states"].shape + assert not torch.allclose(first, second) + + @torch.no_grad() + def test_24b_variant(self): + model = ( + self.model_class( + **self.get_init_dict(), duplicate_channels=True, gated_linear_unit=True, x_rescale_factor=0.1 + ) + .to(torch_device) + .eval() + ) + inputs = self.get_dummy_inputs() + output = model(**inputs).sample + assert output.shape == inputs["hidden_states"].shape + assert torch.isfinite(output).all() + + @torch.no_grad() + def test_explicit_kv_ranges(self): + model = self.model_class(**self.get_init_dict()).to(torch_device).eval() + inputs = self.get_dummy_inputs() + inputs["timestep"] = inputs["timestep"][:, None].expand(-1, 2) + expected = model(**inputs).sample + actual = model(**inputs, kv_ranges=((0, 8), (0, 16))).sample + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + local = model(**inputs, kv_ranges=((0, 8), (8, 16))).sample + assert not torch.allclose(local[:, :, 2:], expected[:, :, 2:]) + with pytest.raises(ValueError, match="range"): + model(**inputs, kv_ranges=((0, 8), (8, 17))) + + +class TestMagiTransformerMemory(MagiTransformerTesterConfig, MemoryTesterMixin): + pass + + +class TestMagiTransformerTorchCompile(MagiTransformerTesterConfig, TorchCompileTesterMixin): + @pytest.fixture(autouse=True) + def capture_packed_text_shapes(self): + with torch._dynamo.config.patch(capture_dynamic_output_shape_ops=True): + yield + + @property + def different_shapes_for_compilation(self): + return [(4, 4), (4, 8), (8, 8)] + + def get_dummy_inputs(self, height=4, width=4): + inputs = super().get_dummy_inputs() + inputs["hidden_states"] = randn_tensor((2, 4, 4, height, width), generator=self.generator, device=torch_device) + inputs["encoder_attention_mask"] = torch.tensor([[True, False] * 4] * 2, device=torch_device) + return inputs + + +class TestMagiTransformerTraining(MagiTransformerTesterConfig, TrainingTesterMixin): + def test_gradient_checkpointing_is_applied(self): + super().test_gradient_checkpointing_is_applied(expected_set={"MagiTransformer3DModel"}) + + +class TestMagiTransformerAttention(MagiTransformerTesterConfig, AttentionTesterMixin): + @pytest.mark.skipif( + not is_flash_attn_available() or torch_device != "cuda", reason="Requires CUDA FlashAttention." + ) + @pytest.mark.parametrize("backend", ["flash", "flash_varlen"]) + @torch.no_grad() + def test_flash_backend(self, tmp_path, backend): + with attention_backend("native"): + model = self.model_class(**self.get_init_dict()).eval() + model.save_pretrained(tmp_path) + model = self.model_class.from_pretrained(tmp_path, torch_dtype=torch.bfloat16).to(torch_device) + model.set_attention_backend(backend) + inputs = self.get_dummy_inputs() + if backend == "flash_varlen": + inputs["encoder_attention_mask"] = torch.tensor([[True, False] * 4] * 2, device=torch_device) + expected = model(**inputs).sample + assert expected.shape == inputs["hidden_states"].shape + assert torch.isfinite(expected).all() + if backend == "flash_varlen": + inputs["encoder_hidden_states"][:, 1::2] += 100 + torch.testing.assert_close(model(**inputs).sample, expected, atol=0, rtol=0) diff --git a/tests/modular_pipelines/magi/__init__.py b/tests/modular_pipelines/magi/__init__.py new file mode 100644 index 000000000000..ef08599ad40f --- /dev/null +++ b/tests/modular_pipelines/magi/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/tests/modular_pipelines/magi/test_magi_denoise.py b/tests/modular_pipelines/magi/test_magi_denoise.py new file mode 100644 index 000000000000..118875556dc1 --- /dev/null +++ b/tests/modular_pipelines/magi/test_magi_denoise.py @@ -0,0 +1,240 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +import torch + +from diffusers import MagiClassifierFreeGuidance, MagiDenoiseStep, MagiEulerScheduler, MagiTransformer3DModel + + +class TestMagiDenoise: + def make_pipeline(self, **model_kwargs): + torch.manual_seed(0) + config = { + "in_channels": 4, + "out_channels": 4, + "num_layers": 2, + "num_attention_heads": 2, + "num_key_value_heads": 1, + "attention_head_dim": 32, + "ffn_dim": 96, + "condition_dim": 16, + "caption_channels": 16, + "caption_max_length": 8, + "frequency_embedding_size": 16, + } + config.update(model_kwargs) + pipe = MagiDenoiseStep().init_pipeline() + pipe.update_components( + transformer=MagiTransformer3DModel(**config).eval(), + scheduler=MagiEulerScheduler(), + guider=MagiClassifierFreeGuidance(), + ) + pipe.load_components() + return pipe + + def inputs(self, batch=1, chunks=3): + generator = torch.Generator().manual_seed(12) + return { + "latents": torch.randn(batch, 4, chunks * 2, 4, 4, generator=generator), + "prompt_embeds": torch.randn(batch, chunks, 3, 16, generator=generator), + "prompt_attention_mask": torch.ones(batch, chunks, 3, dtype=torch.bool), + "negative_prompt_embeds": torch.randn(batch, 3, 16, generator=generator), + "negative_prompt_attention_mask": torch.tensor([[True, False, True]]).expand(batch, -1), + "chunk_width": 2, + "window_size": 2, + "num_inference_steps": 4, + } + + @pytest.mark.parametrize("prefix_chunks", [0, 1]) + def test_cpu_cache_preserves_denoising(self, prefix_chunks): + pipe = self.make_pipeline() + device = "cuda" if torch.cuda.is_available() else "cpu" + pipe.to(device) + inputs = {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in self.inputs(chunks=4).items()} + if prefix_chunks: + inputs["prefix_latents"] = inputs["latents"][:, :, :2].clone() + expected = pipe(**inputs, output=["latents", "clean_kv_cache"]) + actual = pipe(**inputs, cache_device="cpu", output=["latents", "clean_kv_cache"]) + torch.testing.assert_close(actual["latents"], expected["latents"], atol=0, rtol=0) + for full_pair, cpu_pair in zip(expected["clean_kv_cache"], actual["clean_kv_cache"]): + for full, compact in zip(full_pair, cpu_pair): + assert compact.device.type == "cpu" + torch.testing.assert_close(compact.to(full), full, atol=0, rtol=0) + assert compact.untyped_storage().nbytes() == compact.numel() * compact.element_size() + + def test_window_branches_and_clean_cache(self): + pipe = self.make_pipeline() + inputs = self.inputs() + original = {k: v.clone() for k, v in inputs.items() if isinstance(v, torch.Tensor)} + calls = [] + + def capture(module, args, kwargs): + calls.append({k: v.clone() if isinstance(v, torch.Tensor) else v for k, v in kwargs.items()}) + + projected_rows = [] + projection_handle = pipe.transformer.transformer_blocks[ + 0 + ].self_attention.linear_kv_xattn.register_forward_pre_hook( + lambda module, args: projected_rows.append(tuple(args[0].shape)) + ) + handle = pipe.transformer.register_forward_pre_hook(capture, with_kwargs=True) + output = pipe(**inputs, output=["latents", "clean_kv_cache", "completed_chunks"]) + handle.remove() + projection_handle.remove() + assert [shape[0] for shape in projected_rows[:3]] == [3, 2, 2] + assert all(len(shape) == 2 for shape in projected_rows) + assert len(calls) == 24 + assert output["completed_chunks"] == [0, 1, 2] + assert output["latents"].dtype == torch.float32 + assert not output["latents"].requires_grad + assert [x["hidden_states"].shape[2] for x in calls[::3]] == [2, 2, 4, 4, 6, 4, 4, 2] + assert [i for i, x in enumerate(calls) if x["use_cache"]] == [13, 19] + for i in range(0, len(calls), 3): + text, prefix, unconditional = calls[i : i + 3] + assert not text["caption_dropout_mask"].any() + assert prefix["caption_dropout_mask"].all() + assert unconditional["caption_dropout_mask"].all() + assert unconditional["caption_dropout_mask"].numel() == 1 + assert text["kv_cache"] is prefix["kv_cache"] + assert unconditional["kv_cache"] is None + assert unconditional["kv_ranges"] is None + assert unconditional["hidden_states"].shape[2] == 2 + assert unconditional["timestep"].shape[1] == 1 + for name, value in original.items(): + torch.testing.assert_close(inputs[name], value, atol=0, rtol=0) + with torch.no_grad(): + clean = pipe.transformer( + output["latents"][:, :, :4], + inputs["negative_prompt_embeds"], + torch.full((1, 2), 0.9999), + encoder_attention_mask=inputs["negative_prompt_attention_mask"], + caption_dropout_mask=torch.ones(1, dtype=torch.bool), + kv_ranges=((0, 8), (8, 16)), + use_cache=True, + ).kv_cache + for actual_pair, expected_pair in zip(output["clean_kv_cache"], clean): + for actual, expected in zip(actual_pair, expected_pair): + torch.testing.assert_close(actual, expected, atol=1e-5, rtol=1e-5) + assert actual.shape[1] == 16 + assert actual.untyped_storage().nbytes() == actual.numel() * actual.element_size() + + def test_pipeline_roundtrip(self, tmp_path): + from diffusers import ModularPipeline + + pipe = self.make_pipeline() + inputs = self.inputs() + expected = pipe(**inputs, output="latents") + pipe.save_pretrained(str(tmp_path), overwrite_modular_index=True) + restored = ModularPipeline.from_pretrained(str(tmp_path)) + restored.load_components() + actual = restored(**inputs, output="latents") + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + def test_repeat_run_resets_state(self): + pipe = self.make_pipeline() + inputs = self.inputs() + first = pipe(**inputs, output="latents") + second = pipe(**inputs, output="latents") + torch.testing.assert_close(first, second, atol=0, rtol=0) + assert pipe.scheduler.step_index is None + + def test_full_chunk_prefix(self): + pipe = self.make_pipeline() + inputs = self.inputs(chunks=4) + inputs["prefix_latents"] = torch.randn(1, 4, 4, 4, 4) + prefix = inputs["prefix_latents"].clone() + output = pipe(**inputs, output=["latents", "clean_kv_cache", "completed_chunks"]) + torch.testing.assert_close(output["latents"][:, :, :4], prefix, atol=0, rtol=0) + torch.testing.assert_close(inputs["prefix_latents"], prefix, atol=0, rtol=0) + assert output["completed_chunks"] == [0, 1, 2, 3] + assert output["clean_kv_cache"][0][0].shape[1] == 24 + + def test_batch_matches_individual(self): + pipe = self.make_pipeline() + inputs = self.inputs(batch=2) + batched = pipe(**inputs, output="latents") + for i in range(2): + single = {k: v[i : i + 1] if isinstance(v, torch.Tensor) else v for k, v in inputs.items()} + expected = pipe(**single, output="latents") + torch.testing.assert_close(batched[i : i + 1], expected, atol=2e-5, rtol=2e-5) + + def test_single_chunk_and_window_larger_than_video(self): + pipe = self.make_pipeline() + inputs = self.inputs(chunks=1) + inputs["window_size"] = 4 + output = pipe(**inputs, output=["latents", "clean_kv_cache", "completed_chunks"]) + assert output["clean_kv_cache"] is None + assert output["completed_chunks"] == [0] + + def test_shared_text_matches_per_chunk(self): + pipe = self.make_pipeline() + inputs = self.inputs() + inputs["prompt_embeds"] = inputs["prompt_embeds"][:, 0] + inputs["prompt_attention_mask"] = inputs["prompt_attention_mask"][:, 0] + expected = pipe(**inputs, output="latents") + inputs["prompt_embeds"] = inputs["prompt_embeds"][:, None].expand(-1, 3, -1, -1) + inputs["prompt_attention_mask"] = inputs["prompt_attention_mask"][:, None].expand(-1, 3, -1) + actual = pipe(**inputs, output="latents") + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + def test_invalid_inputs(self): + pipe = self.make_pipeline() + cases = [ + {"num_inference_steps": 6}, + {"window_size": 0}, + {"noise2clean_kvrange": ()}, + {"clean_chunk_kvrange": -1}, + {"clean_t": 1.1}, + {"chunk_width": 4}, + {"prefix_latents": torch.zeros(1, 4, 1, 4, 4)}, + {"prefix_latents": torch.zeros(1, 4, 6, 4, 4)}, + {"negative_prompt_attention_mask": torch.zeros(1, 3, dtype=torch.bool)}, + ] + for changed in cases: + with pytest.raises(ValueError): + pipe(**(self.inputs() | changed), output="latents") + with pytest.raises(ValueError, match="base models only"): + self.make_pipeline(distilled=True)(**self.inputs()) + + +class TestMagiGuidance: + def test_formula_and_thresholds(self): + guider = MagiClassifierFreeGuidance(prefix_scales=(1, 2, 3, 4, 5), text_scales=(10, 20, 30, 40, 50)) + times = torch.tensor([[0.0, 0.0217 - 2e-7, 0.0217, 0.1, 0.3, 0.999]]) + guider.set_state(step=0, num_inference_steps=64, timestep=times) + cond, prefix, uncond = [torch.full((1, 4, 12, 2, 2), x) for x in (3.0, 2.0, 1.0)] + actual = guider.forward(cond, prefix, uncond).pred + expected = ( + torch.tensor([12.0, 12.0, 23.0, 34.0, 45.0, 56.0]) + .repeat_interleave(2)[None, None, :, None, None] + .expand_as(actual) + ) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + def test_config_roundtrip(self, tmp_path): + guider = MagiClassifierFreeGuidance(text_scales=(4, 3, 2, 1, 0)) + guider.save_pretrained(str(tmp_path)) + loaded = MagiClassifierFreeGuidance.from_pretrained(str(tmp_path)) + assert tuple(loaded.config.text_scales) == (4, 3, 2, 1, 0) + assert loaded.num_conditions == 3 + + def test_invalid_configuration(self): + for kwargs in ( + {"prefix_scales": (1,)}, + {"timestep_thresholds": (0, 0.1, 0.1, 0.3, 1)}, + {"timestep_thresholds": (-1, 0.1, 0.2, 0.3, 1)}, + ): + with pytest.raises(ValueError): + MagiClassifierFreeGuidance(**kwargs) diff --git a/tests/modular_pipelines/magi/test_modular_pipeline_magi.py b/tests/modular_pipelines/magi/test_modular_pipeline_magi.py new file mode 100644 index 000000000000..41aa90feca13 --- /dev/null +++ b/tests/modular_pipelines/magi/test_modular_pipeline_magi.py @@ -0,0 +1,321 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json + +import numpy as np +import pytest +import torch +from tokenizers import Tokenizer +from tokenizers.models import WordLevel +from tokenizers.pre_tokenizers import Whitespace +from transformers import PreTrainedTokenizerFast, T5Config, T5EncoderModel + +from diffusers import ( + AutoencoderKLMagi, + MagiClassifierFreeGuidance, + MagiEulerScheduler, + MagiModularPipeline, + MagiTextConditioningModel, + MagiTextToVideoBlocks, + MagiTransformer3DModel, + ModularPipeline, +) + +from ..testing_utils import ( + BaseModularPipelineTesterConfig, + ModularLoadingTesterMixin, + ModularMemoryTesterMixin, + ModularPipelineTesterMixin, +) +from .testing_utils import MagiGuiderTesterMixin + + +@pytest.fixture(scope="module") +def tiny_magi_path(tmp_path_factory): + torch.manual_seed(0) + tokenizer = Tokenizer(WordLevel({"[PAD]": 0, "[UNK]": 1, "a": 2, "cat": 3, "runs": 4}, unk_token="[UNK]")) + tokenizer.pre_tokenizer = Whitespace() + tokenizer = PreTrainedTokenizerFast(tokenizer_object=tokenizer, pad_token="[PAD]", unk_token="[UNK]") + pipe = MagiTextToVideoBlocks().init_pipeline() + conditioning = MagiTextConditioningModel(caption_channels=16, caption_max_length=8, null_token_length=4) + with torch.no_grad(): + conditioning.null_embedding.weight.normal_() + conditioning.special_embedding.weight.normal_() + torch.manual_seed(0) + text_encoder = T5EncoderModel(T5Config(vocab_size=5, d_model=16, d_ff=32, d_kv=8, num_heads=2, num_layers=1)) + torch.manual_seed(0) + transformer = MagiTransformer3DModel( + in_channels=4, + out_channels=4, + num_layers=1, + num_attention_heads=2, + num_key_value_heads=1, + attention_head_dim=32, + ffn_dim=96, + condition_dim=16, + caption_channels=16, + caption_max_length=8, + frequency_embedding_size=16, + ) + torch.manual_seed(0) + vae = AutoencoderKLMagi( + latent_channels=4, + embed_dim=32, + num_layers=1, + num_attention_heads=4, + mlp_ratio=2, + patch_size=2, + patch_length=4, + sample_size=8, + sample_frames=8, + ) + pipe.update_components( + tokenizer=tokenizer, + text_encoder=text_encoder, + text_conditioning=conditioning, + transformer=transformer, + vae=vae, + scheduler=MagiEulerScheduler(), + guider=MagiClassifierFreeGuidance(), + ) + pipe.load_components() + for component in pipe.components.values(): + if isinstance(component, torch.nn.Module): + component.eval() + path = str(tmp_path_factory.mktemp("tiny-magi")) + pipe.save_pretrained(path, overwrite_modular_index=True) + return path + + +class MagiPipelineTesterConfig(BaseModularPipelineTesterConfig): + pipeline_class = MagiModularPipeline + pipeline_blocks_class = MagiTextToVideoBlocks + params = frozenset(["prompt", "height", "width", "num_frames"]) + batch_params = frozenset(["prompt"]) + output_name = "videos" + + @pytest.fixture(scope="class", autouse=True) + @classmethod + def model_path(cls, tiny_magi_path): + cls.pretrained_model_name_or_path = tiny_magi_path + + def get_dummy_inputs(self, seed=0): + return { + "prompt": "a cat runs", + "height": 8, + "width": 8, + "num_frames": 16, + "chunk_width": 2, + "window_size": 2, + "num_inference_steps": 4, + "max_sequence_length": 8, + "clean_caption": False, + "output_type": "pt", + "generator": self.get_generator(seed), + } + + +class TestMagiPipelineFast(MagiPipelineTesterConfig, ModularPipelineTesterMixin): + def test_convert_sharded_t5(self, tmp_path): + from scripts.convert_magi_to_diffusers import load_t5 + + torch.manual_seed(0) + model = T5EncoderModel(T5Config(vocab_size=5, d_model=16, d_ff=32, d_kv=8, num_heads=2, num_layers=1)) + model.config.save_pretrained(tmp_path) + tensors = list(model.state_dict().items()) + mapping = {} + for index in range(2): + name = f"pytorch_model-{index + 1:05d}-of-00002.bin" + shard = dict(tensors[index::2]) + torch.save(shard, tmp_path / name) + mapping.update(dict.fromkeys(shard, name)) + (tmp_path / "pytorch_model.bin.index.json").write_text(json.dumps({"weight_map": mapping})) + restored = load_t5(tmp_path) + for name, tensor in model.state_dict().items(): + torch.testing.assert_close(tensor, restored.state_dict()[name], atol=0, rtol=0) + assert not list(tmp_path.glob("*.safetensors")) + + def test_repeat_and_roundtrip(self, tmp_path): + pipe = self.get_pipeline() + expected = self.run_pipe(pipe) + torch.testing.assert_close(self.run_pipe(pipe), expected, rtol=0, atol=0) + pipe.save_pretrained(str(tmp_path), overwrite_modular_index=True) + restored = ModularPipeline.from_pretrained(str(tmp_path)) + restored.load_components() + torch.testing.assert_close(self.run_pipe(restored), expected, rtol=0, atol=0) + assert expected.shape == (1, 16, 3, 8, 8) + + def test_chunk_conditioning_and_decode(self): + pipe = self.get_pipeline() + inputs = self.get_dummy_inputs() + output = pipe( + **inputs, + output=[ + "videos", + "latents", + "prompt_embeds", + "prompt_attention_mask", + "negative_prompt_embeds", + "negative_prompt_attention_mask", + "text_embeds", + ], + ) + text = output["prompt_embeds"] + torch.testing.assert_close(text[0, :, 0], pipe.text_conditioning.special_embedding.weight[[2, 1]]) + torch.testing.assert_close(text[0, :, 1], pipe.text_conditioning.special_embedding.weight[0].expand(2, -1)) + torch.testing.assert_close(text[:, :, 2:], output["text_embeds"][:, None, :6].expand(-1, 2, -1, -1)) + torch.testing.assert_close(output["negative_prompt_embeds"][0], pipe.text_conditioning.null_embedding.weight) + assert output["negative_prompt_attention_mask"].sum().item() == 4 + chunks = [] + with torch.no_grad(): + for chunk in output["latents"].split(2, dim=2): + chunks.append(pipe.vae.decode(chunk / 0.18215).sample) + expected = pipe.video_processor.postprocess_video(torch.cat(chunks, dim=2), output_type="pt") + torch.testing.assert_close(output["videos"], expected, rtol=0, atol=0) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA.") + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) + @torch.no_grad() + def test_low_precision_decode_matches_autocast(self, dtype): + from diffusers.modular_pipelines.magi.decoders import MagiVaeDecoderStep + + pipe = MagiVaeDecoderStep().init_pipeline(self.pretrained_model_name_or_path) + pipe.load_components(dtype=dtype) + pipe.to("cuda") + torch.manual_seed(0) + latents = torch.randn(1, 4, 4, 8, 8, device="cuda") + actual = pipe(latents=latents, chunk_width=2, output_type="pt", output="videos") + chunks = [] + with torch.autocast("cuda", dtype=dtype): + for chunk in latents.split(2, dim=2): + chunks.append(pipe.vae.decode((chunk / 0.18215).to(dtype), num_frames=8).sample) + expected = pipe.video_processor.postprocess_video(torch.cat(chunks, dim=2).float(), output_type="pt") + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + def test_single_latent_frame_is_video(self): + pipe = self.get_pipeline() + video = self.run_pipe(pipe, num_frames=4, chunk_width=1) + assert video.shape[1] == 4 + + def test_round_up_frames(self): + pipe = self.get_pipeline() + video = self.run_pipe(pipe, num_frames=12) + assert video.shape[1] == 16 + + def test_caption_cleaning(self): + pytest.importorskip("ftfy") + pytest.importorskip("bs4") + pipe = self.get_pipeline() + first = self.run_pipe(pipe, prompt="

A CAT runs

https://example.com @someone", clean_caption=True) + second = self.run_pipe(pipe, prompt="a cat runs", clean_caption=False) + torch.testing.assert_close(first, second, atol=0, rtol=0) + + def test_duration_saturates(self): + pipe = self.get_pipeline() + prepare = pipe.blocks.sub_blocks["prepare_latents"].init_pipeline(self.pretrained_model_name_or_path) + prepare.load_components() + features = torch.randn(1, 8, 16) + result = prepare( + text_embeds=features, + text_attention_mask=torch.ones(1, 8, dtype=torch.bool), + num_frames=80, + height=8, + width=8, + chunk_width=2, + output=["prompt_embeds", "latents"], + ) + torch.testing.assert_close( + result["prompt_embeds"][0, :, 0], + prepare.text_conditioning.special_embedding.weight[[8, 8, 8, 7, 6, 5, 4, 3, 2, 1]], + ) + assert result["latents"].shape[2] == 20 + + def test_output_formats_and_latent(self): + pipe = self.get_pipeline() + latent = self.run_pipe(pipe, output_type="latent") + assert latent.shape == (1, 4, 4, 4, 4) + pt = self.run_pipe(pipe) + array = self.run_pipe(pipe, output_type="np") + np.testing.assert_allclose(array, pt.permute(0, 1, 3, 4, 2).numpy(), rtol=0, atol=0) + pil = self.run_pipe(pipe, output_type="pil") + assert len(pil) == 1 and len(pil[0]) == 16 + assert pil[0][0].size == (8, 8) + + @pytest.mark.parametrize( + "changed", + [ + {"prompt": []}, + {"height": 7}, + {"num_frames": 7}, + {"num_images_per_prompt": 0}, + {"latents": torch.zeros(1)}, + {"prefix_latents": torch.zeros(1)}, + {"max_sequence_length": 7}, + {"output_type": "invalid"}, + ], + ) + def test_invalid_input(self, changed): + with pytest.raises(ValueError): + self.run_pipe(self.get_pipeline(), **changed) + + def test_pipeline_converter(self, tiny_magi_path, tmp_path): + from scripts.convert_magi_to_diffusers import convert_pipeline + + pipe = MagiModularPipeline.from_pretrained(tiny_magi_path) + pipe.load_components() + transformer_config = dict(pipe.transformer.config) + transformer_config["caption_max_length"] = 64 + transformer = MagiTransformer3DModel.from_config(transformer_config) + transformer.save_pretrained(str(tmp_path / "transformer")) + pipe.text_encoder.save_pretrained(str(tmp_path / "t5")) + pipe.tokenizer.save_pretrained(str(tmp_path / "t5")) + other = np.random.default_rng(0).normal(size=(100, 16)).astype(np.float32) + np.savez(tmp_path / "special.npz", other_tokens=other) + converted = convert_pipeline(transformer, pipe.vae, str(tmp_path / "t5"), tmp_path / "special.npz") + expected = torch.from_numpy(other[[1, *range(7, 15)]].astype(np.float16)).float() + torch.testing.assert_close(converted.text_conditioning.special_embedding.weight, expected, atol=0, rtol=0) + torch.testing.assert_close( + converted.text_conditioning.null_embedding.weight, + transformer.condition_embedder.y_embedder.null_caption_embedding, + atol=0, + rtol=0, + ) + output = converted( + prompt="a cat", + height=8, + width=8, + num_frames=8, + chunk_width=2, + num_inference_steps=4, + window_size=2, + max_sequence_length=64, + clean_caption=False, + output_type="pt", + output="videos", + ) + assert output.shape == (1, 8, 3, 8, 8) + assert output.isfinite().all() + + +class TestMagiPipelineLoading(MagiPipelineTesterConfig, ModularLoadingTesterMixin): + pass + + +class TestMagiPipelineMemory(MagiPipelineTesterConfig, ModularMemoryTesterMixin): + pass + + +class TestMagiPipelineGuider(MagiPipelineTesterConfig, MagiGuiderTesterMixin): + pass diff --git a/tests/modular_pipelines/magi/test_modular_pipeline_magi_prefix.py b/tests/modular_pipelines/magi/test_modular_pipeline_magi_prefix.py new file mode 100644 index 000000000000..69b0f8118552 --- /dev/null +++ b/tests/modular_pipelines/magi/test_modular_pipeline_magi_prefix.py @@ -0,0 +1,242 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +import torch + +from diffusers import MagiImageToVideoBlocks, MagiModularPipeline, MagiVideoToVideoBlocks, ModularPipeline +from diffusers.modular_pipelines.magi.decoders import MagiPrefixVaeDecoderStep +from diffusers.modular_pipelines.magi.encoders import MagiImageVaeEncoderStep, MagiVideoVaeEncoderStep + +from ..testing_utils import ( + BaseModularPipelineTesterConfig, + ModularLoadingTesterMixin, + ModularMemoryTesterMixin, + ModularPipelineTesterMixin, +) +from .test_modular_pipeline_magi import tiny_magi_path # noqa: F401 +from .testing_utils import MagiGuiderTesterMixin + + +@pytest.fixture(scope="module") +def tiny_prefix_paths(request, tmp_path_factory): + base_path = request.getfixturevalue("tiny_magi_path") + paths = {} + for name, blocks in [("image", MagiImageToVideoBlocks), ("video", MagiVideoToVideoBlocks)]: + pipe = blocks().init_pipeline(base_path) + pipe.load_components() + path = str(tmp_path_factory.mktemp(f"tiny-magi-{name}")) + pipe.save_pretrained(path, overwrite_modular_index=True) + paths[name] = path + return paths + + +class MagiImagePipelineTesterConfig(BaseModularPipelineTesterConfig): + pipeline_class = MagiModularPipeline + pipeline_blocks_class = MagiImageToVideoBlocks + params = frozenset(["prompt", "image", "height", "width", "num_frames"]) + batch_params = frozenset(["prompt"]) + output_name = "videos" + workflow = "image" + + @pytest.fixture(scope="class", autouse=True) + @classmethod + def model_path(cls, tiny_prefix_paths): + cls.pretrained_model_name_or_path = tiny_prefix_paths[cls.workflow] + + def get_dummy_inputs(self, seed=0): + pixels = torch.arange(3 * 8 * 8).reshape(1, 3, 8, 8).to(torch.uint8) + return { + "prompt": "a cat runs", + "image": pixels, + "height": 8, + "width": 8, + "num_frames": 16, + "chunk_width": 2, + "window_size": 2, + "num_inference_steps": 4, + "max_sequence_length": 8, + "clean_caption": False, + "output_type": "pt", + "generator": self.get_generator(seed), + } + + +class TestMagiImagePipelineFast(MagiImagePipelineTesterConfig, ModularPipelineTesterMixin): + def test_image_prefix_and_roundtrip(self, tmp_path): + pipe = self.get_pipeline() + inputs = self.get_dummy_inputs() + result = pipe(**inputs, output=["videos", "latents", "conditioning_latents", "completed_chunks"]) + assert result["videos"].shape == (1, 24, 3, 8, 8) + assert result["conditioning_latents"].shape == (1, 4, 1, 4, 4) + assert result["completed_chunks"] == [0, 1, 2] + assert not torch.equal(result["latents"][:, :, :1], result["conditioning_latents"]) + pipe.save_pretrained(str(tmp_path), overwrite_modular_index=True) + restored = ModularPipeline.from_pretrained(str(tmp_path)) + restored.load_components() + actual = restored(**self.get_dummy_inputs(), output="videos") + torch.testing.assert_close(actual, result["videos"], atol=0, rtol=0) + + def test_prefix_batch_expansion(self): + pipe = self.get_pipeline() + inputs = self.get_dummy_inputs() + pixels = torch.cat([inputs["image"], 255 - inputs["image"]]) + original = pixels.clone() + inputs.update(prompt=["a cat", "a cat runs"], image=pixels, num_images_per_prompt=2) + result = pipe(**inputs, output=["conditioning_latents", "videos"]) + prefix = result["conditioning_latents"] + assert prefix.shape[0] == 4 and result["videos"].shape[0] == 4 + torch.testing.assert_close(prefix[0], prefix[1], atol=0, rtol=0) + torch.testing.assert_close(prefix[2], prefix[3], atol=0, rtol=0) + assert not torch.equal(prefix[0], prefix[2]) + torch.testing.assert_close(pixels, original, atol=0, rtol=0) + + def test_prefix_reinjected_for_every_branch(self): + pipe = self.get_pipeline() + calls = [] + + def observe(module, args, kwargs): + calls.append(kwargs["hidden_states"].detach().cpu().clone()) + + handle = pipe.transformer.register_forward_pre_hook(observe, with_kwargs=True) + try: + result = pipe(**self.get_dummy_inputs(), output=["conditioning_latents", "latents"]) + finally: + handle.remove() + prefix = result["conditioning_latents"].cpu() + for states in calls[:12]: + torch.testing.assert_close(states[:1, :, :1], prefix, atol=0, rtol=0) + assert not torch.equal(result["latents"][:, :, :1].cpu(), prefix) + + @pytest.mark.parametrize("bad", [torch.zeros(1, 3, 8, 8), torch.zeros(3, 8, 8, dtype=torch.uint8)]) + def test_invalid_image(self, bad): + with pytest.raises(ValueError): + self.run_pipe(self.get_pipeline(), image=bad) + + @torch.no_grad() + def test_encoder_matches_posterior_mode(self): + pipe = MagiImageVaeEncoderStep().init_pipeline(self.pretrained_model_name_or_path) + pipe.load_components() + image = self.get_dummy_inputs()["image"] + actual = pipe(image=image, output="conditioning_latents") + expected = pipe.vae.encode(image.unsqueeze(2).float() / 127.5 - 1).latent_dist.mode() * 0.18215 + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + +class TestMagiImagePipelineLoading(MagiImagePipelineTesterConfig, ModularLoadingTesterMixin): + pass + + +class TestMagiImagePipelineMemory(MagiImagePipelineTesterConfig, ModularMemoryTesterMixin): + pass + + +class MagiVideoPipelineTesterConfig(MagiImagePipelineTesterConfig): + pipeline_blocks_class = MagiVideoToVideoBlocks + params = frozenset(["prompt", "video", "height", "width", "num_frames"]) + workflow = "video" + + def get_dummy_inputs(self, seed=0): + inputs = super().get_dummy_inputs(seed) + image = inputs.pop("image") + inputs["video"] = image.unsqueeze(2).repeat(1, 1, 12, 1, 1) + return inputs + + +class TestMagiVideoPipelineFast(MagiVideoPipelineTesterConfig, ModularPipelineTesterMixin): + @pytest.mark.parametrize("prefix_frames, expected_frames", [(8, 16), (12, 17)]) + def test_prefix_duration_and_trim(self, prefix_frames, expected_frames): + pipe = self.get_pipeline() + inputs = self.get_dummy_inputs() + inputs["video"] = inputs["video"][:, :, :prefix_frames] + result = pipe(**inputs, output=["latents", "videos", "conditioning_latents", "prompt_embeds"]) + prefix = result["conditioning_latents"] + full_length = prefix.shape[2] // 2 * 2 + torch.testing.assert_close(result["latents"][:, :, :full_length], prefix[:, :, :full_length], atol=0, rtol=0) + assert result["videos"].shape == (1, expected_frames, 3, 8, 8) + generated = result["latents"].shape[2] // 2 - prefix.shape[2] // 2 + expected_duration = pipe.text_conditioning.special_embedding.weight[generated] + torch.testing.assert_close(result["prompt_embeds"][0, prefix.shape[2] // 2, 0], expected_duration) + decoder = MagiPrefixVaeDecoderStep().init_pipeline(self.pretrained_model_name_or_path) + decoder.load_components() + suffix = decoder( + latents=result["latents"], + conditioning_latents=prefix, + chunk_width=2, + output_type="latent", + output="videos", + ) + torch.testing.assert_close(suffix, result["latents"][:, :, prefix.shape[2] :], atol=0, rtol=0) + chunks = [] + with torch.no_grad(): + for start in range(0, result["latents"].shape[2], 2): + if start + 2 <= prefix.shape[2]: + continue + chunk = result["latents"][:, :, max(start, prefix.shape[2]) : start + 2] + chunks.append(pipe.vae.decode(chunk / 0.18215).sample) + expected = pipe.video_processor.postprocess_video(torch.cat(chunks, dim=2), output_type="pt") + torch.testing.assert_close(result["videos"], expected, atol=0, rtol=0) + + def test_tiled_single_position_suffix(self): + decoder = MagiPrefixVaeDecoderStep().init_pipeline(self.pretrained_model_name_or_path) + decoder.load_components() + decoder.vae.enable_tiling(tile_sample_min_length=12) + latents = torch.randn(1, 4, 18, 4, 4) + video = decoder( + latents=latents, conditioning_latents=latents[:, :, :8], chunk_width=6, output_type="pt", output="videos" + ) + assert video.shape == (1, 37, 3, 8, 8) + + def test_partial_prefix_reinjected(self): + pipe = self.get_pipeline() + calls = [] + + def observe(module, args, kwargs): + calls.append(kwargs["hidden_states"].detach().cpu().clone()) + + handle = pipe.transformer.register_forward_pre_hook(observe, with_kwargs=True) + try: + result = pipe(**self.get_dummy_inputs(), output=["conditioning_latents", "latents"]) + finally: + handle.remove() + prefix = result["conditioning_latents"].cpu() + torch.testing.assert_close(calls[0], prefix[:, :, :2], atol=0, rtol=0) + for states in calls[1:13]: + torch.testing.assert_close(states[:1, :, :1], prefix[:, :, 2:3], atol=0, rtol=0) + assert not torch.equal(result["latents"][:, :, 2:3].cpu(), prefix[:, :, 2:3]) + + @torch.no_grad() + def test_encoder_reusable(self): + pipe = MagiVideoVaeEncoderStep().init_pipeline(self.pretrained_model_name_or_path) + pipe.load_components() + video = self.get_dummy_inputs()["video"] + actual = pipe(video=video, output="conditioning_latents") + expected = pipe.vae.encode(video.float() / 127.5 - 1).latent_dist.mode() * 0.18215 + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + +class TestMagiVideoPipelineLoading(MagiVideoPipelineTesterConfig, ModularLoadingTesterMixin): + pass + + +class TestMagiVideoPipelineMemory(MagiVideoPipelineTesterConfig, ModularMemoryTesterMixin): + pass + + +class TestMagiImagePipelineGuider(MagiImagePipelineTesterConfig, MagiGuiderTesterMixin): + pass + + +class TestMagiVideoPipelineGuider(MagiVideoPipelineTesterConfig, MagiGuiderTesterMixin): + pass diff --git a/tests/modular_pipelines/magi/testing_utils.py b/tests/modular_pipelines/magi/testing_utils.py new file mode 100644 index 000000000000..68ca58ceee5d --- /dev/null +++ b/tests/modular_pipelines/magi/testing_utils.py @@ -0,0 +1,57 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +import torch + +from diffusers import MagiClassifierFreeGuidance + +from ...testing_utils import torch_device +from ..testing_utils import ModularGuiderTesterMixin + + +class MagiGuiderTesterMixin(ModularGuiderTesterMixin): + def test_guider_cfg(self): + pipe = self.get_pipeline().to(torch_device) + pipe.update_components( + guider=MagiClassifierFreeGuidance(timestep_thresholds=(0.0,), prefix_scales=(1.0,), text_scales=(1.0,)) + ) + conditional = pipe(**self.get_dummy_inputs(), output="latents") + pipe.update_components(guider=MagiClassifierFreeGuidance()) + guided = pipe(**self.get_dummy_inputs(), output=["latents", "completed_chunks", "clean_kv_cache"]) + pipe.guider.disable() + disabled = pipe(**self.get_dummy_inputs(), output=["latents", "completed_chunks", "clean_kv_cache"]) + torch.testing.assert_close(disabled["latents"], conditional, atol=0, rtol=0) + assert not torch.allclose(guided["latents"], disabled["latents"]) + assert guided["completed_chunks"] == disabled["completed_chunks"] + assert guided["clean_kv_cache"] is not None and disabled["clean_kv_cache"] is not None + pipe.guider.enable() + restored = pipe(**self.get_dummy_inputs(), output="latents") + torch.testing.assert_close(restored, guided["latents"], atol=0, rtol=0) + + @pytest.mark.parametrize("prefix_scale,text_scale", [(2.0, 1.0), (1.0, 3.0)]) + def test_guidance_scales(self, prefix_scale, text_scale): + pipe = self.get_pipeline().to(torch_device) + pipe.update_components( + guider=MagiClassifierFreeGuidance(timestep_thresholds=(0.0,), prefix_scales=(1.0,), text_scales=(1.0,)) + ) + expected = pipe(**self.get_dummy_inputs(), output="latents") + pipe.update_components( + guider=MagiClassifierFreeGuidance( + timestep_thresholds=(0.0,), prefix_scales=(prefix_scale,), text_scales=(text_scale,) + ) + ) + actual = pipe(**self.get_dummy_inputs(), output="latents") + assert actual.shape == expected.shape + assert not torch.allclose(actual, expected) diff --git a/tests/schedulers/test_scheduler_magi_euler.py b/tests/schedulers/test_scheduler_magi_euler.py new file mode 100644 index 000000000000..87ee5d5df6e5 --- /dev/null +++ b/tests/schedulers/test_scheduler_magi_euler.py @@ -0,0 +1,179 @@ +# Copyright 2025 SandAI and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +import torch + +from diffusers import MagiEulerScheduler + + +class TestMagiEulerScheduler: + @pytest.mark.parametrize("steps", [1, 4, 12, 16, 64]) + @pytest.mark.parametrize("time_schedule", ["sd3", "square", "piecewise", "linear"]) + def test_time_grid(self, steps, time_schedule): + scheduler = MagiEulerScheduler(time_schedule=time_schedule) + scheduler.set_timesteps(steps) + grid = scheduler.timestep_schedule + assert grid.dtype == torch.float32 + assert scheduler.timesteps.shape == (steps,) + assert grid.shape == (steps + 1,) + assert grid[0] == 0 + torch.testing.assert_close(grid[-1], torch.tensor(1.0), atol=2e-7, rtol=0) + assert (grid.diff() > 0).all() + torch.testing.assert_close(scheduler.timesteps, grid[:-1], atol=0, rtol=0) + + @pytest.mark.parametrize( + "mode,first", [("8,16,16", [0, 0.125, 0.1875, 0.25]), ("16,16,8", [0, 0.0625, 0.125, 0.25])] + ) + def test_twelve_step_shortcut(self, mode, first): + scheduler = MagiEulerScheduler(time_schedule="linear", shortcut_mode=mode) + scheduler.set_timesteps(12) + torch.testing.assert_close(scheduler.timesteps[:4], torch.tensor(first), atol=0, rtol=0) + + def test_sd3_operation_order_and_endpoint(self): + scheduler = MagiEulerScheduler() + scheduler.set_timesteps(64) + squared = torch.linspace(0, 1, 65) ** 2 + expected = (1 / 3) * squared / (1 + (1 / 3 - 1) * squared) + torch.testing.assert_close(scheduler.timestep_schedule, expected, atol=0, rtol=0) + assert scheduler.timestep_schedule[-1] > 1 + + @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16, torch.float16]) + def test_fp32_update(self, dtype): + scheduler = MagiEulerScheduler(time_schedule="linear") + scheduler.set_timesteps(4) + sample = torch.randn((2, 4, 4, 4, 6), generator=torch.Generator().manual_seed(0)).to(dtype) + velocity = torch.randn(sample.shape, generator=torch.Generator().manual_seed(1)).to(dtype) + original = sample.clone() + output = scheduler.step(velocity, 0.25, sample, next_timestep=0.5) + assert output.prev_sample.dtype == torch.float32 + torch.testing.assert_close(output.prev_sample, sample.float() + velocity.float() * 0.25, atol=0, rtol=0) + torch.testing.assert_close(sample, original, atol=0, rtol=0) + assert scheduler.step_index is None + torch.testing.assert_close( + scheduler.step(velocity, 0.25, sample, next_timestep=0.5, return_dict=False)[0], + output.prev_sample, + atol=0, + rtol=0, + ) + + def test_per_batch_chunk_update(self): + scheduler = MagiEulerScheduler() + scheduler.set_timesteps(64) + sample = torch.randn((2, 4, 6, 4, 6), generator=torch.Generator().manual_seed(0)) + velocity = sample.sin() + before = torch.tensor([[0.4, 0.2, 0.0], [0.6, 0.4, 0.2]]) + after = torch.tensor([[0.5, 0.3, 0.1], [0.7, 0.5, 0.3]]) + expected = torch.cat( + [ + sample[:, :, index * 2 : (index + 1) * 2] + + velocity[:, :, index * 2 : (index + 1) * 2] + * (after[:, index] - before[:, index])[:, None, None, None, None] + for index in range(3) + ], + dim=2, + ) + actual = scheduler.step(velocity, before, sample, next_timestep=after).prev_sample + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + assert scheduler.step_index is None + + def test_zero_velocity_and_zero_interval(self): + scheduler = MagiEulerScheduler() + scheduler.set_timesteps(4) + sample = torch.randn(1, 4, 4, 4, 4) + torch.testing.assert_close(scheduler.step(torch.zeros_like(sample), 0.2, sample, 0.9).prev_sample, sample) + torch.testing.assert_close(scheduler.step(sample, 0.2, sample, 0.2).prev_sample, sample) + + def test_full_trajectory_and_reset(self): + scheduler = MagiEulerScheduler() + scheduler.set_timesteps(64) + initial = torch.randn((1, 4, 4, 4, 6), generator=torch.Generator().manual_seed(0)) + sample = initial.clone() + expected = initial.clone() + for index, timestep in enumerate(scheduler.timesteps): + velocity = sample.sin() * 0.2 + timestep + expected_velocity = expected.sin() * 0.2 + timestep + expected = expected + expected_velocity * (scheduler.timestep_schedule[index + 1] - timestep) + sample = scheduler.step(velocity, timestep, sample).prev_sample + torch.testing.assert_close(sample, expected, atol=0, rtol=0) + assert scheduler.step_index == index + 1 + with pytest.raises(ValueError, match="complete"): + scheduler.step(velocity, scheduler.timesteps[-1], sample) + scheduler.set_timesteps(64) + assert scheduler.step_index is None + for timestep in scheduler.timesteps: + initial = scheduler.step(initial.sin() * 0.2 + timestep, timestep, initial).prev_sample + torch.testing.assert_close(initial, sample, atol=0, rtol=0) + + def test_out_of_order_step(self): + scheduler = MagiEulerScheduler() + scheduler.set_timesteps(4) + sample = torch.zeros(1, 1, 1, 1, 1) + with pytest.raises(ValueError, match="scalar from"): + scheduler.step(sample, 0.123, sample) + scheduler.step(sample, scheduler.timesteps[0], sample) + with pytest.raises(ValueError, match="next sequential"): + scheduler.step(sample, scheduler.timesteps[0], sample) + assert scheduler.step_index == 1 + + @pytest.mark.parametrize( + "config", [{"shift": 0}, {"shift": float("nan")}, {"time_schedule": "unknown"}, {"shortcut_mode": "unknown"}] + ) + def test_invalid_config(self, config): + with pytest.raises(ValueError): + MagiEulerScheduler(**config) + + @pytest.mark.parametrize("steps", [0, -1, 1.5, True]) + def test_invalid_step_count(self, steps): + with pytest.raises(ValueError, match="positive integer"): + MagiEulerScheduler().set_timesteps(steps) + + @pytest.mark.parametrize( + "before,after", [(0.5, 0.4), (-0.1, 0.1), (0.9, 1.1), (float("nan"), 1), (0, float("inf"))] + ) + def test_invalid_endpoints(self, before, after): + scheduler = MagiEulerScheduler() + scheduler.set_timesteps(4) + with pytest.raises(ValueError, match="Timesteps"): + scheduler.step(torch.ones(1), before, torch.ones(1), after) + + def test_invalid_shapes_and_uninitialized(self): + scheduler = MagiEulerScheduler() + sample = torch.zeros(2, 4, 4, 4, 4) + with pytest.raises(ValueError, match="set_timesteps"): + scheduler.step(sample, 0.0, sample) + scheduler.set_timesteps(4) + with pytest.raises(ValueError, match="same shape"): + scheduler.step(sample[:1], 0.0, sample) + with pytest.raises(ValueError, match="require next_timestep"): + scheduler.step(sample, torch.zeros(2), sample) + for shape in [(3,), (3, 2), (1, 1, 1), (0,)]: + with pytest.raises(ValueError, match="Chunk times"): + scheduler.step(sample, torch.zeros(shape), sample, torch.ones(shape)) + + @pytest.mark.parametrize("time_schedule", ["sd3", "square", "piecewise", "linear"]) + def test_save_load(self, tmp_path, time_schedule): + scheduler = MagiEulerScheduler(shift=5.0, time_schedule=time_schedule, shortcut_mode="16,16,8") + scheduler.set_timesteps(12) + scheduler.save_pretrained(tmp_path) + restored = MagiEulerScheduler.from_pretrained(tmp_path) + restored.set_timesteps(12) + torch.testing.assert_close(restored.timestep_schedule, scheduler.timestep_schedule, atol=0, rtol=0) + sample = torch.ones(1, 2, 4, 4, 4) + torch.testing.assert_close( + restored.step(sample, restored.timesteps[0], sample).prev_sample, + scheduler.step(sample, scheduler.timesteps[0], sample).prev_sample, + atol=0, + rtol=0, + )