diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index d611c0311..f3d22aa76 100644 --- a/docs/models/supported_models.md +++ b/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` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md index 9cd407d37..e7f33306e 100644 --- a/docs/user_guide/diffusion_features.md +++ b/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 diff --git a/examples/offline_inference/magi_human/README.md b/examples/offline_inference/magi_human/README.md new file mode 100644 index 000000000..2b89093d9 --- /dev/null +++ b/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" + `````` diff --git a/examples/offline_inference/magi_human/end2end.py b/examples/offline_inference/magi_human/end2end.py new file mode 100644 index 000000000..39451ccc4 --- /dev/null +++ b/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" + ": " + "\"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" + "" + ) + + 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() diff --git a/requirements/common.txt b/requirements/common.txt index 138a61ed2..89eaac32b 100644 --- a/requirements/common.txt +++ b/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 diff --git a/tests/e2e/offline_inference/test_magi_human.py b/tests/e2e/offline_inference/test_magi_human.py new file mode 100644 index 000000000..6211fdafc --- /dev/null +++ b/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" + ": " + "\"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" + "" + ) + + 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() diff --git a/vllm_omni/diffusion/models/magi_human/__init__.py b/vllm_omni/diffusion/models/magi_human/__init__.py new file mode 100644 index 000000000..988131360 --- /dev/null +++ b/vllm_omni/diffusion/models/magi_human/__init__.py @@ -0,0 +1 @@ +# SPDX-License-Identifier: Apache-2.0 diff --git a/vllm_omni/diffusion/models/magi_human/magi_human_dit.py b/vllm_omni/diffusion/models/magi_human/magi_human_dit.py new file mode 100644 index 000000000..491b1b3c4 --- /dev/null +++ b/vllm_omni/diffusion/models/magi_human/magi_human_dit.py @@ -0,0 +1,1624 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 SandAI. All Rights Reserved. +# Ported from daVinci-MagiHuman inference/model/dit/dit_module.py +# Adaptations: removed Ulysses context-parallelism, inlined Modality/VarlenHandler. + +from __future__ import annotations + +import importlib +from collections.abc import Callable +from dataclasses import dataclass, field +from enum import Enum, IntEnum +from typing import TYPE_CHECKING, Any, Literal + +import torch +import torch.nn as nn +from einops import rearrange, repeat +from torch.nn import Parameter +from torch.nn import functional as F +from vllm.distributed import ( + get_tensor_model_parallel_world_size, +) +from vllm.model_executor.layers.linear import ( + ColumnParallelLinear, + QKVParallelLinear, + RowParallelLinear, +) +from vllm.vllm_flash_attn import flash_attn_varlen_func as _vllm_fa_varlen + +try: + from magi_compiler.api import magi_register_custom_op + from magi_compiler.config import CompileConfig +except Exception: + + class CompileConfig: # type: ignore[no-redef] + pass + + def magi_register_custom_op(*args, **kwargs): # type: ignore[no-redef] + def decorator(func): + return func + + return decorator + + +def magi_compile(*args, **kwargs): + """No-op stub — vllm-omni handles execution; magi compilation is skipped.""" + + def decorator(cls_or_fn): + return cls_or_fn + + return decorator + + +# --------------------------------------------------------------------------- +# Inlined from inference/common/sequence_schema.py +# --------------------------------------------------------------------------- +class Modality(IntEnum): + VIDEO = 0 + AUDIO = 1 + TEXT = 2 + + +@dataclass +class VarlenHandler: + cu_seqlens_q: torch.Tensor + cu_seqlens_k: torch.Tensor + max_seqlen_q: int + max_seqlen_k: int + + +def _is_hopper_arch() -> bool: + if not torch.cuda.is_available(): + return False + return torch.cuda.get_device_capability()[0] == 9 + + +# --------------------------------------------------------------------------- +# FFA handler for local / flex attention +# --------------------------------------------------------------------------- +@dataclass +class FFAHandler: + q_ranges: torch.Tensor + k_ranges: torch.Tensor + max_seqlen_q: int + max_seqlen_k: int + attn_type_map: torch.Tensor + softmax_scale: float + + +# --------------------------------------------------------------------------- +# Activation helpers +# --------------------------------------------------------------------------- +class MLPActivationType(Enum): + SWIGLU7 = "swiglu7" + GELU7 = "gelu7" + + +def swiglu7(x, alpha: float = 1.702, limit: float = 7.0, out_dtype: torch.dtype | None = None): + out_dtype = x.dtype if out_dtype is None else out_dtype + x = x.to(torch.float32) + x_glu, x_linear = x[..., ::2], x[..., 1::2] + x_glu = x_glu.clamp(min=None, max=limit) + x_linear = x_linear.clamp(min=-limit, max=limit) + out_glu = x_glu * torch.sigmoid(alpha * x_glu) + return (out_glu * (x_linear + 1)).to(out_dtype) + + +def gelu7(x, alpha: float = 1.702, limit: float = 7.0, out_dtype: torch.dtype | None = None): + out_dtype = x.dtype if out_dtype is None else out_dtype + x = x.to(torch.float32) + x_glu = x.clamp(min=None, max=limit) + out_glu = x_glu * torch.sigmoid(alpha * x_glu) + return out_glu.to(out_dtype) + + +def create_activation_func(activation_type: MLPActivationType) -> Callable: + match activation_type: + case MLPActivationType.SWIGLU7: + return swiglu7 + case MLPActivationType.GELU7: + return gelu7 + case _: + raise ValueError(f"Unknown activation type: {activation_type}") + + +# --------------------------------------------------------------------------- +# Modality dispatcher (permutation helper) +# --------------------------------------------------------------------------- +class ModalityDispatcher: + permuted_modality_mapping: torch.Tensor + group_size: torch.Tensor + group_size_cpu: list[int] + num_modalities: int + + def __init__(self, modality_mapping: torch.Tensor, num_modalities: int): + self.modality_mapping = modality_mapping + self.num_modalities = num_modalities + self.permuted_modality_mapping = self._precompute_permute_mapping(modality_mapping) + self.group_size = torch.bincount(self.permuted_modality_mapping, minlength=num_modalities).to(torch.int32) + self.group_size_cpu: list[int] = [int(x) for x in self.group_size.to("cpu").tolist()] + + def _precompute_permute_mapping(self, modality_mapping): + self.permute_mapping = torch.argsort(modality_mapping) + self.inv_permute_mapping = torch.argsort(self.permute_mapping) + return modality_mapping[self.permute_mapping] + + def dispatch(self, x: torch.Tensor) -> list[torch.Tensor]: + return list(torch.split(x, self.group_size_cpu, dim=0)) + + def undispatch(self, *processed_groups: list[torch.Tensor]) -> torch.Tensor: + return torch.cat(processed_groups, dim=0) + + @staticmethod + def permute(x: torch.Tensor, permute_mapping: torch.Tensor) -> torch.Tensor: + return x[permute_mapping] + + @staticmethod + def inv_permute(x: torch.Tensor, inv_permute_mapping: torch.Tensor) -> torch.Tensor: + return x[inv_permute_mapping] + + +# --------------------------------------------------------------------------- +# Positional / rotary embedding helpers +# --------------------------------------------------------------------------- +def freq_bands( + num_bands: int, temperature: float = 10000.0, step: int = 2, device: torch.device | None = None +) -> torch.Tensor: + exp = torch.arange(0, num_bands, step, dtype=torch.int64, device=device).to(torch.float32) / num_bands + return 1.0 / (temperature**exp) + + +def rotate_half(x, interleaved=False): + if not interleaved: + x1, x2 = x.chunk(2, dim=-1) + return torch.cat((-x2, x1), dim=-1) + else: + x1, x2 = x[..., ::2], x[..., 1::2] + return rearrange(torch.stack((-x2, x1), dim=-1), "... d two -> ... (d two)", two=2) + + +def apply_rotary_emb_torch(x, cos, sin, interleaved=False): + ro_dim = cos.shape[-1] * 2 + assert ro_dim <= x.shape[-1] + cos = repeat(cos, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)") + sin = repeat(sin, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)") + return torch.cat([x[..., :ro_dim] * cos + rotate_half(x[..., :ro_dim], interleaved) * sin, x[..., ro_dim:]], dim=-1) + + +# --------------------------------------------------------------------------- +# Fourier positional embedding +# --------------------------------------------------------------------------- +class ElementWiseFourierEmbed(nn.Module): + def __init__( + self, + dim: int, + max_res: int = 224, + temperature: float = 10000.0, + in_pixels: bool = True, + linear_bands: bool = False, + learnable: bool = False, + device: torch.device = torch.device("cpu"), + dtype: torch.dtype = torch.float32, + ): + super().__init__() + self.dim = dim + self.in_pixels = in_pixels + self.learnable = learnable + self.temperature = temperature + self.max_res = max_res + self.linear_bands = linear_bands + self.device = device + self.dtype = dtype + bands = self.get_default_bands() + self.bands = nn.Parameter(bands, requires_grad=self.learnable) + + def forward(self, coords: torch.Tensor) -> torch.Tensor: + coords_xyz = coords[:, :3] + sizes = coords[:, 3:6] + refs = coords[:, 6:9] + + scales = (refs - 1) / (sizes - 1) + scales[(refs == 1) & (sizes == 1)] = 1 + assert not scales.isnan().any(), "scales has nan" + assert not scales.isinf().any(), "scales has inf" + + centers = (sizes - 1) / 2 + centers[:, 0] = 0 + coords_xyz = coords_xyz - centers + + bands = self.bands.to(coords.device, coords.dtype) + proj = coords_xyz.unsqueeze(-1) * scales.unsqueeze(-1) * bands + sin_proj = proj.sin() + cos_proj = proj.cos() + return torch.cat((sin_proj, cos_proj), dim=1).flatten(1) + + def reset_parameters(self): + self.bands.copy_(self.get_default_bands()) + + def get_default_bands(self): + if self.in_pixels: + raise NotImplementedError("in_pixels are not implemented yet") + return freq_bands(self.dim // 8, temperature=self.temperature, step=1, device=self.device).to(self.dtype) + + +# --------------------------------------------------------------------------- +# Multi-modality RMSNorm +# --------------------------------------------------------------------------- +class MultiModalityRMSNorm(nn.Module): + __constants__ = ["dim", "eps", "num_modality"] + + def __init__(self, dim: int, eps: float = 1e-6, device: torch.device | None = None, num_modality: int = 1): + super().__init__() + self.dim = dim + self.eps = eps + self.num_modality = num_modality + self.weight = nn.Parameter(torch.zeros(dim * num_modality, device=device, dtype=torch.float32)) + if num_modality > 1: + self.forward = self.forward_multi_experts + else: + self.forward = self.forward_single_expert + self.reset_parameters() + + def reset_parameters(self): + nn.init.zeros_(self.weight) + + def rms(self, x: torch.Tensor) -> torch.Tensor: + t = x.float() + return t * torch.rsqrt(torch.mean(t**2, dim=-1, keepdim=True) + self.eps) + + def forward_multi_experts(self, x: torch.Tensor, modality_dispatcher: ModalityDispatcher) -> torch.Tensor: + original_dtype = x.dtype + t = self.rms(x) + weight_chunked = self.weight.chunk(self.num_modality, dim=0) + t_list = modality_dispatcher.dispatch(t) + for i in range(self.num_modality): + t_list[i] = t_list[i] * (weight_chunked[i] + 1) + t = modality_dispatcher.undispatch(*t_list) + return t.to(original_dtype) + + def forward_single_expert( + self, x: torch.Tensor, modality_dispatcher: ModalityDispatcher | None = None + ) -> torch.Tensor: + t, original_dtype = x.float(), x.dtype + t = t * torch.rsqrt(torch.mean(t**2, dim=-1, keepdim=True) + self.eps) + return (t * (self.weight + 1)).to(original_dtype) + + +# --------------------------------------------------------------------------- +# Linear layers with bf16 compute and MoE dispatch +# --------------------------------------------------------------------------- +class _BF16ComputeLinear(torch.autograd.Function): + @staticmethod + def forward( + ctx, + input: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + output_dtype: torch.dtype | None, + compute_dtype: torch.dtype = torch.bfloat16, + ): + input_cast = input.to(compute_dtype) + weight_cast = weight.to(compute_dtype) + output = torch.matmul(input_cast, weight_cast.t()) + if bias is not None: + output = output + bias.to(compute_dtype) + return output.to(output_dtype) + + +class BaseLinear(nn.Module): + __constants__ = ["in_features", "out_features", "num_layers", "num_experts"] + + def __init__( + self, in_features, out_features, num_layers_for_initialization, num_experts, bias=True, device=None, dtype=None + ): + super().__init__() + factory_kwargs = {"device": device, "dtype": torch.bfloat16} + self.in_features = in_features + self.out_features = out_features + self.num_layers_for_initialization = num_layers_for_initialization + self.num_experts = num_experts + self.use_bias = bias + self.weight = Parameter(torch.empty((out_features * num_experts, in_features), **factory_kwargs)) + if bias: + self.bias = Parameter(torch.empty(out_features * num_experts, **factory_kwargs)) + else: + self.register_parameter("bias", None) + + def forward( + self, + input: torch.Tensor, + output_dtype: torch.dtype | None = None, + modality_dispatcher: ModalityDispatcher | None = None, + ) -> torch.Tensor: + output_dtype = input.dtype if output_dtype is None else output_dtype + return _BF16ComputeLinear.apply(input, self.weight, self.bias, output_dtype, torch.bfloat16) + + +class NativeMoELinear(BaseLinear): + def forward( + self, + input: torch.Tensor, + output_dtype: torch.dtype | None = None, + modality_dispatcher: ModalityDispatcher | None = None, + ) -> torch.Tensor: + output_dtype = input.dtype if output_dtype is None else output_dtype + input_list = modality_dispatcher.dispatch(input) # type: ignore + weight_chunked = self.weight.chunk(self.num_experts, dim=0) + if self.bias is not None: + bias_chunked = self.bias.chunk(self.num_experts, dim=0) + for i in range(self.num_experts): + input_list[i] = _BF16ComputeLinear.apply( + input_list[i], + weight_chunked[i], + bias_chunked[i] if self.bias is not None else None, + output_dtype, + torch.bfloat16, + ) + return modality_dispatcher.undispatch(*input_list) # type: ignore + + +def create_linear( + in_features, out_features, num_layers=1, num_experts=1, bias=True, device=None, dtype=None +) -> BaseLinear | NativeMoELinear: + if num_experts == 1: + return BaseLinear(in_features, out_features, num_layers, num_experts, bias, device, dtype) + else: + return NativeMoELinear(in_features, out_features, num_layers, num_experts, bias, device, dtype) + + +# --------------------------------------------------------------------------- +# MoE TP parallel linear wrappers: per-expert vLLM parallel layers +# --------------------------------------------------------------------------- +class MoEQKVParallelLinear(nn.Module): + """Per-expert QKVParallelLinear with modality dispatch. + + Wraps ``num_experts`` independent QKVParallelLinear instances. + Forward: dispatch tokens by modality → per-expert QKV matmul (TP-sharded) + → undispatch. + """ + + def __init__( + self, + hidden_size: int, + head_size: int, + total_num_heads: int, + total_num_kv_heads: int, + num_experts: int, + bias: bool = False, + ): + super().__init__() + self.num_experts = num_experts + self.experts = nn.ModuleList( + [ + QKVParallelLinear( + hidden_size=hidden_size, + head_size=head_size, + total_num_heads=total_num_heads, + total_num_kv_heads=total_num_kv_heads, + bias=bias, + return_bias=False, + ) + for _ in range(num_experts) + ] + ) + # Expose per-rank head info from the first expert (all are identical). + self.num_heads = self.experts[0].num_heads + self.num_kv_heads = self.experts[0].num_kv_heads + self.head_size = head_size + + def forward( + self, + x: torch.Tensor, + modality_dispatcher: ModalityDispatcher, + ) -> torch.Tensor: + x_list = modality_dispatcher.dispatch(x) + out_list: list[torch.Tensor] = [] + for i in range(self.num_experts): + out = self.experts[i](x_list[i]) + out_list.append(out) + return modality_dispatcher.undispatch(*out_list) + + +class MoEColumnParallelLinear(nn.Module): + """Per-expert ColumnParallelLinear with modality dispatch. + + Forward: dispatch → per-expert column-parallel matmul → undispatch. + Output stays TP-local (no gather). + """ + + def __init__( + self, + input_size: int, + output_size: int, + num_experts: int, + bias: bool = False, + ): + super().__init__() + self.num_experts = num_experts + self.experts = nn.ModuleList( + [ + ColumnParallelLinear( + input_size=input_size, + output_size=output_size, + bias=bias, + gather_output=False, + return_bias=False, + ) + for _ in range(num_experts) + ] + ) + + def forward( + self, + x: torch.Tensor, + modality_dispatcher: ModalityDispatcher, + ) -> torch.Tensor: + x_list = modality_dispatcher.dispatch(x) + out_list: list[torch.Tensor] = [] + for i in range(self.num_experts): + out = self.experts[i](x_list[i]) + out_list.append(out) + return modality_dispatcher.undispatch(*out_list) + + +class MoERowParallelLinear(nn.Module): + """Per-expert RowParallelLinear with modality dispatch. + + Forward: dispatch → per-expert row-parallel matmul (includes all-reduce) + → undispatch. + """ + + def __init__( + self, + input_size: int, + output_size: int, + num_experts: int, + bias: bool = False, + ): + super().__init__() + self.num_experts = num_experts + self.experts = nn.ModuleList( + [ + RowParallelLinear( + input_size=input_size, + output_size=output_size, + bias=bias, + input_is_parallel=True, + return_bias=False, + ) + for _ in range(num_experts) + ] + ) + + def forward( + self, + x: torch.Tensor, + modality_dispatcher: ModalityDispatcher, + ) -> torch.Tensor: + x_list = modality_dispatcher.dispatch(x) + out_list: list[torch.Tensor] = [] + for i in range(self.num_experts): + out = self.experts[i](x_list[i]) + out_list.append(out) + return modality_dispatcher.undispatch(*out_list) + + +def validate_magi_human_tp_constraints( + *, + hidden_size: int, + num_heads_q: int, + num_heads_kv: int, + tensor_parallel_size: int, +) -> None: + """Validate MagiHuman TP divisibility constraints. + + Both shared layers (num_modality == 1) and MoE layers (num_modality == 3) + support TP via vLLM's parallel linear layers (QKVParallelLinear / + ColumnParallelLinear / RowParallelLinear). MoE layers use per-expert + parallel layers with modality dispatch. + + Supported tp_sizes given default config (hidden=5120, heads_q=40, kv=8): 1, 2, 4. + """ + tp = tensor_parallel_size + if tp <= 1: + return + errors: list[str] = [] + if num_heads_q % tp != 0: + errors.append(f"num_heads_q ({num_heads_q}) must be divisible by tensor_parallel_size ({tp})") + if num_heads_kv % tp != 0: + errors.append(f"num_heads_kv ({num_heads_kv}) must be divisible by tensor_parallel_size ({tp})") + # SWIGLU layers use intermediate = int(hidden * 8/3) // 4 * 4 + intermediate_swiglu = int(hidden_size * 4 * 2 / 3) // 4 * 4 + if intermediate_swiglu % tp != 0: + errors.append( + f"swiglu intermediate_size ({intermediate_swiglu}) must be divisible by " + f"tensor_parallel_size ({tp}). Supported tp values: 1, 2, 4" + ) + # GELU7 MoE layers use intermediate = hidden * 4 + intermediate_gelu = hidden_size * 4 + if intermediate_gelu % tp != 0: + errors.append(f"gelu intermediate_size ({intermediate_gelu}) must be divisible by tensor_parallel_size ({tp})") + if errors: + raise ValueError("MagiHuman TP constraint violations:\n" + "\n".join(f" - {e}" for e in errors)) + + +# --------------------------------------------------------------------------- +# Flash attention (no context-parallelism) — uses vllm's flash attention +# --------------------------------------------------------------------------- + +HAS_MAGI_ATTENTION = importlib.util.find_spec("magi_attention") is not None + + +def _fa_varlen_simple( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, +) -> torch.Tensor: + had_batch = query.ndim == 4 + if had_batch: + query = query.squeeze(0) + key = key.squeeze(0) + value = value.squeeze(0) + seq_len = query.shape[0] + cu_seqlens = torch.tensor([0, seq_len], dtype=torch.int32, device=query.device) + out = _vllm_fa_varlen( + q=query, + k=key, + v=value, + cu_seqlens_q=cu_seqlens, + cu_seqlens_k=cu_seqlens, + max_seqlen_q=seq_len, + max_seqlen_k=seq_len, + ) + if had_batch: + out = out.unsqueeze(0) + return out + + +@magi_register_custom_op(name="infra::flash_attn_func", is_subgraph_boundary=True) +def flash_attn_func(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor) -> torch.Tensor: + return _fa_varlen_simple(query, key, value) + + +def _split_q_range_with_no_overlap( + q_ranges: torch.Tensor, k_ranges: torch.Tensor +) -> tuple[list[list[int]], list[list[list[int]]]]: + range_boundary = torch.unique(q_ranges, sorted=True).tolist() + candidates = [[start, end, []] for start, end in zip(range_boundary[:-1], range_boundary[1:])] + q_ranges = q_ranges.tolist() + k_ranges = k_ranges.tolist() + for q_range, k_range in zip(q_ranges, k_ranges): + q_start, q_end = q_range + for q_range_cand in candidates: + if q_start <= q_range_cand[0] and q_range_cand[1] <= q_end: + q_range_cand[2].append(k_range) + q_ranges_out = [] + k_ranges_out = [] + for q_range_cand in candidates: + if len(q_range_cand[2]) > 0: + q_ranges_out.append(q_range_cand[0:2]) + k_ranges_out.append(q_range_cand[2]) + return q_ranges_out, k_ranges_out + + +def _flash_attn_with_correction( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + q_ranges: list[list[int]], + k_range_list: list[list[list[int]]], +): + output = torch.zeros_like(query) + output_lse = torch.zeros((query.shape[0], query.shape[1]), dtype=torch.float32, device=query.device) + + for q_range, k_ranges in zip(q_ranges, k_range_list): + q_start, q_end = q_range + q_chunk = query[q_start:q_end] + q_len = q_chunk.shape[0] + + # Concatenate all k_ranges into a single key/value block, then run one + # flash-attention call. This avoids the need to merge per-chunk LSEs. + k_parts = [key[ks:ke] for ks, ke in k_ranges] + v_parts = [value[ks:ke] for ks, ke in k_ranges] + k_combined = torch.cat(k_parts, dim=0) if len(k_parts) > 1 else k_parts[0] + v_combined = torch.cat(v_parts, dim=0) if len(v_parts) > 1 else v_parts[0] + k_len = k_combined.shape[0] + + cu_q = torch.tensor([0, q_len], dtype=torch.int32, device=query.device) + cu_k = torch.tensor([0, k_len], dtype=torch.int32, device=query.device) + qo_out = _vllm_fa_varlen( + q=q_chunk, + k=k_combined, + v=v_combined, + cu_seqlens_q=cu_q, + cu_seqlens_k=cu_k, + max_seqlen_q=q_len, + max_seqlen_k=k_len, + ) + output[q_start:q_end] = qo_out + return output, output_lse + + +def _flex_flash_attn_func_infer_output_meta( + query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, q_ranges: torch.Tensor, k_ranges: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor]: + output = torch.empty_like(query) + output_lse = torch.empty((query.shape[0], query.shape[1]), dtype=torch.float32, device=query.device) + return output, output_lse + + +@magi_register_custom_op( + name="infra::flex_flash_attn_func", + mutates_args=(), + infer_output_meta_fn=_flex_flash_attn_func_infer_output_meta, + is_subgraph_boundary=True, +) +def flex_flash_attn_func( + query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, q_ranges: torch.Tensor, k_ranges: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor]: + if HAS_MAGI_ATTENTION and _is_hopper_arch(): + from magi_attention.api import flex_flash_attn_func as magi_flex_flash_attn_func + + return magi_flex_flash_attn_func(query, key, value, q_ranges, k_ranges) + else: + q_ranges_split, k_range_list = _split_q_range_with_no_overlap(q_ranges, k_ranges) + return _flash_attn_with_correction(query, key, value, q_ranges_split, k_range_list) + + +def flash_attn_no_cp(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: + q, k, v = q.to(torch.bfloat16), k.to(torch.bfloat16), v.to(torch.bfloat16) + return flash_attn_func(q, k, v).squeeze(0) + + +def flex_flash_attn_no_cp( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + q_ranges: torch.Tensor, + k_ranges: torch.Tensor, +) -> torch.Tensor: + q, k, v = q.to(torch.bfloat16).squeeze(0), k.to(torch.bfloat16).squeeze(0), v.to(torch.bfloat16).squeeze(0) + out, _ = flex_flash_attn_func(q, k, v, q_ranges=q_ranges, k_ranges=k_ranges) + return out + + +# --------------------------------------------------------------------------- +# Attention module (no context-parallelism) +# --------------------------------------------------------------------------- +@dataclass +class AttentionConfig: + hidden_size: int + num_heads_q: int + num_heads_kv: int + head_dim: int + params_dtype: torch.dtype + checkpoint_qk_layernorm_rope: bool + num_modality: int + num_layers: int + use_local_attn: bool = False + enable_attn_gating: bool = False + + +class Attention(torch.nn.Module): + config: AttentionConfig + + def __init__(self, config: AttentionConfig): + super().__init__() + self.config = config + self.pre_norm = MultiModalityRMSNorm(config.hidden_size, eps=1e-6, num_modality=config.num_modality) + self.gating_size = config.num_heads_q if config.enable_attn_gating else 0 + + # Both shared blocks (num_modality == 1) and MoE blocks (num_modality > 1) + # use vLLM's parallel linear layers for TP support. + # MoE blocks wrap per-expert parallel layers with modality dispatch. + if config.num_modality == 1: + # QKVParallelLinear handles GQA head-sharding for any tp_size. + # The combined checkpoint weight [Q, K, V, G] is split during + # load_weights: Q+K+V → linear_qkv, G → linear_gating. + self.linear_qkv = QKVParallelLinear( + hidden_size=config.hidden_size, + head_size=config.head_dim, + total_num_heads=config.num_heads_q, + total_num_kv_heads=config.num_heads_kv, + bias=False, + return_bias=False, + ) + self.linear_proj = RowParallelLinear( + input_size=config.num_heads_q * config.head_dim, + output_size=config.hidden_size, + bias=False, + input_is_parallel=True, + return_bias=False, + ) + if config.enable_attn_gating: + self.linear_gating = ColumnParallelLinear( + input_size=config.hidden_size, + output_size=config.num_heads_q, + bias=False, + gather_output=False, + return_bias=False, + ) + else: + self.linear_gating = None + else: + # MoE blocks: per-expert TP-sharded parallel layers. + self.linear_qkv = MoEQKVParallelLinear( + hidden_size=config.hidden_size, + head_size=config.head_dim, + total_num_heads=config.num_heads_q, + total_num_kv_heads=config.num_heads_kv, + num_experts=config.num_modality, + bias=False, + ) + self.linear_proj = MoERowParallelLinear( + input_size=config.num_heads_q * config.head_dim, + output_size=config.hidden_size, + num_experts=config.num_modality, + bias=False, + ) + if config.enable_attn_gating: + self.linear_gating = MoEColumnParallelLinear( + input_size=config.hidden_size, + output_size=config.num_heads_q, + num_experts=config.num_modality, + bias=False, + ) + else: + self.linear_gating = None + + self.q_norm = MultiModalityRMSNorm(config.head_dim, num_modality=config.num_modality) + self.k_norm = MultiModalityRMSNorm(config.head_dim, num_modality=config.num_modality) + + # q_size / kv_size reflect the per-rank head count when tp > 1. + # Both shared and MoE QKV layers expose .num_heads / .num_kv_heads. + if config.num_modality == 1: + self.q_size = self.linear_qkv.num_heads * config.head_dim + self.kv_size = self.linear_qkv.num_kv_heads * config.head_dim + self._local_heads_q = self.linear_qkv.num_heads + self._local_heads_kv = self.linear_qkv.num_kv_heads + else: + self.q_size = self.linear_qkv.num_heads * config.head_dim + self.kv_size = self.linear_qkv.num_kv_heads * config.head_dim + self._local_heads_q = self.linear_qkv.num_heads + self._local_heads_kv = self.linear_qkv.num_kv_heads + + def forward( + self, + hidden_states: torch.Tensor, + rope: torch.Tensor, + permute_mapping: torch.Tensor, + inv_permute_mapping: torch.Tensor, + varlen_handler: VarlenHandler, + local_attn_handler: FFAHandler | None, + modality_dispatcher: ModalityDispatcher, + ) -> torch.Tensor: + hidden_states = self.pre_norm(hidden_states, modality_dispatcher=modality_dispatcher).to(torch.bfloat16) + + if self.config.num_modality == 1: + # vLLM parallel layers with return_bias=False return a single tensor. + qkv = self.linear_qkv(hidden_states).to(torch.float32) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + if self.linear_gating is not None: + g = self.linear_gating(hidden_states).to(torch.float32) + else: + g = hidden_states.new_empty(hidden_states.shape[0], 0) + else: + # MoE TP path: per-expert QKV parallel layers. + qkv = self.linear_qkv(hidden_states, modality_dispatcher=modality_dispatcher).to(torch.float32) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + if self.linear_gating is not None: + g = self.linear_gating(hidden_states, modality_dispatcher=modality_dispatcher).to(torch.float32) + else: + g = hidden_states.new_empty(hidden_states.shape[0], 0) + + q = q.view(-1, self._local_heads_q, self.config.head_dim) + k = k.view(-1, self._local_heads_kv, self.config.head_dim) + v = v.view(-1, self._local_heads_kv, self.config.head_dim) + g = g.view(k.shape[0], self._local_heads_q, -1) + + q = self.q_norm(q, modality_dispatcher=modality_dispatcher) + k = self.k_norm(k, modality_dispatcher=modality_dispatcher) + + q = ModalityDispatcher.inv_permute(q, inv_permute_mapping).unsqueeze(0) + k = ModalityDispatcher.inv_permute(k, inv_permute_mapping).unsqueeze(0) + v = ModalityDispatcher.inv_permute(v, inv_permute_mapping).unsqueeze(0) + + sin_emb, cos_emb = rope.tensor_split(2, -1) + q = apply_rotary_emb_torch(q, cos_emb, sin_emb) + k = apply_rotary_emb_torch(k, cos_emb, sin_emb) + + if self.config.use_local_attn and local_attn_handler is not None: + self_attn_out = flex_flash_attn_no_cp(q, k, v, local_attn_handler.q_ranges, local_attn_handler.k_ranges) + else: + self_attn_out = flash_attn_no_cp(q, k, v) + self_attn_out = ModalityDispatcher.permute(self_attn_out, permute_mapping) + + if self.config.enable_attn_gating: + self_attn_out = self_attn_out * torch.sigmoid(g) + + self_attn_out = self_attn_out.view(-1, self._local_heads_q * self.config.head_dim).to(torch.bfloat16) + if self.config.num_modality == 1: + return self.linear_proj(self_attn_out) + return self.linear_proj(self_attn_out, modality_dispatcher=modality_dispatcher) + + +# --------------------------------------------------------------------------- +# MLP module +# --------------------------------------------------------------------------- +@dataclass +class MLPConfig: + hidden_size: int + intermediate_size: int + activation_type: MLPActivationType + params_dtype: torch.dtype + num_modality: int = 1 + num_layers: int = 1 + gated_act: bool = False + + +class MLP(torch.nn.Module): + config: MLPConfig + + def __init__(self, config: MLPConfig): + super().__init__() + num_experts = config.num_modality + self.pre_norm = MultiModalityRMSNorm(config.hidden_size, num_modality=config.num_modality) + intermediate_size_up = config.intermediate_size * 2 if config.gated_act else config.intermediate_size + + # Both shared blocks (num_experts == 1) and MoE blocks (num_experts > 1) + # use vLLM's parallel linear layers for TP support. + if num_experts == 1: + # ColumnParallelLinear shards the output dim uniformly. For + # SWIGLU7 the interleaved [up0, gate0, up1, gate1, ...] format + # is preserved within each rank's contiguous slice, so swiglu7 + # (which uses x[..., ::2] / x[..., 1::2]) still works correctly. + self.up_gate_proj = ColumnParallelLinear( + input_size=config.hidden_size, + output_size=intermediate_size_up, + bias=False, + gather_output=False, + return_bias=False, + ) + self.down_proj = RowParallelLinear( + input_size=config.intermediate_size, + output_size=config.hidden_size, + bias=False, + input_is_parallel=True, + return_bias=False, + ) + else: + # MoE blocks: per-expert TP-sharded parallel layers. + self.up_gate_proj = MoEColumnParallelLinear( + input_size=config.hidden_size, + output_size=intermediate_size_up, + num_experts=num_experts, + bias=False, + ) + self.down_proj = MoERowParallelLinear( + input_size=config.intermediate_size, + output_size=config.hidden_size, + num_experts=num_experts, + bias=False, + ) + self.activation_func = create_activation_func(config.activation_type) + + def forward(self, x: torch.Tensor, modality_dispatcher: ModalityDispatcher) -> torch.Tensor: + x = self.pre_norm(x, modality_dispatcher=modality_dispatcher).to(torch.bfloat16) + if isinstance(self.up_gate_proj, ColumnParallelLinear): + x = self.up_gate_proj(x).to(torch.float32) + x = self.activation_func(x).to(torch.bfloat16) + return self.down_proj(x).to(torch.float32) + # MoE TP path: per-expert column/row parallel layers. + x = self.up_gate_proj(x, modality_dispatcher=modality_dispatcher).to(torch.float32) + x = self.activation_func(x).to(torch.bfloat16) + x = self.down_proj(x, modality_dispatcher=modality_dispatcher).to(torch.float32) + return x + + +# --------------------------------------------------------------------------- +# Adapter (per-modality embedders + RoPE) +# --------------------------------------------------------------------------- +@dataclass +class AdapterConfig: + hidden_size: int + num_attention_heads: int + text_in_channels: int + video_in_channels: int + audio_in_channels: int + params_dtype: torch.dtype + + +class Adapter(torch.nn.Module): + config: AdapterConfig + + def __init__(self, config: AdapterConfig): + super().__init__() + self.config = config + self.video_embedder = nn.Linear(config.video_in_channels, config.hidden_size, bias=True, dtype=torch.float32) + self.text_embedder = nn.Linear(config.text_in_channels, config.hidden_size, bias=True, dtype=torch.float32) + self.audio_embedder = nn.Linear(config.audio_in_channels, config.hidden_size, bias=True, dtype=torch.float32) + self.rope = ElementWiseFourierEmbed( + config.hidden_size // config.num_attention_heads, in_pixels=False, learnable=False + ) + + def forward(self, x, coords_mapping, video_mask, audio_mask, text_mask): + rope = self.rope(coords_mapping) + + text_input = x[text_mask, : self.config.text_in_channels] + audio_input = x[audio_mask, : self.config.audio_in_channels] + video_input = x[video_mask, : self.config.video_in_channels] + + text_out = self.text_embedder(text_input) + audio_out = self.audio_embedder(audio_input) + video_out = self.video_embedder(video_input) + + output_x = torch.zeros(x.shape[0], self.config.hidden_size, device=x.device, dtype=x.dtype) + output_x[text_mask] = text_out + output_x[audio_mask] = audio_out + output_x[video_mask] = video_out + return output_x, rope + + +# --------------------------------------------------------------------------- +# Transformer layer (no CP) +# --------------------------------------------------------------------------- +class TransFormerLayer(torch.nn.Module): + def __init__(self, config: Any, layer_idx: int): + super().__init__() + num_modality = 3 if layer_idx in config.mm_layers else 1 + use_local_attn = layer_idx in config.local_attn_layers + self.post_norm = layer_idx in config.post_norm_layers + attention_config = AttentionConfig( + hidden_size=config.hidden_size, + num_heads_q=config.num_heads_q, + num_heads_kv=config.num_heads_kv, + head_dim=config.head_dim, + params_dtype=config.params_dtype, + checkpoint_qk_layernorm_rope=config.checkpoint_qk_layernorm_rope, + num_modality=num_modality, + num_layers=config.num_layers, + use_local_attn=use_local_attn, + enable_attn_gating=config.enable_attn_gating, + ) + self.attention: Attention = Attention(attention_config) + + activation_type = MLPActivationType.GELU7 if layer_idx in config.gelu7_layers else MLPActivationType.SWIGLU7 + if activation_type == MLPActivationType.SWIGLU7: + gated_act = True + intermediate_size = int(config.hidden_size * 4 * 2 / 3) // 4 * 4 + else: + gated_act = False + intermediate_size = config.hidden_size * 4 + mlp_config = MLPConfig( + hidden_size=config.hidden_size, + intermediate_size=intermediate_size, + activation_type=activation_type, + params_dtype=config.params_dtype, + num_modality=num_modality, + num_layers=config.num_layers, + gated_act=gated_act, + ) + self.mlp: MLP = MLP(mlp_config) + if self.post_norm: + self.attn_post_norm = MultiModalityRMSNorm(config.hidden_size, num_modality=num_modality) + self.mlp_post_norm = MultiModalityRMSNorm(config.hidden_size, num_modality=num_modality) + + def forward( + self, + hidden_states: torch.Tensor, + rope: torch.Tensor, + permute_mapping: torch.Tensor, + inv_permute_mapping: torch.Tensor, + varlen_handler: VarlenHandler, + local_attn_handler: FFAHandler | None, + modality_dispatcher: ModalityDispatcher, + ) -> torch.Tensor: + attn_out = self.attention( + hidden_states, + rope, + permute_mapping, + inv_permute_mapping, + varlen_handler, + local_attn_handler, + modality_dispatcher, + ) + if self.post_norm: + attn_out = self.attn_post_norm(attn_out, modality_dispatcher=modality_dispatcher) + hidden_states = hidden_states + attn_out + + mlp_out = self.mlp(hidden_states, modality_dispatcher) + if self.post_norm: + mlp_out = self.mlp_post_norm(mlp_out, modality_dispatcher=modality_dispatcher) + hidden_states = hidden_states + mlp_out + return hidden_states + + +# --------------------------------------------------------------------------- +# TransformerBlock with magi_compile +# --------------------------------------------------------------------------- +is_base_model = True + + +def config_patch(compile_config: CompileConfig) -> CompileConfig: + global is_base_model + if is_base_model: + is_base_model = False + else: + compile_config.offload_config.gpu_resident_weight_ratio = 0.0 + return compile_config + + +@magi_compile( + config_patch=config_patch, dynamic_arg_dims={"x": 0, "rope": 0, "permute_mapping": 0, "inv_permute_mapping": 0} +) +class TransformerBlock(torch.nn.Module): + def __init__(self, model_config: Any): + super().__init__() + self.layers: list[TransFormerLayer] = nn.ModuleList() + for layer_idx in range(model_config.num_layers): + self.layers.append(TransFormerLayer(model_config, layer_idx)) + + def forward( + self, + x: torch.Tensor, + rope: torch.Tensor, + permute_mapping: torch.Tensor, + inv_permute_mapping: torch.Tensor, + varlen_handler: VarlenHandler, + local_attn_handler: FFAHandler | None, + modality_dispatcher: ModalityDispatcher, + ) -> torch.Tensor: + for layer in self.layers: + x = layer( + x, rope, permute_mapping, inv_permute_mapping, varlen_handler, local_attn_handler, modality_dispatcher + ) + return x + + +# --------------------------------------------------------------------------- +# Internal config for TransformerBlock / DiTModel construction +# --------------------------------------------------------------------------- +@dataclass +class TransformerConfig: + hidden_size: int + video_in_channels: int + audio_in_channels: int + text_in_channels: int + params_dtype: torch.dtype + post_process_dtype: torch.dtype + + +# --------------------------------------------------------------------------- +# DiTModel (no context-parallelism) +# --------------------------------------------------------------------------- +class DiTModel(torch.nn.Module): + config: TransformerConfig + _layerwise_offload_blocks_attr = "blocks" + + @property + def blocks(self) -> nn.ModuleList: + return self.block.layers + + def __init__(self, model_config: Any): + super().__init__() + validate_magi_human_tp_constraints( + hidden_size=model_config.hidden_size, + num_heads_q=model_config.hidden_size // model_config.head_dim, + num_heads_kv=model_config.num_query_groups, + tensor_parallel_size=get_tensor_model_parallel_world_size(), + ) + self.config = TransformerConfig( + hidden_size=model_config.hidden_size, + video_in_channels=model_config.video_in_channels, + audio_in_channels=model_config.audio_in_channels, + text_in_channels=model_config.text_in_channels, + params_dtype=model_config.params_dtype, + post_process_dtype=torch.float32, + ) + adapter_config = AdapterConfig( + hidden_size=model_config.hidden_size, + num_attention_heads=model_config.num_heads_q, + text_in_channels=model_config.text_in_channels, + video_in_channels=model_config.video_in_channels, + audio_in_channels=model_config.audio_in_channels, + params_dtype=torch.float32, + ) + self.adapter: Adapter = Adapter(adapter_config) + self.block: TransformerBlock = TransformerBlock(model_config=model_config) + self.final_norm_video = MultiModalityRMSNorm(self.config.hidden_size) + self.final_norm_audio = MultiModalityRMSNorm(self.config.hidden_size) + self.final_linear_video = nn.Linear( + self.config.hidden_size, self.config.video_in_channels, bias=False, dtype=torch.float32 + ) + self.final_linear_audio = nn.Linear( + self.config.hidden_size, self.config.audio_in_channels, bias=False, dtype=torch.float32 + ) + + def forward( + self, + x: torch.Tensor, + coords_mapping: torch.Tensor, + modality_mapping: torch.Tensor, + varlen_handler: VarlenHandler, + local_attn_handler: FFAHandler | None, + ): + modality_dispatcher = ModalityDispatcher(modality_mapping, 3) + permute_mapping = modality_dispatcher.permute_mapping + inv_permute_mapping = modality_dispatcher.inv_permute_mapping + video_mask = modality_mapping == Modality.VIDEO + audio_mask = modality_mapping == Modality.AUDIO + text_mask = modality_mapping == Modality.TEXT + + x, rope = self.adapter(x, coords_mapping, video_mask, audio_mask, text_mask) + + x = x.to(self.config.params_dtype) + x = ModalityDispatcher.permute(x, permute_mapping) + + x = self.block( + x, + rope, + permute_mapping=permute_mapping, + inv_permute_mapping=inv_permute_mapping, + varlen_handler=varlen_handler, + local_attn_handler=local_attn_handler, + modality_dispatcher=modality_dispatcher, + ) + + x = ModalityDispatcher.inv_permute(x, inv_permute_mapping) + + x_video = x[video_mask].to(self.final_norm_video.weight.dtype) + x_video = self.final_norm_video(x_video) + x_video = self.final_linear_video(x_video) + + x_audio = x[audio_mask].to(self.final_norm_audio.weight.dtype) + x_audio = self.final_norm_audio(x_audio) + x_audio = self.final_linear_audio(x_audio) + + x_out = torch.zeros( + x.shape[0], + max(self.config.video_in_channels, self.config.audio_in_channels), + device=x.device, + dtype=x.dtype, + ) + x_out[video_mask, : self.config.video_in_channels] = x_video + x_out[audio_mask, : self.config.audio_in_channels] = x_audio + + return x_out + + +# --------------------------------------------------------------------------- +# Public config dataclass for building DiTModel from JSON +# --------------------------------------------------------------------------- +@dataclass +class MagiHumanDiTConfig: + num_layers: int = 40 + hidden_size: int = 5120 + head_dim: int = 128 + num_query_groups: int = 8 + video_in_channels: int = 48 * 4 + audio_in_channels: int = 64 + text_in_channels: int = 3584 + checkpoint_qk_layernorm_rope: bool = False + params_dtype: torch.dtype = torch.float32 + mm_layers: list = field(default_factory=lambda: [0, 1, 2, 3, 36, 37, 38, 39]) + local_attn_layers: list = field(default_factory=list) + enable_attn_gating: bool = True + gelu7_layers: list = field(default_factory=lambda: [0, 1, 2, 3]) + post_norm_layers: list = field(default_factory=list) + + def __post_init__(self): + self.num_heads_q = self.hidden_size // self.head_dim + self.num_heads_kv = self.num_query_groups + + +if TYPE_CHECKING: + from .pipeline_magi_human import EvalInput + + +# =========================================================================== +# Data proxy (ported from daVinci-MagiHuman inference/pipeline/data_proxy.py) +# =========================================================================== +def _unfold_3d( + x: torch.Tensor, + kernel_size: tuple[int, int, int], + stride: tuple[int, int, int], +) -> torch.Tensor: + """Pure-PyTorch 3D unfold matching UnfoldAnd behavior. + + After N unfold ops the shape is (batch, C, oD, oH, oW, kD, kH, kW). + UnfoldAnd permutes kernel dims next to channel before reshape so that the + col_dim axis is ordered as (C, kD, kH, kW) -- matching F.unfold semantics. + Without this permute, .view() interleaves spatial and kernel positions. + + Args: + x: (N, C, D, H, W) + kernel_size: (kD, kH, kW) + stride: (sD, sH, sW) + Returns: + (N, C*kD*kH*kW, L) where L = product of output spatial dims. + """ + ndim = len(kernel_size) + for d in range(ndim): + x = x.unfold(d + 2, kernel_size[d], stride[d]) + # x: (N, C, oD, oH, oW, kD, kH, kW) + # Permute to (N, C, kD, kH, kW, oD, oH, oW) so that view groups correctly + perm = [0, 1] + list(range(ndim + 2, 2 * ndim + 2)) + list(range(2, ndim + 2)) + x = x.permute(*perm).contiguous() + + batch_size = x.shape[0] + col_dim = 1 + for i in range(1, ndim + 2): + col_dim *= x.shape[i] + spatial = 1 + for i in range(ndim + 2, 2 * ndim + 2): + spatial *= x.shape[i] + return x.view(batch_size, col_dim, spatial) + + +def calc_local_qk_range( + num_video_tokens, + num_audio_and_txt_tokens, + num_frames, + frame_receptive_field, +): + token_per_frame = num_video_tokens // num_frames + total_tokens = num_video_tokens + num_audio_and_txt_tokens + + q_range_list = [] + k_range_list = [] + for i in range(num_frames): + q_range_list.append(torch.tensor([i * token_per_frame, (i + 1) * token_per_frame])) + k_range_list.append( + torch.tensor( + [ + (i - frame_receptive_field) * token_per_frame, + (i + frame_receptive_field + 1) * token_per_frame, + ] + ) + ) + local_q_range = torch.stack(q_range_list, dim=0) + local_k_range = torch.stack(k_range_list, dim=0) + + local_k_range[local_k_range < 0] = 0 + local_k_range[local_k_range > num_video_tokens] = num_video_tokens + + video_q_range = torch.tensor([[0, num_video_tokens]]) + video_k_range = torch.tensor([[num_video_tokens, num_video_tokens + num_audio_and_txt_tokens]]) + + at_q_ranges = torch.tensor([[num_video_tokens, total_tokens]]) + at_k_ranges = torch.tensor([[0, total_tokens]]) + + q_ranges = ( + torch.cat([local_q_range, video_q_range, at_q_ranges], dim=0).to(torch.int32).to("cuda", non_blocking=True) + ) + k_ranges = ( + torch.cat([local_k_range, video_k_range, at_k_ranges], dim=0).to(torch.int32).to("cuda", non_blocking=True) + ) + return q_ranges, k_ranges + + +def calc_local_attn_ffa_handler( + num_video_tokens, + num_audio_and_txt_tokens, + num_frames, + frame_receptive_field, +): + q_ranges, k_ranges = calc_local_qk_range( + num_video_tokens, + num_audio_and_txt_tokens, + num_frames, + frame_receptive_field, + ) + total = num_video_tokens + num_audio_and_txt_tokens + return FFAHandler( + q_ranges=q_ranges, + k_ranges=k_ranges, + max_seqlen_q=total, + max_seqlen_k=total, + attn_type_map=torch.zeros([q_ranges.shape[0]], device="cuda", dtype=torch.int32), + softmax_scale=None, + ) + + +def get_coords( + shape: list[int], + ref_feat_shape: list[int], + offset_thw: list[int] | None = None, + device: torch.device = torch.device("cpu"), + dtype: torch.dtype = torch.float32, +): + if offset_thw is None: + offset_thw = [0, 0, 0] + ori_t, ori_h, ori_w = shape + ref_t, ref_h, ref_w = ref_feat_shape + + offset_t, offset_h, offset_w = offset_thw + time_rng = torch.arange(ori_t, device=device, dtype=dtype) + offset_t + height_rng = torch.arange(ori_h, device=device, dtype=dtype) + offset_h + width_rng = torch.arange(ori_w, device=device, dtype=dtype) + offset_w + + time_grid, height_grid, width_grid = torch.meshgrid( + time_rng, + height_rng, + width_rng, + indexing="ij", + ) + coords_flat = torch.stack([time_grid, height_grid, width_grid], dim=-1).reshape(-1, 3) + + meta = torch.tensor( + [ori_t, ori_h, ori_w, ref_t, ref_h, ref_w], + device=device, + dtype=dtype, + ) + meta_expanded = meta.expand(coords_flat.size(0), -1) + return torch.cat([coords_flat, meta_expanded], dim=-1) + + +@dataclass +class SingleData: + video_x_t: torch.Tensor + audio_x_t: torch.Tensor + audio_feat_len: int + txt_feat: torch.Tensor + txt_feat_len: int + t: int + h: int + w: int + patch_size: int + t_patch_size: int + spatial_rope_interpolation: Literal["inter", "extra"] + ref_audio_offset: int + text_offset: int + coords_style: Literal["v1", "v2"] = "v1" + + def __post_init__(self): + self.video_token_num = self.video_x_t.shape[0] + self.audio_x_t = self.audio_x_t[: self.audio_feat_len] + self.txt_feat = self.txt_feat[: self.txt_feat_len] + self.video_channel = self.video_x_t.shape[-1] + self.audio_channel = self.audio_x_t.shape[-1] + self.txt_channel = self.txt_feat.shape[-1] + + @property + def device(self): + return self.video_x_t.device + + @property + def default_dtype(self): + return self.video_x_t.dtype + + @property + def total_token_num(self): + return self.video_token_num + self.audio_feat_len + self.txt_feat_len + + @property + def token_sequence(self): + tensors = [self.video_x_t, self.audio_x_t, self.txt_feat] + max_channel = max(t.shape[-1] for t in tensors) + padded = [F.pad(t, (0, max_channel - t.shape[-1])) for t in tensors] + return torch.cat(padded, dim=0) + + @property + def modality_mapping(self): + v_map = torch.full((self.video_token_num,), Modality.VIDEO, dtype=torch.int64, device=self.device) + a_map = torch.full((self.audio_feat_len,), Modality.AUDIO, dtype=torch.int64, device=self.device) + t_map = torch.full((self.txt_feat_len,), Modality.TEXT, dtype=torch.int64, device=self.device) + return torch.cat([v_map, a_map, t_map], dim=0) + + def default_coords(self, shape, ref_feat_shape, offset_thw=None): + if offset_thw is None: + offset_thw = [0, 0, 0] + return get_coords( + shape=shape, + ref_feat_shape=ref_feat_shape, + offset_thw=offset_thw, + device=self.device, + dtype=self.default_dtype, + ) + + @property + def coords_mapping(self): + if self.spatial_rope_interpolation == "inter": + video_ref_feat_shape = (self.t // self.t_patch_size, 32, 32) + else: + video_ref_feat_shape = ( + self.t // self.t_patch_size, + self.h // self.patch_size, + self.w // self.patch_size, + ) + + video_coords = self.default_coords( + shape=( + self.t // self.t_patch_size, + self.h // self.patch_size, + self.w // self.patch_size, + ), + ref_feat_shape=video_ref_feat_shape, + ) + + if self.coords_style == "v1": + audio_coords = self.default_coords( + shape=(self.audio_feat_len, 1, 1), + ref_feat_shape=(self.t // self.t_patch_size, 1, 1), + ) + text_coords = self.default_coords( + shape=(self.txt_feat_len, 1, 1), + ref_feat_shape=(2, 1, 1), + offset_thw=[self.text_offset, 0, 0], + ) + elif self.coords_style == "v2": + magic_audio_ref_t = (self.audio_feat_len - 1) // 4 + 1 + audio_coords = self.default_coords( + shape=(self.audio_feat_len, 1, 1), + ref_feat_shape=(magic_audio_ref_t // self.t_patch_size, 1, 1), + ) + text_coords = self.default_coords( + shape=(self.txt_feat_len, 1, 1), + ref_feat_shape=(1, 1, 1), + offset_thw=[-self.txt_feat_len, 0, 0], + ) + else: + raise ValueError(f"Unknown coords_style: {self.coords_style}") + + return torch.cat([video_coords, audio_coords, text_coords], dim=0) + + def depack_token_sequence(self, token_sequence): + video_x_t = token_sequence[: self.video_token_num, : self.video_channel] + video_x_t = rearrange( + video_x_t, + "(T H W) (pT pH pW C) -> C (T pT) (H pH) (W pW)", + H=self.h // self.patch_size, + W=self.w // self.patch_size, + pT=self.t_patch_size, + pH=self.patch_size, + pW=self.patch_size, + ).contiguous() + audio_x_t = token_sequence[ + self.video_token_num : self.video_token_num + self.audio_feat_len, + : self.audio_channel, + ] + return video_x_t, audio_x_t + + +@dataclass +class SimplePackedData: + items: list[SingleData] + + @property + def token_sequence(self): + return torch.cat([item.token_sequence for item in self.items], dim=0) + + @property + def modality_mapping(self): + return torch.cat([item.modality_mapping for item in self.items], dim=0) + + @property + def coords_mapping(self): + return torch.cat([item.coords_mapping for item in self.items], dim=0) + + @property + def total_token_num(self): + return sum(item.total_token_num for item in self.items) + + def __getitem__(self, index): + return self.items[index] + + @property + def cu_seqlen(self): + cu = torch.cumsum( + torch.tensor([item.total_token_num for item in self.items]), + dim=0, + ) + return F.pad(cu, (1, 0)) + + @property + def max_seqlen(self): + return torch.tensor(max(item.total_token_num for item in self.items)) + + def depack_token_sequence(self, token_sequence): + video_list, audio_list = [], [] + parts = torch.split( + token_sequence, + [item.total_token_num for item in self.items], + dim=0, + ) + for item, part in zip(self.items, parts): + v, a = item.depack_token_sequence(part) + video_list.append(v) + audio_list.append(a) + return torch.stack(video_list, dim=0), torch.stack(audio_list, dim=0) + + +class MagiDataProxy: + def __init__( + self, + patch_size: int = 2, + t_patch_size: int = 1, + frame_receptive_field: int = 11, + spatial_rope_interpolation: str = "extra", + ref_audio_offset: int = 1000, + text_offset: int = 0, + coords_style: str = "v2", + ): + self.patch_size = patch_size + self.t_patch_size = t_patch_size + self.frame_receptive_field = frame_receptive_field + self.spatial_rope_interpolation = spatial_rope_interpolation + self.ref_audio_offset = ref_audio_offset + self.text_offset = text_offset + self.coords_style = coords_style + self._kernel = (t_patch_size, patch_size, patch_size) + self._stride = (t_patch_size, patch_size, patch_size) + self._saved_data: dict[str, Any] = {} + + def saved_for_output(self, **kwargs): + self._saved_data.update(kwargs) + + def get_saved_data(self, key: str): + return self._saved_data[key] + + def img2tokens(self, x_t: torch.Tensor): + x_t_unfolded = _unfold_3d(x_t, self._kernel, self._stride) + return rearrange( + x_t_unfolded, + "N col_dim num_tokens -> N num_tokens col_dim", + ).contiguous() + + def process_input(self, transported_data: EvalInput): + batch_size, _, t, h, w = transported_data.x_t.shape + x_t = self.img2tokens(transported_data.x_t) + audio_x_t = transported_data.audio_x_t.contiguous() + text_in = transported_data.txt_feat.contiguous() + + simple_packed_data = SimplePackedData(items=[]) + for i in range(batch_size): + single_data = SingleData( + video_x_t=x_t[i], + audio_x_t=audio_x_t[i], + audio_feat_len=transported_data.audio_feat_len[i], + txt_feat=text_in[i], + txt_feat_len=transported_data.txt_feat_len[i], + t=t, + h=h, + w=w, + patch_size=self.patch_size, + t_patch_size=self.t_patch_size, + spatial_rope_interpolation=self.spatial_rope_interpolation, + ref_audio_offset=self.ref_audio_offset, + text_offset=self.text_offset, + coords_style=self.coords_style, + ) + simple_packed_data.items.append(single_data) + + if self.frame_receptive_field != -1: + assert batch_size == 1, "local attention only supports batch size 1" + local_attn_handler = calc_local_attn_ffa_handler( + num_video_tokens=simple_packed_data[0].video_token_num, + num_audio_and_txt_tokens=(simple_packed_data[0].audio_feat_len + simple_packed_data[0].txt_feat_len), + num_frames=t, + frame_receptive_field=self.frame_receptive_field, + ) + if isinstance(local_attn_handler.max_seqlen_k, torch.Tensor): + local_attn_handler.max_seqlen_k = local_attn_handler.max_seqlen_k.item() + if isinstance(local_attn_handler.max_seqlen_q, torch.Tensor): + local_attn_handler.max_seqlen_q = local_attn_handler.max_seqlen_q.item() + else: + local_attn_handler = None + + varlen_handler = VarlenHandler( + cu_seqlens_q=simple_packed_data.cu_seqlen.to(torch.int32).cuda(), + cu_seqlens_k=simple_packed_data.cu_seqlen.to(torch.int32).cuda(), + max_seqlen_q=simple_packed_data.max_seqlen.to(torch.int32).cuda(), + max_seqlen_k=simple_packed_data.max_seqlen.to(torch.int32).cuda(), + ) + + self.saved_for_output(simple_packed_data=simple_packed_data) + + x = simple_packed_data.token_sequence + coords_mapping = simple_packed_data.coords_mapping + modality_mapping = simple_packed_data.modality_mapping + return (x, coords_mapping, modality_mapping, varlen_handler, local_attn_handler) + + def process_output(self, x: torch.Tensor): + simple_packed_data: SimplePackedData = self.get_saved_data("simple_packed_data") + return simple_packed_data.depack_token_sequence(x) diff --git a/vllm_omni/diffusion/models/magi_human/pipeline_magi_human.py b/vllm_omni/diffusion/models/magi_human/pipeline_magi_human.py new file mode 100644 index 000000000..9e6efcad3 --- /dev/null +++ b/vllm_omni/diffusion/models/magi_human/pipeline_magi_human.py @@ -0,0 +1,2277 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 SandAI. All Rights Reserved. +# Ported from daVinci-MagiHuman inference/pipeline/video_generate.py +# Adapted for vllm-omni: single-GPU, diffusers VAE, configurable dit_subfolder. + +from __future__ import annotations + +import json +import logging +import math +import os +from collections.abc import Iterable +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Literal + +import numpy as np +import torch +import torch.nn as nn +import whisper +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.schedulers.scheduling_utils import ( + KarrasDiffusionSchedulers, + SchedulerMixin, + SchedulerOutput, +) +from diffusers.utils import deprecate, load_image +from diffusers.utils.torch_utils import randn_tensor +from diffusers.video_processor import VideoProcessor +from einops import rearrange +from PIL import Image +from safetensors.torch import load_file +from torch.nn import functional as F +from torch.nn.utils import weight_norm +from transformers import AutoTokenizer +from transformers.models.t5gemma import T5GemmaEncoderModel +from vllm.distributed import ( + get_tensor_model_parallel_world_size, +) +from vllm.model_executor.model_loader.weight_utils import default_weight_loader + +from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import ( + DistributedAutoencoderKLWan, +) +from vllm_omni.diffusion.model_loader.diffusers_loader import ( + DiffusersPipelineLoader, +) +from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin +from vllm_omni.diffusion.models.t5_encoder.t5_gemma_encoder import T5GemmaEncoderModelTP +from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import ( + DiffusionPipelineProfilerMixin, +) +from vllm_omni.diffusion.request import OmniDiffusionRequest + +from .magi_human_dit import ( + DiTModel, + FFAHandler, + MagiHumanDiTConfig, + Modality, + VarlenHandler, +) + +logger = logging.getLogger(__name__) + + +# =========================================================================== +# Scheduler (ported from daVinci-MagiHuman inference/pipeline/scheduler_unipc.py) +# =========================================================================== +class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin): + _compatibles = [e.name for e in KarrasDiffusionSchedulers] + order = 1 + + @register_to_config + def __init__( + self, + num_train_timesteps: int = 1000, + solver_order: int = 2, + prediction_type: str = "flow_prediction", + shift: float = 1.0, + use_dynamic_shifting=False, + thresholding: bool = False, + dynamic_thresholding_ratio: float = 0.995, + sample_max_value: float = 1.0, + predict_x0: bool = True, + solver_type: str = "bh2", + lower_order_final: bool = True, + disable_corrector: list[int] = [], + solver_p: SchedulerMixin = None, + timestep_spacing: str = "linspace", + steps_offset: int = 0, + final_sigmas_type: str | None = "zero", + ): + if solver_type not in ["bh1", "bh2"]: + if solver_type in ["midpoint", "heun", "logrho"]: + self.register_to_config(solver_type="bh2") + else: + raise NotImplementedError(f"{solver_type} is not implemented for {self.__class__}") + + self.predict_x0 = predict_x0 + self.num_inference_steps = None + alphas = np.linspace(1, 1 / num_train_timesteps, num_train_timesteps)[::-1].copy() + sigmas = 1.0 - alphas + sigmas = torch.from_numpy(sigmas).to(dtype=torch.float32) + + if not use_dynamic_shifting: + sigmas = shift * sigmas / (1 + (shift - 1) * sigmas) + + self.sigmas = sigmas + self.timesteps = sigmas * num_train_timesteps + + self.model_outputs = [None] * solver_order + self.timestep_list = [None] * solver_order + self.lower_order_nums = 0 + self.disable_corrector = disable_corrector + self.solver_p = solver_p + self.last_sample = None + self._step_index: int | None = None + self._begin_index: int | None = None + + self.sigmas = self.sigmas.to("cpu") + self.sigma_min = self.sigmas[-1].item() + self.sigma_max = self.sigmas[0].item() + + @property + def step_index(self): + return self._step_index + + @property + def begin_index(self): + return self._begin_index + + def set_begin_index(self, begin_index: int = 0): + self._begin_index = begin_index + + def set_timesteps( + self, + num_inference_steps: int | None = None, + device: str | torch.device = None, + sigmas: list[float] | None = None, + mu: float | None | None = None, + shift: float | None | None = None, + ): + if self.config.use_dynamic_shifting and mu is None: + raise ValueError(" you have to pass a value for `mu` when `use_dynamic_shifting` is set to be `True`") + + if sigmas is None: + sigmas = np.linspace(self.sigma_max, self.sigma_min, num_inference_steps + 1).copy()[:-1] + + if self.config.use_dynamic_shifting: + sigmas = self.time_shift(mu, 1.0, sigmas) + else: + if shift is None: + shift = self.config.shift + sigmas = shift * sigmas / (1 + (shift - 1) * sigmas) + + if self.config.final_sigmas_type == "sigma_min": + sigma_last = ((1 - self.alphas_cumprod[0]) / self.alphas_cumprod[0]) ** 0.5 + elif self.config.final_sigmas_type == "zero": + sigma_last = 0 + else: + raise ValueError( + f"`final_sigmas_type` must be one of 'zero', or 'sigma_min', but got {self.config.final_sigmas_type}" + ) + + timesteps = sigmas * self.config.num_train_timesteps + sigmas = np.concatenate([sigmas, [sigma_last]]).astype(np.float32) + + self.sigmas = torch.from_numpy(sigmas) + self.timesteps = torch.from_numpy(timesteps).to(device=device, dtype=torch.int64) + + self.num_inference_steps = len(timesteps) + + self.model_outputs = [None] * self.config.solver_order + self.lower_order_nums = 0 + self.last_sample = None + if self.solver_p: + self.solver_p.set_timesteps(self.num_inference_steps, device=device) + + self._step_index = None + self._begin_index = None + self.sigmas = self.sigmas.to("cpu") + + def _threshold_sample(self, sample: torch.Tensor) -> torch.Tensor: + dtype = sample.dtype + batch_size, channels, *remaining_dims = sample.shape + + if dtype not in (torch.float32, torch.float64): + sample = sample.float() + + sample = sample.reshape(batch_size, channels * np.prod(remaining_dims)) + abs_sample = sample.abs() + s = torch.quantile(abs_sample, self.config.dynamic_thresholding_ratio, dim=1) + s = torch.clamp(s, min=1, max=self.config.sample_max_value) + s = s.unsqueeze(1) + sample = torch.clamp(sample, -s, s) / s + sample = sample.reshape(batch_size, channels, *remaining_dims) + return sample.to(dtype) + + def _sigma_to_t(self, sigma): + return sigma * self.config.num_train_timesteps + + def _sigma_to_alpha_sigma_t(self, sigma): + return 1 - sigma, sigma + + def time_shift(self, mu: float, sigma: float, t: torch.Tensor): + return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma) + + def convert_model_output( + self, model_output: torch.Tensor, *args, sample: torch.Tensor = None, **kwargs + ) -> torch.Tensor: + timestep = args[0] if len(args) > 0 else kwargs.pop("timestep", None) + if sample is None: + if len(args) > 1: + sample = args[1] + else: + raise ValueError("missing `sample` as a required keyword argument") + if timestep is not None: + deprecate( + "timesteps", + "1.0.0", + "Passing `timesteps` is deprecated and has no effect as model output " + "conversion is now handled via an internal counter `self.step_index`", + ) + + sigma = self.sigmas[self.step_index] + alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma) + + if self.predict_x0: + if self.config.prediction_type == "flow_prediction": + sigma_t = self.sigmas[self.step_index] + x0_pred = sample - sigma_t * model_output + else: + raise ValueError( + f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`," + " `v_prediction` or `flow_prediction` for the UniPCMultistepScheduler." + ) + if self.config.thresholding: + x0_pred = self._threshold_sample(x0_pred) + return x0_pred + else: + if self.config.prediction_type == "flow_prediction": + sigma_t = self.sigmas[self.step_index] + epsilon = sample - (1 - sigma_t) * model_output + else: + raise ValueError( + f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`," + " `v_prediction` or `flow_prediction` for the UniPCMultistepScheduler." + ) + if self.config.thresholding: + sigma_t = self.sigmas[self.step_index] + x0_pred = sample - sigma_t * model_output + x0_pred = self._threshold_sample(x0_pred) + epsilon = model_output + x0_pred + return epsilon + + def multistep_uni_p_bh_update( + self, + model_output: torch.Tensor, + *args, + sample: torch.Tensor | None = None, + order: int | None = None, + **kwargs, + ) -> torch.Tensor: + prev_timestep = args[0] if len(args) > 0 else kwargs.pop("prev_timestep", None) + if sample is None: + if len(args) > 1: + sample = args[1] + else: + raise ValueError(" missing `sample` as a required keyword argument") + if order is None: + if len(args) > 2: + order = args[2] + else: + raise ValueError(" missing `order` as a required keyword argument") + if prev_timestep is not None: + deprecate("prev_timestep", "1.0.0", "Passing `prev_timestep` is deprecated and has no effect.") + + model_output_list = self.model_outputs + s0 = self.timestep_list[-1] + m0 = model_output_list[-1] + x = sample + + if self.solver_p: + return self.solver_p.step(model_output, s0, x).prev_sample + + sigma_t, sigma_s0 = self.sigmas[self.step_index + 1], self.sigmas[self.step_index] + alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t) + alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0) + + lambda_t = torch.log(alpha_t) - torch.log(sigma_t) + lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0) + h = lambda_t - lambda_s0 + device = sample.device + + rks = [] + D1s: list[Any] | None = [] + for i in range(1, order): + si = self.step_index - i + mi = model_output_list[-(i + 1)] + alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si]) + lambda_si = torch.log(alpha_si) - torch.log(sigma_si) + rk = (lambda_si - lambda_s0) / h + rks.append(rk) + D1s.append((mi - m0) / rk) + + rks.append(1.0) + rks = torch.tensor(rks, device=device) + + R = [] + b = [] + hh = -h if self.predict_x0 else h + h_phi_1 = torch.expm1(hh) + h_phi_k = h_phi_1 / hh - 1 + factorial_i = 1 + + if self.config.solver_type == "bh1": + B_h = hh + elif self.config.solver_type == "bh2": + B_h = torch.expm1(hh) + else: + raise NotImplementedError() + + for i in range(1, order + 1): + R.append(torch.pow(rks, i - 1)) + b.append(h_phi_k * factorial_i / B_h) + factorial_i *= i + 1 + h_phi_k = h_phi_k / hh - 1 / factorial_i + + R = torch.stack(R) + b = torch.tensor(b, device=device) + + if len(D1s) > 0: + D1s = torch.stack(D1s, dim=1) + if order == 2: + rhos_p = torch.tensor([0.5], dtype=x.dtype, device=device) + else: + rhos_p = torch.linalg.solve(R[:-1, :-1], b[:-1]).to(device).to(x.dtype) + else: + D1s = None + + if self.predict_x0: + x_t_ = sigma_t / sigma_s0 * x - alpha_t * h_phi_1 * m0 + pred_res = torch.einsum("k,bkc...->bc...", rhos_p, D1s) if D1s is not None else 0 + x_t = x_t_ - alpha_t * B_h * pred_res + else: + x_t_ = alpha_t / alpha_s0 * x - sigma_t * h_phi_1 * m0 + pred_res = torch.einsum("k,bkc...->bc...", rhos_p, D1s) if D1s is not None else 0 + x_t = x_t_ - sigma_t * B_h * pred_res + + return x_t.to(x.dtype) + + def multistep_uni_c_bh_update( + self, + this_model_output: torch.Tensor, + *args, + last_sample: torch.Tensor = None, + this_sample: torch.Tensor = None, + order: int | None = None, + **kwargs, + ) -> torch.Tensor: + this_timestep = args[0] if len(args) > 0 else kwargs.pop("this_timestep", None) + if last_sample is None: + if len(args) > 1: + last_sample = args[1] + else: + raise ValueError(" missing`last_sample` as a required keyword argument") + if this_sample is None: + if len(args) > 2: + this_sample = args[2] + else: + raise ValueError(" missing`this_sample` as a required keyword argument") + if order is None: + if len(args) > 3: + order = args[3] + else: + raise ValueError(" missing`order` as a required keyword argument") + if this_timestep is not None: + deprecate("this_timestep", "1.0.0", "Passing `this_timestep` is deprecated and has no effect.") + + model_output_list = self.model_outputs + m0 = model_output_list[-1] + x = last_sample + x_t = this_sample + model_t = this_model_output + + sigma_t, sigma_s0 = self.sigmas[self.step_index], self.sigmas[self.step_index - 1] + alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t) + alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0) + + lambda_t = torch.log(alpha_t) - torch.log(sigma_t) + lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0) + h = lambda_t - lambda_s0 + device = this_sample.device + + rks = [] + D1s: list[Any] | None = [] + for i in range(1, order): + si = self.step_index - (i + 1) + mi = model_output_list[-(i + 1)] + alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si]) + lambda_si = torch.log(alpha_si) - torch.log(sigma_si) + rk = (lambda_si - lambda_s0) / h + rks.append(rk) + D1s.append((mi - m0) / rk) + + rks.append(1.0) + rks = torch.tensor(rks, device=device) + + R = [] + b = [] + hh = -h if self.predict_x0 else h + h_phi_1 = torch.expm1(hh) + h_phi_k = h_phi_1 / hh - 1 + factorial_i = 1 + + if self.config.solver_type == "bh1": + B_h = hh + elif self.config.solver_type == "bh2": + B_h = torch.expm1(hh) + else: + raise NotImplementedError() + + for i in range(1, order + 1): + R.append(torch.pow(rks, i - 1)) + b.append(h_phi_k * factorial_i / B_h) + factorial_i *= i + 1 + h_phi_k = h_phi_k / hh - 1 / factorial_i + + R = torch.stack(R) + b = torch.tensor(b, device=device) + + if len(D1s) > 0: + D1s = torch.stack(D1s, dim=1) + else: + D1s = None + + if order == 1: + rhos_c = torch.tensor([0.5], dtype=x.dtype, device=device) + else: + rhos_c = torch.linalg.solve(R, b).to(device).to(x.dtype) + + if self.predict_x0: + x_t_ = sigma_t / sigma_s0 * x - alpha_t * h_phi_1 * m0 + corr_res = torch.einsum("k,bkc...->bc...", rhos_c[:-1], D1s) if D1s is not None else 0 + D1_t = model_t - m0 + x_t = x_t_ - alpha_t * B_h * (corr_res + rhos_c[-1] * D1_t) + else: + x_t_ = alpha_t / alpha_s0 * x - sigma_t * h_phi_1 * m0 + corr_res = torch.einsum("k,bkc...->bc...", rhos_c[:-1], D1s) if D1s is not None else 0 + D1_t = model_t - m0 + x_t = x_t_ - sigma_t * B_h * (corr_res + rhos_c[-1] * D1_t) + return x_t.to(x.dtype) + + def index_for_timestep(self, timestep, schedule_timesteps=None): + if schedule_timesteps is None: + schedule_timesteps = self.timesteps + indices = (schedule_timesteps == timestep).nonzero() + pos = 1 if len(indices) > 1 else 0 + return indices[pos].item() + + def _init_step_index(self, timestep): + if self.begin_index is None: + if isinstance(timestep, torch.Tensor): + timestep = timestep.to(self.timesteps.device) + self._step_index = self.index_for_timestep(timestep) + else: + self._step_index = self._begin_index + + def step( + self, + model_output: torch.Tensor, + timestep: int | torch.Tensor, + sample: torch.Tensor, + return_dict: bool = True, + generator=None, + ) -> SchedulerOutput | tuple: + if self.num_inference_steps is None: + raise ValueError( + "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler" + ) + + if self.step_index is None: + self._init_step_index(timestep) + + use_corrector = ( + self.step_index > 0 and self.step_index - 1 not in self.disable_corrector and self.last_sample is not None + ) + + model_output_convert = self.convert_model_output(model_output, sample=sample) + if use_corrector: + sample = self.multistep_uni_c_bh_update( + this_model_output=model_output_convert, + last_sample=self.last_sample, + this_sample=sample, + order=self.this_order, + ) + + for i in range(self.config.solver_order - 1): + self.model_outputs[i] = self.model_outputs[i + 1] + self.timestep_list[i] = self.timestep_list[i + 1] + + self.model_outputs[-1] = model_output_convert + self.timestep_list[-1] = timestep + + if self.config.lower_order_final: + this_order = min(self.config.solver_order, len(self.timesteps) - self.step_index) + else: + this_order = self.config.solver_order + + self.this_order = min(this_order, self.lower_order_nums + 1) + assert self.this_order > 0 + + self.last_sample = sample + prev_sample = self.multistep_uni_p_bh_update(model_output=model_output, sample=sample, order=self.this_order) + + if self.lower_order_nums < self.config.solver_order: + self.lower_order_nums += 1 + + self._step_index += 1 + + if not return_dict: + return (prev_sample,) + return SchedulerOutput(prev_sample=prev_sample) + + def step_ddim( + self, + velocity: torch.FloatTensor, + t: int, + curr_state: torch.FloatTensor, + prev_state: torch.FloatTensor | None = None, + generator: torch.Generator | None = None, + ): + device = curr_state.device + curr_t = self.sigmas[t] + prev_t = self.sigmas[t + 1] + variance_noise = randn_tensor(curr_state.shape, generator=generator, device=device, dtype=curr_state.dtype) + cur_clean_ = curr_state - curr_t * velocity + return prev_t * variance_noise + (1 - prev_t) * cur_clean_ + + def step_sde( + self, + velocity: torch.FloatTensor, + t: int, + curr_state: torch.FloatTensor, + noise_theta: float = 1.0, + prev_state: torch.FloatTensor | None = None, + generator: torch.Generator | None = None, + ): + device = curr_state.device + curr_t = self.sigmas[t] + prev_t = self.sigmas[t + 1] + cos = torch.cos(torch.tensor(noise_theta) * torch.pi / 2).to(device) + sin = torch.sin(torch.tensor(noise_theta) * torch.pi / 2).to(device) + prev_sample_mean = (1 - prev_t + prev_t * cos) * (curr_state - curr_t * velocity) + prev_t * cos * velocity + std_dev_t = prev_t * sin + std_dev_t = torch.ones((1, 1)).to(curr_state) * std_dev_t + if prev_state is None: + variance_noise = randn_tensor(curr_state.shape, generator=generator, device=device, dtype=curr_state.dtype) + prev_state = prev_sample_mean + std_dev_t * variance_noise + else: + prev_state = prev_sample_mean + (prev_state - prev_sample_mean.detach()) + return prev_state + + def scale_model_input(self, sample: torch.Tensor, *args, **kwargs) -> torch.Tensor: + return sample + + def add_noise( + self, original_samples: torch.Tensor, noise: torch.Tensor, timesteps: torch.IntTensor + ) -> torch.Tensor: + sigmas = self.sigmas.to(device=original_samples.device, dtype=original_samples.dtype) + if original_samples.device.type == "mps" and torch.is_floating_point(timesteps): + schedule_timesteps = self.timesteps.to(original_samples.device, dtype=torch.float32) + timesteps = timesteps.to(original_samples.device, dtype=torch.float32) + else: + schedule_timesteps = self.timesteps.to(original_samples.device) + timesteps = timesteps.to(original_samples.device) + + if self.begin_index is None: + step_indices = [self.index_for_timestep(t, schedule_timesteps) for t in timesteps] + elif self.step_index is not None: + step_indices = [self.step_index] * timesteps.shape[0] + else: + step_indices = [self.begin_index] * timesteps.shape[0] + + sigma = sigmas[step_indices].flatten() + while len(sigma.shape) < len(original_samples.shape): + sigma = sigma.unsqueeze(-1) + + alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma) + return alpha_t * original_samples + sigma_t * noise + + def __len__(self): + return self.config.num_train_timesteps + + +# =========================================================================== +# Audio VAE (ported from daVinci-MagiHuman inference/model/sa_audio/) +# =========================================================================== +def _snake_beta(x, alpha, beta): + return x + (1.0 / (beta + 1e-9)) * torch.pow(torch.sin(x * alpha), 2) + + +class _SnakeBeta(nn.Module): + def __init__(self, in_features: int, alpha: float = 1.0, alpha_trainable: bool = True, alpha_logscale: bool = True): + super().__init__() + self.alpha_logscale = alpha_logscale + if self.alpha_logscale: + self.alpha = nn.Parameter(torch.zeros(in_features) * alpha) + self.beta = nn.Parameter(torch.zeros(in_features) * alpha) + else: + self.alpha = nn.Parameter(torch.ones(in_features) * alpha) + self.beta = nn.Parameter(torch.ones(in_features) * alpha) + self.alpha.requires_grad = alpha_trainable + self.beta.requires_grad = alpha_trainable + + def forward(self, x): + alpha = self.alpha.unsqueeze(0).unsqueeze(-1) + beta = self.beta.unsqueeze(0).unsqueeze(-1) + if self.alpha_logscale: + alpha = torch.exp(alpha) + beta = torch.exp(beta) + return _snake_beta(x, alpha, beta) + + +def _vae_sample(mean, scale): + stdev = F.softplus(scale) + 1e-4 + var = stdev * stdev + logvar = torch.log(var) + latents = torch.randn_like(mean) * stdev + mean + kl = (mean * mean + var - logvar - 1).sum(1).mean() + return latents, kl + + +class _VAEBottleneck(nn.Module): + def encode(self, x, return_info=False, **kwargs): + info = {} + mean, scale = x.chunk(2, dim=1) + x, kl = _vae_sample(mean, scale) + info["kl"] = kl + return (x, info) if return_info else x + + def decode(self, x): + return x + + +def _WNConv1d(*args, **kwargs): + return weight_norm(nn.Conv1d(*args, **kwargs)) + + +def _WNConvTranspose1d(*args, **kwargs): + return weight_norm(nn.ConvTranspose1d(*args, **kwargs)) + + +def _checkpoint(function, *args, **kwargs): + kwargs.setdefault("use_reentrant", False) + return torch.utils.checkpoint.checkpoint(function, *args, **kwargs) + + +def _get_activation(activation: Literal["elu", "snake", "none"], antialias: bool = False, channels=None) -> nn.Module: + if antialias: + raise NotImplementedError("antialias activation not supported") + if activation == "elu": + return nn.ELU() + if activation == "snake": + return _SnakeBeta(channels) + if activation == "none": + return nn.Identity() + raise ValueError(f"Unknown activation {activation}") + + +class _ResidualUnit(nn.Module): + def __init__(self, in_channels, out_channels, dilation, use_snake=False, antialias_activation=False): + super().__init__() + padding = (dilation * (7 - 1)) // 2 + self.layers = nn.Sequential( + _get_activation("snake" if use_snake else "elu", antialias=antialias_activation, channels=out_channels), + _WNConv1d(in_channels, out_channels, kernel_size=7, dilation=dilation, padding=padding), + _get_activation("snake" if use_snake else "elu", antialias=antialias_activation, channels=out_channels), + _WNConv1d(out_channels, out_channels, kernel_size=1), + ) + + def forward(self, x): + return (_checkpoint(self.layers, x) if self.training else self.layers(x)) + x + + +class _EncoderBlock(nn.Module): + def __init__(self, in_channels, out_channels, stride, use_snake=False, antialias_activation=False): + super().__init__() + self.layers = nn.Sequential( + _ResidualUnit(in_channels, in_channels, 1, use_snake=use_snake), + _ResidualUnit(in_channels, in_channels, 3, use_snake=use_snake), + _ResidualUnit(in_channels, in_channels, 9, use_snake=use_snake), + _get_activation("snake" if use_snake else "elu", antialias=antialias_activation, channels=in_channels), + _WNConv1d(in_channels, out_channels, kernel_size=2 * stride, stride=stride, padding=math.ceil(stride / 2)), + ) + + def forward(self, x): + return self.layers(x) + + +class _DecoderBlock(nn.Module): + def __init__( + self, in_channels, out_channels, stride, use_snake=False, antialias_activation=False, use_nearest_upsample=False + ): + super().__init__() + if use_nearest_upsample: + upsample_layer = nn.Sequential( + nn.Upsample(scale_factor=stride, mode="nearest"), + _WNConv1d(in_channels, out_channels, kernel_size=2 * stride, stride=1, bias=False, padding="same"), + ) + else: + upsample_layer = _WNConvTranspose1d( + in_channels, out_channels, kernel_size=2 * stride, stride=stride, padding=math.ceil(stride / 2) + ) + self.layers = nn.Sequential( + _get_activation("snake" if use_snake else "elu", antialias=antialias_activation, channels=in_channels), + upsample_layer, + _ResidualUnit(out_channels, out_channels, 1, use_snake=use_snake), + _ResidualUnit(out_channels, out_channels, 3, use_snake=use_snake), + _ResidualUnit(out_channels, out_channels, 9, use_snake=use_snake), + ) + + def forward(self, x): + return self.layers(x) + + +class _OobleckEncoder(nn.Module): + def __init__( + self, + in_channels=2, + channels=128, + latent_dim=32, + c_mults=[1, 2, 4, 8], + strides=[2, 4, 8, 8], + use_snake=False, + antialias_activation=False, + ): + super().__init__() + c_mults = [1] + c_mults + depth = len(c_mults) + layers = [_WNConv1d(in_channels, c_mults[0] * channels, kernel_size=7, padding=3)] + for i in range(depth - 1): + layers.append( + _EncoderBlock(c_mults[i] * channels, c_mults[i + 1] * channels, strides[i], use_snake=use_snake) + ) + layers.extend( + [ + _get_activation( + "snake" if use_snake else "elu", antialias=antialias_activation, channels=c_mults[-1] * channels + ), + _WNConv1d(c_mults[-1] * channels, latent_dim, kernel_size=3, padding=1), + ] + ) + self.layers = nn.Sequential(*layers) + + def forward(self, x): + return self.layers(x) + + +class _OobleckDecoder(nn.Module): + def __init__( + self, + out_channels=2, + channels=128, + latent_dim=32, + c_mults=[1, 2, 4, 8], + strides=[2, 4, 8, 8], + use_snake=False, + antialias_activation=False, + use_nearest_upsample=False, + final_tanh=True, + ): + super().__init__() + c_mults = [1] + c_mults + depth = len(c_mults) + layers = [_WNConv1d(latent_dim, c_mults[-1] * channels, kernel_size=7, padding=3)] + for i in range(depth - 1, 0, -1): + layers.append( + _DecoderBlock( + c_mults[i] * channels, + c_mults[i - 1] * channels, + strides[i - 1], + use_snake=use_snake, + antialias_activation=antialias_activation, + use_nearest_upsample=use_nearest_upsample, + ) + ) + layers.extend( + [ + _get_activation( + "snake" if use_snake else "elu", antialias=antialias_activation, channels=c_mults[0] * channels + ), + _WNConv1d(c_mults[0] * channels, out_channels, kernel_size=7, padding=3, bias=False), + nn.Tanh() if final_tanh else nn.Identity(), + ] + ) + self.layers = nn.Sequential(*layers) + + def forward(self, x): + return self.layers(x) + + +class _AudioAutoencoder(nn.Module): + def __init__( + self, + encoder, + decoder, + latent_dim, + downsampling_ratio, + sample_rate, + io_channels=2, + bottleneck=None, + in_channels=None, + out_channels=None, + soft_clip=False, + ): + super().__init__() + self.downsampling_ratio = downsampling_ratio + self.sample_rate = sample_rate + self.latent_dim = latent_dim + self.io_channels = io_channels + self.in_channels = in_channels if in_channels is not None else io_channels + self.out_channels = out_channels if out_channels is not None else io_channels + self.bottleneck = bottleneck + self.encoder = encoder + self.decoder = decoder + self.soft_clip = soft_clip + + def encode(self, audio, skip_bottleneck=False, return_info=False, **kwargs): + info = {} + latents = self.encoder(audio) + info["pre_bottleneck_latents"] = latents + if self.bottleneck is not None and not skip_bottleneck: + latents, bottleneck_info = self.bottleneck.encode(latents, return_info=True, **kwargs) + info.update(bottleneck_info) + return (latents, info) if return_info else latents + + def decode(self, latents, skip_bottleneck=False, **kwargs): + if self.bottleneck is not None and not skip_bottleneck: + latents = self.bottleneck.decode(latents) + decoded = self.decoder(latents, **kwargs) + if self.soft_clip: + decoded = torch.tanh(decoded) + return decoded + + +def _create_encoder_from_config(cfg: dict[str, Any]): + assert cfg.get("type") == "oobleck", f"Only 'oobleck' encoder supported, got: {cfg.get('type')}" + enc = _OobleckEncoder(**cfg["config"]) + if not cfg.get("requires_grad", True): + for p in enc.parameters(): + p.requires_grad = False + return enc + + +def _create_decoder_from_config(cfg: dict[str, Any]): + assert cfg.get("type") == "oobleck", f"Only 'oobleck' decoder supported, got: {cfg.get('type')}" + dec = _OobleckDecoder(**cfg["config"]) + if not cfg.get("requires_grad", True): + for p in dec.parameters(): + p.requires_grad = False + return dec + + +def _create_bottleneck_from_config(cfg: dict[str, Any]): + assert cfg.get("type") == "vae", f"Only 'vae' bottleneck supported, got: {cfg.get('type')}" + bn = _VAEBottleneck() + if not cfg.get("requires_grad", True): + for p in bn.parameters(): + p.requires_grad = False + return bn + + +def _create_autoencoder_from_config(config: dict[str, Any]): + ae_config = config["model"] + if ae_config.get("pretransform") is not None: + raise NotImplementedError("Nested pretransform not supported") + encoder = _create_encoder_from_config(ae_config["encoder"]) + decoder = _create_decoder_from_config(ae_config["decoder"]) + bottleneck_cfg = ae_config.get("bottleneck") + bottleneck = _create_bottleneck_from_config(bottleneck_cfg) if bottleneck_cfg else None + return _AudioAutoencoder( + encoder=encoder, + decoder=decoder, + latent_dim=ae_config["latent_dim"], + downsampling_ratio=ae_config["downsampling_ratio"], + sample_rate=config["sample_rate"], + io_channels=ae_config["io_channels"], + bottleneck=bottleneck, + in_channels=ae_config.get("in_channels"), + out_channels=ae_config.get("out_channels"), + soft_clip=ae_config["decoder"].get("soft_clip", False), + ) + + +class SAAudioFeatureExtractor: + def __init__(self, device, model_path): + self.device = device + self.vae_model, self.sample_rate = self._load_vae(model_path) + self.resampler = None + + def _load_vae(self, model_path): + if not (isinstance(model_path, str) and Path(model_path).is_dir()): + raise ValueError("model_path must be a local directory") + + model_config_path = os.path.join(model_path, "model_config.json") + with open(model_config_path) as f: + full_config = json.load(f) + + vae_config = full_config["model"]["pretransform"]["config"] + sample_rate = full_config["sample_rate"] + + autoencoder_config = { + "model_type": "autoencoder", + "sample_rate": sample_rate, + "model": vae_config, + } + vae_model = _create_autoencoder_from_config(autoencoder_config) + + weights_path = Path(model_path) / "model.safetensors" + if not weights_path.exists(): + raise FileNotFoundError(f"Weight file does not exist: {weights_path}") + + full_state_dict = load_file(weights_path, device=str(self.device)) + vae_state_dict = {} + for key, value in full_state_dict.items(): + if key.startswith("pretransform.model."): + vae_state_dict[key[len("pretransform.model.") :]] = value + + model_keys = set(vae_model.state_dict().keys()) + vae_keys = set(vae_state_dict.keys()) + missing = model_keys - vae_keys + extra = vae_keys - model_keys + if missing: + logger.warning("Audio VAE missing keys (%d): %s", len(missing), list(missing)[:5]) + if extra: + logger.warning("Audio VAE unexpected keys (%d): %s", len(extra), list(extra)[:5]) + + vae_model.load_state_dict(vae_state_dict) + vae_model.to(self.device) + return vae_model, sample_rate + + def decode(self, latents): + with torch.no_grad(): + return self.vae_model.decode(latents) + + def encode(self, waveform): + with torch.no_grad(): + return self.vae_model.encode(waveform) + + +# =========================================================================== +# Audio utilities (ported from daVinci-MagiHuman inference/pipeline/video_process.py) +# =========================================================================== +_SAMPLE_RATE = 51200 +_AUDIO_CHUNK_DURATION = 29 +_OVERLAP_RATIO = 0.5 + + +def _merge_overlapping_vae_features(audio_feats: list[torch.Tensor], overlap_ratio: float = 0.5) -> torch.Tensor | None: + if not audio_feats: + return None + if len(audio_feats) == 1: + return audio_feats[0] + + batch_size, total_frames, feature_dim = audio_feats[0].shape + overlap_frames = int(total_frames * overlap_ratio) + step_frames = total_frames - overlap_frames + final_length = (len(audio_feats) - 1) * step_frames + total_frames + output_feat = torch.zeros( + batch_size, final_length, feature_dim, device=audio_feats[0].device, dtype=audio_feats[0].dtype + ) + + for block_idx, current_feat in enumerate(audio_feats): + output_start = block_idx * step_frames + if block_idx == 0: + output_feat[:, output_start : output_start + total_frames, :] = current_feat + continue + + non_overlap_start = output_start + overlap_frames + non_overlap_end = output_start + total_frames + output_feat[:, non_overlap_start:non_overlap_end, :] = current_feat[:, overlap_frames:, :] + + for frame_idx in range(overlap_frames): + output_pos = output_start + frame_idx + prev_weight = (overlap_frames - frame_idx) / overlap_frames + curr_weight = frame_idx / overlap_frames + output_feat[:, output_pos, :] = ( + prev_weight * output_feat[:, output_pos, :] + curr_weight * current_feat[:, frame_idx, :] + ) + return output_feat + + +def load_audio_and_encode(audio_vae, audio_path: str, seconds: int | None = None) -> torch.Tensor: + """Load audio from file and encode to latent space using the Stable Audio VAE.""" + audio_full = whisper.load_audio(audio_path, sr=_SAMPLE_RATE) + if seconds is not None: + audio_full = audio_full[: min(int(seconds * _SAMPLE_RATE), audio_full.shape[0])] + total_samples = audio_full.shape[0] + + window_size = int(_AUDIO_CHUNK_DURATION * _SAMPLE_RATE) + step_size = int(window_size * (1 - _OVERLAP_RATIO)) + if total_samples <= window_size: + audio = torch.from_numpy(audio_full).cuda() + audio = audio.unsqueeze(0).expand(2, -1) + return audio_vae.vae_model.encode(audio) + + encoded_chunks = [] + latent_to_audio_ratio = None + for offset_start in range(0, total_samples, step_size): + offset_end = min(offset_start + window_size, total_samples) + chunk = whisper.pad_or_trim(audio_full[offset_start:offset_end], length=window_size) + chunk_tensor = torch.from_numpy(chunk).cuda().unsqueeze(0).expand(2, -1) + encoded_chunk = audio_vae.vae_model.encode(chunk_tensor) + + if latent_to_audio_ratio is None: + latent_to_audio_ratio = encoded_chunk.shape[-1] / window_size + + encoded_chunks.append(encoded_chunk.permute(0, 2, 1)) + if offset_end >= total_samples: + break + + final_feat = _merge_overlapping_vae_features(encoded_chunks, overlap_ratio=_OVERLAP_RATIO).permute(0, 2, 1) + final_target_len = math.ceil(total_samples * latent_to_audio_ratio) + return final_feat[:, :, :final_target_len] + + +# =========================================================================== +# Data proxy (ported from daVinci-MagiHuman inference/pipeline/data_proxy.py) +# =========================================================================== +def _unfold_3d(x: torch.Tensor, kernel_size: tuple[int, int, int], stride: tuple[int, int, int]) -> torch.Tensor: + """Pure-PyTorch 3D unfold matching UnfoldAnd behavior. + + After N unfold ops the shape is (batch, C, oD, oH, oW, kD, kH, kW). + UnfoldAnd permutes kernel dims next to channel before reshape so that the + col_dim axis is ordered as (C, kD, kH, kW) -- matching F.unfold semantics. + Without this permute, .view() interleaves spatial and kernel positions. + + Args: + x: (N, C, D, H, W) + kernel_size: (kD, kH, kW) + stride: (sD, sH, sW) + Returns: + (N, C*kD*kH*kW, L) where L = product of output spatial dims. + """ + ndim = len(kernel_size) + for d in range(ndim): + x = x.unfold(d + 2, kernel_size[d], stride[d]) + perm = [0, 1] + list(range(ndim + 2, 2 * ndim + 2)) + list(range(2, ndim + 2)) + x = x.permute(*perm).contiguous() + + batch_size = x.shape[0] + col_dim = 1 + for i in range(1, ndim + 2): + col_dim *= x.shape[i] + spatial = 1 + for i in range(ndim + 2, 2 * ndim + 2): + spatial *= x.shape[i] + return x.view(batch_size, col_dim, spatial) + + +def _calc_local_qk_range(num_video_tokens, num_audio_and_txt_tokens, num_frames, frame_receptive_field): + token_per_frame = num_video_tokens // num_frames + total_tokens = num_video_tokens + num_audio_and_txt_tokens + + q_range_list = [] + k_range_list = [] + for i in range(num_frames): + q_range_list.append(torch.tensor([i * token_per_frame, (i + 1) * token_per_frame])) + k_range_list.append( + torch.tensor( + [ + (i - frame_receptive_field) * token_per_frame, + (i + frame_receptive_field + 1) * token_per_frame, + ] + ) + ) + local_q_range = torch.stack(q_range_list, dim=0) + local_k_range = torch.stack(k_range_list, dim=0) + + local_k_range[local_k_range < 0] = 0 + local_k_range[local_k_range > num_video_tokens] = num_video_tokens + + video_q_range = torch.tensor([[0, num_video_tokens]]) + video_k_range = torch.tensor([[num_video_tokens, num_video_tokens + num_audio_and_txt_tokens]]) + + at_q_ranges = torch.tensor([[num_video_tokens, total_tokens]]) + at_k_ranges = torch.tensor([[0, total_tokens]]) + + q_ranges = ( + torch.cat([local_q_range, video_q_range, at_q_ranges], dim=0).to(torch.int32).to("cuda", non_blocking=True) + ) + k_ranges = ( + torch.cat([local_k_range, video_k_range, at_k_ranges], dim=0).to(torch.int32).to("cuda", non_blocking=True) + ) + return q_ranges, k_ranges + + +def _calc_local_attn_ffa_handler(num_video_tokens, num_audio_and_txt_tokens, num_frames, frame_receptive_field): + q_ranges, k_ranges = _calc_local_qk_range( + num_video_tokens, num_audio_and_txt_tokens, num_frames, frame_receptive_field + ) + total = num_video_tokens + num_audio_and_txt_tokens + return FFAHandler( + q_ranges=q_ranges, + k_ranges=k_ranges, + max_seqlen_q=total, + max_seqlen_k=total, + attn_type_map=torch.zeros([q_ranges.shape[0]], device="cuda", dtype=torch.int32), + softmax_scale=None, + ) + + +def _get_coords( + shape: list[int], + ref_feat_shape: list[int], + offset_thw: list[int] | None = None, + device: torch.device = torch.device("cpu"), + dtype: torch.dtype = torch.float32, +): + if offset_thw is None: + offset_thw = [0, 0, 0] + ori_t, ori_h, ori_w = shape + ref_t, ref_h, ref_w = ref_feat_shape + + offset_t, offset_h, offset_w = offset_thw + time_rng = torch.arange(ori_t, device=device, dtype=dtype) + offset_t + height_rng = torch.arange(ori_h, device=device, dtype=dtype) + offset_h + width_rng = torch.arange(ori_w, device=device, dtype=dtype) + offset_w + + time_grid, height_grid, width_grid = torch.meshgrid(time_rng, height_rng, width_rng, indexing="ij") + coords_flat = torch.stack([time_grid, height_grid, width_grid], dim=-1).reshape(-1, 3) + + meta = torch.tensor([ori_t, ori_h, ori_w, ref_t, ref_h, ref_w], device=device, dtype=dtype) + meta_expanded = meta.expand(coords_flat.size(0), -1) + return torch.cat([coords_flat, meta_expanded], dim=-1) + + +@dataclass +class _SingleData: + video_x_t: torch.Tensor + audio_x_t: torch.Tensor + audio_feat_len: int + txt_feat: torch.Tensor + txt_feat_len: int + t: int + h: int + w: int + patch_size: int + t_patch_size: int + spatial_rope_interpolation: Literal["inter", "extra"] + ref_audio_offset: int + text_offset: int + coords_style: Literal["v1", "v2"] = "v1" + + def __post_init__(self): + self.video_token_num = self.video_x_t.shape[0] + self.audio_x_t = self.audio_x_t[: self.audio_feat_len] + self.txt_feat = self.txt_feat[: self.txt_feat_len] + self.video_channel = self.video_x_t.shape[-1] + self.audio_channel = self.audio_x_t.shape[-1] + self.txt_channel = self.txt_feat.shape[-1] + + @property + def device(self): + return self.video_x_t.device + + @property + def default_dtype(self): + return self.video_x_t.dtype + + @property + def total_token_num(self): + return self.video_token_num + self.audio_feat_len + self.txt_feat_len + + @property + def token_sequence(self): + tensors = [self.video_x_t, self.audio_x_t, self.txt_feat] + max_channel = max(t.shape[-1] for t in tensors) + padded = [F.pad(t, (0, max_channel - t.shape[-1])) for t in tensors] + return torch.cat(padded, dim=0) + + @property + def modality_mapping(self): + v_map = torch.full((self.video_token_num,), Modality.VIDEO, dtype=torch.int64, device=self.device) + a_map = torch.full((self.audio_feat_len,), Modality.AUDIO, dtype=torch.int64, device=self.device) + t_map = torch.full((self.txt_feat_len,), Modality.TEXT, dtype=torch.int64, device=self.device) + return torch.cat([v_map, a_map, t_map], dim=0) + + def _default_coords(self, shape, ref_feat_shape, offset_thw=None): + if offset_thw is None: + offset_thw = [0, 0, 0] + return _get_coords( + shape=shape, + ref_feat_shape=ref_feat_shape, + offset_thw=offset_thw, + device=self.device, + dtype=self.default_dtype, + ) + + @property + def coords_mapping(self): + if self.spatial_rope_interpolation == "inter": + video_ref_feat_shape = (self.t // self.t_patch_size, 32, 32) + else: + video_ref_feat_shape = (self.t // self.t_patch_size, self.h // self.patch_size, self.w // self.patch_size) + + video_coords = self._default_coords( + shape=(self.t // self.t_patch_size, self.h // self.patch_size, self.w // self.patch_size), + ref_feat_shape=video_ref_feat_shape, + ) + + if self.coords_style == "v1": + audio_coords = self._default_coords( + shape=(self.audio_feat_len, 1, 1), + ref_feat_shape=(self.t // self.t_patch_size, 1, 1), + ) + text_coords = self._default_coords( + shape=(self.txt_feat_len, 1, 1), + ref_feat_shape=(2, 1, 1), + offset_thw=[self.text_offset, 0, 0], + ) + elif self.coords_style == "v2": + magic_audio_ref_t = (self.audio_feat_len - 1) // 4 + 1 + audio_coords = self._default_coords( + shape=(self.audio_feat_len, 1, 1), + ref_feat_shape=(magic_audio_ref_t // self.t_patch_size, 1, 1), + ) + text_coords = self._default_coords( + shape=(self.txt_feat_len, 1, 1), + ref_feat_shape=(1, 1, 1), + offset_thw=[-self.txt_feat_len, 0, 0], + ) + else: + raise ValueError(f"Unknown coords_style: {self.coords_style}") + + return torch.cat([video_coords, audio_coords, text_coords], dim=0) + + def depack_token_sequence(self, token_sequence): + video_x_t = token_sequence[: self.video_token_num, : self.video_channel] + video_x_t = rearrange( + video_x_t, + "(T H W) (pT pH pW C) -> C (T pT) (H pH) (W pW)", + H=self.h // self.patch_size, + W=self.w // self.patch_size, + pT=self.t_patch_size, + pH=self.patch_size, + pW=self.patch_size, + ).contiguous() + audio_x_t = token_sequence[ + self.video_token_num : self.video_token_num + self.audio_feat_len, : self.audio_channel + ] + return video_x_t, audio_x_t + + +@dataclass +class _SimplePackedData: + items: list[_SingleData] + + @property + def token_sequence(self): + return torch.cat([item.token_sequence for item in self.items], dim=0) + + @property + def modality_mapping(self): + return torch.cat([item.modality_mapping for item in self.items], dim=0) + + @property + def coords_mapping(self): + return torch.cat([item.coords_mapping for item in self.items], dim=0) + + @property + def total_token_num(self): + return sum(item.total_token_num for item in self.items) + + def __getitem__(self, index): + return self.items[index] + + @property + def cu_seqlen(self): + cu = torch.cumsum(torch.tensor([item.total_token_num for item in self.items]), dim=0) + return F.pad(cu, (1, 0)) + + @property + def max_seqlen(self): + return torch.tensor(max(item.total_token_num for item in self.items)) + + def depack_token_sequence(self, token_sequence): + video_list, audio_list = [], [] + parts = torch.split(token_sequence, [item.total_token_num for item in self.items], dim=0) + for item, part in zip(self.items, parts): + v, a = item.depack_token_sequence(part) + video_list.append(v) + audio_list.append(a) + return torch.stack(video_list, dim=0), torch.stack(audio_list, dim=0) + + +class MagiDataProxy: + def __init__( + self, + patch_size: int = 2, + t_patch_size: int = 1, + frame_receptive_field: int = 11, + spatial_rope_interpolation: str = "extra", + ref_audio_offset: int = 1000, + text_offset: int = 0, + coords_style: str = "v2", + ): + self.patch_size = patch_size + self.t_patch_size = t_patch_size + self.frame_receptive_field = frame_receptive_field + self.spatial_rope_interpolation = spatial_rope_interpolation + self.ref_audio_offset = ref_audio_offset + self.text_offset = text_offset + self.coords_style = coords_style + self._kernel = (t_patch_size, patch_size, patch_size) + self._stride = (t_patch_size, patch_size, patch_size) + self._saved_data: dict[str, Any] = {} + + def saved_for_output(self, **kwargs): + self._saved_data.update(kwargs) + + def get_saved_data(self, key: str): + return self._saved_data[key] + + def img2tokens(self, x_t: torch.Tensor): + x_t_unfolded = _unfold_3d(x_t, self._kernel, self._stride) + return rearrange(x_t_unfolded, "N col_dim num_tokens -> N num_tokens col_dim").contiguous() + + def process_input(self, transported_data: EvalInput): + batch_size, _, t, h, w = transported_data.x_t.shape + x_t = self.img2tokens(transported_data.x_t) + audio_x_t = transported_data.audio_x_t.contiguous() + text_in = transported_data.txt_feat.contiguous() + + simple_packed_data = _SimplePackedData(items=[]) + for i in range(batch_size): + single_data = _SingleData( + video_x_t=x_t[i], + audio_x_t=audio_x_t[i], + audio_feat_len=transported_data.audio_feat_len[i], + txt_feat=text_in[i], + txt_feat_len=transported_data.txt_feat_len[i], + t=t, + h=h, + w=w, + patch_size=self.patch_size, + t_patch_size=self.t_patch_size, + spatial_rope_interpolation=self.spatial_rope_interpolation, + ref_audio_offset=self.ref_audio_offset, + text_offset=self.text_offset, + coords_style=self.coords_style, + ) + simple_packed_data.items.append(single_data) + + if self.frame_receptive_field != -1: + assert batch_size == 1, "local attention only supports batch size 1" + local_attn_handler = _calc_local_attn_ffa_handler( + num_video_tokens=simple_packed_data[0].video_token_num, + num_audio_and_txt_tokens=simple_packed_data[0].audio_feat_len + simple_packed_data[0].txt_feat_len, + num_frames=t, + frame_receptive_field=self.frame_receptive_field, + ) + if isinstance(local_attn_handler.max_seqlen_k, torch.Tensor): + local_attn_handler.max_seqlen_k = local_attn_handler.max_seqlen_k.item() + if isinstance(local_attn_handler.max_seqlen_q, torch.Tensor): + local_attn_handler.max_seqlen_q = local_attn_handler.max_seqlen_q.item() + else: + local_attn_handler = None + + varlen_handler = VarlenHandler( + cu_seqlens_q=simple_packed_data.cu_seqlen.to(torch.int32).cuda(), + cu_seqlens_k=simple_packed_data.cu_seqlen.to(torch.int32).cuda(), + max_seqlen_q=simple_packed_data.max_seqlen.to(torch.int32).cuda(), + max_seqlen_k=simple_packed_data.max_seqlen.to(torch.int32).cuda(), + ) + + self.saved_for_output(simple_packed_data=simple_packed_data) + + x = simple_packed_data.token_sequence + coords_mapping = simple_packed_data.coords_mapping + modality_mapping = simple_packed_data.modality_mapping + return (x, coords_mapping, modality_mapping, varlen_handler, local_attn_handler) + + def process_output(self, x: torch.Tensor): + simple_packed_data: _SimplePackedData = self.get_saved_data("simple_packed_data") + return simple_packed_data.depack_token_sequence(x) + + +# =========================================================================== +# Pipeline helpers +# =========================================================================== +@dataclass +class EvalInput: + x_t: torch.Tensor + audio_x_t: torch.Tensor + audio_feat_len: torch.Tensor | list[int] + txt_feat: torch.Tensor + txt_feat_len: torch.Tensor | list[int] + + +class _T5GemmaEncoder: + def __init__(self, model_path: str, device: str, weight_dtype: torch.dtype, subfolder: str | None = None): + from vllm.distributed import get_tensor_model_parallel_world_size + + self.device = device + hf_kwargs: dict[str, Any] = {} + if subfolder is not None: + hf_kwargs["subfolder"] = subfolder + self.tokenizer = AutoTokenizer.from_pretrained(model_path, **hf_kwargs) + + tp_size = get_tensor_model_parallel_world_size() + if tp_size > 1: + from transformers.models.t5gemma.configuration_t5gemma import T5GemmaConfig + + config = T5GemmaConfig.from_pretrained(model_path, **hf_kwargs) + # The config we need is the encoder config + config_encoder = config.encoder + # Propagate some outer config values + config_encoder.vocab_size = config.vocab_size + config_encoder.rms_norm_eps = getattr(config, "rms_norm_eps", config_encoder.rms_norm_eps) + self.model = T5GemmaEncoderModelTP(config_encoder).to(device).to(weight_dtype) + self.is_tp = True + else: + self.model = T5GemmaEncoderModel.from_pretrained( + model_path, is_encoder_decoder=False, dtype=weight_dtype, **hf_kwargs + ).to(device) + self.is_tp = False + + @torch.inference_mode() + def encode(self, prompt: str) -> torch.Tensor: + inputs = self.tokenizer([prompt], return_tensors="pt").to(self.device) + outputs = self.model(**inputs) + + if self.is_tp: + # T5GemmaEncoderModelTP just returns the hidden states tensor + return outputs.half() + else: + # HF model returns BaseModelOutput + return outputs["last_hidden_state"].half() + + +def _pad_or_trim(tensor: torch.Tensor, target_size: int, dim: int, pad_value: float = 0.0) -> tuple[torch.Tensor, int]: + current_size = tensor.size(dim) + if current_size < target_size: + padding_amount = target_size - current_size + padding_tuple = [0] * (2 * tensor.dim()) + padding_dim_index = tensor.dim() - 1 - dim + padding_tuple[2 * padding_dim_index + 1] = padding_amount + return F.pad(tensor, tuple(padding_tuple), "constant", pad_value), current_size + slicing = [slice(None)] * tensor.dim() + slicing[dim] = slice(0, target_size) + return tensor[tuple(slicing)], target_size + + +def _get_padded_t5_gemma_embedding( + prompt: str, + encoder: _T5GemmaEncoder, + target_length: int, +) -> tuple[torch.Tensor, int]: + txt_feat = encoder.encode(prompt) + txt_feat, original_len = _pad_or_trim(txt_feat, target_size=target_length, dim=1) + return txt_feat.to(torch.float32), original_len + + +def _resizecrop(img: Image.Image, target_height: int, target_width: int) -> Image.Image: + """Centre-crop resize keeping aspect ratio then letterbox to target.""" + pil_image = img.convert("RGB") + original_width, original_height = pil_image.size + scale_x = target_width / original_width + scale_y = target_height / original_height + scale = max(scale_x, scale_y) + new_width = int(round(original_width * scale)) + new_height = int(round(original_height * scale)) + resized_image = pil_image.resize((new_width, new_height), Image.LANCZOS) + left = (new_width - target_width) // 2 + top = (new_height - target_height) // 2 + return resized_image.crop((left, top, left + target_width, top + target_height)) + + +class ZeroSNRDDPMDiscretization: + """ZeroSNR DDPM sigma schedule, ported from daVinci-MagiHuman. + Used to compute sigma values for SR noise injection. + """ + + def __init__( + self, + linear_start: float = 0.00085, + linear_end: float = 0.0120, + num_timesteps: int = 1000, + shift_scale: float = 1.0, + keep_start: bool = False, + post_shift: bool = False, + ): + from functools import partial + + if keep_start and not post_shift: + linear_start = linear_start / (shift_scale + (1 - shift_scale) * linear_start) + self.num_timesteps = num_timesteps + betas = torch.linspace(linear_start**0.5, linear_end**0.5, num_timesteps, dtype=torch.float64) ** 2 + alphas = 1.0 - betas.cpu().numpy() + self.alphas_cumprod = np.cumprod(alphas, axis=0) + self.to_torch = partial(torch.tensor, dtype=torch.float32) + if not post_shift: + self.alphas_cumprod = self.alphas_cumprod / (shift_scale + (1 - shift_scale) * self.alphas_cumprod) + self.post_shift = post_shift + self.shift_scale = shift_scale + + def __call__( + self, + n: int, + do_append_zero: bool = True, + device: str = "cpu", + flip: bool = False, + return_idx: bool = False, + ): + from functools import partial + + if n < self.num_timesteps: + timesteps = np.linspace(self.num_timesteps - 1, 0, n, endpoint=False).astype(int)[::-1] + alphas_cumprod = self.alphas_cumprod[timesteps] + elif n == self.num_timesteps: + alphas_cumprod = self.alphas_cumprod + else: + raise ValueError(f"n={n} > num_timesteps={self.num_timesteps}") + + to_torch = partial(torch.tensor, dtype=torch.float32, device=device) + alphas_cumprod = to_torch(alphas_cumprod) + alphas_cumprod_sqrt = alphas_cumprod.sqrt() + alphas_cumprod_sqrt_0 = alphas_cumprod_sqrt[0].clone() + alphas_cumprod_sqrt_T = alphas_cumprod_sqrt[-1].clone() + alphas_cumprod_sqrt -= alphas_cumprod_sqrt_T + alphas_cumprod_sqrt *= alphas_cumprod_sqrt_0 / (alphas_cumprod_sqrt_0 - alphas_cumprod_sqrt_T) + + if self.post_shift: + alphas_cumprod_sqrt = ( + alphas_cumprod_sqrt**2 / (self.shift_scale + (1 - self.shift_scale) * alphas_cumprod_sqrt**2) + ) ** 0.5 + + sigmas = torch.flip(alphas_cumprod_sqrt, (0,)) + sigmas = torch.cat([sigmas, sigmas.new_zeros([1])]) if do_append_zero else sigmas + if return_idx: + return sigmas if not flip else torch.flip(sigmas, (0,)), timesteps + return sigmas if not flip else torch.flip(sigmas, (0,)) + + +def _schedule_latent_step( + *, + video_scheduler: FlowUniPCMultistepScheduler, + audio_scheduler: FlowUniPCMultistepScheduler, + latent_video: torch.Tensor, + latent_audio: torch.Tensor, + t, + idx: int, + steps, + v_cfg_video: torch.Tensor, + v_cfg_audio: torch.Tensor, + is_a2v: bool, + cfg_number: int, + using_sde_flag: bool, + use_sr_model: bool = False, +): + # Fast DDIM path for cfg_number==1, only used during the BR stage + if cfg_number == 1 and not use_sr_model: + latent_video = video_scheduler.step_ddim(v_cfg_video, idx, latent_video) + latent_audio = audio_scheduler.step_ddim(v_cfg_audio, idx, latent_audio) + return latent_video, latent_audio + + if using_sde_flag: + if use_sr_model: + # SR stage with SDE: only update video, keep audio unchanged + latent_video = video_scheduler.step(v_cfg_video, t, latent_video, return_dict=False)[0] + return latent_video, latent_audio + if idx < int(len(steps) * (3 / 4)): + noise_theta = 1.0 if (idx + 1) % 2 == 0 else 0.0 + else: + noise_theta = 1.0 if idx % 3 == 0 else 0.0 + latent_video = video_scheduler.step_sde(v_cfg_video, idx, latent_video, noise_theta=noise_theta) + if not is_a2v: + latent_audio = audio_scheduler.step_sde(v_cfg_audio, idx, latent_audio, noise_theta=noise_theta) + return latent_video, latent_audio + + latent_video = video_scheduler.step(v_cfg_video, t, latent_video, return_dict=False)[0] + # Do not update audio latent during the SR stage + if not is_a2v and not use_sr_model: + latent_audio = audio_scheduler.step(v_cfg_audio, t, latent_audio, return_dict=False)[0] + return latent_video, latent_audio + + +_NEGATIVE_PROMPT = ( + "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, " + "overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, " + "poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, " + "still picture, messy background, three legs, many people in the background, walking backwards" + ", low quality, worst quality, poor quality, noise, background noise, hiss, hum, buzz, crackle, static, " + "compression artifacts, MP3 artifacts, digital clipping, distortion, muffled, muddy, unclear, echo, " + "reverb, room echo, over-reverberated, hollow sound, distant, washed out, harsh, shrill, piercing, " + "grating, tinny, thin sound, boomy, bass-heavy, flat EQ, over-compressed, abrupt cut, jarring transition, " + "sudden silence, looping artifact, music, instrumental, sirens, alarms, crowd noise, unrelated sound " + "effects, chaotic, disorganized, messy, cheap sound" + ", emotionless, flat delivery, deadpan, lifeless, apathetic, robotic, mechanical, monotone, flat " + "intonation, undynamic, boring, reading from a script, AI voice, synthetic, text-to-speech, TTS, " + "insincere, fake emotion, exaggerated, overly dramatic, melodramatic, cheesy, cringey, hesitant, " + "unconfident, tired, weak voice, stuttering, stammering, mumbling, slurred speech, mispronounced, " + "bad articulation, lisp, vocal fry, creaky voice, mouth clicks, lip smacks, wet mouth sounds, heavy " + "breathing, audible inhales, plosives, p-pops, coughing, clearing throat, sneezing, speaking too fast, " + "rushed, speaking too slow, dragged out, unnatural pauses, awkward silence, choppy, disjointed, multiple " + "speakers, two voices, background talking, out of tune, off-key, autotune artifacts" +) + + +# =========================================================================== +# Pre/post process funcs (registered in registry) +# =========================================================================== +def get_magi_human_pre_process_func(*args, **kwargs): + def pre_process(request: OmniDiffusionRequest): + return request + + return pre_process + + +def get_magi_human_post_process_func(*args, **kwargs): + def post_process(output): + if isinstance(output, tuple) and len(output) == 2: + video, audio = output + return {"video": video, "audio": audio} + return output + + return post_process + + +# =========================================================================== +# HF Hub / local path helpers +# =========================================================================== + + +def _load_json(model_path: str, filename: str, local_files_only: bool = True) -> dict: + """Load a JSON config file from a local path or HuggingFace Hub repo.""" + if local_files_only: + path = os.path.join(model_path, *filename.split("/")) + with open(path) as f: + return json.load(f) + else: + from huggingface_hub import hf_hub_download + + cached = hf_hub_download(repo_id=model_path, filename=filename) + with open(cached) as f: + return json.load(f) + + +def _resolve_subdir( + model_path: str, + subfolder: str, + local_files_only: bool = True, + required_files: list[str] | None = None, +) -> str: + """Resolve a model subfolder to a local directory path. + + For HF Hub repos, downloads all ``required_files`` (default: ``["config.json"]``) + into the HF cache and returns the parent directory. + """ + if local_files_only: + return os.path.join(model_path, subfolder) + from huggingface_hub import hf_hub_download + + files = required_files or ["config.json"] + last_cached: str | None = None + for fname in files: + last_cached = hf_hub_download(repo_id=model_path, filename=f"{subfolder}/{fname}") + return os.path.dirname(last_cached) + + +# =========================================================================== +# Main Pipeline +# =========================================================================== +class MagiHumanPipeline(nn.Module, ProgressBarMixin, DiffusionPipelineProfilerMixin): + def __init__(self, od_config: OmniDiffusionConfig, **kwargs): + super().__init__() + model_path = od_config.model + local_files_only = os.path.exists(model_path) + device = f"cuda:{torch.cuda.current_device()}" + self.device_str = device + self.dtype = od_config.dtype or torch.bfloat16 + + model_index = _load_json(model_path, "model_index.json", local_files_only) + eval_cfg = model_index + dp_cfg = model_index.get("data_proxy", {}) + + dit_subfolder = "transformer" + + dit_json = _load_json(model_path, f"{dit_subfolder}/config.json", local_files_only) + dit_model_config = MagiHumanDiTConfig(**dit_json) + + self.dit = DiTModel(dit_model_config) + self.dit.eval() + + self.vae = DistributedAutoencoderKLWan.from_pretrained(model_path, subfolder="vae") + self.vae.to(device) + self.vae.eval() + vae_cfg = _load_json(model_path, "vae/config.json", local_files_only) + self.vae_latent_mean = torch.tensor(vae_cfg["latents_mean"], dtype=torch.float32) + self.vae_latent_std = torch.tensor(vae_cfg["latents_std"], dtype=torch.float32) + + self.audio_vae = SAAudioFeatureExtractor( + device=device, + model_path=_resolve_subdir( + model_path, + "audio_vae", + local_files_only, + required_files=["config.json", "model_config.json", "model.safetensors"], + ), + ) + + logger.info("Loading T5Gemma text encoder from %s (subfolder=text_encoder)", model_path) + if local_files_only: + txt_enc_path = os.path.join(model_path, "text_encoder") + txt_enc_subfolder = None + else: + txt_enc_path = model_path + txt_enc_subfolder = "text_encoder" + self.text_encoder = _T5GemmaEncoder( + model_path=txt_enc_path, + device=device, + weight_dtype=self.dtype, + subfolder=txt_enc_subfolder, + ) + + self.data_proxy = MagiDataProxy( + patch_size=dp_cfg.get("patch_size", 2), + t_patch_size=dp_cfg.get("t_patch_size", 1), + frame_receptive_field=dp_cfg.get("frame_receptive_field", 11), + spatial_rope_interpolation=dp_cfg.get("spatial_rope_interpolation", "extra"), + ref_audio_offset=dp_cfg.get("ref_audio_offset", 1000), + text_offset=dp_cfg.get("text_offset", 0), + coords_style=dp_cfg.get("coords_style", "v2"), + ) + # SR DataProxy forces v1 coordinate style (consistent with the original) + self.sr_data_proxy = MagiDataProxy( + patch_size=dp_cfg.get("patch_size", 2), + t_patch_size=dp_cfg.get("t_patch_size", 1), + frame_receptive_field=dp_cfg.get("frame_receptive_field", 11), + spatial_rope_interpolation=dp_cfg.get("spatial_rope_interpolation", "extra"), + ref_audio_offset=dp_cfg.get("ref_audio_offset", 1000), + text_offset=dp_cfg.get("text_offset", 0), + coords_style="v1", + ) + + self.fps = eval_cfg.get("fps", 25) + self.num_inference_steps_default = eval_cfg.get("num_inference_steps", 32) + self.video_txt_guidance_scale = eval_cfg.get("video_txt_guidance_scale", 5.0) + self.audio_txt_guidance_scale = eval_cfg.get("audio_txt_guidance_scale", 5.0) + self.shift = eval_cfg.get("shift", 5.0) + self.cfg_number = eval_cfg.get("cfg_number", 2) + self.use_cfg_trick = eval_cfg.get("use_cfg_trick", True) + self.cfg_trick_start_frame = eval_cfg.get("cfg_trick_start_frame", 13) + self.cfg_trick_value = eval_cfg.get("cfg_trick_value", 2.0) + self.using_sde_flag = eval_cfg.get("using_sde_flag", False) + self.t5_gemma_target_length = eval_cfg.get("t5_gemma_target_length", 640) + self.vae_stride = eval_cfg.get("vae_stride", [4, 16, 16]) + self.z_dim = eval_cfg.get("z_dim", 48) + self.patch_size = eval_cfg.get("patch_size", [1, 2, 2]) + # SR-specific hyperparameters + self.sr_num_inference_steps_default = eval_cfg.get("sr_num_inference_steps", 5) + self.sr_cfg_number = eval_cfg.get("sr_cfg_number", 2) + self.sr_video_txt_guidance_scale = eval_cfg.get("sr_video_txt_guidance_scale", 3.5) + self.noise_value = eval_cfg.get("noise_value", 220) + self.sr_audio_noise_scale = eval_cfg.get("sr_audio_noise_scale", 0.7) + # ZeroSNR sigma schedule for SR noise injection (flip=True, high to low) + self.zerosnr_sigmas = ZeroSNRDDPMDiscretization()(1000, do_append_zero=False, flip=True) + + self.context_null, self.original_context_null_len = _get_padded_t5_gemma_embedding( + _NEGATIVE_PROMPT, + self.text_encoder, + self.t5_gemma_target_length, + ) + self.video_processor = VideoProcessor(vae_scale_factor=16) + + # SR DiT model (loaded from the sr/ subdirectory) + sr_dit_subfolder = "sr" + sr_dit_json = _load_json(model_path, f"{sr_dit_subfolder}/config.json", local_files_only) + sr_dit_model_config = MagiHumanDiTConfig(**sr_dit_json) + self.sr_dit = DiTModel(sr_dit_model_config) + self.sr_dit.eval() + + self.weights_sources = [ + DiffusersPipelineLoader.ComponentSource( + model_or_path=model_path, + subfolder=dit_subfolder, + revision=None, + prefix="dit.", + fall_back_to_pt=True, + ), + DiffusersPipelineLoader.ComponentSource( + model_or_path=model_path, + subfolder=sr_dit_subfolder, + revision=None, + prefix="sr_dit.", + fall_back_to_pt=True, + ), + ] + if getattr(self.text_encoder, "is_tp", False): + self.weights_sources.append( + DiffusersPipelineLoader.ComponentSource( + model_or_path=model_path, + subfolder="text_encoder", + revision=None, + prefix="text_encoder.", + fall_back_to_pt=True, + ), + ) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + # Weight loading for MagiHuman DiT with TP support. + # + # The checkpoint stores weights with these naming patterns: + # - attention.linear_qkv.weight: fused [Q, K, V, G] for shared layers, + # or stacked per-expert [expert0_Q|K|V|G, expert1_..., expert2_...] for MoE. + # - attention.linear_proj.weight: single for shared, stacked per-expert for MoE. + # - mlp.up_gate_proj.weight / mlp.down_proj.weight: similarly stacked for MoE. + # + # The model now uses per-expert vLLM parallel layers for MoE blocks: + # attention.linear_qkv.experts.{i}.weight (QKVParallelLinear per expert) + # attention.linear_gating.experts.{i}.weight (ColumnParallelLinear per expert) + # attention.linear_proj.experts.{i}.weight (RowParallelLinear per expert) + # mlp.up_gate_proj.experts.{i}.weight (ColumnParallelLinear per expert) + # mlp.down_proj.experts.{i}.weight (RowParallelLinear per expert) + # + # Shared layers keep the same naming (no .experts.). + params_dict = dict(self.named_parameters()) + modules_dict = dict(self.named_modules()) + loaded_params: set[str] = set() + + for name, loaded_weight in weights: + # ── Text Encoder weights ── + if name.startswith("text_encoder."): + if getattr(self.text_encoder, "is_tp", False): + # Strip "text_encoder." prefix for the T5Gemma TP model + # The T5GemmaEncoderModelTP load_weights handles the "encoder." prefix itself + sub_name = name[len("text_encoder.") :] + loaded_params.update( + f"text_encoder.{k}" for k in self.text_encoder.model.load_weights([(sub_name, loaded_weight)]) + ) + else: + loaded_params.add(name) + continue + + # ── Shared attention QKV + Gating split ── + # Checkpoint: attention.linear_qkv.weight = [Q, K, V, G] fused. + # Model: attention.linear_qkv.weight (QKVParallelLinear) + attention.linear_gating.weight. + if "attention.linear_qkv.weight" in name: + gating_name = name.replace("attention.linear_qkv.weight", "attention.linear_gating.weight") + # Check if this is a shared layer (direct param exists, no .experts.) + if name in params_dict and gating_name in params_dict: + qkv_param = params_dict[name] + gating_param = params_dict[gating_name] + + mod_path = name[: -len(".weight")] + qkv_mod = modules_dict.get(mod_path) + if qkv_mod is not None and hasattr(qkv_mod, "total_num_heads"): + total_heads_q = qkv_mod.total_num_heads + total_heads_kv = qkv_mod.total_num_kv_heads + head_dim = qkv_mod.head_size + else: + head_dim = 128 + tp_size = get_tensor_model_parallel_world_size() + total_heads_q = gating_param.data.shape[0] * tp_size + total_heads_kv = (loaded_weight.shape[0] - total_heads_q * head_dim - total_heads_q) // ( + 2 * head_dim + ) + + q_size = total_heads_q * head_dim + kv_size = total_heads_kv * head_dim + + q_w = loaded_weight[:q_size] + k_w = loaded_weight[q_size : q_size + kv_size] + v_w = loaded_weight[q_size + kv_size : q_size + 2 * kv_size] + g_w = loaded_weight[q_size + 2 * kv_size :] + + qkv_loader = getattr(qkv_param, "weight_loader", default_weight_loader) + qkv_loader(qkv_param, q_w, "q") + qkv_loader(qkv_param, k_w, "k") + qkv_loader(qkv_param, v_w, "v") + + gating_loader = getattr(gating_param, "weight_loader", default_weight_loader) + gating_loader(gating_param, g_w) + + loaded_params.add(name) + loaded_params.add(gating_name) + continue + + # ── MoE attention QKV + Gating split ── + # Checkpoint: attention.linear_qkv.weight = stacked [expert0_QKVG, expert1_QKVG, ...]. + # Model: attention.linear_qkv.experts.{i}.weight (QKVParallelLinear per expert) + # + attention.linear_gating.experts.{i}.weight (ColumnParallelLinear per expert). + expert0_name = name.replace("attention.linear_qkv.weight", "attention.linear_qkv.experts.0.weight") + if expert0_name in params_dict: + # Determine num_experts by checking which expert indices exist. + moe_qkv_mod_path = name[: -len(".weight")] + moe_qkv_mod = modules_dict.get(moe_qkv_mod_path) + num_experts = moe_qkv_mod.num_experts if moe_qkv_mod is not None else 3 + + # Get head info from the first expert's QKVParallelLinear. + expert0_mod_path = name.replace("attention.linear_qkv.weight", "attention.linear_qkv.experts.0") + expert0_mod = modules_dict.get(expert0_mod_path) + if expert0_mod is not None and hasattr(expert0_mod, "total_num_heads"): + total_heads_q = expert0_mod.total_num_heads + total_heads_kv = expert0_mod.total_num_kv_heads + head_dim = expert0_mod.head_size + else: + head_dim = 128 + # Infer from checkpoint weight shape. + # We'll get exact sizes from model config below. + total_heads_q = 40 # fallback for default config + total_heads_kv = 8 + + q_size = total_heads_q * head_dim + kv_size = total_heads_kv * head_dim + # Check if gating is present. + gating_expert0_name = name.replace( + "attention.linear_qkv.weight", "attention.linear_gating.experts.0.weight" + ) + has_gating = gating_expert0_name in params_dict + + # Split stacked checkpoint weight into per-expert chunks. + expert_weights = loaded_weight.chunk(num_experts, dim=0) + + for i in range(num_experts): + expert_w = expert_weights[i] + # Each expert chunk: [Q, K, V, G (optional)]. + q_w = expert_w[:q_size] + k_w = expert_w[q_size : q_size + kv_size] + v_w = expert_w[q_size + kv_size : q_size + 2 * kv_size] + + expert_param_name = name.replace( + "attention.linear_qkv.weight", + f"attention.linear_qkv.experts.{i}.weight", + ) + expert_param = params_dict[expert_param_name] + expert_loader = getattr(expert_param, "weight_loader", default_weight_loader) + expert_loader(expert_param, q_w, "q") + expert_loader(expert_param, k_w, "k") + expert_loader(expert_param, v_w, "v") + loaded_params.add(expert_param_name) + + if has_gating: + g_w = expert_w[q_size + 2 * kv_size :] + gating_param_name = name.replace( + "attention.linear_qkv.weight", + f"attention.linear_gating.experts.{i}.weight", + ) + gating_param = params_dict[gating_param_name] + gating_loader = getattr(gating_param, "weight_loader", default_weight_loader) + gating_loader(gating_param, g_w) + loaded_params.add(gating_param_name) + continue + + # ── MoE stacked weight splitting for proj / MLP layers ── + # Checkpoint: x.y.weight (stacked [expert0, expert1, ...]). + # Model: x.y.experts.{i}.weight. + if name not in params_dict: + # Check if this is a stacked MoE weight by looking for .experts.0. + base, _, suffix = name.rpartition(".") + expert0_name = f"{base}.experts.0.{suffix}" if base else None + if expert0_name and expert0_name in params_dict: + # Determine num_experts. + moe_mod = modules_dict.get(base) + num_experts = getattr(moe_mod, "num_experts", 3) if moe_mod is not None else 3 + + # Split stacked weight into per-expert chunks. + expert_weights = loaded_weight.chunk(num_experts, dim=0) + for i in range(num_experts): + expert_param_name = f"{base}.experts.{i}.{suffix}" + if expert_param_name not in params_dict: + continue + expert_param = params_dict[expert_param_name] + expert_loader = getattr(expert_param, "weight_loader", default_weight_loader) + expert_loader(expert_param, expert_weights[i]) + loaded_params.add(expert_param_name) + continue + # Truly unknown weight — skip. + continue + + # ── Standard weight loading (shared layers + non-MoE params) ── + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + loaded_params.add(name) + + if getattr(self.text_encoder, "is_tp", False): + self.context_null, self.original_context_null_len = _get_padded_t5_gemma_embedding( + _NEGATIVE_PROMPT, + self.text_encoder, + self.t5_gemma_target_length, + ) + + return loaded_params + + def _dit_forward(self, eval_input: EvalInput) -> tuple[torch.Tensor, torch.Tensor]: + packed = self.data_proxy.process_input(eval_input) + noise_pred = self.dit(*packed) + return self.data_proxy.process_output(noise_pred) + + def _sr_dit_forward(self, eval_input: EvalInput) -> tuple[torch.Tensor, torch.Tensor]: + """SR stage uses sr_data_proxy (coords_style=v1) and sr_dit model.""" + packed = self.sr_data_proxy.process_input(eval_input) + noise_pred = self.sr_dit(*packed) + return self.sr_data_proxy.process_output(noise_pred) + + @torch.inference_mode() + def _evaluate_with_latent( + self, + context: torch.Tensor, + original_context_len: int, + latent_image: torch.Tensor | None, + latent_video: torch.Tensor, + latent_audio: torch.Tensor, + num_inference_steps: int, + is_a2v: bool = False, + use_sr_model: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor]: + # Select cfg_number and guidance_scale based on BR/SR stage + cfg_number = self.sr_cfg_number if use_sr_model else self.cfg_number + video_guidance = self.sr_video_txt_guidance_scale if use_sr_model else self.video_txt_guidance_scale + forward_fn = self._sr_dit_forward if use_sr_model else self._dit_forward + + video_scheduler = FlowUniPCMultistepScheduler() + audio_scheduler = FlowUniPCMultistepScheduler() + video_scheduler.set_timesteps(num_inference_steps, device=self.device_str, shift=self.shift) + audio_scheduler.set_timesteps(num_inference_steps, device=self.device_str, shift=self.shift) + timesteps = video_scheduler.timesteps + + latent_length = latent_video.shape[2] + cfg_trick_guidance = ( + torch.tensor(video_guidance, device=self.device_str).expand(1, 1, latent_length, 1, 1).clone() + ) + if self.use_cfg_trick: + cfg_trick_guidance[:, :, : self.cfg_trick_start_frame] = min(self.cfg_trick_value, video_guidance) + + with self.progress_bar(total=len(timesteps)) as pbar: + for idx, t in enumerate(timesteps): + if latent_image is not None: + latent_video[:, :, :1] = latent_image[:, :, :1] + + # Reduce guidance when t<=500 during BR stage (original behavior) + cur_video_guidance = video_guidance if (use_sr_model or t > 500) else 2.0 + + eval_input_cond = EvalInput( + x_t=latent_video, + audio_x_t=latent_audio, + audio_feat_len=[latent_audio.shape[1]], + txt_feat=context, + txt_feat_len=[original_context_len], + ) + + v_cond_video, v_cond_audio = forward_fn(eval_input_cond) + + if cfg_number == 1: + v_cfg_video = v_cond_video + v_cfg_audio = v_cond_audio + elif cfg_number == 2: + eval_input_uncond = EvalInput( + x_t=latent_video, + audio_x_t=latent_audio, + audio_feat_len=[latent_audio.shape[1]], + txt_feat=self.context_null, + txt_feat_len=[self.original_context_null_len], + ) + v_uncond_video, v_uncond_audio = forward_fn(eval_input_uncond) + v_cfg_video = v_uncond_video + cur_video_guidance * (v_cond_video - v_uncond_video) + v_cfg_audio = v_uncond_audio + self.audio_txt_guidance_scale * (v_cond_audio - v_uncond_audio) + else: + raise ValueError(f"Invalid cfg_number: {cfg_number}") + + latent_video, latent_audio = _schedule_latent_step( + video_scheduler=video_scheduler, + audio_scheduler=audio_scheduler, + latent_video=latent_video, + latent_audio=latent_audio, + t=t, + idx=idx, + steps=timesteps, + v_cfg_video=v_cfg_video, + v_cfg_audio=v_cfg_audio, + is_a2v=is_a2v, + cfg_number=cfg_number, + using_sde_flag=self.using_sde_flag, + use_sr_model=use_sr_model, + ) + + pbar.update() + + if latent_image is not None: + latent_video[:, :, :1] = latent_image[:, :, :1] + return latent_video, latent_audio + + def _encode_image(self, image: Image.Image, height: int, width: int) -> torch.Tensor: + image = load_image(image) + image = _resizecrop(image, height, width) + image = self.video_processor.preprocess(image, height=height, width=width) + image = image.to(device=self.device_str, dtype=self.dtype).unsqueeze(2) + vae_out = self.vae.encode(image) + if hasattr(vae_out, "latent_dist"): + return vae_out.latent_dist.mode().to(torch.float32) + return vae_out.to(torch.float32) + + def _decode_video(self, latent: torch.Tensor) -> list[np.ndarray]: + mean = self.vae_latent_mean.to(latent.device, dtype=latent.dtype).view(1, -1, 1, 1, 1) + std = self.vae_latent_std.to(latent.device, dtype=latent.dtype).view(1, -1, 1, 1, 1) + latent = latent * std + mean + + videos = self.vae.decode(latent.to(self.dtype)) + if hasattr(videos, "sample"): + videos = videos.sample + videos.mul_(0.5).add_(0.5).clamp_(0, 1) + videos = [v.float().cpu().permute(1, 2, 3, 0) * 255 for v in videos] + return [v.numpy().astype(np.uint8) for v in videos] + + def _decode_audio(self, latent_audio: torch.Tensor) -> np.ndarray: + latent_audio = latent_audio.squeeze(0).to(self.dtype) + audio_output = self.audio_vae.decode(latent_audio.T) + audio_np = audio_output.squeeze(0).T.float().cpu().numpy() + target_len = int(audio_np.shape[0] * 441 / 512) + from scipy.signal import resample + + return resample(audio_np, target_len) + + @torch.inference_mode() + def forward( + self, + req: OmniDiffusionRequest, + prompt: str | None = None, + height: int = 256, + width: int = 448, + num_inference_steps: int | None = None, + seconds: int = 10, + seed: int | None = None, + image_path: str | None = None, + audio_path: str | None = None, + **kwargs, + ) -> DiffusionOutput: + if len(req.prompts) >= 1: + p = req.prompts[0] + prompt = p if isinstance(p, str) else p.get("prompt", prompt) + if not isinstance(p, str): + image_path = p.get("image_path", image_path) + audio_path = p.get("audio_path", audio_path) + if prompt is None: + raise ValueError("prompt is required") + + height = req.sampling_params.height or height + width = req.sampling_params.width or width + seed = req.sampling_params.seed if req.sampling_params.seed is not None else seed + num_steps = req.sampling_params.num_inference_steps or num_inference_steps or self.num_inference_steps_default + sr_height: int | None = None + sr_width: int | None = None + sr_num_steps: int | None = None + if hasattr(req.sampling_params, "extra_args") and req.sampling_params.extra_args: + seconds = req.sampling_params.extra_args.get("seconds", seconds) + audio_path = req.sampling_params.extra_args.get("audio_path", audio_path) + image_path = req.sampling_params.extra_args.get("image_path", image_path) + sr_height = req.sampling_params.extra_args.get("sr_height", None) + sr_width = req.sampling_params.extra_args.get("sr_width", None) + sr_num_steps = req.sampling_params.extra_args.get("sr_num_inference_steps", None) + + device = self.device_str + + br_latent_height = height // self.vae_stride[1] // self.patch_size[1] * self.patch_size[1] + br_latent_width = width // self.vae_stride[2] // self.patch_size[2] * self.patch_size[2] + br_height = br_latent_height * self.vae_stride[1] + br_width = br_latent_width * self.vae_stride[2] + + if seed is not None: + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + + if audio_path is not None: + latent_audio = load_audio_and_encode(self.audio_vae, audio_path, seconds) + latent_audio = latent_audio.permute(0, 2, 1) + num_frames = latent_audio.shape[1] + is_a2v = True + else: + num_frames = seconds * self.fps + 1 + latent_audio = torch.randn(1, num_frames, 64, dtype=torch.float32, device=device) + is_a2v = False + + latent_length = (num_frames - 1) // 4 + 1 + latent_video = torch.randn( + 1, + self.z_dim, + latent_length, + br_latent_height, + br_latent_width, + dtype=torch.float32, + device=device, + ) + + context, original_context_len = _get_padded_t5_gemma_embedding( + prompt, + self.text_encoder, + self.t5_gemma_target_length, + ) + + if image_path is not None: + br_image = self._encode_image(load_image(image_path), br_height, br_width) + else: + br_image = None + + # ── BR stage ───────────────────────────────────────────────────────── + br_latent_video, br_latent_audio = self._evaluate_with_latent( + context, + original_context_len, + br_image, + latent_video.clone(), + latent_audio.clone(), + num_steps, + is_a2v, + use_sr_model=False, + ) + + # ── SR stage (optional, triggered when sr_height/sr_width are provided) ── + if sr_height is not None and sr_width is not None: + sr_latent_height = sr_height // self.vae_stride[1] // self.patch_size[1] * self.patch_size[1] + sr_latent_width = sr_width // self.vae_stride[2] // self.patch_size[2] * self.patch_size[2] + sr_height = sr_latent_height * self.vae_stride[1] + sr_width = sr_latent_width * self.vae_stride[2] + + # Image condition (at SR resolution) + if image_path is not None: + sr_image = self._encode_image(load_image(image_path), sr_height, sr_width) + else: + sr_image = None + + # Trilinear interpolation of BR latent to SR resolution + sr_latent_video = torch.nn.functional.interpolate( + br_latent_video, + size=(latent_length, sr_latent_height, sr_latent_width), + mode="trilinear", + align_corners=True, + ) + + # Noise injection: sigma-weighted blend (noise_value indexes the ZeroSNR sigma schedule) + if self.noise_value != 0: + noise = torch.randn_like(sr_latent_video) + sigma = self.zerosnr_sigmas.to(sr_latent_video.device)[self.noise_value] + sr_latent_video = sr_latent_video * sigma + noise * (1 - sigma**2) ** 0.5 + + # Audio: blend with noise (noised version used during SR inference; final audio keeps BR result) + sr_latent_audio = torch.randn_like(br_latent_audio) * self.sr_audio_noise_scale + br_latent_audio * ( + 1 - self.sr_audio_noise_scale + ) + + torch.cuda.empty_cache() + sr_steps = sr_num_steps or self.sr_num_inference_steps_default + final_latent_video, _ = self._evaluate_with_latent( + context, + original_context_len, + sr_image, + sr_latent_video.clone(), + sr_latent_audio.clone(), + sr_steps, + is_a2v, + use_sr_model=True, + ) + # SR stage does not update audio; keep the BR result + final_latent_video = final_latent_video + final_latent_audio = br_latent_audio + else: + final_latent_video = br_latent_video + final_latent_audio = br_latent_audio + + torch.cuda.empty_cache() + videos_np = self._decode_video(final_latent_video) + torch.cuda.empty_cache() + audio_np = self._decode_audio(final_latent_audio) + + return DiffusionOutput(output=(videos_np, audio_np)) diff --git a/vllm_omni/diffusion/models/t5_encoder/t5_gemma_encoder.py b/vllm_omni/diffusion/models/t5_encoder/t5_gemma_encoder.py new file mode 100644 index 000000000..eca4267fa --- /dev/null +++ b/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 diff --git a/vllm_omni/diffusion/offloader/module_collector.py b/vllm_omni/diffusion/offloader/module_collector.py index 307ca53a8..d9d21b939 100644 --- a/vllm_omni/diffusion/offloader/module_collector.py +++ b/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: diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index c1f48137e..97bc7fa29 100644 --- a/vllm_omni/diffusion/registry.py +++ b/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", } diff --git a/vllm_omni/diffusion/utils/media_utils.py b/vllm_omni/diffusion/utils/media_utils.py new file mode 100644 index 000000000..ee1f8116f --- /dev/null +++ b/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()