Thoroughly explore /home/ubuntu/daVinci-MagiHuman/inference/ directory. I need to understand:
1. The full file structure
2. How the inference pipeline works (entry.py and pipeline code)
3. How models are loaded
4. Where GPU/attention backend is selected (flash attention, sm80 vs sm90)
5. Any batching or concurrency support
6. The full flow from entry.py to final video output

Read all key Python files. Be very thorough.


I'll thoroughly explore the inference directory to understand the complete pipeline. Let me start with a comprehensive scan of the file structure and then examine the key files.


Bash

find /home/ubuntu/daVinci-MagiHuman/inference -type f -name "*.py" | head -50


Now let me get a complete directory structure overview:


tree -L 3 /home/ubuntu/daVinci-MagiHuman/inference --dirsfirst 2>/dev/null || find /home/ubuntu/daVinci-MagiHuman/inference -type d | sort


/home/ubuntu/daVinci-MagiHuman/inference
/home/ubuntu/daVinci-MagiHuman/inference/common
/home/ubuntu/daVinci-MagiHuman/inference/infra
/home/ubuntu/daVinci-MagiHuman/inference/infra/checkpoint
/home/ubuntu/daVinci-MagiHuman/inference/infra/distributed
/home/ubuntu/daVinci-MagiHuman/inference/infra/parallelism
/home/ubuntu/daVinci-MagiHuman/inference/model
/home/ubuntu/daVinci-MagiHuman/inference/model/dit
/home/ubuntu/daVinci-MagiHuman/inference/model/sa_audio
/home/ubuntu/daVinci-MagiHuman/inference/model/t5_gemma
/home/ubuntu/daVinci-MagiHuman/inference/model/turbo_vaed
/home/ubuntu/daVinci-MagiHuman/inference/model/vae2_2
/home/ubuntu/daVinci-MagiHuman/inference/pipeline
/home/ubuntu/daVinci-MagiHuman/inference/utils


