     1→# Copyright (c) 2026 SandAI. All Rights Reserved.
     2→#
     3→# Licensed under the Apache License, Version 2.0 (the "License");
     4→# you may not use this file except in compliance with the License.
     5→# You may obtain a copy of the License at
     6→#
     7→#     http://www.apache.org/licenses/LICENSE-2.0
     8→#
     9→# Unless required by applicable law or agreed to in writing, software
    10→# distributed under the License is distributed on an "AS IS" BASIS,
    11→# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12→# See the License for the specific language governing permissions and
    13→# limitations under the License.
    14→
    15→import torch
    16→
    17→from inference.common import parse_config
    18→from inference.infra.distributed import get_dp_rank, initialize_distributed
    19→from inference.utils import print_rank_0, set_random_seed
    20→
    21→
    22→def initialize_infra():
    23→    assert torch.cuda.is_available(), "Infra requires CUDA environment."
    24→
    25→    # Initialize distributed environment
    26→    initialize_distributed()
    27→
    28→    # Initialize config
    29→    config = parse_config(verbose=True)
    30→
    31→    # Initialize random seed
    32→    set_random_seed(config.engine_config.seed + 10 * get_dp_rank())
    33→
    34→    print_rank_0("Infra successfully initialized")
    35→
    36→
    37→__all__ = ["initialize_infra"]
    38→


Now let me check the common init and arch detection:


     1→# Copyright (c) 2026 SandAI. All Rights Reserved.
     2→#
     3→# Licensed under the Apache License, Version 2.0 (the "License");
     4→# you may not use this file except in compliance with the License.
     5→# You may obtain a copy of the License at
     6→#
     7→#     http://www.apache.org/licenses/LICENSE-2.0
     8→#
     9→# Unless required by applicable law or agreed to in writing, software
    10→# distributed under the License is distributed on an "AS IS" BASIS,
    11→# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12→# See the License for the specific language governing permissions and
    13→# limitations under the License.
    14→
    15→from .arch import get_arch_memory, is_hopper_arch
    16→from .config import (
    17→    DataProxyConfig,
    18→    EngineConfig,
    19→    EvaluationConfig,
    20→    parse_config,
    21→)
    22→from .cpu_offload_wrapper import CPUOffloadWrapper
    23→from .sequence_schema import Modality, VarlenHandler
    24→
    25→__all__ = [
    26→    # arch
    27→    "get_arch_memory",
    28→    "is_hopper_arch",
    29→    # config
    30→    "EngineConfig",
    31→    "DataProxyConfig",
    32→    "EvaluationConfig",
    33→    "parse_config",
    34→    # cpu offload wrapper
    35→    "CPUOffloadWrapper",
    36→    # sequence schema
    37→    "Modality",
    38→    "VarlenHandler",
    39→]
    40→


     1→# Copyright (c) 2026 SandAI. All Rights Reserved.
     2→#
     3→# Licensed under the Apache License, Version 2.0 (the "License");
     4→# you may not use this file except in compliance with the License.
     5→# You may obtain a copy of the License at
     6→#
     7→#     http://www.apache.org/licenses/LICENSE-2.0
     8→#
     9→# Unless required by applicable law or agreed to in writing, software
    10→# distributed under the License is distributed on an "AS IS" BASIS,
    11→# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12→# See the License for the specific language governing permissions and
    13→# limitations under the License.
    14→
    15→import torch
    16→
    17→
    18→def is_hopper_arch():
    19→    return torch.cuda.get_device_capability()[0] == 9
    20→
    21→
    22→def get_arch_memory(unit: str = "GB"):
    23→    if not torch.cuda.is_available():
    24→        return 0
    25→    total_bytes = torch.cuda.get_device_properties(torch.cuda.current_device()).total_memory
    26→    if unit == "B":
    27→        return float(total_bytes)
    28→    elif unit == "KB":
    29→        return total_bytes / 1024
    30→    elif unit == "MB":
    31→        return total_bytes / 1024 / 1024
    32→    elif unit == "GB":
    33→        return total_bytes / 1024 / 1024 / 1024
    34→    else:
    35→        raise ValueError(f"Invalid unit: {unit}")
    36→


