[Model Support]: Magihuman support (#2301)

Signed-off-by: princepride <wangzhipeng628@gmail.com>
Signed-off-by: 汪志鹏 <wangzhipeng628@gmail.com>
Co-authored-by: Didan Deng <33117903+wtomin@users.noreply.github.com>

docs/models/supported_models.md

@@ -37,6 +37,7 @@ th {
 | `LTX2TwoStagesPipeline` | LTX-2-T2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | |
 | `LTX2ImageToVideoTwoStagesPipeline` | LTX-2-I2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | |
 | `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | ✅︎ | ✅︎ | ✅︎ | |
+| `MagiHumanPipeline` | MagiHuman | `princepride/daVinci-MagiHuman` | ✅︎ | ✅︎ | | |
 | `OvisImagePipeline` | Ovis-Image | `OvisAI/Ovis-Image` | ✅︎ | ✅︎ | | ✅︎ |
 | `LongcatImagePipeline` | LongCat-Image | `meituan-longcat/LongCat-Image` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |

docs/user_guide/diffusion_features.md

@@ -109,6 +109,7 @@ The following tables show which models support each feature:
 | **HunyuanImage3** | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ |
 | **LongCat-Image** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
 | **LongCat-Image-Edit** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
+| **MagiHuman** | ❌ | ❌ | ❌ | ❓ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ |
 | **MammothModa2(T2I)** | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
 | **Nextstep_1(T2I)** | ❓ | ❓ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
 | **OmniGen2** | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
@@ -123,7 +124,7 @@ The following tables show which models support each feature:
 
 > Notes:
 > 1. Nextstep_1(T2I) does not support cache acceleration methods such as TeaCache or Cache-DiT.
-> 2. `Tongyi-MAI/Z-Image-Turbo` is a distilled model with minimal NFEs; CFG-Parallel is not necessary.
+> 2. `Tongyi-MAI/Z-Image-Turbo` and `princepride/daVinci-MagiHuman` are distilled models with minimal NFEs; CFG-Parallel is not necessary.
 
 ### VideoGen
 

examples/offline_inference/magi_human/README.md

@@ -0,0 +1,72 @@
+# MagiHuman Generation
+
+MagiHuman is an advanced, omni-modality model that generates both high-quality video and lip-synced audio from a text prompt.
+
+Because MagiHuman is a very large model featuring a powerful DiT MoE backbone and a ~9B parameter T5Gemma text encoder, it natively supports **Tensor Parallelism (TP)** in vLLM-Omni to run efficiently across multi-GPU setups, reducing device memory bottlenecks.
+
+## Setup
+
+### Install MagiCompiler (recommended)
+
+MagiHuman relies on [MagiCompiler](https://github.com/SandAI-org/MagiCompiler) for custom-op registration used by the DiT attention kernels. While the pipeline can fall back to stub implementations, installing MagiCompiler is **strongly recommended** for correct behaviour.
+
+```bash
+# Clone the repo
+git clone https://github.com/SandAI-org/MagiCompiler.git
+cd MagiCompiler
+
+# System dependencies (optional, for FX graph visualization; Debian/Ubuntu)
+sudo apt update && sudo apt install -y graphviz
+
+# Python dependencies
+pip install -r requirements.txt
+
+# Install MagiCompiler
+pip install .          # end users (recommended)
+# pip install -e .     # developers (editable install)
+```
+
+### Hardware requirements
+
+Ensure your hardware has enough VRAM. For a standard node with 80GB GPUs, running with `--tensor-parallel-size 4` is recommended to shard both the MoE weights and the T5Gemma text encoder across 4 GPUs, reducing the per-GPU peak VRAM overhead significantly (by roughly ~13.5GB per GPU compared to single-device inference).
+
+Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) for further details on allocating memory.
+
+## Run Examples
+
+Get into the example folder:
+```bash
+cd examples/offline_inference/magi_human
+```
+
+### End-to-End Generation (Text to Video+Audio)
+
+Generate a video with synchronized speech natively generated by the model.
+
+```bash
+python end2end.py \
+    --model /proj-tango-pvc/users/zhipeng.wang/workspace/models/daVinci-MagiHuman \
+    --prompt "A young woman with long, wavy golden blonde hair..." \
+    --tensor-parallel-size 4 \
+    --output output_magihuman.mp4
+```
+
+## Common Parameters
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `--model` | *(Required)* | Local model path or HuggingFace ID |
+| `--prompt` | *(built-in demo prompt)* | Highly detailed text prompt dictating visual look and dialogue text |
+| `--tensor-parallel-size` | `4` | Tensor parallelism size (Number of GPUs) |
+| `--height` | `256` | Initial resolution height |
+| `--width` | `448` | Initial resolution width |
+| `--num-inference-steps` | `8` | Denoising steps |
+| `--seed` | `52` | Random seed |
+| `--output` | `output_magihuman.mp4` | Output video with audio path |
+
+## Example materials
+
+??? abstract "end2end.py"
+    ``````py
+    --8<-- "examples/offline_inference/magi_human/end2end.py"
+    ``````

examples/offline_inference/magi_human/end2end.py

@@ -0,0 +1,117 @@
+import argparse
+
+from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes
+from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="End-to-end inference script for MagiHuman.")
+    parser.add_argument("--model", type=str, required=True, help="Path or ID of the MagiHuman model.")
+    parser.add_argument(
+        "--prompt",
+        type=str,
+        default="",
+        help="Text prompt containing visual description, dialogue, and background sound.",
+    )
+    parser.add_argument(
+        "--tensor-parallel-size", "-tp", type=int, default=4, help="Tensor parallel size (number of GPUs)."
+    )
+    parser.add_argument(
+        "--output", type=str, default="output_magihuman.mp4", help="Path to save the generated mp4 file."
+    )
+    parser.add_argument("--height", type=int, default=256, help="Video height.")
+    parser.add_argument("--width", type=int, default=448, help="Video width.")
+    parser.add_argument("--num-inference-steps", type=int, default=8, help="Number of denoising steps.")
+    parser.add_argument("--seed", type=int, default=52, help="Random seed for generation.")
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+
+    print(f"Initializing MagiHuman pipeline with TP={args.tensor_parallel_size}...")
+    omni = Omni(
+        model=args.model,
+        init_timeout=1200,
+        tensor_parallel_size=args.tensor_parallel_size,
+        devices=list(range(args.tensor_parallel_size)),
+    )
+
+    prompt = args.prompt
+    if not prompt:
+        prompt = (
+            "A young woman with long, wavy golden blonde hair and bright blue eyes, "
+            "wearing a fitted ivory silk blouse with a delicate lace collar, sits "
+            "stationary in front of a softly lit, blurred warm-toned interior. Her "
+            "overall disposition is warm, composed, and gently confident. The camera "
+            "holds a static medium close-up, framing her from the shoulders up, "
+            "with shallow depth of field keeping her face in sharp focus. Soft "
+            "directional key light falls from the upper left, casting a gentle "
+            "highlight along her cheekbone and nose bridge. She draws a quiet breath, "
+            "the levator labii superiors relaxing as her lips part. She speaks in "
+            "clear, warm, unhurried American English: "
+            "\"The most beautiful things in life aren't things at all — "
+            "they're moments, feelings, and the people who make you feel truly alive.\" "
+            "Her jaw descends smoothly on each stressed syllable; the orbicularis oris "
+            "shapes each vowel with precision. A faint, genuine smile engages the "
+            "zygomaticus major, lifting her lip corners fractionally. Her brows rest "
+            "in a soft, neutral arch throughout. She maintains steady, forward-facing "
+            "eye contact. Head position remains level; no torso displacement occurs.\n\n"
+            "Dialogue:\n"
+            "<Young blonde woman, American English>: "
+            "\"The most beautiful things in life aren't things at all — "
+            "they're moments, feelings, and the people who make you feel truly alive.\"\n\n"
+            "Background Sound:\n"
+            "<Soft, warm indoor ambience with a faint distant piano melody>"
+        )
+
+    sampling_params = OmniDiffusionSamplingParams(
+        height=args.height,
+        width=args.width,
+        num_inference_steps=args.num_inference_steps,
+        seed=args.seed,
+        extra_args={
+            "seconds": 5,
+            "sr_height": 1080,
+            "sr_width": 1920,
+            "sr_num_inference_steps": 5,
+        },
+    )
+
+    print(f"Generating with prompt: {prompt[:80]}...")
+    outputs = omni.generate(
+        prompts=[prompt],
+        sampling_params_list=[sampling_params],
+    )
+
+    print(f"Generation complete. Output type: {type(outputs)}")
+    if outputs:
+        first = outputs[0]
+
+        if hasattr(first, "images") and first.images:
+            video_frames = first.images[0]
+            print(f"Video frames: shape={video_frames.shape}, dtype={video_frames.dtype}")
+
+            audio_waveform = None
+            if hasattr(first, "multimodal_output") and first.multimodal_output:
+                audio_waveform = first.multimodal_output.get("audio")
+                if audio_waveform is not None:
+                    print(f"Audio waveform: shape={audio_waveform.shape}, dtype={audio_waveform.dtype}")
+
+            video_bytes = mux_video_audio_bytes(
+                video_frames,
+                audio_waveform,
+                fps=25.0,
+                audio_sample_rate=44100,
+            )
+            with open(args.output, "wb") as f:
+                f.write(video_bytes)
+            print(f"Saved MP4 ({len(video_bytes)} bytes) to {args.output}")
+        print("SUCCESS: MagiHuman pipeline generation completed.")
+    else:
+        print("WARNING: No outputs returned.")
+
+
+if __name__ == "__main__":
+    main()

requirements/common.txt

@@ -1,4 +1,5 @@
 # Common dependencies for all platforms
+av>=14.0.0
 omegaconf>=2.3.0
 librosa>=0.11.0
 resampy>=0.4.3

tests/e2e/offline_inference/test_magi_human.py

@@ -0,0 +1,129 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""End-to-end tests for MagiHuman pipeline via vLLM-Omni."""
+
+import io
+
+import av
+import numpy as np
+import pytest
+
+from tests.utils import hardware_test
+from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes
+from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+
+def _validate_mp4(video_bytes: bytes, min_frames: int = 10) -> None:
+    """Validate that the MP4 contains meaningful video and audio tracks."""
+    container = av.open(io.BytesIO(video_bytes))
+
+    v_streams = [s for s in container.streams if s.type == "video"]
+    assert len(v_streams) >= 1, "No video stream found in MP4"
+
+    a_streams = [s for s in container.streams if s.type == "audio"]
+    assert len(a_streams) >= 1, "No audio stream found in MP4"
+
+    v_stream = v_streams[0]
+    assert v_stream.width >= 1080, f"Unexpected video width: {v_stream.width}"
+    assert v_stream.height >= 1056, f"Unexpected video height: {v_stream.height}"
+
+    frame_count = 0
+    for frame in container.decode(video=0):
+        frame_count += 1
+        if frame_count >= min_frames:
+            break
+    assert frame_count >= min_frames, f"Video has only {frame_count} frames (expected >= {min_frames})"
+
+    container.close()
+
+
+@pytest.mark.core_model
+@pytest.mark.advanced_model
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "H100"}, num_cards=2)
+def test_magi_human_e2e(run_level):
+    """End-to-end test for MagiHuman generating video and audio."""
+    if run_level != "advanced_model":
+        pytest.skip("MagiHuman e2e test requires advanced_model run level with real weights.")
+
+    model_path = "princepride/daVinci-MagiHuman"
+
+    omni = Omni(
+        model=model_path,
+        init_timeout=1200,
+        tensor_parallel_size=2,
+    )
+
+    prompt = (
+        "A young woman with long, wavy golden blonde hair and bright blue eyes, "
+        "wearing a fitted ivory silk blouse with a delicate lace collar, sits "
+        "stationary in front of a softly lit, blurred warm-toned interior. Her "
+        "overall disposition is warm, composed, and gently confident. The camera "
+        "holds a static medium close-up, framing her from the shoulders up, "
+        "with shallow depth of field keeping her face in sharp focus. Soft "
+        "directional key light falls from the upper left, casting a gentle "
+        "highlight along her cheekbone and nose bridge. She draws a quiet breath, "
+        "the levator labii superiors relaxing as her lips part. She speaks in "
+        "clear, warm, unhurried American English: "
+        "\"The most beautiful things in life aren't things at all — "
+        "they're moments, feelings, and the people who make you feel truly alive.\" "
+        "Her jaw descends smoothly on each stressed syllable; the orbicularis oris "
+        "shapes each vowel with precision. A faint, genuine smile engages the "
+        "zygomaticus major, lifting her lip corners fractionally. Her brows rest "
+        "in a soft, neutral arch throughout. She maintains steady, forward-facing "
+        "eye contact. Head position remains level; no torso displacement occurs.\n\n"
+        "Dialogue:\n"
+        "<Young blonde woman, American English>: "
+        "\"The most beautiful things in life aren't things at all — "
+        "they're moments, feelings, and the people who make you feel truly alive.\"\n\n"
+        "Background Sound:\n"
+        "<Soft, warm indoor ambience with a faint distant piano melody>"
+    )
+
+    sampling_params = OmniDiffusionSamplingParams(
+        height=256,
+        width=448,
+        num_inference_steps=8,
+        seed=52,
+        extra_args={
+            "seconds": 5,
+            "sr_height": 1080,
+            "sr_width": 1920,
+            "sr_num_inference_steps": 5,
+        },
+    )
+
+    try:
+        outputs = list(
+            omni.generate(
+                prompts=[prompt],
+                sampling_params_list=[sampling_params],
+            )
+        )
+
+        assert len(outputs) > 0, "No outputs returned"
+        first = outputs[0]
+
+        assert hasattr(first, "images") and first.images, "No video frames in output"
+        video_frames = first.images[0]
+        assert isinstance(video_frames, np.ndarray), f"Expected numpy array, got {type(video_frames)}"
+        assert video_frames.ndim == 4, f"Expected 4D array (T,H,W,3), got shape {video_frames.shape}"
+
+        audio_waveform = None
+        if hasattr(first, "multimodal_output") and first.multimodal_output:
+            audio_waveform = first.multimodal_output.get("audio")
+        assert audio_waveform is not None, "No audio waveform in multimodal_output"
+
+        video_bytes = mux_video_audio_bytes(
+            video_frames,
+            audio_waveform,
+            fps=25.0,
+            audio_sample_rate=44100,
+        )
+        assert isinstance(video_bytes, bytes), f"Expected MP4 bytes, got {type(video_bytes)}"
+        assert len(video_bytes) > 1000, f"MP4 too small ({len(video_bytes)} bytes)"
+
+        _validate_mp4(video_bytes)
+    finally:
+        omni.close()

vllm_omni/diffusion/models/magi_human/__init__.py

@@ -0,0 +1 @@
+# SPDX-License-Identifier: Apache-2.0

vllm_omni/diffusion/models/t5_encoder/t5_gemma_encoder.py

@@ -0,0 +1,309 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from __future__ import annotations
+
+from collections.abc import Iterable
+
+import torch
+import torch.nn as nn
+from transformers import PretrainedConfig
+from vllm.config import VllmConfig
+from vllm.distributed import get_tensor_model_parallel_world_size
+from vllm.model_executor.layers.activation import get_act_fn
+from vllm.model_executor.layers.linear import (
+    MergedColumnParallelLinear,
+    QKVParallelLinear,
+    RowParallelLinear,
+)
+from vllm.model_executor.layers.rotary_embedding import get_rope
+from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+
+
+class T5GemmaRMSNorm(nn.Module):
+    def __init__(self, hidden_size: int, eps: float = 1e-6):
+        super().__init__()
+        # Normal RMSNorm but T5Gemma requires (1 + weight)
+        self.weight = nn.Parameter(torch.zeros(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return (hidden_states * (1.0 + self.weight.float())).to(input_dtype)
+
+
+class T5GemmaMLP(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        intermediate_size: int,
+        hidden_act: str,
+    ) -> None:
+        super().__init__()
+        self.gate_up_proj = MergedColumnParallelLinear(
+            input_size=hidden_size,
+            output_sizes=[intermediate_size, intermediate_size],
+            bias=False,
+            gather_output=False,
+        )
+        self.down_proj = RowParallelLinear(
+            input_size=intermediate_size,
+            output_size=hidden_size,
+            bias=False,
+            input_is_parallel=True,
+        )
+        self.act_fn = get_act_fn(hidden_act)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        gate_up, _ = self.gate_up_proj(x)
+        gate, up = gate_up.chunk(2, dim=-1)
+        x = self.act_fn(gate) * up
+        x, _ = self.down_proj(x)
+        return x
+
+
+class T5GemmaAttention(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        num_heads: int,
+        num_kv_heads: int,
+        head_dim: int,
+        max_position_embeddings: int,
+        rope_theta: float,
+        cache_config: VllmConfig | None = None,
+        quant_config: dict | None = None,
+    ) -> None:
+        super().__init__()
+        self.hidden_size = hidden_size
+        tp_size = get_tensor_model_parallel_world_size()
+        self.total_num_heads = num_heads
+        assert self.total_num_heads % tp_size == 0
+        self.num_heads = self.total_num_heads // tp_size
+        self.total_num_kv_heads = num_kv_heads
+        if self.total_num_kv_heads >= tp_size:
+            assert self.total_num_kv_heads % tp_size == 0
+        else:
+            assert tp_size % self.total_num_kv_heads == 0
+        self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
+        self.head_dim = head_dim
+        self.q_size = self.num_heads * self.head_dim
+        self.kv_size = self.num_kv_heads * self.head_dim
+
+        self.qkv_proj = QKVParallelLinear(
+            hidden_size=hidden_size,
+            head_size=self.head_dim,
+            total_num_heads=self.total_num_heads,
+            total_num_kv_heads=self.total_num_kv_heads,
+            bias=False,
+        )
+        self.o_proj = RowParallelLinear(
+            input_size=self.total_num_heads * self.head_dim,
+            output_size=hidden_size,
+            bias=False,
+            input_is_parallel=True,
+        )
+
+        self.rotary_emb = get_rope(
+            self.head_dim,
+            max_position=max_position_embeddings,
+            is_neox_style=True,
+            rope_parameters={"base": rope_theta},
+        )
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        qkv, _ = self.qkv_proj(hidden_states)
+        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+
+        q, k = self.rotary_emb(positions, q, k)
+
+        # Scale Q appropriately. T5Gemma uses query_pre_attn_scalar=256 => 256**-0.5 = 1/16
+        # The standard scaling is head_dim**-0.5. For T5Gemma, head_dim=256.
+        # So we don't need to manually scale if F.scaled_dot_product_attention scales by head_dim.
+        # But we must reshape.
+        batch_size, seq_len, _ = hidden_states.shape
+        q = q.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
+        k = k.view(batch_size, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
+        v = v.view(batch_size, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
+
+        # GQA repeat KV
+        if self.num_kv_heads != self.num_heads:
+            num_repeat = self.num_heads // self.num_kv_heads
+            k = k.repeat_interleave(num_repeat, dim=1)
+            v = v.repeat_interleave(num_repeat, dim=1)
+
+        attn_output = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=attention_mask, dropout_p=0.0)
+        attn_output = attn_output.transpose(1, 2).contiguous()
+        attn_output = attn_output.view(batch_size, seq_len, self.q_size)
+
+        output, _ = self.o_proj(attn_output)
+        return output
+
+
+class T5GemmaEncoderLayer(nn.Module):
+    def __init__(self, config: PretrainedConfig) -> None:
+        super().__init__()
+        self.self_attn = T5GemmaAttention(
+            hidden_size=config.hidden_size,
+            num_heads=config.num_attention_heads,
+            num_kv_heads=config.num_key_value_heads,
+            head_dim=config.head_dim,
+            max_position_embeddings=config.max_position_embeddings,
+            rope_theta=config.rope_theta,
+        )
+        self.mlp = T5GemmaMLP(
+            hidden_size=config.hidden_size,
+            intermediate_size=config.intermediate_size,
+            hidden_act=config.hidden_activation,
+        )
+        self.pre_self_attn_layernorm = T5GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_self_attn_layernorm = T5GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.pre_feedforward_layernorm = T5GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_feedforward_layernorm = T5GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None,
+    ) -> torch.Tensor:
+        # Self Attention
+        residual = hidden_states
+        hidden_states = self.pre_self_attn_layernorm(hidden_states)
+        hidden_states = self.self_attn(
+            positions=positions,
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+        )
+        hidden_states = self.post_self_attn_layernorm(hidden_states)
+        hidden_states = residual + hidden_states
+
+        # Fully Connected
+        residual = hidden_states
+        hidden_states = self.pre_feedforward_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = self.post_feedforward_layernorm(hidden_states)
+        hidden_states = residual + hidden_states
+        return hidden_states
+
+
+class T5GemmaEncoderModelTP(nn.Module):
+    def __init__(self, config: PretrainedConfig) -> None:
+        super().__init__()
+        self.config = config
+        self.vocab_size = config.vocab_size
+
+        self.embed_tokens = VocabParallelEmbedding(
+            config.vocab_size,
+            config.hidden_size,
+        )
+
+        self.layers = nn.ModuleList([T5GemmaEncoderLayer(config) for _ in range(config.num_hidden_layers)])
+        self.norm = T5GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+    @property
+    def dtype(self) -> torch.dtype:
+        return next(self.parameters()).dtype
+
+    @property
+    def device(self) -> torch.device:
+        return next(self.parameters()).device
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        hidden_states = self.embed_tokens(input_ids)
+
+        # Scaling inputs
+        normalizer = torch.tensor(self.config.hidden_size**0.5, dtype=hidden_states.dtype, device=hidden_states.device)
+        hidden_states = hidden_states * normalizer
+
+        # Simple position ids for RoPE
+        batch_size, seq_len = input_ids.shape
+        positions = torch.arange(seq_len, device=input_ids.device, dtype=torch.long).unsqueeze(0).expand(batch_size, -1)
+
+        # Build attention mask: (batch, seq) -> (batch, 1, 1, seq)
+        # Assuming typical bidirectional causal mask handling in HF: T5Gemma uses non-causal encoder.
+        if attention_mask is not None:
+            # HuggingFace expects boolean mask for scaled_dot_product_attention
+            # or additive mask (0 and -inf). Let's use boolean matching FA patterns.
+            # SDPA expects attention_mask to be boolean (True = keep, False = masking)
+            bool_mask = attention_mask.to(torch.bool)
+            extended_mask = bool_mask.unsqueeze(1).unsqueeze(2)  # (B, 1, 1, S)
+        else:
+            extended_mask = None
+
+        for idx, layer in enumerate(self.layers):
+            # T5Gemma has layer_types switching between "sliding_attention" and "full_attention"
+            # However, for text encoder inference, the sequences are typically < max sequence length
+            # and local sliding window only affects very long contexts. For simplicity we use full.
+            hidden_states = layer(
+                positions=positions,
+                hidden_states=hidden_states,
+                attention_mask=extended_mask,
+            )
+
+        hidden_states = self.norm(hidden_states)
+        return hidden_states
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        stacked_params_mapping = [
+            # (param_name, shard_name, shard_id)
+            ("qkv_proj", "q_proj", "q"),
+            ("qkv_proj", "k_proj", "k"),
+            ("qkv_proj", "v_proj", "v"),
+            ("gate_up_proj", "gate_proj", 0),
+            ("gate_up_proj", "up_proj", 1),
+        ]
+
+        params_dict = dict(self.named_parameters())
+        loaded_params: set[str] = set()
+
+        for name, loaded_weight in weights:
+            # HF checkpoint keys may carry a "model." prefix (e.g.
+            # "model.encoder.layers.0...").  Strip it so the rest of the
+            # logic only needs to handle the "encoder.*" namespace.
+            if name.startswith("model."):
+                name = name[len("model.") :]
+
+            if not name.startswith("encoder."):
+                continue
+
+            # Strip "encoder." prefix as this model only wraps the encoder
+            name = name[len("encoder.") :]
+
+            # Map self_attn to self_attn and correct normalization names
+            # HF: layers.0.pre_self_attn_layernorm.weight -> Ours: layers.0.pre_self_attn_layernorm.weight
+
+            lookup_name = name
+            for param_name, weight_name, shard_id in stacked_params_mapping:
+                if f".{weight_name}." not in name:
+                    continue
+                lookup_name = name.replace(f".{weight_name}.", f".{param_name}.")
+                if lookup_name not in params_dict:
+                    continue
+                param = params_dict[lookup_name]
+                weight_loader = param.weight_loader
+                weight_loader(param, loaded_weight, shard_id)
+                break
+            else:
+                if name not in params_dict:
+                    continue
+                param = params_dict[name]
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                weight_loader(param, loaded_weight)
+
+            loaded_params.add("encoder." + name)
+            loaded_params.add("encoder." + lookup_name)
+
+        return loaded_params

vllm_omni/diffusion/offloader/module_collector.py

@@ -21,9 +21,9 @@ class PipelineModules:
 class ModuleDiscovery:
     """Discovers pipeline components for offloading"""
 
-    DIT_ATTRS = ["transformer", "transformer_2", "dit", "language_model", "transformer_blocks"]
+    DIT_ATTRS = ["transformer", "transformer_2", "dit", "sr_dit", "language_model", "transformer_blocks"]
     ENCODER_ATTRS = ["text_encoder", "text_encoder_2", "text_encoder_3", "image_encoder"]
-    VAE_ATTRS = ["vae"]
+    VAE_ATTRS = ["vae", "audio_vae"]
 
     @staticmethod
     def discover(pipeline: nn.Module) -> PipelineModules:

vllm_omni/diffusion/registry.py

@@ -173,6 +173,11 @@ _DIFFUSION_MODELS = {
         "pipeline_hunyuan_video_1_5_i2v",
         "HunyuanVideo15I2VPipeline",
     ),
+    "MagiHumanPipeline": (
+        "magi_human",
+        "pipeline_magi_human",
+        "MagiHumanPipeline",
+    ),
     "OmniVoicePipeline": (
         "omnivoice",
         "pipeline_omnivoice",
@@ -368,6 +373,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = {
     "Flux2Pipeline": "get_flux2_post_process_func",
     "HunyuanVideo15Pipeline": "get_hunyuan_video_15_post_process_func",
     "HunyuanVideo15ImageToVideoPipeline": "get_hunyuan_video_15_i2v_post_process_func",
+    "MagiHumanPipeline": "get_magi_human_post_process_func",
     "OmniVoicePipeline": "get_omnivoice_post_process_func",
 }
 
@@ -387,6 +393,7 @@ _DIFFUSION_PRE_PROCESS_FUNCS = {
     "HeliosPipeline": "get_helios_pre_process_func",
     "HeliosPyramidPipeline": "get_helios_pre_process_func",
     "HunyuanVideo15ImageToVideoPipeline": "get_hunyuan_video_15_i2v_pre_process_func",
+    "MagiHumanPipeline": "get_magi_human_pre_process_func",
 }
 
 

vllm_omni/diffusion/utils/media_utils.py

@@ -0,0 +1,75 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Video/audio muxing utilities using PyAV (no ffmpeg binary dependency)."""
+
+from __future__ import annotations
+
+import io
+from fractions import Fraction
+
+import av
+import numpy as np
+
+
+def mux_video_audio_bytes(
+    video_frames: np.ndarray,
+    audio_waveform: np.ndarray | None = None,
+    *,
+    fps: float = 25.0,
+    audio_sample_rate: int = 44100,
+    video_codec: str = "h264",
+    audio_codec: str = "aac",
+    crf: str = "18",
+) -> bytes:
+    """Mux video frames and optional audio waveform into MP4 bytes.
+
+    Args:
+        video_frames: uint8 array of shape ``(T, H, W, 3)`` (RGB).
+        audio_waveform: float32 array – mono ``(N,)`` or ``(N, C)`` / ``(C, N)``.
+        fps: Video frame rate.
+        audio_sample_rate: Audio sample rate in Hz.
+        video_codec: Video codec name.
+        audio_codec: Audio codec name.
+        crf: Constant rate factor for the video encoder.
+
+    Returns:
+        Raw MP4 bytes ready to be written to disk or streamed.
+    """
+    buf = io.BytesIO()
+    container = av.open(buf, mode="w", format="mp4")
+
+    v_stream = container.add_stream(video_codec, rate=Fraction(fps).limit_denominator(10000))
+    v_stream.width = video_frames.shape[2]
+    v_stream.height = video_frames.shape[1]
+    v_stream.pix_fmt = "yuv420p"
+    v_stream.options = {"crf": crf}
+
+    a_stream = None
+    if audio_waveform is not None:
+        samples = audio_waveform.astype(np.float32)
+        if samples.ndim == 1:
+            samples = samples.reshape(1, -1)
+        elif samples.ndim == 2 and samples.shape[0] > samples.shape[1]:
+            samples = samples.T
+        num_channels = samples.shape[0]
+        layout = "stereo" if num_channels >= 2 else "mono"
+        a_stream = container.add_stream(audio_codec, rate=audio_sample_rate)
+        a_stream.layout = layout
+
+    for frame_data in video_frames:
+        frame = av.VideoFrame.from_ndarray(frame_data, format="rgb24")
+        for packet in v_stream.encode(frame):
+            container.mux(packet)
+    for packet in v_stream.encode():
+        container.mux(packet)
+
+    if a_stream is not None and audio_waveform is not None:
+        audio_frame = av.AudioFrame.from_ndarray(samples, format="fltp", layout=layout)
+        audio_frame.sample_rate = audio_sample_rate
+        for packet in a_stream.encode(audio_frame):
+            container.mux(packet)
+        for packet in a_stream.encode():
+            container.mux(packet)
+
+    container.close()
+    return buf.getvalue()