[Model] Add native MAGI-2 Preview diffusion support (#5918)

Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>

examples/offline_inference/image_to_video/README.md

@@ -1,7 +1,7 @@
 # Image-To-Video
 
 This shared example generates videos from images with VACE, Wan2.2, LTX-2,
-HunyuanVideo-1.5, SANA-Video, Cosmos3, and other compatible pipelines.
+HunyuanVideo-1.5, SANA-Video, Cosmos3, MAGI-2, and other compatible pipelines.
 
 - `image_to_video.py`: command-line script for single video generation with advanced options.
 
@@ -27,9 +27,14 @@ This folder provides a unified CLI script for image-to-video generation using vL
 | `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | 480 x 832 | 81 | 50 | 4.0 | Around 20–25 GiB BF16, smallest I2V model |
 | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v` | 480 x 832 | 121 | 50 | 6.0 | Around 100 GiB at default settings; the example enables `--enable-cpu-offload` + VAE tiling/slicing to fit an 80 GiB card |
 | `Lightricks/LTX-2` | 512 x 768 | 121 | 40 | video 3.0 / audio 7.0 | Memory use depends on frame count and tensor parallelism |
+| `sand-ai/MAGI-2-preview` | 512 x 896 | 125 | 100 | Model-fixed | Native four-GPU TP/SP; resident SP4 default; DLO available |
 | `Efficient-Large-Model/SANA-Video_2B_480p_diffusers` | 480 x 832 | 81 | 50 | 6.0 | Native `SanaImageToVideoPipeline`; Wan VAE |
 | `Efficient-Large-Model/SANA-Video_2B_720p_diffusers` | 704 x 1280 | 81 | 50 | 6.0 | Native `SanaImageToVideoPipeline`; LTX-2 Video VAE |
 
+MAGI-2 native Preview setup, four-GPU topology and DLO choices, I2VA commands,
+and eight-GPU validation status are documented in the
+[`MAGI-2 Preview L20X recipe`](../../../recipes/SandAI/MAGI-2-preview-L20X.md).
+
 !!! info
     Peak VRAM: based on basic single-card usage, batch size = 1, without any acceleration/optimization features. Some model weights cannot fit into one card with 80 GiB VRAM, which may need to use CPU offloading.
 
@@ -147,6 +152,9 @@ python image_to_video.py \
 | `--vae-use-tiling` | flag | off | Enable VAE tiling for memory optimization |
 | `--enable-cpu-offload` | flag | off | Enable CPU offloading for diffusion models |
 | `--enable-layerwise-offload` | flag | off | Enable layerwise offloading on DiT modules |
+| `--enable-distributed-layerwise-offload` | flag | off | Enable distributed layerwise offload |
+| `--dlo-use-allgather` / `--dlo-no-use-allgather` | flag | on / off | Select sharded AllGather reconstruction (default) or rank-local streaming |
+| `--dlo-resident-layers` | int | `0` | Leading main-DiT blocks kept device-resident during distributed layerwise offload |
 | `--cfg-parallel-size` | int | `1` | Set to `2` to enable CFG Parallel |
 | `--tensor-parallel-size` | int | `1` | Tensor parallel size (effective for models that support TP, e.g. LTX2) |
 | `--ulysses-degree` | int | `1` | Ulysses sequence parallel degree |
@@ -339,8 +347,6 @@ Key arguments:
 - `--hsdp-shard-size`: Number of GPUs to shard model weights across within each replica group. -1 (default) auto-calculates as world_size / replicate_size.
 - `--hsdp-replicate-size`: Number of replica groups for HSDP. Each replica holds a full sharded copy. Default 1 means pure sharding (no replication).
 
-
-
 > ℹ️ If you encounter OOM errors, try using `--vae-use-slicing` and `--vae-use-tiling` to reduce memory usage.
 
 ## Wan2.1 VACE Conditional Tasks

examples/offline_inference/image_to_video/image_to_video.py

@@ -3,7 +3,7 @@
 
 """
 Image-to-Video generation example using Wan2.2 I2V/TI2V models, LTX2/LTX-2.3,
-HunyuanVideo-1.5, SANA-Video, Cosmos3, or Wan2.1 VACE.
+HunyuanVideo-1.5, SANA-Video, Cosmos3, MAGI-2, or Wan2.1 VACE.
 
 Supports:
 - Wan2.2-I2V-A14B-Diffusers: MoE model with CLIP image encoder
@@ -64,6 +64,7 @@ import numpy as np
 import PIL.Image
 import torch
 
+from vllm_omni.diffusion.data import resolve_model_class_name
 from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
@@ -75,6 +76,7 @@ from vllm_omni.model_extras import (
 from vllm_omni.model_extras import (
     get_extra_body_params,
     get_model_class_name,
+    get_video_generation_defaults,
     should_preserve_reference_image_size,
 )
 from vllm_omni.outputs import OmniRequestOutput
@@ -122,15 +124,15 @@ def parse_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser(
         description=(
             "Generate a video from one or more images "
-            "(Wan2.2, LTX2/LTX-2.3, HunyuanVideo-1.5, SANA-Video, Cosmos3, or Wan2.1 VACE)."
+            "(Wan2.2, LTX2/LTX-2.3, HunyuanVideo-1.5, SANA-Video, Cosmos3, MAGI-2, or Wan2.1 VACE)."
         )
     )
     parser.add_argument(
         "--model",
         default="Wan-AI/Wan2.2-I2V-A14B-Diffusers",
         help=(
-            "Diffusers I2V model ID or local path "
-            "(Wan2.2, LTX2/LTX-2.3, HunyuanVideo-1.5, SANA-Video, Cosmos3, or Wan2.1 VACE)."
+            "I2V model ID or local path "
+            "(Wan2.2, LTX2/LTX-2.3, HunyuanVideo-1.5, SANA-Video, Cosmos3, MAGI-2, or Wan2.1 VACE)."
         ),
     )
     parser.add_argument(
@@ -209,7 +211,7 @@ def parse_args() -> argparse.Namespace:
         help="Diffusion KV-cache quantization skip-layer selector, e.g. '0,1,4-8'.",
     )
     parser.add_argument("--output", type=str, default="i2v_output.mp4", help="Path to save the video (mp4).")
-    parser.add_argument("--fps", type=int, default=None, help="Frames per second for the output video.")
+    parser.add_argument("--fps", type=float, default=None, help="Frames per second for the output video.")
     parser.add_argument(
         "--vae-use-slicing",
         action="store_true",
@@ -230,6 +232,30 @@ def parse_args() -> argparse.Namespace:
         action="store_true",
         help="Enable layerwise (blockwise) offloading on DiT modules.",
     )
+    parser.add_argument(
+        "--enable-distributed-layerwise-offload",
+        action="store_true",
+        help="Enable distributed layerwise offloading with overlapped host-to-device weight streaming.",
+    )
+    parser.add_argument(
+        "--dlo-use-allgather",
+        dest="dlo_use_allgather",
+        action="store_true",
+        default=True,
+        help="Use shard + AllGather weight reconstruction for distributed layerwise offload (default: enabled).",
+    )
+    parser.add_argument(
+        "--dlo-no-use-allgather",
+        dest="dlo_use_allgather",
+        action="store_false",
+        help="Stream standard-loader rank-local weights without DLO sharding or AllGather.",
+    )
+    parser.add_argument(
+        "--dlo-resident-layers",
+        type=int,
+        default=0,
+        help="Number of leading main-DiT blocks to keep device-resident during distributed layerwise offload.",
+    )
     parser.add_argument(
         "--enforce-eager",
         action="store_true",
@@ -391,12 +417,14 @@ def main():
     generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
     model_name = str(args.model).lower() if args.model is not None else ""
     model_class_name = args.model_class_name
-    model_class_name_lower = (model_class_name or "").lower()
+    resolved_model_class_name = model_class_name or resolve_model_class_name(args.model)
+    model_class_name_lower = (resolved_model_class_name or "").lower()
+    video_defaults = get_video_generation_defaults(resolved_model_class_name, args.extra_body)
     is_ltx2_distilled = "distilled" in model_class_name_lower or "distilled" in model_name
     is_ltx23 = "ltx23" in model_class_name_lower or "ltx-2.3" in model_name
     is_ltx2 = is_ltx2_distilled or is_ltx23 or "ltx2" in model_class_name_lower or "ltx-2" in model_name
-    is_sana = model_class_name == "SanaImageToVideoPipeline" or "sana-video" in model_name
-    is_cosmos = "cosmos" in model_name or (model_class_name is not None and "cosmos" in model_class_name.lower())
+    is_sana = resolved_model_class_name == "SanaImageToVideoPipeline" or "sana-video" in model_name
+    is_cosmos = "cosmos" in model_name or "cosmos" in model_class_name_lower
     is_cosmos_edge = is_cosmos and ("edge" in model_name or "edge" in model_class_name_lower)
 
     image = PIL.Image.open(args.image).convert("RGB") if args.image else None
@@ -411,7 +439,17 @@ def main():
 
     # Per-model generation defaults, applied only when the matching flag is omitted.
     # Cosmos3 would otherwise silently inherit the Wan2.2 defaults (wrong size/steps/shift).
-    if is_cosmos_edge:
+    if video_defaults is not None:
+        d_fps, d_guidance, d_num_frames, d_steps, d_flow_shift, d_max_area, d_mod = (
+            video_defaults.fps,
+            video_defaults.guidance_scale,
+            video_defaults.num_frames,
+            video_defaults.num_inference_steps,
+            video_defaults.flow_shift,
+            video_defaults.max_area,
+            video_defaults.dimension_multiple,
+        )
+    elif is_cosmos_edge:
         # Cosmos3-Edge native defaults: 480x832, gs 5.0, flow_shift 3.0 — NOT the Nano/Super
         # 720p / gs 6.0 / flow_shift 10.0 below, which produce degenerate output on Edge.
         d_fps, d_guidance, d_num_frames, d_steps, d_flow_shift, d_max_area, d_mod = (
@@ -477,7 +515,10 @@ def main():
     # Calculate dimensions if not provided (model-aware max area).
     height = args.height
     width = args.width
-    if height is None or width is None:
+    if video_defaults is not None:
+        height = height or video_defaults.height
+        width = width or video_defaults.width
+    elif height is None or width is None:
         calc_height, calc_width = calculate_dimensions(dimension_image, max_area=d_max_area, mod_value=d_mod)
         height = height or calc_height
         width = width or calc_width
@@ -557,6 +598,9 @@ def main():
         cache_config=cache_config,
         enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
         profiler_config=args.profiler_config,
+        enable_distributed_layerwise_offload=args.enable_distributed_layerwise_offload,
+        dlo_use_allgather=args.dlo_use_allgather,
+        dlo_resident_layers=args.dlo_resident_layers,
     )
     if args.deploy_config:
         omni_kwargs["deploy_config"] = args.deploy_config
@@ -613,9 +657,12 @@ def main():
         )
 
     negative_prompt = args.negative_prompt
-    if negative_prompt is None and not is_ltx2:
-        # Preserve the historical empty-prompt behavior for non-LTX examples.
-        negative_prompt = ""
+    if negative_prompt is None:
+        if video_defaults is not None:
+            negative_prompt = video_defaults.default_negative_prompt
+        elif not is_ltx2:
+            # Preserve the historical empty-prompt behavior for non-LTX examples.
+            negative_prompt = ""
     prompt_dict = build_image_to_video_prompt(
         prompt=args.prompt,
         negative_prompt=negative_prompt,
@@ -642,6 +689,7 @@ def main():
         boundary_ratio=args.boundary_ratio,
         num_inference_steps=num_inference_steps,
         num_frames=num_frames,
+        fps=fps,
         frame_rate=frame_rate,
         extra_args=sampling_extra_args,
     )

examples/offline_inference/text_to_video/text_to_video.md

@@ -14,9 +14,14 @@ A unified script for text-to-video generation. Supports multiple models with mod
 | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | 720x1280 | 121 | 50 | 6.0 | FP8 + VAE tiling required |
 | `nvidia/Cosmos3-Nano` | 720x1280 | 189 | 35 | 6.0 | ~46 GiB (peak, 720p) |
 | `BestWishYsh/Helios-Base` / `Helios-Mid` / `Helios-Distilled` | 384x640 | 99 | 50 | 5.0 / 5.0 / 1.0 | — |
+| `sand-ai/MAGI-2-preview` | 512x896 | 125 | 100 | Model-fixed | Native four-GPU TP/SP; resident SP4 default; DLO available |
 | `Efficient-Large-Model/SANA-Video_2B_480p_diffusers` | 480x832 | 81 | 50 | 6.0 | BF16 DiT + FP32 Wan VAE wrapper |
 | `Efficient-Large-Model/SANA-Video_2B_720p_diffusers` | 704x1280 | 81 | 50 | 6.0 | BF16 DiT + LTX-2 Video VAE wrapper |
 
+MAGI-2 native Preview setup, four-GPU topology and DLO choices, request
+constraints, and eight-GPU validation status are documented in the
+[`MAGI-2 Preview L20X recipe`](../../../recipes/SandAI/MAGI-2-preview-L20X.md).
+
 ## Local CLI Usage
 
 ### Wan2.2 (default)
@@ -264,6 +269,9 @@ python text_to_video.py \
 - `--tensor-parallel-size`: tensor parallel size (effective for models that support TP, e.g. LTX2).
 - `--enable-cpu-offload`: enable CPU offloading for diffusion models.
 - `--enable-layerwise-offload`: enable layerwise offloading on DiT modules.
+- `--enable-distributed-layerwise-offload`: enable distributed layerwise offload.
+- `--dlo-use-allgather` / `--dlo-no-use-allgather`: use sharded AllGather reconstruction (the default) or stream rank-local weights without AllGather.
+- `--dlo-resident-layers`: keep this many leading main-DiT blocks device-resident during distributed layerwise offload (default: `0`).
 - `--frame-rate`: generation FPS for pipelines that require it (e.g., LTX2).
 - `--audio-sample-rate`: fallback audio sample rate when the pipeline returns audio.
 - `--quantization`: quantization method (such as `fp8` for FP8).

examples/offline_inference/text_to_video/text_to_video.py

@@ -10,12 +10,18 @@ from typing import Any
 import numpy as np
 import torch
 
+from vllm_omni.diffusion.data import resolve_model_class_name
 from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.lora.request import LoRARequest
 from vllm_omni.lora.utils import stable_lora_int_id
-from vllm_omni.model_extras import get_extra_body_params, get_model_class_name, get_output_tensor_range
+from vllm_omni.model_extras import (
+    get_extra_body_params,
+    get_model_class_name,
+    get_output_tensor_range,
+    get_video_generation_defaults,
+)
 from vllm_omni.outputs import OmniRequestOutput
 from vllm_omni.platforms import current_omni_platform
 
@@ -135,9 +141,17 @@ _MODEL_PRESETS = {
 }
 
 
-def _detect_preset(model: str, model_class_name: str | None = None) -> dict:
+def _detect_preset(
+    model: str,
+    model_class_name: str | None = None,
+    extra_body: dict[str, Any] | None = None,
+) -> dict:
     model_lower = model.lower()
-    class_lower = (model_class_name or "").lower()
+    resolved_model_class_name = model_class_name or resolve_model_class_name(model)
+    class_lower = (resolved_model_class_name or "").lower()
+    video_defaults = get_video_generation_defaults(resolved_model_class_name, extra_body)
+    if video_defaults is not None:
+        return video_defaults.cli_defaults()
     if "sana-video" in model_lower or "sana_video" in model_lower or "sanavideo" in class_lower:
         return _MODEL_PRESETS["sana_720p" if "720p" in model_lower else "sana_480p"]
     if "lingbot" in model_lower or "lingbotvideo" in class_lower:
@@ -221,7 +235,8 @@ def parse_extra_body(value: str) -> dict[str, Any]:
 def parse_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser(
         description="Generate a video from a text prompt. "
-        "Supports Wan2.2, HunyuanVideo-1.5, Helios, LingBot-Video, and other text-to-video models."
+        "Supports Wan2.2, HunyuanVideo-1.5, Helios, LingBot-Video, MAGI-2, "
+        "and other text-to-video models."
     )
     parser.add_argument(
         "--model",
@@ -288,7 +303,7 @@ def parse_args() -> argparse.Namespace:
         type=str,
         default=None,
         choices=["cache_dit"],
-        help="Cache backend for acceleration (Wan2.2). Default: None.",
+        help="Cache backend for supported diffusion pipelines. Default: None.",
     )
     parser.add_argument(
         "--enable-cache-dit-summary",
@@ -296,7 +311,7 @@ def parse_args() -> argparse.Namespace:
         help="Enable cache-dit summary logging after diffusion forward passes.",
     )
     parser.add_argument("--output", type=str, default=None, help="Output path (mp4). Default: model-specific.")
-    parser.add_argument("--fps", type=int, default=None, help="Frames per second for the output video.")
+    parser.add_argument("--fps", type=float, default=None, help="Frames per second for the output video.")
     parser.add_argument(
         "--vae-use-slicing",
         action="store_true",
@@ -322,6 +337,30 @@ def parse_args() -> argparse.Namespace:
         action="store_true",
         help="Enable layerwise (blockwise) offloading on DiT modules.",
     )
+    parser.add_argument(
+        "--enable-distributed-layerwise-offload",
+        action="store_true",
+        help="Enable distributed layerwise offloading with overlapped host-to-device weight streaming.",
+    )
+    parser.add_argument(
+        "--dlo-use-allgather",
+        dest="dlo_use_allgather",
+        action="store_true",
+        default=True,
+        help="Use shard + AllGather weight reconstruction for distributed layerwise offload (default: enabled).",
+    )
+    parser.add_argument(
+        "--dlo-no-use-allgather",
+        dest="dlo_use_allgather",
+        action="store_false",
+        help="Stream standard-loader rank-local weights without DLO sharding or AllGather.",
+    )
+    parser.add_argument(
+        "--dlo-resident-layers",
+        type=int,
+        default=0,
+        help="Number of leading main-DiT blocks to keep device-resident during distributed layerwise offload.",
+    )
     parser.add_argument(
         "--audio-sample-rate",
         type=int,
@@ -464,15 +503,17 @@ def _extract_peak_memory_mb(result: Any) -> float:
 def main():
     args = parse_args()
     model_class_name = args.model_class_name
+    resolved_model_class_name = model_class_name or resolve_model_class_name(args.model)
+    video_defaults = get_video_generation_defaults(resolved_model_class_name, args.extra_body)
 
-    preset = _detect_preset(args.model, model_class_name)
+    preset = _detect_preset(args.model, resolved_model_class_name, args.extra_body)
     for key, default_val in preset.items():
         if getattr(args, key.replace("-", "_"), None) is None:
             setattr(args, key.replace("-", "_"), default_val)
     model_class_name = args.model_class_name
 
     generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
-    # Cache-dit config (Wan2.2 only)
+    # Shared Cache-DiT defaults for pipelines that declare an adapter.
     cache_config = None
     if args.cache_backend == "cache_dit":
         cache_config = {
@@ -501,6 +542,9 @@ def main():
         cfg_parallel_size=args.cfg_parallel_size,
         tensor_parallel_size=args.tensor_parallel_size,
         vae_patch_parallel_size=args.vae_patch_parallel_size,
+        use_hsdp=args.use_hsdp,
+        hsdp_shard_size=args.hsdp_shard_size,
+        hsdp_replicate_size=args.hsdp_replicate_size,
         pipeline_parallel_size=args.pipeline_parallel_size,
         enable_expert_parallel=args.enable_expert_parallel,
         enforce_eager=args.enforce_eager,
@@ -509,6 +553,9 @@ def main():
         cache_config=cache_config,
         enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
         profiler_config=args.profiler_config,
+        enable_distributed_layerwise_offload=args.enable_distributed_layerwise_offload,
+        dlo_use_allgather=args.dlo_use_allgather,
+        dlo_resident_layers=args.dlo_resident_layers,
     )
     if args.deploy_config:
         omni_kwargs["deploy_config"] = args.deploy_config
@@ -572,6 +619,8 @@ def main():
         )
 
     negative_prompt = args.negative_prompt
+    if negative_prompt is None and video_defaults is not None:
+        negative_prompt = video_defaults.default_negative_prompt
     if negative_prompt is None and all(preset is not _MODEL_PRESETS[name] for name in ("lingbot", "ltx2", "ltx23")):
         # Preserve the historical empty-prompt behavior for non-LTX examples.
         negative_prompt = ""
@@ -591,7 +640,10 @@ def main():
         num_frames=args.num_frames,
         extra_args=extra_args,
     )
-    if args.frame_rate is not None:
+    if video_defaults is not None:
+        sampling_kwargs["fps"] = args.fps
+        sampling_kwargs["frame_rate"] = args.frame_rate or args.fps
+    elif args.frame_rate is not None:
         sampling_kwargs["frame_rate"] = args.frame_rate
     if args.guidance_scale_high is not None:
         sampling_kwargs["guidance_scale_2"] = args.guidance_scale_high