Now let me examine the video processing and data proxy:


     1→# Copyright (c) 2026 SandAI. All Rights Reserved.
     2→#
     3→# Licensed under the Apache License, Version 2.0 (the "License");
     4→# you may not use this file except in compliance with the License.
     5→# You may obtain a copy of the License at
     6→#
     7→#     http://www.apache.org/licenses/LICENSE-2.0
     8→#
     9→# Unless required by applicable law or agreed to in writing, software
    10→# distributed under the License is distributed on an "AS IS" BASIS,
    11→# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12→# See the License for the specific language governing permissions and
    13→# limitations under the License.
    14→
    15→import math
    16→import os
    17→import subprocess
    18→from typing import Optional
    19→
    20→import numpy as np
    21→import torch
    22→import whisper
    23→from einops import rearrange
    24→from PIL import Image
    25→from scipy.signal import resample
    26→from torch.nn import functional as F
    27→
    28→from inference.utils import print_rank_0
    29→
    30→
    31→def merge_video_and_audio(video_path: str, audio_path: str, save_path: str):
    32→    # Merge video with audio and keep the shortest stream.
    33→    cmd = [
    34→        "ffmpeg",
    35→        "-i",
    36→        video_path,
    37→        "-i",
    38→        audio_path,
    39→        "-map",
    40→        "0:v:0",
    41→        "-map",
    42→        "1:a:0",
    43→        "-c:v",
    44→        "copy",
    45→        "-c:a",
    46→        "aac",
    47→        "-y",
    48→        save_path,
    49→        "-loglevel",
    50→        "error",
    51→    ]
    52→    try:
    53→        subprocess.run(cmd, check=True)
    54→        os.remove(video_path)
    55→        os.remove(audio_path)
    56→    except subprocess.CalledProcessError as e:
    57→        print_rank_0(f"ffmpeg failed: {e}")
    58→
    59→
    60→def upsample_video(video_np: np.ndarray, width: int, height: int, upsample_mode: str = "bilinear") -> np.ndarray:
    61→    """
    62→    Upsample video NumPy array to specified resolution.
    63→
    64→    This function assumes the input NumPy array is the result of VAE decoding,
    65→    with data type uint8 and dimension order (T, H, W, C).
    66→
    67→    Args:
    68→        video_np (np.ndarray): Input video array with shape (T, H, W, C),
    69→                               data type uint8.
    70→        width (int): Target width.
    71→        height (int): Target height.
    72→        upsample_mode (str): Upsampling mode. Supports "bilinear", "nearest", "bicubic".
    73→                            Defaults to "bilinear".
    74→
    75→    Returns:
    76→        np.ndarray: Upsampled video array with shape (T, height, width, C),
    77→                    data type uint8.
    78→    """
    79→    assert upsample_mode in ["bilinear", "nearest", "bicubic"], "Supported upsample modes: bilinear, nearest, bicubic"
    80→
    81→     # 1. Convert NumPy array to PyTorch tensor
    82→    video_tensor = torch.from_numpy(video_np)
    83→
    84→    # 2. Convert from uint8 to float32 and normalize to [0, 1]
    85→    #    F.interpolate works better on floating point numbers
    86→    if video_tensor.dtype == torch.uint8:
    87→        video_tensor = video_tensor.float() / 255.0
    88→
    89→    # 3. Adjust dimension order to match F.interpolate requirements (T, H, W, C) -> (C, T, H, W)
    90→    video_tensor = rearrange(video_tensor, "t h w c -> c t h w")
    91→
    92→    # 4. Use F.interpolate for upsampling
    93→    #    Note: interpolate operates on spatial dimensions (H, W), so size=(height, width)
    94→    upsampled_tensor = F.interpolate(
    95→        video_tensor,
    96→        size=(height, width),
    97→        mode=upsample_mode,
    98→        align_corners=False if upsample_mode in ["bilinear", "bicubic"] else None,
    99→    )
   100→
   101→    # 5. Adjust dimension order back (C, T, H, W) -> (T, H, W, C)
   102→    upsampled_tensor = rearrange(upsampled_tensor, "c t h w -> t h w c")
   103→
   104→    # 6. Convert data from [0, 1] range back to [0, 255] and convert to uint8
   105→    upsampled_tensor = (upsampled_tensor.clamp(0, 1) * 255).byte()
   106→
   107→    # 7. Convert PyTorch tensor back to NumPy array
   108→    return upsampled_tensor.numpy()
   109→
   110→
   111→def resizecrop(image: Image.Image, th: int, tw: int) -> Image.Image:
   112→    w, h = image.size
   113→    if w == tw and h == th:
   114→        return image
   115→    if h / w > th / tw:
   116→        new_w = int(w)
   117→        new_h = int(new_w * th / tw)
   118→    else:
   119→        new_h = int(h)
   120→        new_w = int(new_h * tw / th)
   121→    left = (w - new_w) / 2
   122→    top = (h - new_h) / 2
   123→    right = (w + new_w) / 2
   124→    bottom = (h + new_h) / 2
   125→    return image.crop((left, top, right, bottom))
   126→
   127→
   128→def resample_audio_sinc(audio: torch.Tensor, time_stretching: float):
   129→    print_rank_0(f"before resample audio: {audio.shape}")
   130→    new_length = int(audio.shape[0] * time_stretching)
   131→    audio = resample(audio, new_length)
   132→    print_rank_0(f"after resample audio: {audio.shape}")
   133→    return audio
   134→
   135→
   136→def merge_overlapping_vae_features(audio_feats, overlap_ratio=0.5):
   137→    if not audio_feats:
   138→        return None
   139→    if len(audio_feats) == 1:
   140→        return audio_feats[0]
   141→
   142→    batch_size, total_frames, feature_dim = audio_feats[0].shape
   143→    overlap_frames = int(total_frames * overlap_ratio)
   144→    step_frames = total_frames - overlap_frames
   145→    final_length = (len(audio_feats) - 1) * step_frames + total_frames
   146→    output_feat = torch.zeros(batch_size, final_length, feature_dim, device=audio_feats[0].device, dtype=audio_feats[0].dtype)
   147→
   148→    for block_idx, current_feat in enumerate(audio_feats):
   149→        output_start = block_idx * step_frames
   150→        if block_idx == 0:


     1→# Copyright (c) 2026 SandAI. All Rights Reserved.
     2→#
     3→# Licensed under the Apache License, Version 2.0 (the "License");
     4→# you may not use this file except in compliance with the License.
     5→# You may obtain a copy of the License at
     6→#
     7→#     http://www.apache.org/licenses/LICENSE-2.0
     8→#
     9→# Unless required by applicable law or agreed to in writing, software
    10→# distributed under the License is distributed on an "AS IS" BASIS,
    11→# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12→# See the License for the specific language governing permissions and
    13→# limitations under the License.
    14→
    15→import math
    16→from dataclasses import dataclass
    17→from enum import IntEnum
    18→from typing import Any, Literal, Optional, TYPE_CHECKING
    19→
    20→import torch
    21→from einops import rearrange
    22→from inference.common import DataProxyConfig, Modality, VarlenHandler
    23→from inference.model.dit.dit_module import FFAHandler
    24→from torch.nn import functional as F
    25→from unfoldNd import UnfoldNd
    26→
    27→if TYPE_CHECKING:
    28→    from inference.pipeline.video_generate import EvalInput
    29→
    30→
    31→def calc_local_qk_range(num_video_tokens, num_audio_and_txt_tokens, num_frames, frame_receptive_field):
    32→    token_per_frame = num_video_tokens // num_frames
    33→    total_tokens = num_video_tokens + num_audio_and_txt_tokens
    34→
    35→    q_range_list = []
    36→    k_range_list = []
    37→
    38→    for i in range(num_frames):
    39→        local_q_range = torch.tensor([i * token_per_frame, (i + 1) * token_per_frame])
    40→        local_k_range = torch.tensor(
    41→            [(i - frame_receptive_field) * token_per_frame, (i + frame_receptive_field + 1) * token_per_frame]
    42→        )
    43→
    44→        q_range_list.append(local_q_range)
    45→        k_range_list.append(local_k_range)
    46→    local_q_range = torch.stack(q_range_list, dim=0)
    47→    local_k_range = torch.stack(k_range_list, dim=0)
    48→
    49→    local_k_range[local_k_range < 0] = 0
    50→    local_k_range[local_k_range > num_video_tokens] = num_video_tokens
    51→
    52→    video_q_range = torch.tensor([[0, num_video_tokens]])
    53→    video_k_range = torch.tensor([[num_video_tokens, num_video_tokens + num_audio_and_txt_tokens]])
    54→
    55→    at_q_ranges = torch.tensor([[num_video_tokens, total_tokens]])
    56→    at_k_ranges = torch.tensor([[0, total_tokens]])
    57→
    58→    q_ranges = torch.cat([local_q_range, video_q_range, at_q_ranges], dim=0).to(torch.int32).to("cuda", non_blocking=True)
    59→    k_ranges = torch.cat([local_k_range, video_k_range, at_k_ranges], dim=0).to(torch.int32).to("cuda", non_blocking=True)
    60→
    61→    return (q_ranges, k_ranges)
    62→
    63→
    64→def calc_local_attn_ffa_handler(num_video_tokens, num_audio_and_txt_tokens, num_frames, frame_receptive_field):
    65→    q_ranges, k_ranges = calc_local_qk_range(num_video_tokens, num_audio_and_txt_tokens, num_frames, frame_receptive_field)
    66→    max_seqlen_q = num_video_tokens + num_audio_and_txt_tokens
    67→    max_seqlen_k = num_video_tokens + num_audio_and_txt_tokens
    68→    attn_type_map = torch.zeros([q_ranges.shape[0]], device="cuda", dtype=torch.int32)
    69→    softmax_scale = None
    70→
    71→    ffa_handler = FFAHandler(
    72→        q_ranges=q_ranges,
    73→        k_ranges=k_ranges,
    74→        max_seqlen_q=max_seqlen_q,
    75→        max_seqlen_k=max_seqlen_k,
    76→        attn_type_map=attn_type_map,
    77→        softmax_scale=softmax_scale,
    78→    )
    79→    return ffa_handler
    80→
    81→
    82→def get_coords(
    83→    shape: list[int],
    84→    ref_feat_shape: list[int],
    85→    offset_thw: list[int] = [0, 0, 0],
    86→    device: torch.device = torch.device("cpu"),
    87→    dtype: torch.dtype = torch.float32,
    88→):
    89→    """
    90→    Generate feature-grid coordinates and corresponding original/reference size metadata.
    91→    Args:
    92→        feat_shape: [T, H, W] original feature-map shape
    93→        ref_feat_shape: [T_ref, H_ref, W_ref] reference feature-map shape
    94→        device: device for coordinate tensors
    95→    Returns:
    96→        coords: tensor shape (T*H*W, 9), containing (t, h, w, T, H, W, ref_T, ref_H, ref_W)
    97→    """
    98→    ori_t, ori_h, ori_w = shape
    99→    ref_t, ref_h, ref_w = ref_feat_shape
   100→