Now let me read the key entry point and pipeline files:


     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 argparse
    16→import sys
    17→
    18→from inference.common import parse_config
    19→from inference.infra import initialize_infra
    20→from inference.model.dit import get_dit
    21→from inference.utils import print_rank_0
    22→
    23→try:
    24→    from .pipeline import MagiPipeline
    25→except ImportError:
    26→    # Keep compatibility when entry.py is executed as a script path.
    27→    from inference.pipeline import MagiPipeline
    28→
    29→
    30→def parse_arguments():
    31→    parser = argparse.ArgumentParser(description="Run DiT pipeline with unified offline entry.")
    32→    parser.add_argument("--prompt", type=str)
    33→    parser.add_argument("--save_path_prefix", type=str, help="Path prefix for saving outputs.")
    34→    parser.add_argument("--output_path", type=str, help="Alias of --save_path_prefix for MAGI-style CLI.")
    35→
    36→    parser.add_argument("--image_path", type=str, help="Path to image for i2v mode.")
    37→    parser.add_argument(
    38→        "--audio_path", type=str, default=None, help="Path to optional audio for lipsync mode; omit to use i2v or t2v"
    39→    )
    40→
    41→    # Optional runtime controls; forwarded to pipeline methods when provided.
    42→    parser.add_argument("--seed", type=int)
    43→    parser.add_argument("--seconds", type=int)
    44→    parser.add_argument("--br_width", type=int)
    45→    parser.add_argument("--br_height", type=int)
    46→    parser.add_argument("--sr_width", type=int)
    47→    parser.add_argument("--sr_height", type=int)
    48→    parser.add_argument("--output_width", type=int)
    49→    parser.add_argument("--output_height", type=int)
    50→    parser.add_argument("--upsample_mode", type=str)
    51→    args, _ = parser.parse_known_args()
    52→    return args
    53→
    54→
    55→def main():
    56→    args = parse_arguments()
    57→    config = parse_config()
    58→    model = get_dit(config.arch_config, config.engine_config)
    59→    pipeline = MagiPipeline(model, config.evaluation_config)
    60→    save_path_prefix = args.save_path_prefix or args.output_path
    61→    if not save_path_prefix:
    62→        print_rank_0("Error: --save_path_prefix (or --output_path) is required.")
    63→        sys.exit(1)
    64→
    65→    optional_kwargs = {
    66→        "seed": args.seed,
    67→        "seconds": args.seconds,
    68→        "br_width": args.br_width,
    69→        "br_height": args.br_height,
    70→        "sr_width": args.sr_width,
    71→        "sr_height": args.sr_height,
    72→        "output_width": args.output_width,
    73→        "output_height": args.output_height,
    74→        "upsample_mode": args.upsample_mode,
    75→    }
    76→    optional_kwargs = {k: v for k, v in optional_kwargs.items() if v is not None and v is not False}
    77→
    78→    prompt = args.prompt
    79→    image_path = args.image_path
    80→    audio_path = args.audio_path
    81→
    82→    if not prompt:
    83→        print_rank_0("Error: --prompt is required.")
    84→        sys.exit(1)
    85→    if not image_path:
    86→        print_rank_0("Error: --image_path is required.")
    87→        sys.exit(1)
    88→
    89→    pipeline.run_offline(
    90→        prompt=prompt, image=image_path, audio=audio_path, save_path_prefix=save_path_prefix, **optional_kwargs
    91→    )
    92→
    93→
    94→if __name__ == "__main__":
    95→    initialize_infra()
    96→    main()
    97→


     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 os
    16→import random
    17→from typing import Optional, Union
    18→
    19→import imageio
    20→import soundfile as sf
    21→import torch
    22→from PIL import Image
    23→
    24→from inference.common import EvaluationConfig, parse_config
    25→from inference.model.dit import get_dit
    26→from inference.model.dit import DiTModel
    27→from .video_generate import MagiEvaluator
    28→from .video_process import merge_video_and_audio, upsample_video
    29→
    30→
    31→class MagiPipeline:
    32→    """Pipeline facade for inference."""
    33→
    34→    def __init__(self, model: DiTModel, evaluation_config: EvaluationConfig, device: str = "cuda"):
    35→        self.model = model
    36→        self.evaluation_config = evaluation_config
    37→        config = parse_config()
    38→        if evaluation_config.use_sr_model:
    39→            config.engine_config.load = evaluation_config.sr_model_path
    40→            sr_model = get_dit(config.sr_arch_config, config.engine_config)
    41→        else:
    42→            sr_model = None
    43→        self.evaluator = MagiEvaluator(model, sr_model, evaluation_config, device)
    44→
    45→    def _validate_offline_request(
    46→        self,
    47→        prompt: str,
    48→        save_path_prefix: str,
    49→    ):
    50→        if not prompt or not prompt.strip():
    51→            raise ValueError("`prompt` must be a non-empty string.")
    52→        if not save_path_prefix or not save_path_prefix.strip():
    53→            raise ValueError("`save_path_prefix` must be a non-empty string.")
    54→
    55→    def run_offline(
    56→        self,
    57→        prompt: str,
    58→        image: Union[str, Image.Image, None],
    59→        audio: Optional[str],
    60→        save_path_prefix: str,
    61→        seed: int = 42,
    62→        seconds: int = 4,
    63→        br_width: int = 480,
    64→        br_height: int = 272,
    65→        sr_width: Optional[int] = None,
    66→        sr_height: Optional[int] = None,
    67→        output_width: Optional[int] = None,
    68→        output_height: Optional[int] = None,
    69→        upsample_mode: Optional[str] = None,
    70→    ):
    71→        self._validate_offline_request(prompt=prompt, save_path_prefix=save_path_prefix)
    72→
    73→        if self.evaluator.sr_model is not None:
    74→            save_path = f"{save_path_prefix}_{seconds}s_{br_width}x{br_height}_{sr_width}x{sr_height}.mp4"
    75→        else:
    76→            save_path = f"{save_path_prefix}_{seconds}s_{br_width}x{br_height}.mp4"
    77→
    78→        with torch.random.fork_rng(devices=[torch.cuda.current_device()]):
    79→            torch.random.manual_seed(seed)
    80→            video_np, audio_np = self.evaluator.evaluate(
    81→                prompt,
    82→                image,
    83→                audio,
    84→                seconds=seconds,
    85→                br_width=br_width,
    86→                br_height=br_height,
    87→                sr_width=sr_width,
    88→                sr_height=sr_height,
    89→                br_num_inference_steps=self.evaluation_config.num_inference_steps,
    90→                sr_num_inference_steps=self.evaluation_config.sr_num_inference_steps,
    91→            )
    92→
    93→        if output_width is not None and output_height is not None:
    94→            video_np = upsample_video(video_np, output_width, output_height, upsample_mode)
    95→
    96→        if torch.distributed.get_rank() == torch.distributed.get_world_size() - 1:
    97→            saving_name = f"{prompt.replace(' ', '_')[:10]}"
    98→            audio_path = saving_name + str(random.randint(0, 1000000)) + ".wav"
    99→            video_path = saving_name + str(random.randint(0, 1000000)) + ".mp4"
   100→            sf.write(audio_path, audio_np, self.evaluator.audio_vae.sample_rate)
   101→            imageio.mimwrite(video_path, video_np, fps=self.evaluation_config.fps, quality=8, output_params=["-loglevel", "error"])
   102→            assert os.path.exists(video_path)
   103→            merge_video_and_audio(video_path, audio_path, save_path)
   104→
   105→        if torch.distributed.is_initialized():
   106→            torch.distributed.barrier()
   107→        return save_path
   108→
   109→


     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 copy
    16→import math
    17→import os
    18→import shutil
    19→from abc import abstractmethod
    20→from dataclasses import dataclass
    21→from functools import partial
    22→from typing import List, Optional, Tuple, Union
    23→
    24→import numpy as np
    25→import torch
    26→from diffusers.utils import load_image
    27→from diffusers.video_processor import VideoProcessor
    28→from PIL import Image
    29→from tqdm import tqdm
    30→
    31→from inference.common import CPUOffloadWrapper, EvaluationConfig, get_arch_memory
    32→from inference.infra.distributed import get_cp_group
    33→from inference.model.dit import DiTModel
    34→from inference.model.sa_audio import SAAudioFeatureExtractor
    35→from inference.model.turbo_vaed import TurboVAED, get_turbo_vaed
    36→from inference.model.vae2_2 import Wan2_2_VAE, get_vae2_2
    37→from inference.utils import env_is_true, event_path_timer, print_mem_info_rank_0, print_rank_0, print_rank_last
    38→from .prompt_process import get_padded_t5_gemma_embedding, pad_or_trim
    39→from .scheduler_unipc import FlowUniPCMultistepScheduler
    40→from .data_proxy import MagiDataProxy
    41→from .video_process import load_audio_and_encode, resample_audio_sinc, resizecrop
    42→
    43→
    44→def schedule_latent_step(
    45→    *,
    46→    video_scheduler,
    47→    audio_scheduler,
    48→    latent_video: torch.Tensor,
    49→    latent_audio: torch.Tensor,
    50→    t,
    51→    idx: int,
    52→    steps: int,
    53→    v_cfg_video: torch.Tensor,
    54→    v_cfg_audio: torch.Tensor,
    55→    is_a2v: bool,
    56→    cfg_number: int,
    57→    use_sr_model: bool,
    58→    using_sde_flag: bool,
    59→):
    60→    if cfg_number == 1 and (not use_sr_model):
    61→        latent_video = video_scheduler.step_ddim(v_cfg_video, idx, latent_video)
    62→        latent_audio = audio_scheduler.step_ddim(v_cfg_audio, idx, latent_audio)
    63→        return latent_video, latent_audio
    64→
    65→    if using_sde_flag:
    66→        print_rank_0("Using sde scheduler")
    67→        if use_sr_model:
    68→            latent_video = video_scheduler.step(v_cfg_video, t, latent_video, return_dict=False)[0]
    69→            return latent_video, latent_audio
    70→
    71→        if idx < int(steps * (3 / 4)):
    72→            noise_theta = 1.0 if (idx + 1) % 2 == 0 else 0.0
    73→        else:
    74→            noise_theta = 1.0 if idx % 3 == 0 else 0.0
    75→
    76→        latent_video = video_scheduler.step_sde(v_cfg_video, idx, latent_video, noise_theta=noise_theta)
    77→        if not is_a2v:
    78→            latent_audio = audio_scheduler.step_sde(v_cfg_audio, idx, latent_audio, noise_theta=noise_theta)
    79→        return latent_video, latent_audio
    80→
    81→    latent_video = video_scheduler.step(v_cfg_video, t, latent_video, return_dict=False)[0]
    82→    if not is_a2v and not use_sr_model:
    83→        latent_audio = audio_scheduler.step(v_cfg_audio, t, latent_audio, return_dict=False)[0]
    84→    return latent_video, latent_audio
    85→
    86→
    87→@dataclass
    88→class EvalInput:
    89→    x_t: torch.Tensor
    90→    audio_x_t: torch.Tensor
    91→    audio_feat_len: torch.Tensor | list[int]
    92→    txt_feat: torch.Tensor
    93→    txt_feat_len: torch.Tensor | list[int]
    94→
    95→
    96→class ZeroSNRDDPMDiscretization:
    97→    def __init__(
    98→        self,
    99→        linear_start=0.00085,
   100→        linear_end=0.0120,
   101→        num_timesteps=1000,
   102→        shift_scale=1.0,  # noise schedule t_n -> t_m: logSNR(t_m) = logSNR(t_n) - log(shift_scale)
   103→        keep_start=False,
   104→        post_shift=False,
   105→    ):
   106→        if keep_start and not post_shift:
   107→            linear_start = linear_start / (shift_scale + (1 - shift_scale) * linear_start)
   108→        self.num_timesteps = num_timesteps
   109→        betas = torch.linspace(linear_start**0.5, linear_end**0.5, num_timesteps, dtype=torch.float64) ** 2
   110→        betas = betas.numpy()
   111→        alphas = 1.0 - betas
   112→        self.alphas_cumprod = np.cumprod(alphas, axis=0)
   113→        self.to_torch = partial(torch.tensor, dtype=torch.float32)
   114→
   115→        # SNR shift
   116→        if not post_shift:
   117→            self.alphas_cumprod = self.alphas_cumprod / (shift_scale + (1 - shift_scale) * self.alphas_cumprod)
   118→
   119→        self.post_shift = post_shift
   120→        self.shift_scale = shift_scale
   121→
   122→    def __call__(self, n, do_append_zero=True, device="cpu", flip=False, return_idx=False):
   123→        if return_idx:
   124→            sigmas, idx = self.get_sigmas(n, device=device, return_idx=return_idx)
   125→        else:
   126→            sigmas = self.get_sigmas(n, device=device, return_idx=return_idx)
   127→        sigmas = torch.cat([sigmas, sigmas.new_zeros([1])]) if do_append_zero else sigmas
   128→        if return_idx:
   129→            return sigmas if not flip else torch.flip(sigmas, (0,)), idx
   130→        else:
   131→            return sigmas if not flip else torch.flip(sigmas, (0,))
   132→
   133→    def get_sigmas(self, n, device="cpu", return_idx=False):
   134→        if n < self.num_timesteps:
   135→            timesteps = np.linspace(self.num_timesteps - 1, 0, n, endpoint=False).astype(int)[::-1]
   136→            alphas_cumprod = self.alphas_cumprod[timesteps]
   137→        elif n == self.num_timesteps:
   138→            alphas_cumprod = self.alphas_cumprod
   139→        else:
   140→            raise ValueError
   141→
   142→        to_torch = partial(torch.tensor, dtype=torch.float32, device=device)
   143→        alphas_cumprod = to_torch(alphas_cumprod)
   144→        alphas_cumprod_sqrt = alphas_cumprod.sqrt()
   145→        alphas_cumprod_sqrt_0 = alphas_cumprod_sqrt[0].clone()
   146→        alphas_cumprod_sqrt_T = alphas_cumprod_sqrt[-1].clone()
   147→
   148→        alphas_cumprod_sqrt -= alphas_cumprod_sqrt_T
   149→        alphas_cumprod_sqrt *= alphas_cumprod_sqrt_0 / (alphas_cumprod_sqrt_0 - alphas_cumprod_sqrt_T)
   150→
   151→        if self.post_shift:
   152→            alphas_cumprod_sqrt = (
   153→                alphas_cumprod_sqrt**2 / (self.shift_scale + (1 - self.shift_scale) * alphas_cumprod_sqrt**2)
   154→            ) ** 0.5
   155→
   156→        if return_idx:
   157→            return torch.flip(alphas_cumprod_sqrt, (0,)), timesteps
   158→        else:
   159→            return torch.flip(alphas_cumprod_sqrt, (0,))
   160→
   161→
   162→class MagiEvaluator:
   163→    def __init__(
   164→        self,
   165→        model: DiTModel,
   166→        sr_model: Optional[DiTModel],
   167→        config: EvaluationConfig,
   168→        device: str = "cuda",
   169→        weight_dtype: torch.dtype = torch.bfloat16,
   170→    ):
   171→        device = f"cuda:{torch.cuda.current_device()}"
   172→
   173→        self.model = model
   174→        self.model.eval()
   175→        self.sr_model = sr_model
   176→        if self.sr_model is not None:
   177→            self.sr_model.eval()
   178→            if env_is_true("CPU_OFFLOAD") and env_is_true("SR2_1080"):
   179→                self.model = self.model.to(torch.device("cpu"))
   180→                self.sr_model = self.sr_model.to(torch.device("cpu"))
   181→        self.device = device
   182→        self.config = config
   183→        self.dtype = weight_dtype
   184→        self.data_proxy = MagiDataProxy(config.data_proxy_config)
   185→        sr_data_proxy_config = copy.deepcopy(config.data_proxy_config)
   186→        sr_data_proxy_config.coords_style = "v1"
   187→        self.sr_data_proxy = MagiDataProxy(sr_data_proxy_config)
   188→        self.vae_stride = config.vae_stride
   189→        self.z_dim = config.z_dim
   190→        self.patch_size = config.patch_size
   191→
   192→        self.sr_video_txt_guidance_scale = config.sr_video_txt_guidance_scale
   193→        self.video_txt_guidance_scale = config.video_txt_guidance_scale
   194→        self.audio_txt_guidance_scale = config.audio_txt_guidance_scale
   195→        self.noise_value = config.noise_value
   196→        self.shift = config.shift
   197→        self.fps = config.fps
   198→        self.use_cfg_trick = config.use_cfg_trick
   199→        self.cfg_trick_start_frame = config.cfg_trick_start_frame
   200→        self.cfg_trick_value = config.cfg_trick_value
   201→        self.using_sde_flag = config.using_sde_flag
   202→
   203→        print_mem_info_rank_0("Begin init MagiEvaluator")
   204→
   205→        vae_model_path = os.path.join(config.vae_model_path, "Wan2.2_VAE.pth")
   206→        self.vae: Wan2_2_VAE = CPUOffloadWrapper(
   207→            get_vae2_2(vae_model_path, self.device, weight_dtype=weight_dtype), is_cpu_offload=get_arch_memory() <= 48
   208→        )
   209→        if config.use_turbo_vae:
   210→            self.turbo_vae: TurboVAED = CPUOffloadWrapper(
   211→                get_turbo_vaed(config.student_config_path, config.student_ckpt_path, self.device, weight_dtype=weight_dtype),
   212→                is_cpu_offload=get_arch_memory() <= 48,
   213→            )
   214→
   215→        print_mem_info_rank_0("After init video vae")
   216→        print_rank_0(f"vae loaded from {vae_model_path}")
   217→        self.video_processor = VideoProcessor(vae_scale_factor=16)
   218→        self.audio_vae = SAAudioFeatureExtractor(device=self.device, model_path=config.audio_model_path)
   219→        self.sigmas = ZeroSNRDDPMDiscretization()(1000, do_append_zero=False, flip=True)
   220→        print_mem_info_rank_0("After init audio vae")
   221→
   222→        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"  # noqa: E501
   223→        negative_prompt += ", 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"
   224→        negative_prompt += ", 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"
   225→
   226→        txt_encoder_device = "cpu" if get_arch_memory() <= 48 else self.device
   227→        self.txt_model_path = config.txt_model_path
   228→
   229→        self.context_null, self.original_context_null_len = get_padded_t5_gemma_embedding(
   230→            negative_prompt,
   231→            self.txt_model_path,
   232→            txt_encoder_device,
   233→            self.dtype,
   234→            config.t5_gemma_target_length,
   235→        )
   236→        print_mem_info_rank_0("After init t5 gamma")
   237→
   238→    def forward(self, eval_input: EvalInput, use_sr_model: bool = False):
   239→        if use_sr_model:
   240→            eval_input = self.sr_data_proxy.process_input(eval_input)
   241→            noise_pred = self.sr_model(*eval_input)
   242→            noise_pred = self.sr_data_proxy.process_output(noise_pred)
   243→        else:
   244→            eval_input = self.data_proxy.process_input(eval_input)
   245→            noise_pred = self.model(*eval_input)
   246→            noise_pred = self.data_proxy.process_output(noise_pred)
   247→        return noise_pred
   248→
   249→    @torch.inference_mode()
   250→    def evaluate(
   251→        self,
   252→        prompt: str,
   253→        image: Optional[Image.Image],
   254→        audio_path: Optional[str],
   255→        seconds: int,
   256→        br_width: int,
   257→        br_height: int,
   258→        sr_width: Optional[int],
   259→        sr_height: Optional[int],
   260→        br_num_inference_steps: int,
   261→        sr_num_inference_steps: int,
   262→    ):
   263→        event_path_timer().reset()
   264→        event_path_timer().synced_record("Step1: Prepare Latent Features")
   265→        br_latent_height = br_height // self.vae_stride[1] // self.patch_size[1] * self.patch_size[1]
   266→        br_latent_width = br_width // self.vae_stride[2] // self.patch_size[2] * self.patch_size[2]
   267→        br_height = br_latent_height * self.vae_stride[1]
   268→        br_width = br_latent_width * self.vae_stride[2]
   269→
   270→        # init latent
   271→        if audio_path is not None:
   272→            latent_audio = load_audio_and_encode(self.audio_vae, audio_path, seconds)
   273→            latent_audio = latent_audio.permute(0, 2, 1)
   274→            num_frames = latent_audio.shape[1]
   275→            is_a2v = True
   276→            print_rank_0(f"Using provided audio, latent_audio: {latent_audio.shape}")
   277→        else:
   278→            num_frames = seconds * self.fps + 1
   279→            latent_audio = torch.randn(1, num_frames, 64, dtype=torch.float32, device=self.device)
   280→            is_a2v = False
   281→            print_rank_0(f"Using random audio, latent_audio: {latent_audio.shape}")
   282→        latent_length = (num_frames - 1) // 4 + 1
   283→        latent_video = torch.randn(
   284→            1, self.z_dim, latent_length, br_latent_height, br_latent_width, dtype=torch.float32, device=self.device
   285→        )
   286→
   287→        context, original_context_len = get_padded_t5_gemma_embedding(
   288→            prompt, self.txt_model_path, self.device, self.dtype, self.config.t5_gemma_target_length
   289→        )
   290→
   291→        event_path_timer().synced_record("Step2: Encode Image for Basic Resolution")
   292→        if image is not None:
   293→            br_image = self.encode_image(image, br_height, br_width)
   294→        else:
   295→            br_image = None
   296→        event_path_timer().synced_record("Step3: Basic Resolution Evaluation")
   297→
   298→        if env_is_true("CPU_OFFLOAD") and env_is_true("SR2_1080"):
   299→            self.model = self.model.to(self.device)
   300→        br_latent_video, br_latent_audio = self.evaluate_with_latent(
   301→            context,
   302→            original_context_len,
   303→            br_image,
   304→            latent_video.clone(),
   305→            latent_audio.clone(),
   306→            br_num_inference_steps,
   307→            is_a2v,
   308→            use_sr_model=False,
   309→        )
   310→        if env_is_true("CPU_OFFLOAD") and env_is_true("SR2_1080"):
   311→            self.model = self.model.to(torch.device("cpu"))
   312→
   313→        if sr_width is not None and sr_height is not None and self.sr_model is not None:
   314→            event_path_timer().synced_record("Step4: Encode Image for Super Resolution")
   315→            sr_latent_height = sr_height // self.vae_stride[1] // self.patch_size[1] * self.patch_size[1]
   316→            sr_latent_width = sr_width // self.vae_stride[2] // self.patch_size[2] * self.patch_size[2]
   317→            sr_height = sr_latent_height * self.vae_stride[1]
   318→            sr_width = sr_latent_width * self.vae_stride[2]
   319→            if image is not None:
   320→                sr_image = self.encode_image(image, sr_height, sr_width)
   321→            else:
   322→                sr_image = None
   323→            latent_video = torch.nn.functional.interpolate(
   324→                br_latent_video, size=(latent_length, sr_latent_height, sr_latent_width), mode="trilinear", align_corners=True
   325→            )
   326→            if self.noise_value != 0:
   327→                noise = torch.randn_like(latent_video, device=latent_video.device)
   328→                sigmas = self.sigmas.to(latent_video.device)
   329→                sigma = sigmas[self.noise_value]
   330→                latent_video = latent_video * sigma + noise * (1 - sigma**2) ** 0.5
   331→            event_path_timer().synced_record("Step5: Super Resolution Evaluation")
   332→            print_mem_info_rank_0("Before super resolution evaluation")
   333→            latent_audio = br_latent_audio.clone()
   334→            br_latent_audio = torch.randn_like(
   335→                br_latent_audio, device=br_latent_audio.device
   336→            ) * self.config.sr_audio_noise_scale + br_latent_audio * (1 - self.config.sr_audio_noise_scale)
   337→
   338→            if env_is_true("CPU_OFFLOAD") and env_is_true("SR2_1080"):
   339→                self.sr_model = self.sr_model.to(self.device)
   340→            latent_video, _ = self.evaluate_with_latent(
   341→                context,
   342→                original_context_len,
   343→                sr_image,
   344→                latent_video.clone(),
   345→                br_latent_audio.clone(),
   346→                sr_num_inference_steps,
   347→                is_a2v,
   348→                use_sr_model=True,
   349→            )
   350→            if env_is_true("CPU_OFFLOAD") and env_is_true("SR2_1080"):
   351→                self.sr_model = self.sr_model.to(torch.device("cpu"))
   352→        else:
   353→            latent_video = br_latent_video
   354→            latent_audio = br_latent_audio
   355→
   356→        event_path_timer().synced_record("Step6: Decode Video", print_fn=print_rank_last)
   357→        result = self.post_process(latent_video, latent_audio)
   358→        event_path_timer().synced_record("Step8: Post Process", print_fn=print_rank_last)
   359→        return result
   360→
   361→    def schedule(
   362→        self,
   363→        video_scheduler,
   364→        audio_scheduler,
   365→        latent_video,
   366→        latent_audio,
   367→        t,
   368→        idx,
   369→        steps,
   370→        v_cfg_video,
   371→        v_cfg_audio,
   372→        is_a2v,
   373→        cfg_number,
   374→        use_sr_model=False,
   375→    ):
   376→        return schedule_latent_step(
   377→            video_scheduler=video_scheduler,
   378→            audio_scheduler=audio_scheduler,
   379→            latent_video=latent_video,
   380→            latent_audio=latent_audio,
   381→            t=t,
   382→            idx=idx,
   383→            steps=steps,
   384→            v_cfg_video=v_cfg_video,
   385→            v_cfg_audio=v_cfg_audio,
   386→            is_a2v=is_a2v,
   387→            cfg_number=cfg_number,
   388→            use_sr_model=use_sr_model,
   389→            using_sde_flag=self.config.using_sde_flag,
   390→        )
   391→
   392→    @torch.inference_mode()
   393→    def evaluate_with_latent(
   394→        self,
   395→        context: torch.Tensor,
   396→        original_context_len: int,
   397→        latent_image: Optional[torch.Tensor],
   398→        latent_video: torch.Tensor,
   399→        latent_audio: torch.Tensor,
   400→        num_inference_steps: int,
   401→        is_a2v: bool = False,
   402→        use_sr_model: bool = False,
   403→    ) -> Tuple[torch.Tensor, torch.Tensor]:
   404→        video_scheduler = FlowUniPCMultistepScheduler()
   405→        audio_scheduler = FlowUniPCMultistepScheduler()
   406→        video_scheduler.set_timesteps(num_inference_steps, device=self.device, shift=self.shift)
   407→        audio_scheduler.set_timesteps(num_inference_steps, device=self.device, shift=self.shift)
   408→        timesteps = video_scheduler.timesteps
   409→
   410→        # a inference trick to aviod over exposure in the I2V evaluation
   411→        latent_length = latent_video.shape[2]
   412→        sr_video_txt_guidance_scale = (
   413→            torch.tensor(self.sr_video_txt_guidance_scale, device=self.device).expand(1, 1, latent_length, 1, 1).clone()
   414→        )
   415→        if self.use_cfg_trick:
   416→            sr_video_txt_guidance_scale[:, :, : self.cfg_trick_start_frame] = min(
   417→                self.cfg_trick_value, self.sr_video_txt_guidance_scale
   418→            )
   419→
   420→        # forward
   421→        for idx, t in enumerate(
   422→            tqdm(timesteps, disable=torch.distributed.get_rank() != torch.distributed.get_world_size() - 1)
   423→        ):
   424→            if latent_image is not None:
   425→                latent_video[:, :, :1] = latent_image[:, :, :1]
   426→            video_txt_guidance_scale = self.video_txt_guidance_scale if t > 500 else 2.0
   427→            eval_input_cond = EvalInput(
   428→                x_t=latent_video,
   429→                audio_x_t=latent_audio,
   430→                audio_feat_len=[latent_audio.shape[1]],
   431→                txt_feat=context,
   432→                txt_feat_len=[original_context_len],
   433→            )   # txt + audio
   434→            v_output = self.forward(eval_input_cond, use_sr_model=use_sr_model)
   435→            v_cond_video = v_output[0]
   436→            v_cond_audio = v_output[1]
   437→
   438→            cfg_number = self.config.sr_cfg_number if use_sr_model else self.config.cfg_number
   439→            if cfg_number == 1:
   440→                v_cfg_video = v_cond_video
   441→                v_cfg_audio = v_cond_audio
   442→            elif cfg_number == 2:
   443→                eval_input_uncond = EvalInput(
   444→                    x_t=latent_video,
   445→                    audio_x_t=latent_audio,
   446→                    audio_feat_len=[latent_audio.shape[1]],
   447→                    txt_feat=self.context_null,
   448→                    txt_feat_len=[self.original_context_null_len],
   449→                )
   450→                v_output_uncond = self.forward(eval_input_uncond, use_sr_model=use_sr_model)
   451→                v_uncond_video = v_output_uncond[0]
   452→                v_uncond_audio = v_output_uncond[1]
   453→                if use_sr_model:
   454→                    v_cfg_video = v_uncond_video + sr_video_txt_guidance_scale * (v_cond_video - v_uncond_video)
   455→                else:
   456→                    v_cfg_video = v_uncond_video + video_txt_guidance_scale * (v_cond_video - v_uncond_video)
   457→                v_cfg_audio = v_uncond_audio + self.audio_txt_guidance_scale * (v_cond_audio - v_uncond_audio)
   458→            else:
   459→                raise ValueError(f"Invalid cfg_number: {cfg_number}")
   460→
   461→            latent_video, latent_audio = self.schedule(
   462→                video_scheduler,
   463→                audio_scheduler,
   464→                latent_video,
   465→                latent_audio,
   466→                t,
   467→                idx,
   468→                timesteps,
   469→                v_cfg_video,
   470→                v_cfg_audio,
   471→                is_a2v,
   472→                cfg_number,
   473→                use_sr_model,
   474→            )
   475→
   476→        print_rank_0(f"latent_video: {latent_video.shape}, latent_audio: {latent_audio.shape}")
   477→        if latent_image is not None:
   478→            latent_video[:, :, :1] = latent_image[:, :, :1]
   479→        return latent_video, latent_audio
   480→
   481→    def encode_image(self, image: Image.Image, height: int, width: int):
   482→        image = load_image(image)
   483→        image = resizecrop(image, height, width)
   484→        image = self.video_processor.preprocess(image, height=height, width=width)
   485→        image = image.to(device=self.device, dtype=self.dtype).unsqueeze(2)
   486→        image = self.vae.encode(image).to(torch.float32)
   487→        return image
   488→
   489→    def decode_video(self, latent: torch.Tensor, group: torch.distributed.ProcessGroup = None):
   490→        if self.config.use_turbo_vae:
   491→            is_memory_limited = env_is_true("CPU_OFFLOAD") and env_is_true("SR2_1080")
   492→            videos = self.turbo_vae.decode(latent.to(self.dtype), output_offload=is_memory_limited).float()
   493→        else:
   494→            videos = self.vae.decode(latent.squeeze(0).to(self.dtype), group=group)
   495→        if videos is None:
   496→            return None
   497→        videos.mul_(0.5).add_(0.5).clamp_(0, 1)
   498→        videos = [video.cpu() for video in videos]
   499→        videos = [video.permute(1, 2, 3, 0) * 255 for video in videos]
   500→        videos = [video.numpy().astype(np.uint8) for video in videos]
   501→        return videos
   502→
   503→    def post_process(self, latent_video: torch.Tensor, latent_audio: torch.Tensor):
   504→        torch.cuda.empty_cache()
   505→        # CTHW -> THWC
   506→        videos_np = self.decode_video(latent_video, group=get_cp_group())
   507→        torch.cuda.empty_cache()
   508→        event_path_timer().synced_record("Step7: Decode Audio", print_fn=print_rank_last)
   509→
   510→        if torch.distributed.get_rank() == torch.distributed.get_world_size() - 1:
   511→            video_np = videos_np[0]
   512→
   513→            latent_audio = latent_audio.squeeze(0)
   514→            audio_output = self.audio_vae.decode(latent_audio.T)
   515→            audio_output_np = audio_output.squeeze(0).T.cpu().numpy()
   516→            audio_output_np = resample_audio_sinc(audio_output_np, 441 / 512)
   517→
   518→            return video_np, audio_output_np
   519→        else:
   520→            return None, None
   521→
   522→
   523→


Now let me read the model loading and configuration files:


     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 io
    16→import json
    17→import os
    18→import subprocess
    19→from concurrent.futures import ThreadPoolExecutor
    20→
    21→from inference.common import EngineConfig
    22→from inference.utils import print_rank_0
    23→from safetensors.torch import load as load_from_bytes
    24→from safetensors.torch import load_file
    25→from tqdm.auto import tqdm
    26→
    27→
    28→def _load_shard(shard_path, param_names, num_threads=None):
    29→    zstd_path = shard_path + ".zst"
    30→    if os.path.exists(zstd_path):
    31→        cmd = ["zstd", "-d"]
    32→        if num_threads:
    33→            cmd.extend(["-T", str(num_threads)])  # set parallelism
    34→
    35→        process = subprocess.Popen(cmd + ["-c", zstd_path], stdout=subprocess.PIPE, stderr=subprocess.PIPE, bufsize=-1)
    36→
    37→        decompressed_data = process.stdout.read()
    38→        while True:
    39→            new_data = process.stdout.read()
    40→            if not new_data:
    41→                break
    42→            decompressed_data += new_data
    43→        process.stdout.close()
    44→
    45→        retcode = process.wait()
    46→        if retcode != 0:
    47→            raise RuntimeError(f"Decompression failed: {process.stderr.read().decode()}")
    48→
    49→        buffer = io.BytesIO(decompressed_data)
    50→        weights = load_from_bytes(buffer.getvalue())
    51→        buffer.close()
    52→    else:
    53→        weights = load_file(shard_path)
    54→
    55→    return {name: weights[name] for name in param_names}
    56→
    57→
    58→def load_sharded_safetensors_parallel_with_progress(checkpoint_dir):
    59→    index_path = os.path.join(checkpoint_dir, "model.safetensors.index.json")
    60→    if not os.path.exists(index_path):
    61→        model_file_path = os.path.join(checkpoint_dir, "model.safetensors")
    62→        state_dict = load_file(model_file_path)
    63→        return state_dict
    64→
    65→    with open(index_path, "r") as f:
    66→        index = json.load(f)
    67→
    68→    state_dict = {}
    69→    shard_map = {}
    70→
    71→    # Group parameters by shard file
    72→    for param_name, shard_file in index["weight_map"].items():
    73→        shard_path = os.path.join(checkpoint_dir, shard_file)
    74→        if shard_path not in shard_map:
    75→            shard_map[shard_path] = []
    76→        shard_map[shard_path].append(param_name)
    77→
    78→    # Load shards in parallel with a progress bar
    79→    with ThreadPoolExecutor() as executor:
    80→        futures = {
    81→            executor.submit(_load_shard, shard_path, param_names): shard_path for shard_path, param_names in shard_map.items()
    82→        }
    83→        pbar = tqdm(futures, desc="Loading shards", total=len(futures))
    84→        for future in pbar:
    85→            result = future.result()
    86→            state_dict.update(result)
    87→
    88→    return state_dict
    89→
    90→
    91→def load_model_checkpoint(model, engine_config: EngineConfig):
    92→    print_rank_0("Loading checkpoint with safetensors format from pretrained_folder")
    93→    state_dict = load_sharded_safetensors_parallel_with_progress(engine_config.load)
    94→
    95→    missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False)
    96→    print_rank_0(f"Load Weight Missing Keys: {missing_keys}")
    97→    print_rank_0(f"Load Weight Unexpected Keys: {unexpected_keys}")
    98→    print_rank_0("Load checkpoint successfully")
    99→    return model
   100→


     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 argparse
    16→import copy
    17→import json
    18→import os
    19→import sys
    20→from pathlib import Path
    21→from typing import Literal, Tuple
    22→
    23→import torch
    24→from inference.utils import env_is_true, print_rank_0
    25→from pydantic import BaseModel, ConfigDict, Field, field_serializer, field_validator, model_validator
    26→from pydantic_settings import (
    27→    BaseSettings,
    28→    CliSettingsSource,
    29→    JsonConfigSettingsSource,
    30→    PydanticBaseSettingsSource,
    31→    SettingsConfigDict,
    32→)
    33→
    34→
    35→class EngineConfig(BaseModel):
    36→    # Basic settings
    37→    seed: int = Field(1234, description="Random seed used for python, numpy, pytorch, and cuda.")
    38→    load: str | None = Field(None, description="Directory containing a model checkpoint.")
    39→
    40→    # Parallelism strategy
    41→    distributed_backend: Literal["nccl", "gloo"] = Field("nccl", description="Distributed backend. Choices: ['nccl', 'gloo'].")
    42→    distributed_timeout_minutes: int = Field(10, description="Timeout minutes for torch.distributed.")
    43→    sequence_parallel: bool = Field(False, description="Enable sequence parallel optimization.")
    44→    tp_size: int = Field(1, description="Degree of tensor model parallelism.")
    45→    pp_size: int = Field(1, description="Degree of pipeline model parallelism.")
    46→    cp_size: int = Field(1, description="Degree of context parallelism.")
    47→    dp_size: int = Field(1, description="Degree of data parallelism.")
    48→
    49→
    50→class ModelConfig(BaseModel):
    51→    """Model configuration class defining various parameters for video generation model"""
    52→
    53→    model_config = ConfigDict(arbitrary_types_allowed=True, protected_namespaces=())
    54→
    55→    num_layers: int = Field(default=40, description="Number of Transformer layers")
    56→    hidden_size: int = Field(default=5120, description="Hidden size of the Transformer model")
    57→    head_dim: int = Field(default=128, description="Dimension per attention head")
    58→    num_query_groups: int = Field(default=8, description="Number of query groups for grouped-query attention")
    59→    video_in_channels: int = Field(default=48 * 4, description="Number of video input channels after patch embedding")
    60→    audio_in_channels: int = Field(default=64, description="Number of audio input channels")
    61→    text_in_channels: int = Field(default=3584, description="Number of text input channels")
    62→    checkpoint_qk_layernorm_rope: bool = Field(default=False, description="Enable checkpointing for QK layernorm + RoPE")
    63→    params_dtype: torch.dtype | str = Field(default=torch.float32, description="Parameter dtype")
    64→    tread_config: dict = Field(
    65→        default=dict(
    66→            selection_rate=0.5, start_layer_idx=2, end_layer_idx=25  # after forward of 0, 1  # before forward of 26 27 28 29
    67→        ),
    68→        description="TReAD (Token Routing and Early Drop) configuration",
    69→    )
    70→    mm_layers: list[int] = Field(default=[0, 1, 2, 3, 36, 37, 38, 39], description="Indices of multimodal fusion layers")
    71→    local_attn_layers: list[int] = Field(default=[], description="Indices of local attention layers")
    72→    enable_attn_gating: bool = Field(default=True, description="Enable attention gating")
    73→    activation_type: str = Field(default="swiglu7", description="Activation type")
    74→    gelu7_layers: list[int] = Field(default=[0, 1, 2, 3], description="Indices of gelu7 layers")
    75→
    76→    # Add computed fields
    77→    num_heads_q: int = Field(default=0, description="Number of query heads (calculated from hidden_size // head_dim)")
    78→    num_heads_kv: int = Field(default=0, description="Number of key-value heads (calculated from num_query_groups)")
    79→    post_norm_layers: list[int] = Field(default=[], description="Indices of post norm layers")
    80→
    81→    @field_serializer("params_dtype")
    82→    def serialize_dtype(self, value: torch.dtype | str) -> str:
    83→        return str(value)
    84→
    85→    @field_validator("params_dtype", mode="before")
    86→    @classmethod
    87→    def validate_dtype(cls, value):
    88→        if isinstance(value, torch.dtype):
    89→            return value
    90→        if isinstance(value, str):
    91→            if value == "torch.float32" or value == "float32":
    92→                return torch.float32
    93→            elif value == "torch.float16" or value == "float16":
    94→                return torch.float16
    95→            elif value == "torch.bfloat16" or value == "bfloat16":
    96→                return torch.bfloat16
    97→        raise ValueError(f"Unknown torch.dtype string: '{value}'")
    98→
    99→
   100→class DataProxyConfig(BaseModel):
   101→    t_patch_size: int = Field(default=1, description="Patch size for time dimension")
   102→    patch_size: int = Field(default=2, description="Patch size for spatial dimensions")
   103→    frame_receptive_field: int = Field(default=11, description="Frame receptive field")
   104→    spatial_rope_interpolation: Literal["inter", "extra"] = Field(
   105→        default="extra", description="Spatial rope interpolation method."
   106→    )
   107→    ref_audio_offset: int = Field(default=1000, description="Offset for reference audio.")
   108→    text_offset: int = Field(default=0, description="Offset for text.")
   109→    coords_style: Literal["v1", "v2"] = Field(default="v2", description="Coords style.")
   110→
   111→
   112→class EvaluationConfig(BaseModel):
   113→    """Evaluation configuration class defining parameters for model evaluation and inference"""
   114→
   115→    model_config = ConfigDict(protected_namespaces=())
   116→
   117→    data_proxy_config: DataProxyConfig = Field(default=DataProxyConfig(), description="Data proxy configuration")
   118→
   119→    fps: int = Field(default=25, description="Frames per second for video generation")
   120→    num_inference_steps: int = Field(default=32, description="Number of denoising steps during inference")
   121→    video_txt_guidance_scale: float = Field(default=5.0, description="Video text guidance scale for text conditioning")
   122→    audio_txt_guidance_scale: float = Field(default=5.0, description="Audio text guidance scale for text conditioning")
   123→    txt_encoder_type: Literal["t5_gemma"] = Field(default="t5_gemma", description="Text encoder type.")
   124→    t5_gemma_target_length: int = Field(default=640, description="Target length for T5-Gemma encoder.")
   125→    support_ref_audio: bool = Field(default=True, description="Whether to support the ref_audio feature")
   126→    shift: float = Field(default=5.0, description="Temporal shift parameter for video generation")
   127→    exp_name: str = Field(default="exp_debug", description="Experiment name with evaluation suffix")
   128→    audio_model_path: str = Field(default="", description="Path to the pretrained audio model")
   129→    txt_model_path: str = Field(default="", description="Path to the pretrained txt model")
   130→    vae_model_path: str = Field(default="", description="Path to the pretrained vae model")
   131→    vae_stride: Tuple[int, int, int] = Field(default=(4, 16, 16), description="VAE stride in format (time, height, width)")
   132→    z_dim: int = Field(default=48, description="Dimension of z space.")
   133→    patch_size: Tuple[int, int, int] = Field(default=(1, 2, 2), description="Patch size in format (time, height, width)")
   134→    cfg_number: int = Field(default=2, description="Classifier-free guidance number")
   135→    sr_cfg_number: int = Field(default=2, description="SR Classifier-free guidance number")
   136→
   137→    # flops recording
   138→    enable_flops_recording: bool = Field(default=False, description="Whether to enable flops recording")
   139→
   140→    # super resolution model configuration
   141→    use_sr_model: bool = Field(default=False, description="Whether to use the super resolution model")
   142→    sr_model_path: str = Field(default="", description="Path to the pretrained super resolution model")
   143→    sr_num_inference_steps: int = Field(default=5, description="Number of denoising steps during super resolution inference")
   144→    noise_value: int = Field(default=220, description="Noise value for the super resolution model")
   145→    sr_video_txt_guidance_scale: float = Field(
   146→        default=3.5, description="Super resolution video text guidance scale for text conditioning"
   147→    )
   148→    use_cfg_trick: bool = Field(default=True, description="Whether to use the cfg trick")
   149→    cfg_trick_start_frame: int = Field(default=13, description="Start frame for the cfg trick")
   150→    cfg_trick_value: float = Field(default=2.0, description="Value for the cfg trick")
   151→    using_sde_flag: bool = Field(default=False, description="Whether to use the sde flag")
   152→    sr_audio_noise_scale: float = Field(default=0.7, description="Noise scale for the super resolution audio")
   153→
   154→    # turbo-vae config
   155→    use_turbo_vae: bool = Field(default=True, description="Whether to use the turbo-vae")
   156→    student_config_path: str = Field(default="", description="Path to the student config")
   157→    student_ckpt_path: str = Field(default="", description="Path to the student checkpoint")
   158→
   159→
   160→class MagiPipelineConfig(BaseSettings):
   161→    engine_config: EngineConfig = Field(description="Engine configuration.", default_factory=EngineConfig)
   162→    arch_config: ModelConfig = Field(default=ModelConfig(), description="Model configuration.")
   163→    evaluation_config: EvaluationConfig = Field(default=EvaluationConfig(), description="Evaluation configuration.")
   164→    sr_arch_config: ModelConfig = Field(default=ModelConfig(), description="Super resolution model configuration.")
   165→    model_config = SettingsConfigDict(cli_parse_args=True, cli_ignore_unknown_args=True, cli_implicit_flags=True)
   166→
   167→    @classmethod
   168→    def settings_customise_sources(
   169→        cls,
   170→        settings_cls: type[BaseSettings],
   171→        init_settings: PydanticBaseSettingsSource,
   172→        env_settings: PydanticBaseSettingsSource,
   173→        dotenv_settings: PydanticBaseSettingsSource,
   174→        file_secret_settings: PydanticBaseSettingsSource,
   175→    ) -> tuple[PydanticBaseSettingsSource, ...]:
   176→        parser = argparse.ArgumentParser(allow_abbrev=False)
   177→        parser.add_argument("--config-load-path", type=str, default=None, help="Path to load the config.json from")
   178→        args, _ = parser.parse_known_args()
   179→        config_load_path = args.config_load_path
   180→        sources = [env_settings, CliSettingsSource(settings_cls, cli_parse_args=True, cli_ignore_unknown_args=True)]
   181→        if config_load_path:
   182→            sources.append(JsonConfigSettingsSource(settings_cls, json_file=config_load_path))
   183→
   184→        sources.extend([init_settings, dotenv_settings, file_secret_settings])
   185→        return tuple(sources)
   186→
   187→    def save_to_json(self, json_path: str, indent: int = 4):
   188→        path = Path(json_path)
   189→        path.parent.mkdir(parents=True, exist_ok=True)
   190→        path.write_text(self.__str__(indent=indent))
   191→
   192→    def __str__(self, indent: int = 4):
   193→        data = self.model_dump(mode="json")
   194→        formatted = json.dumps(data, indent=indent, ensure_ascii=False, sort_keys=False)
   195→        class_name = self.__class__.__name__
   196→        return f"{class_name}:\n{formatted}".replace('"', "")
   197→
   198→    def __repr__(self, indent: int = 4):
   199→        return self.__str__(indent=indent)
   200→
   201→    @model_validator(mode="after")
   202→    def validate_engine_config(self):
   203→        world_size = int(os.getenv("WORLD_SIZE", "1"))
   204→        self.engine_config.dp_size = world_size // (
   205→            self.engine_config.tp_size * self.engine_config.pp_size * self.engine_config.cp_size
   206→        )
   207→
   208→        assert world_size % self.engine_config.tp_size == 0
   209→        tp_pp_size = self.engine_config.tp_size * self.engine_config.pp_size
   210→        assert world_size % tp_pp_size == 0
   211→        tp_pp_cp_size = tp_pp_size * self.engine_config.cp_size
   212→        assert world_size % tp_pp_cp_size == 0
   213→        assert world_size == self.engine_config.dp_size * tp_pp_cp_size
   214→
   215→        if self.engine_config.tp_size == 1:
   216→            self.engine_config.sequence_parallel = False
   217→
   218→        return self
   219→
   220→    @model_validator(mode="after")
   221→    def post_override_config(self):
   222→        self.arch_config.num_heads_q = self.arch_config.hidden_size // self.arch_config.head_dim
   223→        self.arch_config.num_heads_kv = self.arch_config.num_query_groups
   224→
   225→        self.sr_arch_config = copy.deepcopy(self.arch_config)
   226→        if env_is_true("SR2_1080"):
   227→            self.sr_arch_config = copy.deepcopy(self.arch_config)
   228→            # fmt: off
   229→            self.sr_arch_config.local_attn_layers = [
   230→                0, 1, 2,
   231→                4, 5, 6,
   232→                8, 9, 10,
   233→                12, 13, 14,
   234→                16, 17, 18,
   235→                20, 21, 22,
   236→                24, 25, 26,
   237→                28, 29, 30,
   238→                32, 33, 34,
   239→                35, 36, 37,
   240→                38, 39,
   241→            ]
   242→            # fmt: on
   243→            self.evaluation_config.sr_video_txt_guidance_scale = 3.5
   244→
   245→        return self
   246→
   247→
   248→def prevent_unsupported_list_syntax():
   249→    """
   250→    Check sys.argv before Pydantic parsing to prevent using unsupported list syntax.
   251→    """
   252→    args = sys.argv[1:]
   253→    for i, arg in enumerate(args):
   254→        if i + 2 < len(args):
   255→            value1, value2 = args[i + 1], args[i + 2]
   256→            if not value1.startswith("-") and not value2.startswith("-"):
   257→                error_msg = (
   258→                    f"\n\nError: Detected list parameter '{arg}' using unsupported command line syntax.\n"
   259→                    f"Error pattern: '{arg} {value1} {value2} ...'\n\n"
   260→                    "Pydantic (or related libraries) do not support passing lists with space-separated multiple values.\n"
   261→                    "Please use one of the following supported formats:\n\n"
   262→                    f"1. JSON style:      {arg} '[{value1},{value2},...]'\n"
   263→                    f"2. Argparse style:  {arg} {value1} {arg} {value2}\n"
   264→                    f"3. Lazy style:      {arg} {value1},{value2}\n"
   265→                )
   266→                raise ValueError(error_msg)
   267→
   268→
   269→def parse_config(verbose: bool = False) -> MagiPipelineConfig:
   270→    parser = argparse.ArgumentParser(description="Load and optionally save config", allow_abbrev=False)
   271→    parser.add_argument("--config-save-path", type=str, default=None, help="Path to save the config.json to")
   272→    args, _ = parser.parse_known_args()
   273→
   274→    prevent_unsupported_list_syntax()
   275→    config = MagiPipelineConfig()
   276→
   277→    if args.config_save_path is not None:
   278→        config.save_to_json(args.config_save_path)
   279→
   280→    if verbose:
   281→        print_rank_0(config)
   282→
   283→    return config
   284→


     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 gc
    16→
    17→import torch
    18→from inference.infra.checkpoint import load_model_checkpoint
    19→from inference.infra.distributed import get_cp_rank, get_pp_rank, get_tp_rank
    20→from inference.utils import print_mem_info_rank_0, print_model_size, print_rank_0
    21→
    22→from .dit_module import DiTModel
    23→
    24→
    25→def get_dit(model_config, engine_config):
    26→    """Build and load DiT model."""
    27→    model = DiTModel(model_config=model_config)
    28→
    29→    print_rank_0("Build dit model successfully")
    30→    print_rank_0(model)
    31→    print_model_size(
    32→        model, prefix=f"(tp, cp, pp) rank ({get_tp_rank()}, {get_cp_rank()}, {get_pp_rank()}): ", print_func=print_rank_0
    33→    )
    34→
    35→    model = load_model_checkpoint(model, engine_config)
    36→    model.cuda(torch.cuda.current_device())
    37→    model.eval()
    38→    print_mem_info_rank_0("Load model successfully")
    39→
    40→    gc.collect()
    41→    torch.cuda.empty_cache()
    42→    return model
    43→