Let me check the prompt processing and scheduler:


     1→# Copyright (c) 2026 SandAI. All Rights Reserved.
     2→#
     3→# Licensed under the Apache License, Version 2.0 (the "License");
     4→# you may not use this file except in compliance with the License.
     5→# You may obtain a copy of the License at
     6→#
     7→#     http://www.apache.org/licenses/LICENSE-2.0
     8→#
     9→# Unless required by applicable law or agreed to in writing, software
    10→# distributed under the License is distributed on an "AS IS" BASIS,
    11→# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12→# See the License for the specific language governing permissions and
    13→# limitations under the License.
    14→
    15→from typing import Tuple
    16→
    17→import torch
    18→from torch.nn import functional as F
    19→
    20→from inference.model.t5_gemma import get_t5_gemma_embedding
    21→
    22→
    23→def pad_or_trim(tensor: torch.Tensor, target_size: int, dim: int, pad_value: float = 0.0) -> Tuple[torch.Tensor, int]:
    24→    """
    25→    Pads or trims a tensor along a specified dimension to reach a target size.
    26→
    27→    Args:
    28→        tensor (torch.Tensor): The input tensor to be processed.
    29→        target_size (int): The desired size for the specified dimension.
    30→        dim (int): The dimension along which to pad or trim.
    31→        pad_value (float, optional): The value used for padding. Defaults to 0.0.
    32→
    33→    Returns:
    34→        torch.Tensor: The resulting tensor with the target size in the specified dimension.
    35→    """
    36→    current_size = tensor.size(dim)
    37→    if current_size < target_size:
    38→        padding_amount = target_size - current_size
    39→        padding_tuple = [0] * (2 * tensor.dim())
    40→        padding_dim_index = tensor.dim() - 1 - dim
    41→        padding_tuple[2 * padding_dim_index + 1] = padding_amount
    42→        return F.pad(tensor, tuple(padding_tuple), "constant", pad_value), current_size
    43→
    44→    slicing = [slice(None)] * tensor.dim()
    45→    slicing[dim] = slice(0, target_size)
    46→    return tensor[tuple(slicing)], target_size
    47→
    48→
    49→def get_padded_t5_gemma_embedding(
    50→    prompt: str,
    51→    model_path: str,
    52→    device: str,
    53→    weight_dtype: torch.dtype,
    54→    target_length: int,
    55→) -> Tuple[torch.Tensor, int]:
    56→    txt_feat = get_t5_gemma_embedding(prompt, model_path, device, weight_dtype)
    57→    txt_feat, original_len = pad_or_trim(txt_feat, target_size=target_length, dim=1)
    58→    return txt_feat.to(torch.float32), original_len
    59→
    60→
    61→


     1→# Copyright (c) 2026 SandAI. All Rights Reserved.
     2→#
     3→# Licensed under the Apache License, Version 2.0 (the "License");
     4→# you may not use this file except in compliance with the License.
     5→# You may obtain a copy of the License at
     6→#
     7→#     http://www.apache.org/licenses/LICENSE-2.0
     8→#
     9→# Unless required by applicable law or agreed to in writing, software
    10→# distributed under the License is distributed on an "AS IS" BASIS,
    11→# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12→# See the License for the specific language governing permissions and
    13→# limitations under the License.
    14→
    15→# Copied from https://github.com/huggingface/diffusers/blob/v0.31.0/src/diffusers/schedulers/scheduling_unipc_multistep.py
    16→# Convert unipc for flow matching
    17→# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
    18→import math
    19→from typing import Any, List, Optional, Tuple, Union
    20→
    21→import numpy as np
    22→import torch
    23→from diffusers.configuration_utils import ConfigMixin, register_to_config
    24→from diffusers.schedulers.scheduling_utils import KarrasDiffusionSchedulers, SchedulerMixin, SchedulerOutput
    25→from diffusers.utils import deprecate
    26→from diffusers.utils.torch_utils import randn_tensor
    27→
    28→
    29→class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin):
    30→    """
    31→    `UniPCMultistepScheduler` is a training-free framework designed for the fast sampling of diffusion models.
    32→
    33→    This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic
    34→    methods the library implements for all schedulers such as loading and saving.
    35→
    36→    Args:
    37→        num_train_timesteps (`int`, defaults to 1000):
    38→            The number of diffusion steps to train the model.
    39→        solver_order (`int`, default `2`):
    40→            The UniPC order which can be any positive integer. The effective order of accuracy is `solver_order + 1`
    41→            due to the UniC. It is recommended to use `solver_order=2` for guided sampling, and `solver_order=3` for
    42→            unconditional sampling.
    43→        prediction_type (`str`, defaults to "flow_prediction"):
    44→            Prediction type of the scheduler function; must be `flow_prediction` for this scheduler, which predicts
    45→            the flow of the diffusion process.
    46→        thresholding (`bool`, defaults to `False`):
    47→            Whether to use the "dynamic thresholding" method. This is unsuitable for latent-space diffusion models such
    48→            as Stable Diffusion.
    49→        dynamic_thresholding_ratio (`float`, defaults to 0.995):
    50→            The ratio for the dynamic thresholding method. Valid only when `thresholding=True`.
    51→        sample_max_value (`float`, defaults to 1.0):
    52→            The threshold value for dynamic thresholding. Valid only when `thresholding=True` and `predict_x0=True`.
    53→        predict_x0 (`bool`, defaults to `True`):
    54→            Whether to use the updating algorithm on the predicted x0.
    55→        solver_type (`str`, default `bh2`):
    56→            Solver type for UniPC. It is recommended to use `bh1` for unconditional sampling when steps < 10, and `bh2`
    57→            otherwise.
    58→        lower_order_final (`bool`, default `True`):
    59→            Whether to use lower-order solvers in the final steps. Only valid for < 15 inference steps. This can
    60→            stabilize the sampling of DPMSolver for steps < 15, especially for steps <= 10.
    61→        disable_corrector (`list`, default `[]`):
    62→            Decides which step to disable the corrector to mitigate the misalignment between `epsilon_theta(x_t, c)`
    63→            and `epsilon_theta(x_t^c, c)` which can influence convergence for a large guidance scale. Corrector is
    64→            usually disabled during the first few steps.
    65→        solver_p (`SchedulerMixin`, default `None`):
    66→            Any other scheduler that if specified, the algorithm becomes `solver_p + UniC`.
    67→        use_karras_sigmas (`bool`, *optional*, defaults to `False`):
    68→            Whether to use Karras sigmas for step sizes in the noise schedule during the sampling process. If `True`,
    69→            the sigmas are determined according to a sequence of noise levels {σi}.
    70→        use_exponential_sigmas (`bool`, *optional*, defaults to `False`):
    71→            Whether to use exponential sigmas for step sizes in the noise schedule during the sampling process.
    72→        timestep_spacing (`str`, defaults to `"linspace"`):
    73→            The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and
    74→            Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information.
    75→        steps_offset (`int`, defaults to 0):
    76→            An offset added to the inference steps, as required by some model families.
    77→        final_sigmas_type (`str`, defaults to `"zero"`):
    78→            The final `sigma` value for the noise schedule during the sampling process. If `"sigma_min"`, the final
    79→            sigma is the same as the last sigma in the training schedule. If `zero`, the final sigma is set to 0.
    80→    """
    81→
    82→    _compatibles = [e.name for e in KarrasDiffusionSchedulers]
    83→    order = 1
    84→
    85→    @register_to_config
    86→    def __init__(
    87→        self,
    88→        num_train_timesteps: int = 1000,
    89→        solver_order: int = 2,
    90→        prediction_type: str = "flow_prediction",
    91→        shift: float = 1.0,
    92→        use_dynamic_shifting=False,
    93→        thresholding: bool = False,
    94→        dynamic_thresholding_ratio: float = 0.995,
    95→        sample_max_value: float = 1.0,
    96→        predict_x0: bool = True,
    97→        solver_type: str = "bh2",
    98→        lower_order_final: bool = True,
    99→        disable_corrector: List[int] = [],
   100→        solver_p: SchedulerMixin = None,


Now let me look at the VAE models and audio model loading:


     1→import gc
     2→
     3→import torch
     4→
     5→from .vae2_2_module import Wan2_2_VAE
     6→
     7→
     8→def get_vae2_2(model_path, device="cuda", weight_dtype=torch.float32) -> Wan2_2_VAE:
     9→    vae = Wan2_2_VAE(vae_pth=model_path).to(device).to(weight_dtype)
    10→    vae.vae.requires_grad_(False)
    11→    vae.vae.eval()
    12→    gc.collect()
    13→    torch.cuda.empty_cache()
    14→    return vae
    15→
    16→
    17→__all__ = ["Wan2_2_VAE", "get_vae2_2"]
    18→


     1→# Copyright (c) 2026 SandAI. All Rights Reserved.
     2→#
     3→# Licensed under the Apache License, Version 2.0 (the "License");
     4→# you may not use this file except in compliance with the License.
     5→# You may obtain a copy of the License at
     6→#
     7→#     http://www.apache.org/licenses/LICENSE-2.0
     8→#
     9→# Unless required by applicable law or agreed to in writing, software
    10→# distributed under the License is distributed on an "AS IS" BASIS,
    11→# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12→# See the License for the specific language governing permissions and
    13→# limitations under the License.
    14→
    15→import json
    16→import os
    17→from pathlib import Path
    18→
    19→import torch
    20→from safetensors.torch import load_file
    21→
    22→# Set env vars for local T5 loading
    23→os.environ["TRANSFORMERS_OFFLINE"] = "1"
    24→os.environ["HF_HUB_OFFLINE"] = "1"
    25→
    26→from .sa_audio_module import create_model_from_config
    27→
    28→from inference.utils import print_rank_0
    29→
    30→
    31→class SAAudioFeatureExtractor:
    32→    """Stable Audio Feature Extractor that loads model once and reuses it."""
    33→
    34→    def __init__(self, device, model_path):
    35→        """Initialize the extractor with model loading."""
    36→        self.device = device
    37→        self.vae_model, self.sample_rate = self._get_vae_only(model_path)
    38→        # self.vae_model.to(self.device).to(torch.bfloat16)
    39→        self.resampler = None  # Will be initialized when needed
    40→
    41→    def _get_vae_only(self, model_path):
    42→        """Load VAE only, skip T5 and diffusion model."""
    43→        if isinstance(model_path, str) and Path(model_path).is_dir():
    44→            try:
    45→                # Read full config
    46→                model_config_path = os.path.join(model_path, "model_config.json")
    47→                with open(model_config_path) as f:
    48→                    full_config = json.load(f)
    49→
    50→                vae_config = full_config["model"]["pretransform"]["config"]
    51→                sample_rate = full_config["sample_rate"]
    52→
    53→                # Rebuild config structure expected by create_autoencoder_from_config
    54→                autoencoder_config = {
    55→                    "model_type": "autoencoder",
    56→                    "sample_rate": sample_rate,  # sample_rate is required
    57→                    "model": vae_config,  # create_autoencoder_from_config expects key "model"
    58→                }
    59→
    60→                vae_model = create_model_from_config(autoencoder_config)
    61→                # Load weights
    62→                weights_path = Path(model_path) / "model.safetensors"
    63→
    64→                if not weights_path.exists():
    65→                    raise FileNotFoundError(f"Weight file does not exist: {weights_path}")
    66→
    67→                # Load full state dict
    68→                full_state_dict = load_file(weights_path, device=str(self.device))
    69→
    70→                # Filter VAE-related weights (prefix: pretransform.model)
    71→                vae_state_dict = {}
    72→                for key, value in full_state_dict.items():
    73→                    if key.startswith("pretransform.model."):
    74→                        vae_key = key[len("pretransform.model.") :]
    75→                        vae_state_dict[vae_key] = value
    76→
    77→                # Check expected model keys
    78→                model_keys = set(vae_model.state_dict().keys())
    79→                vae_keys = set(vae_state_dict.keys())
    80→


Let me check the sequence schema for batching/concurrency info:


     1→# Copyright (c) 2026 SandAI. All Rights Reserved.
     2→#
     3→# Licensed under the Apache License, Version 2.0 (the "License");
     4→# you may not use this file except in compliance with the License.
     5→# You may obtain a copy of the License at
     6→#
     7→#     http://www.apache.org/licenses/LICENSE-2.0
     8→#
     9→# Unless required by applicable law or agreed to in writing, software
    10→# distributed under the License is distributed on an "AS IS" BASIS,
    11→# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12→# See the License for the specific language governing permissions and
    13→# limitations under the License.
    14→
    15→from dataclasses import dataclass
    16→from enum import IntEnum
    17→
    18→import torch
    19→
    20→
    21→class Modality(IntEnum):
    22→    VIDEO = 0
    23→    AUDIO = 1
    24→    TEXT = 2
    25→
    26→
    27→@dataclass
    28→class VarlenHandler:
    29→    cu_seqlens_q: torch.Tensor
    30→    cu_seqlens_k: torch.Tensor
    31→    max_seqlen_q: int
    32→    max_seqlen_k: int
    33→
    34→


     1→# Copyright (c) 2026 SandAI. All Rights Reserved.
     2→#
     3→# Licensed under the Apache License, Version 2.0 (the "License");
     4→# you may not use this file except in compliance with the License.
     5→# You may obtain a copy of the License at
     6→#
     7→#     http://www.apache.org/licenses/LICENSE-2.0
     8→#
     9→# Unless required by applicable law or agreed to in writing, software
    10→# distributed under the License is distributed on an "AS IS" BASIS,
    11→# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12→# See the License for the specific language governing permissions and
    13→# limitations under the License.
    14→
    15→from typing import Any, Callable, Dict, Tuple
    16→
    17→import torch
    18→
    19→
    20→class CPUOffloadWrapper:
    21→    def __init__(self, model: Any, is_cpu_offload: bool = False, is_running_on_gpu: bool = True):
    22→        object.__setattr__(self, "model", model)
    23→        object.__setattr__(self, "is_cpu_offload", is_cpu_offload)
    24→        object.__setattr__(self, "is_running_on_gpu", is_running_on_gpu)
    25→
    26→        cpu_device = torch.device("cpu")
    27→        cuda_device = torch.device("cuda")
    28→        object.__setattr__(self, "cpu_device", cpu_device)
    29→        object.__setattr__(self, "cuda_device", cuda_device)
    30→
    31→        # Initialize placement location
    32→        if is_cpu_offload:
    33→            self.model.to(cpu_device)
    34→        else:
    35→            self.model.to(cuda_device)
    36→
    37→        # Whitelist non-compute methods that shouldn't trigger device hops (pass-through only; no device switch)
    38→        object.__setattr__(
    39→            self,
    40→            "_non_compute_methods",
    41→            {
    42→                "to",
    43→                "cpu",
    44→                "cuda",
    45→                "eval",
    46→                "train",
    47→                "state_dict",
    48→                "load_state_dict",
    49→                "parameters",
    50→                "named_parameters",
    51→                "buffers",
    52→                "named_buffers",
    53→                "modules",
    54→                "named_modules",
    55→                "children",
    56→                "named_children",
    57→                "register_forward_hook",
    58→                "register_forward_pre_hook",
    59→                "register_full_backward_hook",
    60→                "zero_grad",
    61→                "share_memory",
    62→                "half",
    63→                "float",
    64→                "bfloat16",
    65→            },
    66→        )
    67→
    68→    # Get current primary device (for external reads)
    69→    @property
    70→    def device(self) -> torch.device:
    71→        if isinstance(self.model, torch.nn.Module):
    72→            return next(self.model.parameters()).device
    73→        else:
    74→            for k, v in self.model.__dict__.items():
    75→                if isinstance(v, torch.Tensor):
    76→                    return v.device
    77→                elif isinstance(v, torch.nn.Module):
    78→                    return next(v.parameters()).device
    79→            return self.cuda_device
    80→
    81→    def _backup_cpu_state(self) -> Tuple[Dict[str, torch.Tensor], Dict[str, torch.Tensor], Dict[str, Any]]:
    82→        # Backup module parameters and buffers
    83→        module_param_backup = {}
    84→        module_buffer_backup = {}
    85→        other_backup = {}
    86→
    87→        def save_module_state(mod: torch.nn.Module, prefix: str):
    88→            for name, param in mod.named_parameters():
    89→                if param is not None:
    90→                    full_key = prefix + name
    91→                    module_param_backup[full_key] = param.data
    92→            for name, buffer in mod.named_buffers():
    93→                if buffer is not None:
    94→                    full_key = prefix + name
    95→                    module_buffer_backup[full_key] = buffer.data
    96→
    97→        if isinstance(self.model, torch.nn.Module):
    98→            save_module_state(self.model, "")
    99→        else:
   100→            for name, attr_val in self.model.__dict__.items():
   101→                if isinstance(attr_val, torch.nn.Module):
   102→                    save_module_state(attr_val, name + ".")
   103→                elif isinstance(attr_val, torch.Tensor):
   104→                    other_backup[name] = attr_val
   105→
   106→        return module_param_backup, module_buffer_backup, other_backup
   107→
   108→    def _restore_cpu_state(self, backups: Tuple[Dict[str, torch.Tensor], Dict[str, torch.Tensor], Dict[str, Any]]):
   109→        # Restore module parameters and buffers
   110→        module_param_backup, module_buffer_backup, other_backup = backups
   111→
   112→        def restore_module_state(mod: torch.nn.Module, prefix: str):
   113→            for name, param in mod.named_parameters():
   114→                full_key = prefix + name
   115→                if full_key in module_param_backup:
   116→                    param.data = module_param_backup[full_key]
   117→
   118→            for name, buffer in mod.named_buffers():
   119→                full_key = prefix + name
   120→                if full_key in module_buffer_backup:
   121→                    buffer.data = module_buffer_backup[full_key]
   122→
   123→        if isinstance(self.model, torch.nn.Module):
   124→            restore_module_state(self.model, "")
   125→        else:
   126→            for name, attr_val in self.model.__dict__.items():
   127→                if isinstance(attr_val, torch.nn.Module):
   128→                    restore_module_state(attr_val, name + ".")
   129→
   130→        if not isinstance(self.model, torch.nn.Module):
   131→            for name, val in other_backup.items():
   132→                setattr(self.model, name, val)
   133→
   134→    # Unified on/offload executor
   135→    def _run_with_optional_offload(self, func: Callable[..., Any], *args, **kwargs):
   136→        if self.is_cpu_offload and self.is_running_on_gpu:
   137→            backups = self._backup_cpu_state()
   138→            self.model.to(self.cuda_device)
   139→            try:
   140→                return func(*args, **kwargs)
   141→            finally:
   142→                if torch.cuda.is_available():
   143→                    torch.cuda.synchronize()
   144→                self._restore_cpu_state(backups)
   145→        else:
   146→            # Make sure model and args are on the same device
   147→            args = [
   148→                arg.to(self.device) if isinstance(arg, torch.Tensor) and arg.device != self.device else arg for arg in args
   149→            ]
   150→            kwargs = {
   151→                k: v.to(self.device) if isinstance(v, torch.Tensor) and v.device != self.device else v
   152→                for k, v in kwargs.items()
   153→            }
   154→            return func(*args, **kwargs)
   155→
   156→    # Direct call (equivalent to forward)
   157→    def __call__(self, *args, **kwargs):
   158→        return self._run_with_optional_offload(self.model.__call__, *args, **kwargs)
   159→
   160→    # Explicit forward; some code calls model.forward(...)
   161→    def forward(self, *args, **kwargs):
   162→        return self._run_with_optional_offload(self.model.forward, *args, **kwargs)
   163→
   164→    # Key: passthrough all attrs/methods. For callables, wrap with on/offload; for non-compute methods, pass-through only with no device switch.
   165→    def __getattr__(self, name: str):
   166→        # Fetch attribute from the wrapped model first
   167→        attr = getattr(self.model, name)
   168→
   169→        # Wrap methods (except in whitelist)
   170→        if callable(attr) and name not in self._non_compute_methods:
   171→
   172→            def _wrapped(*args, **kwargs):
   173→                return self._run_with_optional_offload(attr, *args, **kwargs)
   174→
   175→            return _wrapped
   176→
   177→        return attr
   178→
   179→    def __dir__(self):
   180→        return sorted(set(list(super().__dir__()) + dir(self.model)))
   181→
   182→    def __setattr__(self, name: str, value: Any):
   183→        raise AttributeError("CPUOffloadWrapper is immutable")
   184→
   185→    def __repr__(self) -> str:
   186→        return f"CPUOffloadWrapper(is_cpu_offload={self.is_cpu_offload}, is_running_on_gpu={self.is_running_on_gpu}, model={repr(self.model)})"
   187→


Now let me check the distributed utilities for parallelism:


     1→# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
     2→# Copyright (c) 2026 SandAI. All Rights Reserved.
     3→#
     4→# Licensed under the Apache License, Version 2.0 (the "License");
     5→# you may not use this file except in compliance with the License.
     6→# You may obtain a copy of the License at
     7→#
     8→#     http://www.apache.org/licenses/LICENSE-2.0
     9→#
    10→# Unless required by applicable law or agreed to in writing, software
    11→# distributed under the License is distributed on an "AS IS" BASIS,
    12→# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13→# See the License for the specific language governing permissions and
    14→# limitations under the License.
    15→
    16→"""Model and data parallel groups."""
    17→
    18→import warnings
    19→from datetime import timedelta
    20→from typing import List, Optional
    21→
    22→import torch
    23→
    24→# Intra-layer model parallel group that the current rank belongs to.
    25→_TENSOR_MODEL_PARALLEL_GROUP = None
    26→# Tensor parallel group information with context parallel combined.
    27→_TENSOR_MODEL_PARALLEL_GROUP_WITH_CP = None
    28→_TENSOR_MODEL_PARALLEL_GLOBAL_RANKS_WITH_CP = None
    29→# Inter-layer model parallel group that the current rank belongs to.
    30→_PIPELINE_MODEL_PARALLEL_GROUP = None
    31→# Model parallel group (both intra- and pipeline) that the current rank belongs to.
    32→_MODEL_PARALLEL_GROUP = None
    33→# Data parallel group that the current rank belongs to.
    34→_DATA_PARALLEL_GROUP = None
    35→# tensor model parallel group and data parallel group combined
    36→# used for fp8 and moe training
    37→_TENSOR_AND_DATA_PARALLEL_GROUP = None
    38→
    39→# A list of global ranks for each pipeline group to ease calculation of the source
    40→# rank when broadcasting from the first or last pipeline stage.
    41→_PIPELINE_GLOBAL_RANKS = None
    42→
    43→# A list of global ranks for each data parallel group to ease calculation of the source
    44→# rank when broadcasting weights from src to all other data parallel ranks
    45→_DATA_PARALLEL_GLOBAL_RANKS = None
    46→
    47→# A list of global ranks for each tensor model parallel group to ease calculation of
    48→# the first local rank in the tensor model parallel group
    49→_TENSOR_MODEL_PARALLEL_GLOBAL_RANKS = None
    50→
    51→# Context parallel group that the current rank belongs to
    52→_CONTEXT_PARALLEL_GROUP = None
    53→# A list of global ranks for each context parallel group to ease calculation of the
    54→# destination rank when exchanging KV/dKV between context parallel_ranks
    55→_CONTEXT_PARALLEL_GLOBAL_RANKS = None
    56→
    57→_CONTEXT_PARALLEL_EXTRA_GROUP = None
    58→
    59→# Data parallel group information with context parallel combined.
    60→_DATA_PARALLEL_GROUP_WITH_CP = None
    61→_DATA_PARALLEL_GLOBAL_RANKS_WITH_CP = None
    62→
    63→# combined parallel group of TP, DP, and CP used for fp8
    64→_TENSOR_AND_DATA_PARALLEL_GROUP_WITH_CP = None
    65→
    66→
    67→def _get_nccl_options(pg_name, nccl_comm_cfgs):
    68→    """Set the NCCL process group options.
    69→
    70→    Args:
    71→        pg_name (str): process group name
    72→        nccl_comm_cfgs (dict): nccl communicator configurations
    73→
    74→    When an option (e.g., max_ctas) is not found in the config, use the NCCL default setting.
    75→    """
    76→    if pg_name in nccl_comm_cfgs:
    77→        nccl_options = torch.distributed.ProcessGroupNCCL.Options()
    78→        nccl_options.config.cga_cluster_size = nccl_comm_cfgs[pg_name].get("cga_cluster_size", 4)
    79→        nccl_options.config.max_ctas = nccl_comm_cfgs[pg_name].get("max_ctas", 32)
    80→        nccl_options.config.min_ctas = nccl_comm_cfgs[pg_name].get("min_ctas", 1)
    81→        return nccl_options
    82→    else:
    83→        return None
    84→
    85→
    86→def generate_masked_orthogonal_rank_groups(world_size: int, parallel_size: List[int], mask: List[bool]) -> List[List[int]]:
    87→    r"""Generate orthogonal parallel groups based on the parallel size and mask.
    88→
    89→    Arguments:
    90→        world_size (int): world size
    91→
    92→        parallel_size (List[int]):
    93→            The parallel size of each orthogonal parallel type. For example, if
    94→            tensor_parallel_size = 2, pipeline_model_parallel_group = 3, data_parallel_size = 4,
    95→            and the parallel mapping order is tp-pp-dp, then the parallel_size = [2, 3, 4].
    96→
    97→        mask (List[bool]):
    98→            The mask controls which parallel methods the generated groups represent. If mask[i] is
    99→            True, it means the generated group contains the i-th parallelism method. For example,
   100→            if parallel_size = [tp_size, pp_size, dp_size], and mask = [True, False , True], then
   101→            the generated group is the `tp-dp` group, if the mask = [False, True, False], then the
   102→            generated group is the `pp` group.
   103→
   104→    Algorithm:
   105→        For orthogonal parallelism, such as tp/dp/pp/cp, the global_rank and
   106→        local_rank satisfy the following equation:
   107→            global_rank = tp_rank + dp_rank * tp_size + pp_rank * tp_size * dp_size (1)
   108→                tp_rank \in [0, tp_size)
   109→                dp_rank \in [0, dp_size)
   110→                pp_rank \in [0, pp_size)
   111→
   112→        If we want to get the `dp_group` (tp_size * pp_size groups of dp_size ranks each.
   113→        For example,  if the gpu size is 8 and order is 'tp-pp-dp', size is '2-2-2', and the
   114→        dp_group here is [[0, 4], [1, 5], [2, 6], [3, 7]].)
   115→        The tp_rank and pp_rank will be combined to form the `dp_group_index`.
   116→            dp_group_index = tp_rank + pp_rank * tp_size (2)
   117→
   118→        So, Given that tp_rank and pp_rank satisfy equation (2), and dp_rank in
   119→        range(0, dp_size), the ranks in dp_group[dp_group_index] satisfies the
   120→        equation (1).
   121→
   122→        This function solve this math problem.
   123→
   124→    For example, if the parallel_size = [tp_size, dp_size, pp_size] = [2, 3, 4],
   125→    and the mask = [False, True, False]. Then,
   126→        dp_group_index(0) = tp_rank(0) + pp_rank(0) * 2
   127→        dp_group_index(1) = tp_rank(1) + pp_rank(0) * 2
   128→        ...
   129→        dp_group_index(7) = tp_rank(1) + pp_rank(3) * 2
   130→
   131→        dp_group[0] = 0 + range(0, 3) * 2 + 0 = [0, 2, 4]
   132→        dp_group[1] = 1 + range(0, 3) * 2 + 0 = [1, 3, 5]
   133→        ...
   134→        dp_group[7] = 1 + range(0, 3) * 2 + 3 * 2 * 3 = [19, 21, 23]
   135→    """
   136→
   137→    def prefix_product(a: List[int], init=1) -> List[int]:
   138→        r = [init]
   139→        for v in a:
   140→            init = init * v
   141→            r.append(init)
   142→        return r
   143→
   144→    def inner_product(a: List[int], b: List[int]) -> int:
   145→        return sum([x * y for x, y in zip(a, b)])
   146→
   147→    def decompose(index, shape, stride=None):
   148→        """
   149→        This function solve the math problem below:
   150→            There is an equation:


Now let me get a comprehensive overview by examining how the full model is constructed. Let me check the DiT model class:


Bash

grep -n "class DiTModel" /home/ubuntu/daVinci-MagiHuman/inference/model/dit/dit_module.py


870:class DiTModel(torch.nn.Module):


   870→class DiTModel(torch.nn.Module):
   871→    config: TransformerConfig
   872→
   873→    def __init__(self, model_config: Any):
   874→        super().__init__()
   875→        self.config = TransformerConfig(
   876→            hidden_size=model_config.hidden_size,
   877→            video_in_channels=model_config.video_in_channels,
   878→            audio_in_channels=model_config.audio_in_channels,
   879→            text_in_channels=model_config.text_in_channels,
   880→            params_dtype=model_config.params_dtype,
   881→            post_process_dtype=torch.float32,
   882→        )
   883→        adapter_config = AdapterConfig(
   884→            hidden_size=model_config.hidden_size,
   885→            num_attention_heads=model_config.num_heads_q,
   886→            text_in_channels=model_config.text_in_channels,
   887→            video_in_channels=model_config.video_in_channels,
   888→            audio_in_channels=model_config.audio_in_channels,
   889→            params_dtype=torch.float32,
   890→        )
   891→        self.adapter: Adapter = Adapter(adapter_config)
   892→        self.block: TransformerBlock = TransformerBlock(model_config=model_config)
   893→        self.final_norm_video = MultiModalityRMSNorm(self.config.hidden_size)
   894→        self.final_norm_audio = MultiModalityRMSNorm(self.config.hidden_size)
   895→        self.final_linear_video = nn.Linear(
   896→            self.config.hidden_size, self.config.video_in_channels, bias=False, dtype=torch.float32
   897→        )
   898→        self.final_linear_audio = nn.Linear(
   899→            self.config.hidden_size, self.config.audio_in_channels, bias=False, dtype=torch.float32
   900→        )
   901→
   902→    def forward(
   903→        self,
   904→        x: torch.Tensor,
   905→        coords_mapping: torch.Tensor,
   906→        modality_mapping: torch.Tensor,
   907→        varlen_handler: VarlenHandler,
   908→        local_attn_handler: FFAHandler,
   909→    ):
   910→        x = ulysses_scheduler().dispatch(x)
   911→        coords_mapping = ulysses_scheduler().dispatch(coords_mapping)
   912→        modality_mapping = ulysses_scheduler().dispatch(modality_mapping)
   913→        cp_split_sizes = ulysses_scheduler().cp_split_sizes
   914→
   915→        modality_dispatcher = ModalityDispatcher(modality_mapping, 3)
   916→        permute_mapping, inv_permute_mapping = modality_dispatcher.permute_mapping, modality_dispatcher.inv_permute_mapping
   917→        video_mask = modality_mapping == Modality.VIDEO
   918→        audio_mask = modality_mapping == Modality.AUDIO
   919→        text_mask = modality_mapping == Modality.TEXT
   920→
   921→        x, rope = self.adapter(x, coords_mapping, video_mask, audio_mask, text_mask)
   922→        x = x.to(self.config.params_dtype)
   923→        x = ModalityDispatcher.permute(x, permute_mapping)
   924→        x = self.block(
   925→            x,
   926→            rope,
   927→            permute_mapping=permute_mapping,
   928→            inv_permute_mapping=inv_permute_mapping,
   929→            varlen_handler=varlen_handler,
   930→            local_attn_handler=local_attn_handler,
   931→            modality_dispatcher=modality_dispatcher,
   932→            cp_split_sizes=cp_split_sizes,
   933→        )
   934→        x = ModalityDispatcher.inv_permute(x, inv_permute_mapping)
   935→
   936→        x_video = x[video_mask].to(self.final_norm_video.weight.dtype)
   937→        x_video = self.final_norm_video(x_video)
   938→        x_video = self.final_linear_video(x_video)
   939→
   940→        x_audio = x[audio_mask].to(self.final_norm_audio.weight.dtype)
   941→        x_audio = self.final_norm_audio(x_audio)
   942→        x_audio = self.final_linear_audio(x_audio)
   943→
   944→        x_out = torch.zeros(
   945→            x.shape[0], max(self.config.video_in_channels, self.config.audio_in_channels), device=x.device, dtype=x.dtype
   946→        )
   947→        x_out[video_mask, : self.config.video_in_channels] = x_video
   948→        x_out[audio_mask, : self.config.audio_in_channels] = x_audio
   949→        x_out = ulysses_scheduler().undispatch(x_out)
   950→        return x_out
   951→


Perfect! Now let me create a comprehensive summary document. Let me check one more key file - the TransformerBlock to understand the full architecture:


grep -n "class TransformerBlock\|class TransformerConfig" /home/ubuntu/daVinci-MagiHuman/inference/model/dit/dit_module.py


828:class TransformerBlock(torch.nn.Module):
861:class TransformerConfig:


   828→class TransformerBlock(torch.nn.Module):
   829→    def __init__(self, model_config: Any):
   830→        super().__init__()
   831→        self.layers: list[TransFormerLayer] = nn.ModuleList()
   832→        for layer_idx in range(model_config.num_layers):
   833→            self.layers.append(TransFormerLayer(model_config, layer_idx))
   834→
   835→    def forward(
   836→        self,
   837→        x: torch.Tensor,
   838→        rope: torch.Tensor,
   839→        permute_mapping: torch.Tensor,
   840→        inv_permute_mapping: torch.Tensor,
   841→        varlen_handler: VarlenHandler,
   842→        local_attn_handler: FFAHandler,
   843→        modality_dispatcher: ModalityDispatcher,
   844→        cp_split_sizes: List[int],
   845→    ) -> torch.Tensor:
   846→        for _, layer in enumerate(self.layers):
   847→            x = layer(
   848→                x,
   849→                rope,
   850→                permute_mapping,
   851→                inv_permute_mapping,
   852→                varlen_handler,
   853→                local_attn_handler,
   854→                modality_dispatcher,
   855→                cp_split_sizes,
   856→            )
   857→        return x
   858→
   859→
   860→@dataclass
   861→class TransformerConfig:
   862→    hidden_size: int
   863→    video_in_channels: int
   864→    audio_in_channels: int
   865→    text_in_channels: int
   866→    params_dtype: torch.dtype
   867→    post_process_dtype: torch.dtype
   868→
   869→
   870→class DiTModel(torch.nn.Module):
   871→    config: TransformerConfig
   872→
   873→    def __init__(self, model_config: Any):
   874→        super().__init__()
   875→        self.config = TransformerConfig(
   876→            hidden_size=model_config.hidden_size,
   877→            video_in_channels=model_config.video_in_channels,


Perfect! Now I have a comprehensive understanding. Let me create a detailed summary report.

## Comprehensive Analysis: daVinci-MagiHuman Inference Pipeline

### 1. FULL FILE STRUCTURE

The inference directory is organized as follows:

```
/home/ubuntu/daVinci-MagiHuman/inference/
├── common/                          # Configuration and utility schemas
│   ├── arch.py                      # GPU architecture detection (Hopper/SM90 vs SM80)
│   ├── config.py                    # Configuration system (Pydantic-based)
│   ├── cpu_offload_wrapper.py       # CPU/GPU memory management wrapper
│   ├── sequence_schema.py           # Modality enum and VarlenHandler
│   └── __init__.py
├── infra/                           # Infrastructure and initialization
│   ├── checkpoint/
│   │   └── load_model_checkpoint.py # Parallel safetensors loading
│   ├── distributed/
│   │   ├── init_dist_env.py         # torch.distributed initialization
│   │   ├── parallel_state.py        # TP/PP/CP/DP group management
│   │   ├── utils.py
│   │   └── __init__.py
│   ├── parallelism/
│   │   ├── ulysses_scheduler.py     # Sequence parallel scheduling
│   │   ├── gather_scatter_primitive.py
│   │   ├── all_to_all_primitive.py
│   │   └── __init__.py
│   └── __init__.py
├── model/                           # Model components
│   ├── dit/                         # Diffusion Transformer
│   │   ├── dit_module.py            # Attention, MLP, DiT architecture
│   │   ├── dit_model.py             # Model instantiation and loading
│   │   └── __init__.py
│   ├── sa_audio/                    # Stable Audio VAE
│   │   ├── sa_audio_module.py
│   │   ├── sa_audio_model.py
│   │   └── __init__.py
│   ├── t5_gemma/                    # T5-Gemma text encoder
│   │   ├── t5_gemma_model.py
│   │   └── __init__.py
│   ├── turbo_vaed/                  # Turbo VAE Decoder
│   │   ├── turbo_vaed_model.py
│   │   ├── turbo_vaed_module.py
│   │   └── __init__.py
│   └── vae2_2/                      # Wan 2.2 Video VAE
│       ├── vae2_2_model.py
│       ├── vae2_2_module.py
│       └── __init__.py
├── pipeline/                        # Main inference pipeline
│   ├── entry.py                     # CLI entry point
│   ├── pipeline.py                  # MagiPipeline facade
│   ├── video_generate.py            # MagiEvaluator (core inference loop)
│   ├── video_process.py             # Video/audio processing utilities
│   ├── prompt_process.py            # Text encoding and padding
│   ├── scheduler_unipc.py           # UniPC diffusion scheduler
│   ├── data_proxy.py                # Input/output marshalling
│   └── __init__.py
└── utils/                           # Utilities
    ├── logger.py
    ├── seed.py
    ├── env.py
    ├── timer.py
    ├── math.py
    └── __init__.py
```

---

### 2. INFERENCE PIPELINE FLOW (entry.py to video output)

#### **Stage 1: Entry Point (`entry.py`)**

```
entry.py (CLI)
├─ parse_arguments()
│  └─ Extract: prompt, image_path, audio_path, save_path_prefix, seed, dimensions
├─ initialize_infra()
│  ├─ Initialize torch.distributed backend (NCCL/Gloo)
│  ├─ Setup TP/PP/CP/DP parallelism groups
│  └─ Set random seed (per rank)
├─ parse_config()
│  └─ Load config from JSON or CLI args (Pydantic-based)
├─ get_dit(config.arch_config, config.engine_config)
│  └─ Load main diffusion transformer model
└─ MagiPipeline(model, config.evaluation_config)
   └─ run_offline(prompt, image, audio, save_path_prefix, **kwargs)
```

#### **Stage 2: Configuration (`config.py`)**

The system uses Pydantic BaseSettings with layered config sources:

**Key Classes:**
- **`EngineConfig`**: Parallelism settings (tp_size, pp_size, cp_size, dp_size), seed, checkpoint path
- **`ModelConfig`**: Model architecture (num_layers=40, hidden_size=5120, head_dim=128, num_query_groups=8)
- **`EvaluationConfig`**: Inference parameters (num_inference_steps=32, sr_num_inference_steps=5, fps=25)
- **`DataProxyConfig`**: Input coordinate system (patch_size, frame_receptive_field)

#### **Stage 3: Model Loading (`dit_model.py` → `load_model_checkpoint.py`)**

```python
get_dit(model_config, engine_config)
├─ DiTModel(model_config)  # Build architecture
├─ load_model_checkpoint(model, engine_config)
│  ├─ Load sharded safetensors in parallel via ThreadPoolExecutor
│  └─ Optional zstd decompression (if .zst files exist)
├─ model.cuda(torch.cuda.current_device())
└─ model.eval()
```

**Checkpoint Loading Features:**
- Parallel safetensors loading with progress bar
- Support for zstd-compressed shards
- Graceful handling of missing/unexpected keys

#### **Stage 4: Core Inference Loop (`video_generate.py` → `MagiEvaluator.evaluate()`)**

```
evaluate(prompt, image, audio_path, seconds, br_width, br_height, ...)
│
├─ Step 1: Prepare Latents
│  ├─ Image encoding (if provided)
│  │  ├─ Load/resize image via resizecrop()
│  │  ├─ Preprocess with VideoProcessor
│  │  └─ VAE encode to latent space
│  ├─ Audio encoding (if audio_path provided)
│  │  ├─ load_audio_and_encode(audio_vae, audio_path)
│  │  └─ Output shape: (1, num_frames, 64)
│  ├─ Text embedding
│  │  ├─ get_padded_t5_gemma_embedding(prompt, model_path, device, dtype, target_length)
│  │  └─ Pad/trim to target_length (default 640)
│  └─ Initialize random latents:
│      ├─ latent_video: (1, z_dim=48, latent_frames, latent_h, latent_w)
│      └─ latent_audio: (1, num_frames, 64)
│
├─ Step 2: Basic Resolution (BR) Inference Loop
│  ├─ evaluate_with_latent(..., use_sr_model=False, num_inference_steps=32)
│  │  ├─ FlowUniPCMultistepScheduler.set_timesteps()
│  │  ├─ for each timestep t in timesteps:
│  │  │  ├─ [Optional] Lock first frame to image latent
│  │  │  ├─ Forward pass (conditional):
│  │  │  │  ├─ EvalInput(x_t, audio_x_t, txt_feat, ...)
│  │  │  │  ├─ forward() → (v_cond_video, v_cond_audio)
│  │  │  ├─ Classifier-Free Guidance (if cfg_number=2):
│  │  │  │  ├─ Forward unconditioned pass → (v_uncond_video, v_uncond_audio)
│  │  │  │  ├─ v_cfg_video = v_uncond + guidance_scale * (v_cond - v_uncond)
│  │  │  │  └─ v_cfg_audio = v_uncond + audio_guidance_scale * (v_cond - v_uncond)
│  │  │  └─ Scheduler step (UniPC or SDE):
│  │  │     ├─ schedule_latent_step(..., scheduler, ...)
│  │  │     └─ Update latent_video, latent_audio
│  │  └─ Return latent_video_br, latent_audio_br
│  └─ [CPU Offload Optional] Move main model to CPU if SR2_1080 mode
│
├─ Step 3: Super Resolution (SR) Inference [if sr_width & sr_height provided]
│  ├─ Interpolate BR latents to SR dimensions
│  ├─ Add noise to interpolated latents (controlled by noise_value)
│  ├─ [CPU Offload Optional] Load SR model to GPU
│  ├─ evaluate_with_latent(..., use_sr_model=True, num_inference_steps=5)
│  │  └─ Similar loop but with sr_model and sr_num_inference_steps
│  └─ [CPU Offload Optional] Move SR model back to CPU
│
├─ Step 4: Decode Video Latents
│  ├─ If use_turbo_vae: use faster TurboVAED decoder
│  ├─ Else: use standard Wan2.2_VAE decoder
│  ├─ Decode and denormalize: output range [0, 255], uint8, shape (T, H, W, C)
│  └─ [CPU Offload] Empty CUDA cache
│
├─ Step 5: Decode Audio Latents
│  ├─ audio_vae.decode(latent_audio.T)
│  ├─ Resample audio via sinc interpolation
│  └─ Output numpy array
│
├─ Step 6: Post-process & Save
│  └─ [Rank sync] Only rank (world_size-1) saves:
│     ├─ Save audio: sf.write(audio_path, audio_np, sample_rate)
│     ├─ Save video: imageio.mimwrite(video_path, video_np, fps=25, quality=8)
│     ├─ Merge: merge_video_and_audio(video_path, audio_path, save_path)
│     └─ Barrier: torch.distributed.barrier()
```

---

### 3. MODEL ARCHITECTURE DETAILS

#### **DiT Model Structure** (`dit_module.py`)

```
DiTModel (main class)
├─ adapter: Adapter
│  ├─ video_embedder: Linear(video_in_channels=192 → hidden_size=5120)
│  ├─ text_embedder: Linear(text_in_channels=3584 → hidden_size=5120)
│  ├─ audio_embedder: Linear(audio_in_channels=64 → hidden_size=5120)
│  └─ rope: ElementWiseFourierEmbed (positional encoding)
│
├─ block: TransformerBlock
│  └─ layers[40]: TransFormerLayer (layer_idx 0-39)
│     ├─ Attention (self-attention)
│     │  ├─ pre_norm: MultiModalityRMSNorm
│     │  ├─ linear_qkv: create_linear() → (Q, K, V, gating)
│     │  ├─ q_norm, k_norm: MultiModalityRMSNorm
│     │  ├─ linear_proj: Output projection
│     │  └─ forward():
│     │     ├─ Normalize input
│     │     ├─ Project to Q, K, V
│     │     ├─ Apply RoPE (rotary position embeddings)
│     │     ├─ Flash Attention or Local Attention (flex_flash_attn)
│     │     ├─ [Optional] Attention gating
│     │     └─ Project output
│     │
│     └─ MLP
│        ├─ pre_norm: MultiModalityRMSNorm
│        ├─ up_gate_proj: expand to intermediate_size
│        ├─ Activation: SwiGLU7 or GELU7 (clamped to [-7, 7])
│        └─ down_proj: project back to hidden_size
│
├─ final_norm_video: MultiModalityRMSNorm
├─ final_norm_audio: MultiModalityRMSNorm
├─ final_linear_video: Linear(hidden_size=5120 → video_in_channels=192)
└─ final_linear_audio: Linear(hidden_size=5120 → audio_in_channels=64)
```

#### **Key Architectural Features:**

1. **Multimodal Expert Layers** (specified by `mm_layers=[0, 1, 2, 3, 36, 37, 38, 39]`)
   - Layers in `mm_layers` use 3 modality experts (Video, Audio, Text)
   - Other layers use single expert
   - ModalityDispatcher groups sequences by modality, processes independently

2. **Local Attention** (specified by `local_attn_layers` in config)
   - For SR2_1080 mode, ~20 layers use local attention with frame receptive field
   - Computed via `calc_local_qk_range()` to limit attention scope per frame

3. **Attention Gating** (`enable_attn_gating=True`)
   - Gating parameter learned per head: `g = sigmoid(g) * attn_output`

4. **Activation Functions:**
   - **SwiGLU7**: `(sigmoid(1.702 * clamp(x_glu, -7, 7)) * (clamp(x_linear, -7, 7) + 1))` (default)
   - **GELU7**: `(1.702 * clamp(x, -7, 7)) * sigmoid(1.702 * clamp(x, -7, 7))` (for early layers)

5. **Post-norm Layers** (empty list by default, but configurable)
   - If layer in `post_norm_layers`, apply norm after residual

---

### 4. GPU/ATTENTION BACKEND SELECTION

#### **Architecture Detection** (`arch.py`)

```python
def is_hopper_arch():
    return torch.cuda.get_device_capability()[0] == 9  # SM90 → True, SM80 → False

def get_arch_memory(unit: str = "GB"):
    return torch.cuda.get_device_properties(torch.cuda.current_device()).total_memory
```

#### **Flash Attention Selection** (`dit_module.py` lines 387-400)

```python
HAS_MAGI_ATTENTION = importlib.util.find_spec("magi_attention") is not None  # Custom optimized attention
HAS_FA3 = importlib.util.find_spec("flash_attn_interface") is not None        # FA3 for Hopper

def flash_attn_func(query, key, value):
    if HAS_FA3 and is_hopper_arch():
        from flash_attn_interface import flash_attn_func as fa3_flash_attn_func
        return fa3_flash_attn_func(query, key, value)  # SM90: FA3 (optimized)
    else:
        from flash_attn.flash_attn_interface import flash_attn_func as fa2_flash_attn_func
        return fa2_flash_attn_func(query, key, value)  # SM80: FA2
```

#### **Local Attention Backend** (`dit_module.py` lines 485-493)

```python
def flex_flash_attn_func(query, key, value, q_ranges, k_ranges):
    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)  # SM90: Optimized
    else:
        return _custom_flex_flash_attn_func(query, key, value, q_ranges, k_ranges)  # Fallback
```

#### **Decision Flow:**

```
SM90 (Hopper) with all libraries:
├─ Global attention → FA3 (via magi_compiler)
├─ Local attention → MAGI Attention (flex_flash_attn)
└─ Uses magi_compile for custom ops

SM80 (Ada) or missing libraries:
├─ Global attention → Flash Attention v2
├─ Local attention → Custom flex_flash_attn fallback
└─ Uses magi_compiler if available
```

**CPU Memory Decision:**
```python
if get_arch_memory() <= 48:  # ≤ 48 GB
    vae = CPUOffloadWrapper(vae, is_cpu_offload=True)  # Offload VAE to CPU
if env_is_true("CPU_OFFLOAD") and env_is_true("SR2_1080"):
    model.to("cpu")  # During SR phase, move main model to CPU
```

---

### 5. PARALLELISM & CONCURRENCY SUPPORT

#### **Distributed Parallelism** (`infra/distributed/parallel_state.py`)

The system supports 4 orthogonal parallelism strategies:

1. **Tensor Parallelism (TP)**: Split weights across GPUs horizontally
   - Each rank processes different attention heads

2. **Pipeline Parallelism (PP)**: Split layers across GPUs vertically
   - Each rank processes different transformer layers sequentially

3. **Context Parallelism (CP)**: Split sequence along time/space dimensions
   - Each rank processes different parts of the token sequence
   - Uses Ulysses scheduler for efficient scatter/gather

4. **Data Parallelism (DP)**: Process different samples on different GPUs
   - Each rank processes different batch elements

**Config Example:**
```python
engine_config = EngineConfig(
    tp_size=1,  # Tensor parallel degree
    pp_size=1,  # Pipeline parallel degree
    cp_size=1,  # Context parallel degree
    dp_size=world_size // (tp_size * pp_size * cp_size)  # Derived
)
```

#### **Sequence Parallelism** (`ulysses_scheduler.py`)

For CP>1, the Ulysses scheduler:
```python
def dispatch(x):
    # Split sequence dimension across CP group ranks
    x_local = x[rank * seq_len_per_rank : (rank+1) * seq_len_per_rank]
    return x_local

def undispatch(x_gathered):
    # Gather all-to-all to reconstruct full sequence
    return torch.cat([x_rank0, x_rank1, ...], dim=seq_dim)
```

#### **Modality-based Batching** (`dit_module.py` lines 83-130)

The ModalityDispatcher provides **online batching by modality**:

```python
class ModalityDispatcher:
    def dispatch(x):
        # Group tokens by modality (Video=0, Audio=1, Text=2)
        # Split into 3 groups and process independently
        [video_tokens, audio_tokens, text_tokens] = split(x)
        return [video_tokens, audio_tokens, text_tokens]

    def undispatch(*groups):
        # Concatenate back in original order via inverse permutation
        return cat([video, audio, text])
```

**Benefits:**
- Efficient expert dispatching (multimodal MLP layers)
- Different computation per modality type
- No wasteful padding

#### **No Explicit Batch Concurrency:**
The current implementation processes **one video per inference call**:
- `batch_size=1` hardcoded in latent initialization
- Single inference request per `run_offline()` call
- Distributed concurrency via TP/PP/CP on single batch

---

### 6. VIDEO OUTPUT GENERATION

#### **VAE Decoding Path:**

```
latent_video (1, 48, frames, latent_h, latent_w)
│
├─ If use_turbo_vae (default):
│  └─ TurboVAED.decode(latent.to(dtype), output_offload=...)
│     └─ Fast approximation of Wan2.2_VAE
│
└─ Else:
   └─ Wan2_2_VAE.decode(latent.squeeze(0).to(dtype), group=cp_group)
      ├─ CP distributed decoding via get_cp_group() barrier
      └─ Returns list of video tensors
│
├─ Denormalization:
│  ├─ output.mul_(0.5).add_(0.5).clamp_(0, 1)  # [-1,1] → [0,1]
│  ├─ Convert to uint8: * 255
│  └─ Permute: CTHW → THWC
│
└─ Output: (frames, height, width, 3) uint8 numpy array
```

#### **Audio Decoding Path:**

```
latent_audio (1, num_frames, 64)
│
└─ audio_vae.decode(latent_audio.squeeze(0).T)
   ├─ Stable Audio VAE inverse transform
   └─ Output: (num_samples,) numpy array
│
└─ Resample if needed:
   └─ resample_audio_sinc(audio, time_stretching_factor)
      ├─ Via scipy.signal.resample()
      └─ Adjust for video frame rate
```

#### **Merging & Output:**

```
save_path = f"{save_path_prefix}_{seconds}s_{br_width}x{br_height}_{sr_width}x{sr_height}.mp4"

1. Save audio WAV: sf.write(audio_path, audio_np, sample_rate=44100)

2. Save video MP4: imageio.mimwrite(video_path, video_np, fps=25, quality=8)
   ├─ FFMPEG encoding
   ├─ H.264 codec (default)
   └─ Quality 8 (high fidelity)

3. Merge: ffmpeg -i video.mp4 -i audio.wav -c:v copy -c:a aac -y output.mp4
   ├─ Copy video stream (no re-encoding)
   ├─ AAC audio encoding
   └─ Synchronize streams (shortest)

4. Cleanup: remove temporary video.mp4 and audio.wav

Output: save_path (final MP4 with synced video+audio)
```

#### **Optional Upsampling:**

```
if output_width and output_height provided:
    video_np = upsample_video(video_np, output_width, output_height, mode='bilinear')
    ├─ Converts uint8 → float32 [0,1]
    ├─ Permutes THWC → CTHW
    ├─ Uses F.interpolate(mode='bilinear', align_corners=False)
    ├─ Permutes back CTHW → THWC
    └─ Converts to uint8 [0,255]
```

---

### 7. KEY CONFIGURATION PARAMETERS

| Parameter | Default | Purpose |
|-----------|---------|---------|
| **Inference Steps** | 32 | BR diffusion steps (num_inference_steps) |
| **SR Steps** | 5 | Super-resolution steps (sr_num_inference_steps) |
| **FPS** | 25 | Frames per second (fps) |
| **Video Guidance Scale** | 5.0 | CFG scale for video generation |
| **Audio Guidance Scale** | 5.0 | CFG scale for audio generation |
| **SR Guidance Scale** | 3.5 | CFG scale for super-resolution |
| **Shift** | 5.0 | Temporal shift in scheduler |
| **T5-Gemma Target Length** | 640 | Padded text embedding length |
| **Noise Value (SR)** | 220 | Noise injection for SR initialization |
| **Use TurboVAE** | True | Use fast VAE decoder |
| **Use SR Model** | False | Enable super-resolution stage |
| **SR Audio Noise Scale** | 0.7 | Noise for audio in SR stage |

---

### 8. SUMMARY: COMPLETE INFERENCE FLOW

```
User Input (CLI)
    ↓
entry.py: Parse args + initialize_infra()
    ↓
Load Config (Pydantic) → EngineConfig, ModelConfig, EvaluationConfig
    ↓
Load Model: get_dit() → safetensors parallel loading → GPU
    ↓
MagiPipeline.run_offline()
    ↓
MagiEvaluator.evaluate()
    ├─ Encode image (if provided) → VAE latent
    ├─ Encode audio (if provided) → Audio VAE latent
    ├─ Encode prompt → T5-Gemma embedding
    ├─ Initialize random latents
    ├─ BR Loop (32 steps):
    │  ├─ Forward pass with TP/PP/CP/DP
    │  ├─ Flash Attention (FA3 if SM90, FA2 if SM80)
    │  ├─ Local attention (if enabled, using magi_attention if SM90)
    │  ├─ CFG denoising
    │  └─ UniPC scheduler step
    ├─ [Optional] SR Loop (5 steps): Repeat above at higher resolution
    ├─ Decode video: TurboVAED or Wan2.2_VAE → (T, H, W, C) uint8
    ├─ Decode audio: Stable Audio VAE → waveform
    └─ Save: imageio.mimwrite + ffmpeg merge → final MP4
    ↓
Distributed Barrier (wait for rank sync)
    ↓
Output: save_path/prompt_*.mp4
```

**Key Files Summary:**

| File | Purpose |
|------|---------|
| `/pipeline/entry.py` | CLI entry point, argument parsing |
| `/pipeline/pipeline.py` | MagiPipeline facade |
| `/pipeline/video_generate.py` | MagiEvaluator core loop |
| `/model/dit/dit_module.py` | Attention, MLP, DiTModel architecture |
| `/infra/checkpoint/load_model_checkpoint.py` | Parallel safetensors loading |
| `/infra/distributed/init_dist_env.py` | torch.distributed setup |
| `/common/arch.py` | GPU detection (SM90 vs SM80) |
| `/common/config.py` | Configuration system |