Now let me examine the DiT module implementation and attention mechanisms:


     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 importlib
    16→from dataclasses import dataclass
    17→from enum import Enum
    18→from typing import Any, Callable, List, Optional, Tuple
    19→
    20→import torch
    21→import torch.nn as nn
    22→from einops import rearrange, repeat
    23→from inference.common import Modality, VarlenHandler, is_hopper_arch
    24→from inference.infra.parallelism import ulysses_scheduler
    25→from magi_compiler import magi_compile
    26→from magi_compiler.api import magi_register_custom_op
    27→from magi_compiler.config import CompileConfig
    28→from torch import Tensor
    29→from torch.nn import Parameter
    30→
    31→
    32→@dataclass
    33→class FFAHandler:
    34→    q_ranges: torch.Tensor
    35→    k_ranges: torch.Tensor
    36→    max_seqlen_q: int
    37→    max_seqlen_k: int
    38→    attn_type_map: torch.Tensor
    39→    softmax_scale: float
    40→
    41→
    42→# Define the MLP activation type
    43→class MLPActivationType(Enum):
    44→    """Enumeration of supported activation functions for MLP"""
    45→
    46→    SWIGLU7 = "swiglu7"
    47→    GELU7 = "gelu7"
    48→
    49→
    50→def swiglu7(x, alpha: float = 1.702, limit: float = 7.0, out_dtype: Optional[torch.dtype] = None):
    51→    out_dtype = x.dtype if out_dtype is None else out_dtype
    52→    x = x.to(torch.float32)
    53→    x_glu, x_linear = x[..., ::2], x[..., 1::2]
    54→    # Clamp the input values
    55→    x_glu = x_glu.clamp(min=None, max=limit)
    56→    x_linear = x_linear.clamp(min=-limit, max=limit)
    57→    out_glu = x_glu * torch.sigmoid(alpha * x_glu)
    58→    # Note we add an extra bias of 1 to the linear layer (from GPT-OSS)
    59→    return (out_glu * (x_linear + 1)).to(out_dtype)
    60→
    61→
    62→def gelu7(x, alpha: float = 1.702, limit: float = 7.0, out_dtype: Optional[torch.dtype] = None):
    63→    out_dtype = x.dtype if out_dtype is None else out_dtype
    64→    x = x.to(torch.float32)
    65→    x_glu = x
    66→    # Clamp the input values
    67→    x_glu = x_glu.clamp(min=None, max=limit)
    68→    out_glu = x_glu * torch.sigmoid(alpha * x_glu)
    69→    # Note we add an extra bias of 1 to the linear layer
    70→    return out_glu.to(out_dtype)
    71→
    72→
    73→def create_activation_func(activation_type: MLPActivationType) -> Callable:
    74→    match activation_type:
    75→        case MLPActivationType.SWIGLU7:
    76→            return swiglu7
    77→        case MLPActivationType.GELU7:
    78→            return gelu7
    79→        case _:
    80→            raise ValueError(f"Unknown activation type: {activation_type}")
    81→
    82→
    83→class ModalityDispatcher:
    84→    permuted_modality_mapping: torch.Tensor
    85→    group_size: torch.Tensor
    86→    group_size_cpu: list[int]
    87→    num_modalities: int
    88→
    89→    def __init__(self, modality_mapping: torch.Tensor, num_modalities: int):
    90→        """
    91→        Initialize dispatcher.
    92→        This runs once during object construction and precomputes all mappings.
    93→        """
    94→        self.modality_mapping = modality_mapping
    95→        self.num_modalities = num_modalities
    96→
    97→        self.permuted_modality_mapping = self._precompute_permute_mapping(modality_mapping)
    98→
    99→        self.group_size = torch.bincount(self.permuted_modality_mapping, minlength=num_modalities).to(torch.int32)
   100→        self.group_size_cpu: list[int] = [int(x) for x in self.group_size.to("cpu").tolist()]
   101→
   102→    def _precompute_permute_mapping(self, modality_mapping):
   103→        # 1. Compute forward and inverse permutation mappings.
   104→        # argsort is an efficient O(N log N) operation.
   105→        self.permute_mapping = torch.argsort(modality_mapping)
   106→        self.inv_permute_mapping = torch.argsort(self.permute_mapping)
   107→
   108→        # 2. Compute group size for each modality.
   109→        # bincount is highly efficient for counting.
   110→        permuted_modality_mapping = modality_mapping[self.permute_mapping]
   111→
   112→        return permuted_modality_mapping
   113→
   114→    def dispatch(self, x: torch.Tensor) -> list[torch.Tensor]:
   115→        grouped_tensors = torch.split(x, self.group_size_cpu, dim=0)
   116→        return list(grouped_tensors)
   117→
   118→    def undispatch(self, *processed_groups: list[torch.Tensor]) -> torch.Tensor:
   119→        return torch.cat(processed_groups, dim=0)
   120→
   121→    @staticmethod
   122→    def permute(x: torch.Tensor, permute_mapping: torch.Tensor) -> torch.Tensor:
   123→        """Apply forward permutation to tensor."""
   124→        return x[permute_mapping]
   125→
   126→    @staticmethod
   127→    def inv_permute(x: torch.Tensor, inv_permute_mapping: torch.Tensor) -> torch.Tensor:
   128→        """Apply inverse permutation to tensor."""
   129→        return x[inv_permute_mapping]
   130→
   131→
   132→def freq_bands(
   133→    num_bands: int, temperature: float = 10000.0, step: int = 2, device: Optional[torch.device] = None
   134→) -> torch.Tensor:
   135→    exp = torch.arange(0, num_bands, step, dtype=torch.int64, device=device).to(torch.float32) / num_bands
   136→    bands = 1.0 / (temperature**exp)
   137→    return bands
   138→
   139→
   140→def rotate_half(x, interleaved=False):
   141→    if not interleaved:
   142→        x1, x2 = x.chunk(2, dim=-1)
   143→        return torch.cat((-x2, x1), dim=-1)
   144→    else:
   145→        x1, x2 = x[..., ::2], x[..., 1::2]
   146→        return rearrange(torch.stack((-x2, x1), dim=-1), "... d two -> ... (d two)", two=2)
   147→
   148→
   149→def apply_rotary_emb_torch(x, cos, sin, interleaved=False):
   150→    """
   151→    x: (batch_size, seqlen, nheads, headdim)
   152→    cos, sin: (seqlen, rotary_dim / 2) or (batch_size, seqlen, rotary_dim / 2)
   153→    """
   154→    ro_dim = cos.shape[-1] * 2
   155→    assert ro_dim <= x.shape[-1]
   156→    cos = repeat(cos, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)")
   157→    sin = repeat(sin, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)")
   158→    return torch.cat([x[..., :ro_dim] * cos + rotate_half(x[..., :ro_dim], interleaved) * sin, x[..., ro_dim:]], dim=-1)
   159→
   160→
   161→class ElementWiseFourierEmbed(nn.Module):
   162→    def __init__(
   163→        self,
   164→        dim: int,
   165→        max_res: int = 224,
   166→        temperature: float = 10000.0,
   167→        in_pixels: bool = True,
   168→        linear_bands: bool = False,
   169→        learnable: bool = False,
   170→        device: torch.device = torch.device("cpu"),
   171→        dtype: torch.dtype = torch.float32,
   172→    ):
   173→        """
   174→        Args:
   175→            dim: Output feature dimension, total channels, must be divisible by 6
   176→            max_res: Max pixel-frequency resolution for pixel-domain bands
   177→            temperature: Temperature in inverse-frequency mode
   178→            in_pixels: True -> pixel-frequency bands, False -> inverse-frequency bands
   179→            linear_bands: Whether pixel-frequency bands are linearly spaced
   180→            learnable: Whether frequency bands are trainable
   181→        """
   182→        super().__init__()
   183→        self.dim = dim
   184→        self.in_pixels = in_pixels
   185→        self.learnable = learnable
   186→        self.temperature = temperature
   187→        self.max_res = max_res
   188→        self.linear_bands = linear_bands
   189→        self.device = device
   190→        self.dtype = dtype
   191→        # Make frequency bands trainable or register as buffer
   192→        bands = self.get_default_bands()
   193→        if self.learnable:
   194→            self.bands = nn.Parameter(bands)
   195→        else:
   196→            self.register_buffer("bands", bands)
   197→
   198→    def forward(self, coords: torch.Tensor) -> torch.Tensor:
   199→        """
   200→        Args:


Let me continue reading the DiT module to find attention implementation:


   200→        Args:
   201→            coords: [L,9], column order (time, row, col, T, H, W, ref_T, ref_H, ref_W)
   202→        Returns:
   203→            emb: [L, dim] element-wise Fourier embedding
   204→        """
   205→        # Use slicing instead of unbind + stack to reduce intermediates
   206→        coords_xyz = coords[:, :3]  # [L,3] -> (t, h, w)
   207→        sizes = coords[:, 3:6]  # [L,3] -> (T, H, W)
   208→        refs = coords[:, 6:9]  # [L,3] -> (ref_T, ref_H, ref_W)
   209→
   210→        # Compute scale factors
   211→        scales = (refs - 1) / (sizes - 1)  # [L,3]
   212→
   213→        # NOTE: if both ref and size are 1, scale is fixed to 1; otherwise invalid
   214→        scales[(refs == 1) & (sizes == 1)] = 1
   215→        assert not scales.isnan().any(), "scales has nan"
   216→        assert not scales.isinf().any(), "scales has inf"
   217→
   218→        # Center alignment: apply to h,w only (not time)
   219→        centers = (sizes - 1) / 2  # [L,3]
   220→        centers[:, 0] = 0  # Do not center the time dimension
   221→        coords_xyz = coords_xyz - centers  # [L,3]
   222→
   223→        # Project to frequency bands in one shot: [L,3,B]
   224→        proj = coords_xyz.unsqueeze(-1) * scales.unsqueeze(-1) * self.bands
   225→
   226→        # Compute sin & cos and concatenate
   227→        sin_proj = proj.sin()  # [L,3,B]
   228→        cos_proj = proj.cos()
   229→
   230→        return torch.cat((sin_proj, cos_proj), dim=1).flatten(1)
   231→
   232→    def reset_parameters(self):
   233→        bands = self.get_default_bands()
   234→        self.bands.copy_(bands)
   235→
   236→    def get_default_bands(self):
   237→        if self.in_pixels:
   238→            raise NotImplementedError("in_pixels are not implemented yet")
   239→        else:
   240→            bands = freq_bands(self.dim // 8, temperature=self.temperature, step=1, device=self.device).to(self.dtype)
   241→        return bands
   242→
   243→
   244→class MultiModalityRMSNorm(nn.Module):
   245→    __constants__ = ["dim", "eps", "num_modality"]
   246→    dim: int
   247→    eps: float
   248→    num_modality: int
   249→
   250→    def __init__(self, dim: int, eps: float = 1e-6, device: torch.device | None = None, num_modality: int = 1):
   251→        super().__init__()
   252→        self.dim = dim
   253→        self.eps = eps
   254→        self.num_modality = num_modality
   255→
   256→        self.weight = torch.nn.Parameter(torch.zeros(dim * num_modality, device=device, dtype=torch.float32))
   257→        if num_modality > 1:
   258→            self.forward = self.forward_multi_experts
   259→        else:
   260→            self.forward = self.forward_single_expert
   261→
   262→        self.reset_parameters()
   263→
   264→    def reset_parameters(self):
   265→        nn.init.zeros_(self.weight)
   266→
   267→    def rms(self, x: torch.Tensor) -> torch.Tensor:
   268→        t, original_dtype = x.float(), x.dtype
   269→        t = t * torch.rsqrt(torch.mean(t**2, dim=-1, keepdim=True) + self.eps)
   270→        return t
   271→
   272→    def forward_multi_experts(self, x: torch.Tensor, modality_dispatcher: ModalityDispatcher) -> torch.Tensor:
   273→        original_dtype = x.dtype
   274→        t = self.rms(x)
   275→
   276→        weight_chunked = self.weight.chunk(self.num_modality, dim=0)
   277→        t_list = modality_dispatcher.dispatch(t)
   278→        for i in range(self.num_modality):
   279→            t_list[i] = t_list[i] * (weight_chunked[i] + 1)
   280→        t = modality_dispatcher.undispatch(*t_list)
   281→
   282→        return t.to(original_dtype)
   283→
   284→    def forward_single_expert(self, x: torch.Tensor, modality_dispatcher: Optional[ModalityDispatcher] = None) -> torch.Tensor:
   285→        t, original_dtype = x.float(), x.dtype
   286→        t = t * torch.rsqrt(torch.mean(t**2, dim=-1, keepdim=True) + self.eps)
   287→        return (t * (self.weight + 1)).to(original_dtype)
   288→
   289→
   290→class _BF16ComputeLinear(torch.autograd.Function):
   291→    @staticmethod
   292→    def forward(
   293→        ctx,
   294→        input: torch.Tensor,
   295→        weight: torch.Tensor,
   296→        bias: Optional[torch.Tensor],
   297→        output_dtype: Optional[torch.dtype],
   298→        compute_dtype: torch.dtype = torch.bfloat16,
   299→    ):
   300→        # Convert input to specified input data type
   301→        input_cast = input.to(compute_dtype)
   302→        # Convert weight to computation data type
   303→        weight_cast = weight.to(compute_dtype)
   304→        # Perform linear operation
   305→        output = torch.matmul(input_cast, weight_cast.t())
   306→
   307→        # Add bias if present
   308→        if bias is not None:
   309→            bias_cast = bias.to(compute_dtype)
   310→            output = output + bias_cast
   311→        else:
   312→            bias_cast = None
   313→
   314→        # Convert output to specified output data type
   315→        return output.to(output_dtype)
   316→
   317→
   318→class BaseLinear(nn.Module):
   319→    __constants__ = ["in_features", "out_features", "num_layers", "num_experts"]
   320→    in_features: int
   321→    out_features: int
   322→    num_layers_for_initialization: int
   323→    num_experts: int
   324→    weight: Tensor
   325→
   326→    def __init__(
   327→        self, in_features, out_features, num_layers_for_initialization, num_experts, bias=True, device=None, dtype=None
   328→    ):
   329→        super().__init__()
   330→        factory_kwargs = {"device": device, "dtype": torch.bfloat16}
   331→        self.in_features = in_features
   332→        self.out_features = out_features
   333→        self.num_layers_for_initialization = num_layers_for_initialization
   334→        self.num_experts = num_experts
   335→        self.use_bias = bias
   336→        self.weight = Parameter(torch.empty((out_features * num_experts, in_features), **factory_kwargs))
   337→        if bias:
   338→            self.bias = Parameter(torch.empty(out_features * num_experts, **factory_kwargs))
   339→        else:
   340→            self.register_parameter("bias", None)
   341→
   342→    def forward(
   343→        self,
   344→        input: torch.Tensor,
   345→        output_dtype: Optional[torch.dtype] = None,
   346→        modality_dispatcher: Optional[ModalityDispatcher] = None,
   347→    ) -> torch.Tensor:
   348→        output_dtype = input.dtype if output_dtype is None else output_dtype
   349→        return _BF16ComputeLinear.apply(input, self.weight, self.bias, output_dtype, torch.bfloat16)
   350→
   351→
   352→class NativeMoELinear(BaseLinear):
   353→    def forward(
   354→        self,
   355→        input: torch.Tensor,
   356→        output_dtype: Optional[torch.dtype] = None,
   357→        modality_dispatcher: Optional[ModalityDispatcher] = None,
   358→    ) -> torch.Tensor:
   359→        output_dtype = input.dtype if output_dtype is None else output_dtype
   360→
   361→        input_list = modality_dispatcher.dispatch(input)  # type: ignore
   362→        weight_chunked = self.weight.chunk(self.num_experts, dim=0)
   363→
   364→        if self.bias is not None:
   365→            bias_chunked = self.bias.chunk(self.num_experts, dim=0)
   366→
   367→        for i in range(self.num_experts):
   368→            input_list[i] = _BF16ComputeLinear.apply(
   369→                input_list[i],
   370→                weight_chunked[i],
   371→                bias_chunked[i] if self.bias is not None else None,
   372→                output_dtype,
   373→                torch.bfloat16,
   374→            )
   375→        return modality_dispatcher.undispatch(*input_list)  # type: ignore
   376→
   377→
   378→def create_linear(
   379→    in_features, out_features, num_layers=1, num_experts=1, bias=True, device=None, dtype=None
   380→) -> BaseLinear | NativeMoELinear:
   381→    if num_experts == 1:
   382→        return BaseLinear(in_features, out_features, num_layers, num_experts, bias, device, dtype)
   383→    else:
   384→        return NativeMoELinear(in_features, out_features, num_layers, num_experts, bias, device, dtype)
   385→
   386→
   387→HAS_MAGI_ATTENTION = importlib.util.find_spec("magi_attention") is not None
   388→HAS_FA3 = importlib.util.find_spec("flash_attn_interface") is not None
   389→
   390→
   391→@magi_register_custom_op(name="infra::flash_attn_func", is_subgraph_boundary=True)
   392→def flash_attn_func(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor) -> torch.Tensor:
   393→    if HAS_FA3 and is_hopper_arch():
   394→        from flash_attn_interface import flash_attn_func as fa3_flash_attn_func
   395→
   396→        return fa3_flash_attn_func(query, key, value)
   397→    else:
   398→        from flash_attn.flash_attn_interface import flash_attn_func as fa2_flash_attn_func
   399→
   400→        return fa2_flash_attn_func(query, key, value)
   401→
   402→
   403→def _split_q_range_with_no_overlap(
   404→    q_ranges: torch.Tensor, k_ranges: torch.Tensor
   405→) -> Tuple[List[List[int]], List[List[List[int]]]]:
   406→    range_boundary = torch.unique(q_ranges, sorted=True).tolist()
   407→    candidates = [[start, end, []] for start, end in zip(range_boundary[:-1], range_boundary[1:])]
   408→    q_ranges = q_ranges.tolist()
   409→    k_ranges = k_ranges.tolist()
   410→    for q_range, k_range in zip(q_ranges, k_ranges):
   411→        q_start, q_end = q_range
   412→        for q_range_cand in candidates:
   413→            if q_start <= q_range_cand[0] and q_range_cand[1] <= q_end:
   414→                q_range_cand[2].append(k_range)
   415→    q_ranges_out = []
   416→    k_ranges_out = []
   417→    for q_range_cand in candidates:
   418→        if len(q_range_cand[2]) > 0:
   419→            q_ranges_out.append(q_range_cand[0:2])
   420→            k_ranges_out.append(q_range_cand[2])
   421→    return q_ranges_out, k_ranges_out
   422→
   423→
   424→def _flash_attn_with_correction(
   425→    query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, q_ranges: List[List[int]], k_range_list: List[List[List[int]]]
   426→):
   427→    output = torch.zeros_like(query)
   428→    output_lse = torch.zeros((query.shape[0], query.shape[1]), dtype=torch.float32, device=query.device)
   429→
   430→    from flash_attn.flash_attn_interface import flash_attn_func
   431→
   432→    for q_range, k_ranges in zip(q_ranges, k_range_list):
   433→        q_start, q_end = q_range
   434→        qo_out, qo_lse = None, None
   435→        for k_range in k_ranges:
   436→            k_start, k_end = k_range
   437→            cur_qo_out, cur_qo_lse, _ = flash_attn_func(
   438→                query[q_start:q_end].unsqueeze(0),
   439→                key[k_start:k_end].unsqueeze(0),
   440→                value[k_start:k_end].unsqueeze(0),
   441→                return_attn_probs=True,
   442→            )
   443→            cur_qo_out, cur_qo_lse = cur_qo_out.squeeze(0), cur_qo_lse.squeeze(0)
   444→
   445→            if qo_out is None:
   446→                qo_out = cur_qo_out
   447→                qo_lse = cur_qo_lse
   448→            else:
   449→                qo_lse[qo_lse == torch.inf] = -torch.inf
   450→                cur_qo_lse[cur_qo_lse == torch.inf] = -torch.inf
   451→                max_lse = torch.max(qo_lse, cur_qo_lse)
   452→                qo_se, cur_qo_se = torch.exp(qo_lse - max_lse), torch.exp(cur_qo_lse - max_lse)
   453→                sum_se = qo_se + cur_qo_se
   454→                qo_scale, cur_qo_scale = qo_se / sum_se, cur_qo_se / sum_se
   455→
   456→                qo_out = qo_out * qo_scale.permute(1, 0).unsqueeze(-1) + cur_qo_out * cur_qo_scale.permute(1, 0).unsqueeze(-1)
   457→                qo_lse = torch.log(sum_se) + max_lse
   458→
   459→        output[q_start:q_end] = qo_out
   460→        output_lse[q_start:q_end, :] = qo_lse.permute(1, 0)
   461→    return output, output_lse
   462→
   463→
   464→def _custom_flex_flash_attn_func(
   465→    query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, q_ranges: torch.Tensor, k_ranges: torch.Tensor, **kwargs
   466→):
   467→    q_ranges, k_range_list = _split_q_range_with_no_overlap(q_ranges, k_ranges)
   468→    return _flash_attn_with_correction(query, key, value, q_ranges, k_range_list)
   469→
   470→
   471→def _flex_flash_attn_func_infer_output_meta(
   472→    query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, q_ranges: torch.Tensor, k_ranges: torch.Tensor
   473→) -> Tuple[torch.Tensor, torch.Tensor]:
   474→    output = torch.empty_like(query)
   475→    output_lse = torch.empty((query.shape[0], query.shape[1]), dtype=torch.float32, device=query.device)
   476→    return output, output_lse
   477→
   478→
   479→@magi_register_custom_op(
   480→    name="infra::flex_flash_attn_func",
   481→    mutates_args=(),
   482→    infer_output_meta_fn=_flex_flash_attn_func_infer_output_meta,
   483→    is_subgraph_boundary=True,
   484→)
   485→def flex_flash_attn_func(
   486→    query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, q_ranges: torch.Tensor, k_ranges: torch.Tensor
   487→) -> Tuple[torch.Tensor, torch.Tensor]:
   488→    if HAS_MAGI_ATTENTION and is_hopper_arch():
   489→        from magi_attention.api import flex_flash_attn_func as magi_flex_flash_attn_func
   490→
   491→        return magi_flex_flash_attn_func(query, key, value, q_ranges, k_ranges)
   492→    else:
   493→        return _custom_flex_flash_attn_func(query, key, value, q_ranges, k_ranges)
   494→
   495→
   496→def _attention_with_cp_infer_output_meta(q: torch.Tensor, *args, **kwargs) -> torch.Tensor:
   497→    return torch.empty_like(q, dtype=torch.bfloat16).squeeze(0)
   498→
   499→


Let me continue and look for the actual attention layer implementation:


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


559:class AttentionConfig:
572:class Attention(torch.nn.Module):


   559→class AttentionConfig:
   560→    hidden_size: int
   561→    num_heads_q: int
   562→    num_heads_kv: int
   563→    head_dim: int
   564→    params_dtype: torch.dtype
   565→    checkpoint_qk_layernorm_rope: bool
   566→    num_modality: int
   567→    num_layers: int
   568→    use_local_attn: bool = False
   569→    enable_attn_gating: bool = False
   570→
   571→
   572→class Attention(torch.nn.Module):
   573→    config: AttentionConfig
   574→
   575→    def __init__(self, config: AttentionConfig):
   576→        super().__init__()
   577→        self.config = config
   578→
   579→        self.pre_norm = MultiModalityRMSNorm(config.hidden_size, eps=1e-6, num_modality=config.num_modality)
   580→        self.gating_size = config.num_heads_q if config.enable_attn_gating else 0
   581→
   582→        self.linear_qkv = create_linear(
   583→            config.hidden_size,
   584→            config.num_heads_q * config.head_dim + config.num_heads_kv * config.head_dim * 2 + self.gating_size,
   585→            num_experts=config.num_modality,
   586→            bias=False,
   587→            dtype=config.params_dtype,
   588→            num_layers=config.num_layers,
   589→        )
   590→        self.linear_proj = create_linear(
   591→            config.num_heads_q * config.head_dim,
   592→            config.hidden_size,
   593→            bias=False,
   594→            num_experts=config.num_modality,
   595→            dtype=config.params_dtype,
   596→            num_layers=config.num_layers,
   597→        )
   598→        self.q_norm = MultiModalityRMSNorm(config.head_dim, num_modality=config.num_modality)
   599→        self.k_norm = MultiModalityRMSNorm(config.head_dim, num_modality=config.num_modality)
   600→
   601→        self.q_size = config.num_heads_q * config.head_dim
   602→        self.kv_size = config.num_heads_kv * config.head_dim
   603→
   604→    def reset_parameters(self):
   605→        if hasattr(self.linear_proj, "reset_parameters_output_layer"):
   606→            self.linear_proj.reset_parameters_output_layer()
   607→
   608→    def forward(
   609→        self,
   610→        hidden_states: torch.Tensor,
   611→        rope: torch.Tensor,
   612→        permute_mapping: torch.Tensor,
   613→        inv_permute_mapping: torch.Tensor,
   614→        varlen_handler: VarlenHandler,
   615→        local_attn_handler: FFAHandler,
   616→        modality_dispatcher: ModalityDispatcher,
   617→        cp_split_sizes: List[int],
   618→    ) -> torch.Tensor:
   619→        hidden_states = self.pre_norm(hidden_states, modality_dispatcher=modality_dispatcher).to(torch.bfloat16)
   620→        qkv: torch.Tensor = self.linear_qkv(hidden_states, modality_dispatcher=modality_dispatcher).to(torch.float32)
   621→
   622→        q, k, v, g = torch.split(qkv, [self.q_size, self.kv_size, self.kv_size, self.gating_size], dim=1)
   623→        q = q.view(-1, self.config.num_heads_q, self.config.head_dim)
   624→        k = k.view(-1, self.config.num_heads_kv, self.config.head_dim)
   625→        v = v.view(-1, self.config.num_heads_kv, self.config.head_dim)
   626→        g = g.view(k.shape[0], self.config.num_heads_q, -1)
   627→
   628→        q = self.q_norm(q, modality_dispatcher=modality_dispatcher)
   629→        k = self.k_norm(k, modality_dispatcher=modality_dispatcher)
   630→
   631→        q = ModalityDispatcher.inv_permute(q, inv_permute_mapping).unsqueeze(0)
   632→        k = ModalityDispatcher.inv_permute(k, inv_permute_mapping).unsqueeze(0)
   633→        v = ModalityDispatcher.inv_permute(v, inv_permute_mapping).unsqueeze(0)
   634→
   635→        sin_emb, cos_emb = rope.tensor_split(2, -1)
   636→        q = apply_rotary_emb_torch(q, cos_emb, sin_emb)
   637→        k = apply_rotary_emb_torch(k, cos_emb, sin_emb)
   638→
   639→        if self.config.use_local_attn:
   640→            self_attn_out = flex_flash_attn_with_cp(
   641→                q, k, v, local_attn_handler.q_ranges, local_attn_handler.k_ranges, cp_split_sizes
   642→            )
   643→        else:
   644→            self_attn_out = flash_attn_with_cp(q, k, v, cp_split_sizes)
   645→        self_attn_out = ModalityDispatcher.permute(self_attn_out, permute_mapping)
   646→
   647→        if self.config.enable_attn_gating:
   648→            self_attn_out = self_attn_out * torch.sigmoid(g)
   649→
   650→        self_attn_out = self_attn_out.view(-1, self.config.num_heads_q * self.config.head_dim).to(torch.bfloat16)
   651→        out = self.linear_proj(self_attn_out, modality_dispatcher=modality_dispatcher)
   652→        return out
   653→
   654→
   655→@dataclass
   656→class MLPConfig:
   657→    hidden_size: int
   658→    intermediate_size: int
   659→    activation_type: MLPActivationType
   660→    params_dtype: torch.dtype
   661→    num_modality: int = 1
   662→    num_layers: int = 1
   663→    gated_act: bool = False
   664→
   665→
   666→class MLP(torch.nn.Module):
   667→    config: MLPConfig
   668→
   669→    def __init__(self, config: MLPConfig):
   670→        super().__init__()
   671→        num_experts = config.num_modality
   672→        self.pre_norm = MultiModalityRMSNorm(config.hidden_size, num_modality=config.num_modality)
   673→        intermediate_size_up = config.intermediate_size * 2 if config.gated_act else config.intermediate_size
   674→
   675→        self.up_gate_proj = create_linear(
   676→            config.hidden_size,
   677→            intermediate_size_up,
   678→            bias=False,
   679→            dtype=config.params_dtype,
   680→            num_layers=config.num_layers,
   681→            num_experts=num_experts,
   682→        )
   683→        self.down_proj = create_linear(
   684→            config.intermediate_size,
   685→            config.hidden_size,
   686→            bias=False,
   687→            dtype=config.params_dtype,
   688→            num_layers=config.num_layers,
   689→            num_experts=num_experts,
   690→        )
   691→        self.activation_func = create_activation_func(config.activation_type)
   692→
   693→    def forward(self, x: torch.Tensor, modality_dispatcher: ModalityDispatcher) -> torch.Tensor:
   694→        x = self.pre_norm(x, modality_dispatcher=modality_dispatcher).to(torch.bfloat16)
   695→        x = self.up_gate_proj(x, modality_dispatcher=modality_dispatcher).to(torch.float32)
   696→        x = self.activation_func(x).to(torch.bfloat16)
   697→        x = self.down_proj(x, modality_dispatcher=modality_dispatcher).to(torch.float32)
   698→        return x
   699→
   700→    def extra_repr(self) -> str:
   701→        return f"{self.up_gate_proj.weight.shape=}, {self.down_proj.weight.shape=}"
   702→
   703→
   704→@dataclass
   705→class AdapterConfig:
   706→    hidden_size: int
   707→    num_attention_heads: int
   708→    text_in_channels: int
   709→    video_in_channels: int
   710→    audio_in_channels: int
   711→    params_dtype: torch.dtype
   712→
   713→
   714→class Adapter(torch.nn.Module):
   715→    config: AdapterConfig
   716→
   717→    def __init__(self, config: AdapterConfig):
   718→        super().__init__()
   719→        self.config = config
   720→        self.video_embedder = nn.Linear(config.video_in_channels, config.hidden_size, bias=True, dtype=torch.float32)
   721→        self.text_embedder = nn.Linear(config.text_in_channels, config.hidden_size, bias=True, dtype=torch.float32)
   722→        self.audio_embedder = nn.Linear(config.audio_in_channels, config.hidden_size, bias=True, dtype=torch.float32)
   723→        self.rope = ElementWiseFourierEmbed(config.hidden_size // config.num_attention_heads, in_pixels=False, learnable=False)
   724→
   725→    def forward(
   726→        self,
   727→        x: torch.Tensor,
   728→        coords_mapping: torch.Tensor,
   729→        video_mask: torch.Tensor,
   730→        audio_mask: torch.Tensor,
   731→        text_mask: torch.Tensor,
   732→    ):
   733→        rope = self.rope(coords_mapping)
   734→        output_x = torch.zeros(x.shape[0], self.config.hidden_size, device=x.device, dtype=x.dtype)
   735→        output_x[text_mask] = self.text_embedder(x[text_mask, : self.config.text_in_channels])
   736→        output_x[audio_mask] = self.audio_embedder(x[audio_mask, : self.config.audio_in_channels])
   737→        output_x[video_mask] = self.video_embedder(x[video_mask, : self.config.video_in_channels])
   738→        return output_x, rope
   739→
   740→
   741→class TransFormerLayer(torch.nn.Module):
   742→    def __init__(self, config: Any, layer_idx: int):
   743→        super().__init__()
   744→        num_modality = 3 if layer_idx in config.mm_layers else 1
   745→        use_local_attn = layer_idx in config.local_attn_layers
   746→        self.post_norm = layer_idx in config.post_norm_layers
   747→        attention_config = AttentionConfig(
   748→            hidden_size=config.hidden_size,
   749→            num_heads_q=config.num_heads_q,
   750→            num_heads_kv=config.num_heads_kv,
   751→            head_dim=config.head_dim,
   752→            params_dtype=config.params_dtype,
   753→            checkpoint_qk_layernorm_rope=config.checkpoint_qk_layernorm_rope,
   754→            num_modality=num_modality,
   755→            num_layers=config.num_layers,
   756→            use_local_attn=use_local_attn,
   757→            enable_attn_gating=config.enable_attn_gating,
   758→        )
   759→        self.attention: Attention = Attention(attention_config)
   760→
   761→        activation_type = MLPActivationType.GELU7 if layer_idx in config.gelu7_layers else MLPActivationType.SWIGLU7
   762→        if activation_type == MLPActivationType.SWIGLU7:
   763→            gated_act = True
   764→            intermediate_size = int(config.hidden_size * 4 * 2 / 3) // 4 * 4
   765→        else:
   766→            gated_act = False
   767→            intermediate_size = config.hidden_size * 4
   768→        mlp_config = MLPConfig(
   769→            hidden_size=config.hidden_size,
   770→            intermediate_size=intermediate_size,
   771→            activation_type=activation_type,
   772→            params_dtype=config.params_dtype,
   773→            num_modality=num_modality,
   774→            num_layers=config.num_layers,
   775→            gated_act=gated_act,
   776→        )
   777→        self.mlp: MLP = MLP(mlp_config)
   778→        if self.post_norm:
   779→            self.attn_post_norm = MultiModalityRMSNorm(config.hidden_size, num_modality=num_modality)
   780→            self.mlp_post_norm = MultiModalityRMSNorm(config.hidden_size, num_modality=num_modality)
   781→
   782→    def forward(
   783→        self,
   784→        hidden_states: torch.Tensor,
   785→        rope: torch.Tensor,
   786→        permute_mapping: torch.Tensor,
   787→        inv_permute_mapping: torch.Tensor,
   788→        varlen_handler: VarlenHandler,
   789→        local_attn_handler: FFAHandler,
   790→        modality_dispatcher: ModalityDispatcher,
   791→        cp_split_sizes: List[int],
   792→    ) -> torch.Tensor:
   793→        attn_out = self.attention(
   794→            hidden_states,
   795→            rope,
   796→            permute_mapping,
   797→            inv_permute_mapping,
   798→            varlen_handler,
   799→            local_attn_handler,
   800→            modality_dispatcher,
   801→            cp_split_sizes,
   802→        )
   803→        if self.post_norm:
   804→            attn_out = self.attn_post_norm(attn_out, modality_dispatcher=modality_dispatcher)
   805→        hidden_states = hidden_states + attn_out
   806→
   807→        mlp_out = self.mlp(hidden_states, modality_dispatcher)
   808→        if self.post_norm:


Now let me look at the infra initialization and distributed setup:


     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 os
    16→from datetime import timedelta
    17→
    18→import torch
    19→
    20→from inference.common import parse_config
    21→
    22→from .parallel_state import initialize_model_parallel, model_parallel_is_initialized
    23→from inference.utils import print_rank_0
    24→
    25→
    26→def initialize_distributed():
    27→    """Initialize torch.distributed and core model parallel."""
    28→    config = parse_config()
    29→
    30→    device_count = torch.cuda.device_count()
    31→    if torch.distributed.is_initialized():
    32→        if torch.distributed.get_rank() == 0:
    33→            print_rank_0("> torch distributed already initialized, skipping initialization ...")
    34→    else:
    35→        rank = int(os.getenv("RANK", "0"))
    36→        world_size = int(os.getenv("WORLD_SIZE", "1"))
    37→        if rank == 0:
    38→            print_rank_0("> initializing torch distributed ...")
    39→        # Manually set the device ids.
    40→        if device_count > 0:
    41→            device = rank % device_count
    42→            torch.cuda.set_device(device)
    43→        # Call the init process
    44→        torch.distributed.init_process_group(
    45→            backend=config.engine_config.distributed_backend,
    46→            world_size=world_size,
    47→            rank=rank,
    48→            timeout=timedelta(minutes=config.engine_config.distributed_timeout_minutes),
    49→        )
    50→
    51→    # Set the tp, pp and dp communicators.
    52→    if device_count > 0:
    53→        if model_parallel_is_initialized():
    54→            return
    55→        initialize_model_parallel(
    56→            tp_size=config.engine_config.tp_size,
    57→            pp_size=config.engine_config.pp_size,
    58→            cp_size=config.engine_config.cp_size,
    59→            nccl_communicator_config_path=None,
    60→            distributed_timeout_minutes=config.engine_config.distributed_timeout_minutes,
    61→            order="tp-cp-pp-dp",
    62→        )
    63→