vllm-omni/vllm_omni/diffusion/models/magi_human/pipeline_magi_human.py

# SPDX-License-Identifier: Apache-2.0
# Copyright (c) 2026 SandAI. All Rights Reserved.
# Ported from daVinci-MagiHuman inference/pipeline/video_generate.py
# Adapted for vllm-omni: single-GPU, diffusers VAE, configurable dit_subfolder.

from __future__ import annotations

import json
import logging
import math
import os
from collections.abc import Iterable
from dataclasses import dataclass
from pathlib import Path
from typing import Any, ClassVar, Literal

import numpy as np
import torch
import torch.nn as nn
import whisper
from diffusers.configuration_utils import ConfigMixin, register_to_config
from diffusers.schedulers.scheduling_utils import (
    KarrasDiffusionSchedulers,
    SchedulerMixin,
    SchedulerOutput,
)
from diffusers.utils import deprecate, load_image
from diffusers.utils.torch_utils import randn_tensor
from diffusers.video_processor import VideoProcessor
from einops import rearrange
from PIL import Image
from safetensors.torch import load_file
from torch.nn import functional as F
from torch.nn.utils import weight_norm
from transformers import AutoTokenizer
from transformers.models.t5gemma import T5GemmaEncoderModel
from vllm.distributed import (
    get_tensor_model_parallel_world_size,
)
from vllm.model_executor.model_loader.weight_utils import default_weight_loader

from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import (
    DistributedAutoencoderKLWan,
)
from vllm_omni.diffusion.model_loader.diffusers_loader import (
    DiffusersPipelineLoader,
)
from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
from vllm_omni.diffusion.models.t5_encoder.t5_gemma_encoder import T5GemmaEncoderModelTP
from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import (
    DiffusionPipelineProfilerMixin,
)
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch

from .magi_human_dit import (
    DiTModel,
    FFAHandler,
    MagiHumanDiTConfig,
    Modality,
    VarlenHandler,
)

logger = logging.getLogger(__name__)


# ===========================================================================
# Scheduler (ported from daVinci-MagiHuman inference/pipeline/scheduler_unipc.py)
# ===========================================================================
class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin):
    _compatibles = [e.name for e in KarrasDiffusionSchedulers]
    order = 1

    @register_to_config
    def __init__(
        self,
        num_train_timesteps: int = 1000,
        solver_order: int = 2,
        prediction_type: str = "flow_prediction",
        shift: float = 1.0,
        use_dynamic_shifting=False,
        thresholding: bool = False,
        dynamic_thresholding_ratio: float = 0.995,
        sample_max_value: float = 1.0,
        predict_x0: bool = True,
        solver_type: str = "bh2",
        lower_order_final: bool = True,
        disable_corrector: list[int] = [],
        solver_p: SchedulerMixin = None,
        timestep_spacing: str = "linspace",
        steps_offset: int = 0,
        final_sigmas_type: str | None = "zero",
    ):
        if solver_type not in ["bh1", "bh2"]:
            if solver_type in ["midpoint", "heun", "logrho"]:
                self.register_to_config(solver_type="bh2")
            else:
                raise NotImplementedError(f"{solver_type} is not implemented for {self.__class__}")

        self.predict_x0 = predict_x0
        self.num_inference_steps = None
        alphas = np.linspace(1, 1 / num_train_timesteps, num_train_timesteps)[::-1].copy()
        sigmas = 1.0 - alphas
        sigmas = torch.from_numpy(sigmas).to(dtype=torch.float32)

        if not use_dynamic_shifting:
            sigmas = shift * sigmas / (1 + (shift - 1) * sigmas)

        self.sigmas = sigmas
        self.timesteps = sigmas * num_train_timesteps

        self.model_outputs = [None] * solver_order
        self.timestep_list = [None] * solver_order
        self.lower_order_nums = 0
        self.disable_corrector = disable_corrector
        self.solver_p = solver_p
        self.last_sample = None
        self._step_index: int | None = None
        self._begin_index: int | None = None

        self.sigmas = self.sigmas.to("cpu")
        self.sigma_min = self.sigmas[-1].item()
        self.sigma_max = self.sigmas[0].item()

    @property
    def step_index(self):
        return self._step_index

    @property
    def begin_index(self):
        return self._begin_index

    def set_begin_index(self, begin_index: int = 0):
        self._begin_index = begin_index

    def set_timesteps(
        self,
        num_inference_steps: int | None = None,
        device: str | torch.device = None,
        sigmas: list[float] | None = None,
        mu: float | None | None = None,
        shift: float | None | None = None,
    ):
        if self.config.use_dynamic_shifting and mu is None:
            raise ValueError(" you have to pass a value for `mu` when `use_dynamic_shifting` is set to be `True`")

        if sigmas is None:
            sigmas = np.linspace(self.sigma_max, self.sigma_min, num_inference_steps + 1).copy()[:-1]

        if self.config.use_dynamic_shifting:
            sigmas = self.time_shift(mu, 1.0, sigmas)
        else:
            if shift is None:
                shift = self.config.shift
            sigmas = shift * sigmas / (1 + (shift - 1) * sigmas)

        if self.config.final_sigmas_type == "sigma_min":
            sigma_last = ((1 - self.alphas_cumprod[0]) / self.alphas_cumprod[0]) ** 0.5
        elif self.config.final_sigmas_type == "zero":
            sigma_last = 0
        else:
            raise ValueError(
                f"`final_sigmas_type` must be one of 'zero', or 'sigma_min', but got {self.config.final_sigmas_type}"
            )

        timesteps = sigmas * self.config.num_train_timesteps
        sigmas = np.concatenate([sigmas, [sigma_last]]).astype(np.float32)

        self.sigmas = torch.from_numpy(sigmas)
        self.timesteps = torch.from_numpy(timesteps).to(device=device, dtype=torch.int64)

        self.num_inference_steps = len(timesteps)

        self.model_outputs = [None] * self.config.solver_order
        self.lower_order_nums = 0
        self.last_sample = None
        if self.solver_p:
            self.solver_p.set_timesteps(self.num_inference_steps, device=device)

        self._step_index = None
        self._begin_index = None
        self.sigmas = self.sigmas.to("cpu")

    def _threshold_sample(self, sample: torch.Tensor) -> torch.Tensor:
        dtype = sample.dtype
        batch_size, channels, *remaining_dims = sample.shape

        if dtype not in (torch.float32, torch.float64):
            sample = sample.float()

        sample = sample.reshape(batch_size, channels * np.prod(remaining_dims))
        abs_sample = sample.abs()
        s = torch.quantile(abs_sample, self.config.dynamic_thresholding_ratio, dim=1)
        s = torch.clamp(s, min=1, max=self.config.sample_max_value)
        s = s.unsqueeze(1)
        sample = torch.clamp(sample, -s, s) / s
        sample = sample.reshape(batch_size, channels, *remaining_dims)
        return sample.to(dtype)

    def _sigma_to_t(self, sigma):
        return sigma * self.config.num_train_timesteps

    def _sigma_to_alpha_sigma_t(self, sigma):
        return 1 - sigma, sigma

    def time_shift(self, mu: float, sigma: float, t: torch.Tensor):
        return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma)

    def convert_model_output(
        self, model_output: torch.Tensor, *args, sample: torch.Tensor = None, **kwargs
    ) -> torch.Tensor:
        timestep = args[0] if len(args) > 0 else kwargs.pop("timestep", None)
        if sample is None:
            if len(args) > 1:
                sample = args[1]
            else:
                raise ValueError("missing `sample` as a required keyword argument")
        if timestep is not None:
            deprecate(
                "timesteps",
                "1.0.0",
                "Passing `timesteps` is deprecated and has no effect as model output "
                "conversion is now handled via an internal counter `self.step_index`",
            )

        sigma = self.sigmas[self.step_index]
        alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma)

        if self.predict_x0:
            if self.config.prediction_type == "flow_prediction":
                sigma_t = self.sigmas[self.step_index]
                x0_pred = sample - sigma_t * model_output
            else:
                raise ValueError(
                    f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`,"
                    " `v_prediction` or `flow_prediction` for the UniPCMultistepScheduler."
                )
            if self.config.thresholding:
                x0_pred = self._threshold_sample(x0_pred)
            return x0_pred
        else:
            if self.config.prediction_type == "flow_prediction":
                sigma_t = self.sigmas[self.step_index]
                epsilon = sample - (1 - sigma_t) * model_output
            else:
                raise ValueError(
                    f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`,"
                    " `v_prediction` or `flow_prediction` for the UniPCMultistepScheduler."
                )
            if self.config.thresholding:
                sigma_t = self.sigmas[self.step_index]
                x0_pred = sample - sigma_t * model_output
                x0_pred = self._threshold_sample(x0_pred)
                epsilon = model_output + x0_pred
            return epsilon

    def multistep_uni_p_bh_update(
        self,
        model_output: torch.Tensor,
        *args,
        sample: torch.Tensor | None = None,
        order: int | None = None,
        **kwargs,
    ) -> torch.Tensor:
        prev_timestep = args[0] if len(args) > 0 else kwargs.pop("prev_timestep", None)
        if sample is None:
            if len(args) > 1:
                sample = args[1]
            else:
                raise ValueError(" missing `sample` as a required keyword argument")
        if order is None:
            if len(args) > 2:
                order = args[2]
            else:
                raise ValueError(" missing `order` as a required keyword argument")
        if prev_timestep is not None:
            deprecate("prev_timestep", "1.0.0", "Passing `prev_timestep` is deprecated and has no effect.")

        model_output_list = self.model_outputs
        s0 = self.timestep_list[-1]
        m0 = model_output_list[-1]
        x = sample

        if self.solver_p:
            return self.solver_p.step(model_output, s0, x).prev_sample

        sigma_t, sigma_s0 = self.sigmas[self.step_index + 1], self.sigmas[self.step_index]
        alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t)
        alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0)

        lambda_t = torch.log(alpha_t) - torch.log(sigma_t)
        lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0)
        h = lambda_t - lambda_s0
        device = sample.device

        rks = []
        D1s: list[Any] | None = []
        for i in range(1, order):
            si = self.step_index - i
            mi = model_output_list[-(i + 1)]
            alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si])
            lambda_si = torch.log(alpha_si) - torch.log(sigma_si)
            rk = (lambda_si - lambda_s0) / h
            rks.append(rk)
            D1s.append((mi - m0) / rk)

        rks.append(1.0)
        rks = torch.tensor(rks, device=device)

        R = []
        b = []
        hh = -h if self.predict_x0 else h
        h_phi_1 = torch.expm1(hh)
        h_phi_k = h_phi_1 / hh - 1
        factorial_i = 1

        if self.config.solver_type == "bh1":
            B_h = hh
        elif self.config.solver_type == "bh2":
            B_h = torch.expm1(hh)
        else:
            raise NotImplementedError()

        for i in range(1, order + 1):
            R.append(torch.pow(rks, i - 1))
            b.append(h_phi_k * factorial_i / B_h)
            factorial_i *= i + 1
            h_phi_k = h_phi_k / hh - 1 / factorial_i

        R = torch.stack(R)
        b = torch.tensor(b, device=device)

        if len(D1s) > 0:
            D1s = torch.stack(D1s, dim=1)
            if order == 2:
                rhos_p = torch.tensor([0.5], dtype=x.dtype, device=device)
            else:
                rhos_p = torch.linalg.solve(R[:-1, :-1], b[:-1]).to(device).to(x.dtype)
        else:
            D1s = None

        if self.predict_x0:
            x_t_ = sigma_t / sigma_s0 * x - alpha_t * h_phi_1 * m0
            pred_res = torch.einsum("k,bkc...->bc...", rhos_p, D1s) if D1s is not None else 0
            x_t = x_t_ - alpha_t * B_h * pred_res
        else:
            x_t_ = alpha_t / alpha_s0 * x - sigma_t * h_phi_1 * m0
            pred_res = torch.einsum("k,bkc...->bc...", rhos_p, D1s) if D1s is not None else 0
            x_t = x_t_ - sigma_t * B_h * pred_res

        return x_t.to(x.dtype)

    def multistep_uni_c_bh_update(
        self,
        this_model_output: torch.Tensor,
        *args,
        last_sample: torch.Tensor = None,
        this_sample: torch.Tensor = None,
        order: int | None = None,
        **kwargs,
    ) -> torch.Tensor:
        this_timestep = args[0] if len(args) > 0 else kwargs.pop("this_timestep", None)
        if last_sample is None:
            if len(args) > 1:
                last_sample = args[1]
            else:
                raise ValueError(" missing`last_sample` as a required keyword argument")
        if this_sample is None:
            if len(args) > 2:
                this_sample = args[2]
            else:
                raise ValueError(" missing`this_sample` as a required keyword argument")
        if order is None:
            if len(args) > 3:
                order = args[3]
            else:
                raise ValueError(" missing`order` as a required keyword argument")
        if this_timestep is not None:
            deprecate("this_timestep", "1.0.0", "Passing `this_timestep` is deprecated and has no effect.")

        model_output_list = self.model_outputs
        m0 = model_output_list[-1]
        x = last_sample
        x_t = this_sample
        model_t = this_model_output

        sigma_t, sigma_s0 = self.sigmas[self.step_index], self.sigmas[self.step_index - 1]
        alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t)
        alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0)

        lambda_t = torch.log(alpha_t) - torch.log(sigma_t)
        lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0)
        h = lambda_t - lambda_s0
        device = this_sample.device

        rks = []
        D1s: list[Any] | None = []
        for i in range(1, order):
            si = self.step_index - (i + 1)
            mi = model_output_list[-(i + 1)]
            alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si])
            lambda_si = torch.log(alpha_si) - torch.log(sigma_si)
            rk = (lambda_si - lambda_s0) / h
            rks.append(rk)
            D1s.append((mi - m0) / rk)

        rks.append(1.0)
        rks = torch.tensor(rks, device=device)

        R = []
        b = []
        hh = -h if self.predict_x0 else h
        h_phi_1 = torch.expm1(hh)
        h_phi_k = h_phi_1 / hh - 1
        factorial_i = 1

        if self.config.solver_type == "bh1":
            B_h = hh
        elif self.config.solver_type == "bh2":
            B_h = torch.expm1(hh)
        else:
            raise NotImplementedError()

        for i in range(1, order + 1):
            R.append(torch.pow(rks, i - 1))
            b.append(h_phi_k * factorial_i / B_h)
            factorial_i *= i + 1
            h_phi_k = h_phi_k / hh - 1 / factorial_i

        R = torch.stack(R)
        b = torch.tensor(b, device=device)

        if len(D1s) > 0:
            D1s = torch.stack(D1s, dim=1)
        else:
            D1s = None

        if order == 1:
            rhos_c = torch.tensor([0.5], dtype=x.dtype, device=device)
        else:
            rhos_c = torch.linalg.solve(R, b).to(device).to(x.dtype)

        if self.predict_x0:
            x_t_ = sigma_t / sigma_s0 * x - alpha_t * h_phi_1 * m0
            corr_res = torch.einsum("k,bkc...->bc...", rhos_c[:-1], D1s) if D1s is not None else 0
            D1_t = model_t - m0
            x_t = x_t_ - alpha_t * B_h * (corr_res + rhos_c[-1] * D1_t)
        else:
            x_t_ = alpha_t / alpha_s0 * x - sigma_t * h_phi_1 * m0
            corr_res = torch.einsum("k,bkc...->bc...", rhos_c[:-1], D1s) if D1s is not None else 0
            D1_t = model_t - m0
            x_t = x_t_ - sigma_t * B_h * (corr_res + rhos_c[-1] * D1_t)
        return x_t.to(x.dtype)

    def index_for_timestep(self, timestep, schedule_timesteps=None):
        if schedule_timesteps is None:
            schedule_timesteps = self.timesteps
        indices = (schedule_timesteps == timestep).nonzero()
        pos = 1 if len(indices) > 1 else 0
        return indices[pos].item()

    def _init_step_index(self, timestep):
        if self.begin_index is None:
            if isinstance(timestep, torch.Tensor):
                timestep = timestep.to(self.timesteps.device)
            self._step_index = self.index_for_timestep(timestep)
        else:
            self._step_index = self._begin_index

    def step(
        self,
        model_output: torch.Tensor,
        timestep: int | torch.Tensor,
        sample: torch.Tensor,
        return_dict: bool = True,
        generator=None,
    ) -> SchedulerOutput | tuple:
        if self.num_inference_steps is None:
            raise ValueError(
                "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler"
            )

        if self.step_index is None:
            self._init_step_index(timestep)

        use_corrector = (
            self.step_index > 0 and self.step_index - 1 not in self.disable_corrector and self.last_sample is not None
        )

        model_output_convert = self.convert_model_output(model_output, sample=sample)
        if use_corrector:
            sample = self.multistep_uni_c_bh_update(
                this_model_output=model_output_convert,
                last_sample=self.last_sample,
                this_sample=sample,
                order=self.this_order,
            )

        for i in range(self.config.solver_order - 1):
            self.model_outputs[i] = self.model_outputs[i + 1]
            self.timestep_list[i] = self.timestep_list[i + 1]

        self.model_outputs[-1] = model_output_convert
        self.timestep_list[-1] = timestep

        if self.config.lower_order_final:
            this_order = min(self.config.solver_order, len(self.timesteps) - self.step_index)
        else:
            this_order = self.config.solver_order

        self.this_order = min(this_order, self.lower_order_nums + 1)
        assert self.this_order > 0

        self.last_sample = sample
        prev_sample = self.multistep_uni_p_bh_update(model_output=model_output, sample=sample, order=self.this_order)

        if self.lower_order_nums < self.config.solver_order:
            self.lower_order_nums += 1

        self._step_index += 1

        if not return_dict:
            return (prev_sample,)
        return SchedulerOutput(prev_sample=prev_sample)

    def step_ddim(
        self,
        velocity: torch.FloatTensor,
        t: int,
        curr_state: torch.FloatTensor,
        prev_state: torch.FloatTensor | None = None,
        generator: torch.Generator | None = None,
    ):
        device = curr_state.device
        curr_t = self.sigmas[t]
        prev_t = self.sigmas[t + 1]
        variance_noise = randn_tensor(curr_state.shape, generator=generator, device=device, dtype=curr_state.dtype)
        cur_clean_ = curr_state - curr_t * velocity
        return prev_t * variance_noise + (1 - prev_t) * cur_clean_

    def step_sde(
        self,
        velocity: torch.FloatTensor,
        t: int,
        curr_state: torch.FloatTensor,
        noise_theta: float = 1.0,
        prev_state: torch.FloatTensor | None = None,
        generator: torch.Generator | None = None,
    ):
        device = curr_state.device
        curr_t = self.sigmas[t]
        prev_t = self.sigmas[t + 1]
        cos = torch.cos(torch.tensor(noise_theta) * torch.pi / 2).to(device)
        sin = torch.sin(torch.tensor(noise_theta) * torch.pi / 2).to(device)
        prev_sample_mean = (1 - prev_t + prev_t * cos) * (curr_state - curr_t * velocity) + prev_t * cos * velocity
        std_dev_t = prev_t * sin
        std_dev_t = torch.ones((1, 1)).to(curr_state) * std_dev_t
        if prev_state is None:
            variance_noise = randn_tensor(curr_state.shape, generator=generator, device=device, dtype=curr_state.dtype)
            prev_state = prev_sample_mean + std_dev_t * variance_noise
        else:
            prev_state = prev_sample_mean + (prev_state - prev_sample_mean.detach())
        return prev_state

    def scale_model_input(self, sample: torch.Tensor, *args, **kwargs) -> torch.Tensor:
        return sample

    def add_noise(
        self, original_samples: torch.Tensor, noise: torch.Tensor, timesteps: torch.IntTensor
    ) -> torch.Tensor:
        sigmas = self.sigmas.to(device=original_samples.device, dtype=original_samples.dtype)
        if original_samples.device.type == "mps" and torch.is_floating_point(timesteps):
            schedule_timesteps = self.timesteps.to(original_samples.device, dtype=torch.float32)
            timesteps = timesteps.to(original_samples.device, dtype=torch.float32)
        else:
            schedule_timesteps = self.timesteps.to(original_samples.device)
            timesteps = timesteps.to(original_samples.device)

        if self.begin_index is None:
            step_indices = [self.index_for_timestep(t, schedule_timesteps) for t in timesteps]
        elif self.step_index is not None:
            step_indices = [self.step_index] * timesteps.shape[0]
        else:
            step_indices = [self.begin_index] * timesteps.shape[0]

        sigma = sigmas[step_indices].flatten()
        while len(sigma.shape) < len(original_samples.shape):
            sigma = sigma.unsqueeze(-1)

        alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma)
        return alpha_t * original_samples + sigma_t * noise

    def __len__(self):
        return self.config.num_train_timesteps


# ===========================================================================
# Audio VAE (ported from daVinci-MagiHuman inference/model/sa_audio/)
# ===========================================================================
def _snake_beta(x, alpha, beta):
    return x + (1.0 / (beta + 1e-9)) * torch.pow(torch.sin(x * alpha), 2)


class _SnakeBeta(nn.Module):
    def __init__(self, in_features: int, alpha: float = 1.0, alpha_trainable: bool = True, alpha_logscale: bool = True):
        super().__init__()
        self.alpha_logscale = alpha_logscale
        if self.alpha_logscale:
            self.alpha = nn.Parameter(torch.zeros(in_features) * alpha)
            self.beta = nn.Parameter(torch.zeros(in_features) * alpha)
        else:
            self.alpha = nn.Parameter(torch.ones(in_features) * alpha)
            self.beta = nn.Parameter(torch.ones(in_features) * alpha)
        self.alpha.requires_grad = alpha_trainable
        self.beta.requires_grad = alpha_trainable

    def forward(self, x):
        alpha = self.alpha.unsqueeze(0).unsqueeze(-1)
        beta = self.beta.unsqueeze(0).unsqueeze(-1)
        if self.alpha_logscale:
            alpha = torch.exp(alpha)
            beta = torch.exp(beta)
        return _snake_beta(x, alpha, beta)


def _vae_sample(mean, scale):
    stdev = F.softplus(scale) + 1e-4
    var = stdev * stdev
    logvar = torch.log(var)
    latents = torch.randn_like(mean) * stdev + mean
    kl = (mean * mean + var - logvar - 1).sum(1).mean()
    return latents, kl


class _VAEBottleneck(nn.Module):
    def encode(self, x, return_info=False, **kwargs):
        info = {}
        mean, scale = x.chunk(2, dim=1)
        x, kl = _vae_sample(mean, scale)
        info["kl"] = kl
        return (x, info) if return_info else x

    def decode(self, x):
        return x


def _WNConv1d(*args, **kwargs):
    return weight_norm(nn.Conv1d(*args, **kwargs))


def _WNConvTranspose1d(*args, **kwargs):
    return weight_norm(nn.ConvTranspose1d(*args, **kwargs))


def _checkpoint(function, *args, **kwargs):
    kwargs.setdefault("use_reentrant", False)
    return torch.utils.checkpoint.checkpoint(function, *args, **kwargs)


def _get_activation(activation: Literal["elu", "snake", "none"], antialias: bool = False, channels=None) -> nn.Module:
    if antialias:
        raise NotImplementedError("antialias activation not supported")
    if activation == "elu":
        return nn.ELU()
    if activation == "snake":
        return _SnakeBeta(channels)
    if activation == "none":
        return nn.Identity()
    raise ValueError(f"Unknown activation {activation}")


class _ResidualUnit(nn.Module):
    def __init__(self, in_channels, out_channels, dilation, use_snake=False, antialias_activation=False):
        super().__init__()
        padding = (dilation * (7 - 1)) // 2
        self.layers = nn.Sequential(
            _get_activation("snake" if use_snake else "elu", antialias=antialias_activation, channels=out_channels),
            _WNConv1d(in_channels, out_channels, kernel_size=7, dilation=dilation, padding=padding),
            _get_activation("snake" if use_snake else "elu", antialias=antialias_activation, channels=out_channels),
            _WNConv1d(out_channels, out_channels, kernel_size=1),
        )

    def forward(self, x):
        return (_checkpoint(self.layers, x) if self.training else self.layers(x)) + x


class _EncoderBlock(nn.Module):
    def __init__(self, in_channels, out_channels, stride, use_snake=False, antialias_activation=False):
        super().__init__()
        self.layers = nn.Sequential(
            _ResidualUnit(in_channels, in_channels, 1, use_snake=use_snake),
            _ResidualUnit(in_channels, in_channels, 3, use_snake=use_snake),
            _ResidualUnit(in_channels, in_channels, 9, use_snake=use_snake),
            _get_activation("snake" if use_snake else "elu", antialias=antialias_activation, channels=in_channels),
            _WNConv1d(in_channels, out_channels, kernel_size=2 * stride, stride=stride, padding=math.ceil(stride / 2)),
        )

    def forward(self, x):
        return self.layers(x)


class _DecoderBlock(nn.Module):
    def __init__(
        self, in_channels, out_channels, stride, use_snake=False, antialias_activation=False, use_nearest_upsample=False
    ):
        super().__init__()
        if use_nearest_upsample:
            upsample_layer = nn.Sequential(
                nn.Upsample(scale_factor=stride, mode="nearest"),
                _WNConv1d(in_channels, out_channels, kernel_size=2 * stride, stride=1, bias=False, padding="same"),
            )
        else:
            upsample_layer = _WNConvTranspose1d(
                in_channels, out_channels, kernel_size=2 * stride, stride=stride, padding=math.ceil(stride / 2)
            )
        self.layers = nn.Sequential(
            _get_activation("snake" if use_snake else "elu", antialias=antialias_activation, channels=in_channels),
            upsample_layer,
            _ResidualUnit(out_channels, out_channels, 1, use_snake=use_snake),
            _ResidualUnit(out_channels, out_channels, 3, use_snake=use_snake),
            _ResidualUnit(out_channels, out_channels, 9, use_snake=use_snake),
        )

    def forward(self, x):
        return self.layers(x)


class _OobleckEncoder(nn.Module):
    def __init__(
        self,
        in_channels=2,
        channels=128,
        latent_dim=32,
        c_mults=[1, 2, 4, 8],
        strides=[2, 4, 8, 8],
        use_snake=False,
        antialias_activation=False,
    ):
        super().__init__()
        c_mults = [1] + c_mults
        depth = len(c_mults)
        layers = [_WNConv1d(in_channels, c_mults[0] * channels, kernel_size=7, padding=3)]
        for i in range(depth - 1):
            layers.append(
                _EncoderBlock(c_mults[i] * channels, c_mults[i + 1] * channels, strides[i], use_snake=use_snake)
            )
        layers.extend(
            [
                _get_activation(
                    "snake" if use_snake else "elu", antialias=antialias_activation, channels=c_mults[-1] * channels
                ),
                _WNConv1d(c_mults[-1] * channels, latent_dim, kernel_size=3, padding=1),
            ]
        )
        self.layers = nn.Sequential(*layers)

    def forward(self, x):
        return self.layers(x)


class _OobleckDecoder(nn.Module):
    def __init__(
        self,
        out_channels=2,
        channels=128,
        latent_dim=32,
        c_mults=[1, 2, 4, 8],
        strides=[2, 4, 8, 8],
        use_snake=False,
        antialias_activation=False,
        use_nearest_upsample=False,
        final_tanh=True,
    ):
        super().__init__()
        c_mults = [1] + c_mults
        depth = len(c_mults)
        layers = [_WNConv1d(latent_dim, c_mults[-1] * channels, kernel_size=7, padding=3)]
        for i in range(depth - 1, 0, -1):
            layers.append(
                _DecoderBlock(
                    c_mults[i] * channels,
                    c_mults[i - 1] * channels,
                    strides[i - 1],
                    use_snake=use_snake,
                    antialias_activation=antialias_activation,
                    use_nearest_upsample=use_nearest_upsample,
                )
            )
        layers.extend(
            [
                _get_activation(
                    "snake" if use_snake else "elu", antialias=antialias_activation, channels=c_mults[0] * channels
                ),
                _WNConv1d(c_mults[0] * channels, out_channels, kernel_size=7, padding=3, bias=False),
                nn.Tanh() if final_tanh else nn.Identity(),
            ]
        )
        self.layers = nn.Sequential(*layers)

    def forward(self, x):
        return self.layers(x)


class _AudioAutoencoder(nn.Module):
    def __init__(
        self,
        encoder,
        decoder,
        latent_dim,
        downsampling_ratio,
        sample_rate,
        io_channels=2,
        bottleneck=None,
        in_channels=None,
        out_channels=None,
        soft_clip=False,
    ):
        super().__init__()
        self.downsampling_ratio = downsampling_ratio
        self.sample_rate = sample_rate
        self.latent_dim = latent_dim
        self.io_channels = io_channels
        self.in_channels = in_channels if in_channels is not None else io_channels
        self.out_channels = out_channels if out_channels is not None else io_channels
        self.bottleneck = bottleneck
        self.encoder = encoder
        self.decoder = decoder
        self.soft_clip = soft_clip

    def encode(self, audio, skip_bottleneck=False, return_info=False, **kwargs):
        info = {}
        latents = self.encoder(audio)
        info["pre_bottleneck_latents"] = latents
        if self.bottleneck is not None and not skip_bottleneck:
            latents, bottleneck_info = self.bottleneck.encode(latents, return_info=True, **kwargs)
            info.update(bottleneck_info)
        return (latents, info) if return_info else latents

    def decode(self, latents, skip_bottleneck=False, **kwargs):
        if self.bottleneck is not None and not skip_bottleneck:
            latents = self.bottleneck.decode(latents)
        decoded = self.decoder(latents, **kwargs)
        if self.soft_clip:
            decoded = torch.tanh(decoded)
        return decoded


def _create_encoder_from_config(cfg: dict[str, Any]):
    assert cfg.get("type") == "oobleck", f"Only 'oobleck' encoder supported, got: {cfg.get('type')}"
    enc = _OobleckEncoder(**cfg["config"])
    if not cfg.get("requires_grad", True):
        for p in enc.parameters():
            p.requires_grad = False
    return enc


def _create_decoder_from_config(cfg: dict[str, Any]):
    assert cfg.get("type") == "oobleck", f"Only 'oobleck' decoder supported, got: {cfg.get('type')}"
    dec = _OobleckDecoder(**cfg["config"])
    if not cfg.get("requires_grad", True):
        for p in dec.parameters():
            p.requires_grad = False
    return dec


def _create_bottleneck_from_config(cfg: dict[str, Any]):
    assert cfg.get("type") == "vae", f"Only 'vae' bottleneck supported, got: {cfg.get('type')}"
    bn = _VAEBottleneck()
    if not cfg.get("requires_grad", True):
        for p in bn.parameters():
            p.requires_grad = False
    return bn


def _create_autoencoder_from_config(config: dict[str, Any]):
    ae_config = config["model"]
    if ae_config.get("pretransform") is not None:
        raise NotImplementedError("Nested pretransform not supported")
    encoder = _create_encoder_from_config(ae_config["encoder"])
    decoder = _create_decoder_from_config(ae_config["decoder"])
    bottleneck_cfg = ae_config.get("bottleneck")
    bottleneck = _create_bottleneck_from_config(bottleneck_cfg) if bottleneck_cfg else None
    return _AudioAutoencoder(
        encoder=encoder,
        decoder=decoder,
        latent_dim=ae_config["latent_dim"],
        downsampling_ratio=ae_config["downsampling_ratio"],
        sample_rate=config["sample_rate"],
        io_channels=ae_config["io_channels"],
        bottleneck=bottleneck,
        in_channels=ae_config.get("in_channels"),
        out_channels=ae_config.get("out_channels"),
        soft_clip=ae_config["decoder"].get("soft_clip", False),
    )


class SAAudioFeatureExtractor:
    def __init__(self, device, model_path):
        self.device = device
        self.vae_model, self.sample_rate = self._load_vae(model_path)
        self.resampler = None

    def _load_vae(self, model_path):
        if not (isinstance(model_path, str) and Path(model_path).is_dir()):
            raise ValueError("model_path must be a local directory")

        model_config_path = os.path.join(model_path, "model_config.json")
        with open(model_config_path) as f:
            full_config = json.load(f)

        vae_config = full_config["model"]["pretransform"]["config"]
        sample_rate = full_config["sample_rate"]

        autoencoder_config = {
            "model_type": "autoencoder",
            "sample_rate": sample_rate,
            "model": vae_config,
        }
        vae_model = _create_autoencoder_from_config(autoencoder_config)

        weights_path = Path(model_path) / "model.safetensors"
        if not weights_path.exists():
            raise FileNotFoundError(f"Weight file does not exist: {weights_path}")

        full_state_dict = load_file(weights_path, device=str(self.device))
        vae_state_dict = {}
        for key, value in full_state_dict.items():
            if key.startswith("pretransform.model."):
                vae_state_dict[key[len("pretransform.model.") :]] = value

        model_keys = set(vae_model.state_dict().keys())
        vae_keys = set(vae_state_dict.keys())
        missing = model_keys - vae_keys
        extra = vae_keys - model_keys
        if missing:
            logger.warning("Audio VAE missing keys (%d): %s", len(missing), list(missing)[:5])
        if extra:
            logger.warning("Audio VAE unexpected keys (%d): %s", len(extra), list(extra)[:5])

        vae_model.load_state_dict(vae_state_dict)
        vae_model.to(self.device)
        return vae_model, sample_rate

    def decode(self, latents):
        with torch.no_grad():
            return self.vae_model.decode(latents)

    def encode(self, waveform):
        with torch.no_grad():
            return self.vae_model.encode(waveform)


# ===========================================================================
# Audio utilities (ported from daVinci-MagiHuman inference/pipeline/video_process.py)
# ===========================================================================
_SAMPLE_RATE = 51200
_AUDIO_CHUNK_DURATION = 29
_OVERLAP_RATIO = 0.5


def _merge_overlapping_vae_features(audio_feats: list[torch.Tensor], overlap_ratio: float = 0.5) -> torch.Tensor | None:
    if not audio_feats:
        return None
    if len(audio_feats) == 1:
        return audio_feats[0]

    batch_size, total_frames, feature_dim = audio_feats[0].shape
    overlap_frames = int(total_frames * overlap_ratio)
    step_frames = total_frames - overlap_frames
    final_length = (len(audio_feats) - 1) * step_frames + total_frames
    output_feat = torch.zeros(
        batch_size, final_length, feature_dim, device=audio_feats[0].device, dtype=audio_feats[0].dtype
    )

    for block_idx, current_feat in enumerate(audio_feats):
        output_start = block_idx * step_frames
        if block_idx == 0:
            output_feat[:, output_start : output_start + total_frames, :] = current_feat
            continue

        non_overlap_start = output_start + overlap_frames
        non_overlap_end = output_start + total_frames
        output_feat[:, non_overlap_start:non_overlap_end, :] = current_feat[:, overlap_frames:, :]

        for frame_idx in range(overlap_frames):
            output_pos = output_start + frame_idx
            prev_weight = (overlap_frames - frame_idx) / overlap_frames
            curr_weight = frame_idx / overlap_frames
            output_feat[:, output_pos, :] = (
                prev_weight * output_feat[:, output_pos, :] + curr_weight * current_feat[:, frame_idx, :]
            )
    return output_feat


def load_audio_and_encode(audio_vae, audio_path: str, seconds: int | None = None) -> torch.Tensor:
    """Load audio from file and encode to latent space using the Stable Audio VAE."""
    audio_full = whisper.load_audio(audio_path, sr=_SAMPLE_RATE)
    if seconds is not None:
        audio_full = audio_full[: min(int(seconds * _SAMPLE_RATE), audio_full.shape[0])]
    total_samples = audio_full.shape[0]

    window_size = int(_AUDIO_CHUNK_DURATION * _SAMPLE_RATE)
    step_size = int(window_size * (1 - _OVERLAP_RATIO))
    if total_samples <= window_size:
        audio = torch.from_numpy(audio_full).cuda()
        audio = audio.unsqueeze(0).expand(2, -1)
        return audio_vae.vae_model.encode(audio)

    encoded_chunks = []
    latent_to_audio_ratio = None
    for offset_start in range(0, total_samples, step_size):
        offset_end = min(offset_start + window_size, total_samples)
        chunk = whisper.pad_or_trim(audio_full[offset_start:offset_end], length=window_size)
        chunk_tensor = torch.from_numpy(chunk).cuda().unsqueeze(0).expand(2, -1)
        encoded_chunk = audio_vae.vae_model.encode(chunk_tensor)

        if latent_to_audio_ratio is None:
            latent_to_audio_ratio = encoded_chunk.shape[-1] / window_size

        encoded_chunks.append(encoded_chunk.permute(0, 2, 1))
        if offset_end >= total_samples:
            break

    final_feat = _merge_overlapping_vae_features(encoded_chunks, overlap_ratio=_OVERLAP_RATIO).permute(0, 2, 1)
    final_target_len = math.ceil(total_samples * latent_to_audio_ratio)
    return final_feat[:, :, :final_target_len]


# ===========================================================================
# Data proxy (ported from daVinci-MagiHuman inference/pipeline/data_proxy.py)
# ===========================================================================
def _unfold_3d(x: torch.Tensor, kernel_size: tuple[int, int, int], stride: tuple[int, int, int]) -> torch.Tensor:
    """Pure-PyTorch 3D unfold matching UnfoldAnd behavior.

    After N unfold ops the shape is (batch, C, oD, oH, oW, kD, kH, kW).
    UnfoldAnd permutes kernel dims next to channel before reshape so that the
    col_dim axis is ordered as (C, kD, kH, kW) -- matching F.unfold semantics.
    Without this permute, .view() interleaves spatial and kernel positions.

    Args:
        x: (N, C, D, H, W)
        kernel_size: (kD, kH, kW)
        stride: (sD, sH, sW)
    Returns:
        (N, C*kD*kH*kW, L) where L = product of output spatial dims.
    """
    ndim = len(kernel_size)
    for d in range(ndim):
        x = x.unfold(d + 2, kernel_size[d], stride[d])
    perm = [0, 1] + list(range(ndim + 2, 2 * ndim + 2)) + list(range(2, ndim + 2))
    x = x.permute(*perm).contiguous()

    batch_size = x.shape[0]
    col_dim = 1
    for i in range(1, ndim + 2):
        col_dim *= x.shape[i]
    spatial = 1
    for i in range(ndim + 2, 2 * ndim + 2):
        spatial *= x.shape[i]
    return x.view(batch_size, col_dim, spatial)


def _calc_local_qk_range(num_video_tokens, num_audio_and_txt_tokens, num_frames, frame_receptive_field):
    token_per_frame = num_video_tokens // num_frames
    total_tokens = num_video_tokens + num_audio_and_txt_tokens

    q_range_list = []
    k_range_list = []
    for i in range(num_frames):
        q_range_list.append(torch.tensor([i * token_per_frame, (i + 1) * token_per_frame]))
        k_range_list.append(
            torch.tensor(
                [
                    (i - frame_receptive_field) * token_per_frame,
                    (i + frame_receptive_field + 1) * token_per_frame,
                ]
            )
        )
    local_q_range = torch.stack(q_range_list, dim=0)
    local_k_range = torch.stack(k_range_list, dim=0)

    local_k_range[local_k_range < 0] = 0
    local_k_range[local_k_range > num_video_tokens] = num_video_tokens

    video_q_range = torch.tensor([[0, num_video_tokens]])
    video_k_range = torch.tensor([[num_video_tokens, num_video_tokens + num_audio_and_txt_tokens]])

    at_q_ranges = torch.tensor([[num_video_tokens, total_tokens]])
    at_k_ranges = torch.tensor([[0, total_tokens]])

    q_ranges = (
        torch.cat([local_q_range, video_q_range, at_q_ranges], dim=0).to(torch.int32).to("cuda", non_blocking=True)
    )
    k_ranges = (
        torch.cat([local_k_range, video_k_range, at_k_ranges], dim=0).to(torch.int32).to("cuda", non_blocking=True)
    )
    return q_ranges, k_ranges


def _calc_local_attn_ffa_handler(num_video_tokens, num_audio_and_txt_tokens, num_frames, frame_receptive_field):
    q_ranges, k_ranges = _calc_local_qk_range(
        num_video_tokens, num_audio_and_txt_tokens, num_frames, frame_receptive_field
    )
    total = num_video_tokens + num_audio_and_txt_tokens
    return FFAHandler(
        q_ranges=q_ranges,
        k_ranges=k_ranges,
        max_seqlen_q=total,
        max_seqlen_k=total,
        attn_type_map=torch.zeros([q_ranges.shape[0]], device="cuda", dtype=torch.int32),
        softmax_scale=None,
    )


def _get_coords(
    shape: list[int],
    ref_feat_shape: list[int],
    offset_thw: list[int] | None = None,
    device: torch.device = torch.device("cpu"),
    dtype: torch.dtype = torch.float32,
):
    if offset_thw is None:
        offset_thw = [0, 0, 0]
    ori_t, ori_h, ori_w = shape
    ref_t, ref_h, ref_w = ref_feat_shape

    offset_t, offset_h, offset_w = offset_thw
    time_rng = torch.arange(ori_t, device=device, dtype=dtype) + offset_t
    height_rng = torch.arange(ori_h, device=device, dtype=dtype) + offset_h
    width_rng = torch.arange(ori_w, device=device, dtype=dtype) + offset_w

    time_grid, height_grid, width_grid = torch.meshgrid(time_rng, height_rng, width_rng, indexing="ij")
    coords_flat = torch.stack([time_grid, height_grid, width_grid], dim=-1).reshape(-1, 3)

    meta = torch.tensor([ori_t, ori_h, ori_w, ref_t, ref_h, ref_w], device=device, dtype=dtype)
    meta_expanded = meta.expand(coords_flat.size(0), -1)
    return torch.cat([coords_flat, meta_expanded], dim=-1)


@dataclass
class _SingleData:
    video_x_t: torch.Tensor
    audio_x_t: torch.Tensor
    audio_feat_len: int
    txt_feat: torch.Tensor
    txt_feat_len: int
    t: int
    h: int
    w: int
    patch_size: int
    t_patch_size: int
    spatial_rope_interpolation: Literal["inter", "extra"]
    ref_audio_offset: int
    text_offset: int
    coords_style: Literal["v1", "v2"] = "v1"

    def __post_init__(self):
        self.video_token_num = self.video_x_t.shape[0]
        self.audio_x_t = self.audio_x_t[: self.audio_feat_len]
        self.txt_feat = self.txt_feat[: self.txt_feat_len]
        self.video_channel = self.video_x_t.shape[-1]
        self.audio_channel = self.audio_x_t.shape[-1]
        self.txt_channel = self.txt_feat.shape[-1]

    @property
    def device(self):
        return self.video_x_t.device

    @property
    def default_dtype(self):
        return self.video_x_t.dtype

    @property
    def total_token_num(self):
        return self.video_token_num + self.audio_feat_len + self.txt_feat_len

    @property
    def token_sequence(self):
        tensors = [self.video_x_t, self.audio_x_t, self.txt_feat]
        max_channel = max(t.shape[-1] for t in tensors)
        padded = [F.pad(t, (0, max_channel - t.shape[-1])) for t in tensors]
        return torch.cat(padded, dim=0)

    @property
    def modality_mapping(self):
        v_map = torch.full((self.video_token_num,), Modality.VIDEO, dtype=torch.int64, device=self.device)
        a_map = torch.full((self.audio_feat_len,), Modality.AUDIO, dtype=torch.int64, device=self.device)
        t_map = torch.full((self.txt_feat_len,), Modality.TEXT, dtype=torch.int64, device=self.device)
        return torch.cat([v_map, a_map, t_map], dim=0)

    def _default_coords(self, shape, ref_feat_shape, offset_thw=None):
        if offset_thw is None:
            offset_thw = [0, 0, 0]
        return _get_coords(
            shape=shape,
            ref_feat_shape=ref_feat_shape,
            offset_thw=offset_thw,
            device=self.device,
            dtype=self.default_dtype,
        )

    @property
    def coords_mapping(self):
        if self.spatial_rope_interpolation == "inter":
            video_ref_feat_shape = (self.t // self.t_patch_size, 32, 32)
        else:
            video_ref_feat_shape = (self.t // self.t_patch_size, self.h // self.patch_size, self.w // self.patch_size)

        video_coords = self._default_coords(
            shape=(self.t // self.t_patch_size, self.h // self.patch_size, self.w // self.patch_size),
            ref_feat_shape=video_ref_feat_shape,
        )

        if self.coords_style == "v1":
            audio_coords = self._default_coords(
                shape=(self.audio_feat_len, 1, 1),
                ref_feat_shape=(self.t // self.t_patch_size, 1, 1),
            )
            text_coords = self._default_coords(
                shape=(self.txt_feat_len, 1, 1),
                ref_feat_shape=(2, 1, 1),
                offset_thw=[self.text_offset, 0, 0],
            )
        elif self.coords_style == "v2":
            magic_audio_ref_t = (self.audio_feat_len - 1) // 4 + 1
            audio_coords = self._default_coords(
                shape=(self.audio_feat_len, 1, 1),
                ref_feat_shape=(magic_audio_ref_t // self.t_patch_size, 1, 1),
            )
            text_coords = self._default_coords(
                shape=(self.txt_feat_len, 1, 1),
                ref_feat_shape=(1, 1, 1),
                offset_thw=[-self.txt_feat_len, 0, 0],
            )
        else:
            raise ValueError(f"Unknown coords_style: {self.coords_style}")

        return torch.cat([video_coords, audio_coords, text_coords], dim=0)

    def depack_token_sequence(self, token_sequence):
        video_x_t = token_sequence[: self.video_token_num, : self.video_channel]
        video_x_t = rearrange(
            video_x_t,
            "(T H W) (pT pH pW C) -> C (T pT) (H pH) (W pW)",
            H=self.h // self.patch_size,
            W=self.w // self.patch_size,
            pT=self.t_patch_size,
            pH=self.patch_size,
            pW=self.patch_size,
        ).contiguous()
        audio_x_t = token_sequence[
            self.video_token_num : self.video_token_num + self.audio_feat_len, : self.audio_channel
        ]
        return video_x_t, audio_x_t


@dataclass
class _SimplePackedData:
    items: list[_SingleData]

    @property
    def token_sequence(self):
        return torch.cat([item.token_sequence for item in self.items], dim=0)

    @property
    def modality_mapping(self):
        return torch.cat([item.modality_mapping for item in self.items], dim=0)

    @property
    def coords_mapping(self):
        return torch.cat([item.coords_mapping for item in self.items], dim=0)

    @property
    def total_token_num(self):
        return sum(item.total_token_num for item in self.items)

    def __getitem__(self, index):
        return self.items[index]

    @property
    def cu_seqlen(self):
        cu = torch.cumsum(torch.tensor([item.total_token_num for item in self.items]), dim=0)
        return F.pad(cu, (1, 0))

    @property
    def max_seqlen(self):
        return torch.tensor(max(item.total_token_num for item in self.items))

    def depack_token_sequence(self, token_sequence):
        video_list, audio_list = [], []
        parts = torch.split(token_sequence, [item.total_token_num for item in self.items], dim=0)
        for item, part in zip(self.items, parts):
            v, a = item.depack_token_sequence(part)
            video_list.append(v)
            audio_list.append(a)
        return torch.stack(video_list, dim=0), torch.stack(audio_list, dim=0)


class MagiDataProxy:
    def __init__(
        self,
        patch_size: int = 2,
        t_patch_size: int = 1,
        frame_receptive_field: int = 11,
        spatial_rope_interpolation: str = "extra",
        ref_audio_offset: int = 1000,
        text_offset: int = 0,
        coords_style: str = "v2",
    ):
        self.patch_size = patch_size
        self.t_patch_size = t_patch_size
        self.frame_receptive_field = frame_receptive_field
        self.spatial_rope_interpolation = spatial_rope_interpolation
        self.ref_audio_offset = ref_audio_offset
        self.text_offset = text_offset
        self.coords_style = coords_style
        self._kernel = (t_patch_size, patch_size, patch_size)
        self._stride = (t_patch_size, patch_size, patch_size)
        self._saved_data: dict[str, Any] = {}

    def saved_for_output(self, **kwargs):
        self._saved_data.update(kwargs)

    def get_saved_data(self, key: str):
        return self._saved_data[key]

    def img2tokens(self, x_t: torch.Tensor):
        x_t_unfolded = _unfold_3d(x_t, self._kernel, self._stride)
        return rearrange(x_t_unfolded, "N col_dim num_tokens -> N num_tokens col_dim").contiguous()

    def process_input(self, transported_data: EvalInput):
        batch_size, _, t, h, w = transported_data.x_t.shape
        x_t = self.img2tokens(transported_data.x_t)
        audio_x_t = transported_data.audio_x_t.contiguous()
        text_in = transported_data.txt_feat.contiguous()

        simple_packed_data = _SimplePackedData(items=[])
        for i in range(batch_size):
            single_data = _SingleData(
                video_x_t=x_t[i],
                audio_x_t=audio_x_t[i],
                audio_feat_len=transported_data.audio_feat_len[i],
                txt_feat=text_in[i],
                txt_feat_len=transported_data.txt_feat_len[i],
                t=t,
                h=h,
                w=w,
                patch_size=self.patch_size,
                t_patch_size=self.t_patch_size,
                spatial_rope_interpolation=self.spatial_rope_interpolation,
                ref_audio_offset=self.ref_audio_offset,
                text_offset=self.text_offset,
                coords_style=self.coords_style,
            )
            simple_packed_data.items.append(single_data)

        if self.frame_receptive_field != -1:
            assert batch_size == 1, "local attention only supports batch size 1"
            local_attn_handler = _calc_local_attn_ffa_handler(
                num_video_tokens=simple_packed_data[0].video_token_num,
                num_audio_and_txt_tokens=simple_packed_data[0].audio_feat_len + simple_packed_data[0].txt_feat_len,
                num_frames=t,
                frame_receptive_field=self.frame_receptive_field,
            )
            if isinstance(local_attn_handler.max_seqlen_k, torch.Tensor):
                local_attn_handler.max_seqlen_k = local_attn_handler.max_seqlen_k.item()
            if isinstance(local_attn_handler.max_seqlen_q, torch.Tensor):
                local_attn_handler.max_seqlen_q = local_attn_handler.max_seqlen_q.item()
        else:
            local_attn_handler = None

        varlen_handler = VarlenHandler(
            cu_seqlens_q=simple_packed_data.cu_seqlen.to(torch.int32).cuda(),
            cu_seqlens_k=simple_packed_data.cu_seqlen.to(torch.int32).cuda(),
            max_seqlen_q=simple_packed_data.max_seqlen.to(torch.int32).cuda(),
            max_seqlen_k=simple_packed_data.max_seqlen.to(torch.int32).cuda(),
        )

        self.saved_for_output(simple_packed_data=simple_packed_data)

        x = simple_packed_data.token_sequence
        coords_mapping = simple_packed_data.coords_mapping
        modality_mapping = simple_packed_data.modality_mapping
        return (x, coords_mapping, modality_mapping, varlen_handler, local_attn_handler)

    def process_output(self, x: torch.Tensor):
        simple_packed_data: _SimplePackedData = self.get_saved_data("simple_packed_data")
        return simple_packed_data.depack_token_sequence(x)


# ===========================================================================
# Pipeline helpers
# ===========================================================================
@dataclass
class EvalInput:
    x_t: torch.Tensor
    audio_x_t: torch.Tensor
    audio_feat_len: torch.Tensor | list[int]
    txt_feat: torch.Tensor
    txt_feat_len: torch.Tensor | list[int]


class _T5GemmaEncoder:
    def __init__(self, model_path: str, device: str, weight_dtype: torch.dtype, subfolder: str | None = None):
        from vllm.distributed import get_tensor_model_parallel_world_size

        self.device = device
        hf_kwargs: dict[str, Any] = {}
        if subfolder is not None:
            hf_kwargs["subfolder"] = subfolder
        self.tokenizer = AutoTokenizer.from_pretrained(model_path, **hf_kwargs)

        tp_size = get_tensor_model_parallel_world_size()
        if tp_size > 1:
            from transformers.models.t5gemma.configuration_t5gemma import T5GemmaConfig

            config = T5GemmaConfig.from_pretrained(model_path, **hf_kwargs)
            # The config we need is the encoder config
            config_encoder = config.encoder
            # Propagate some outer config values
            config_encoder.vocab_size = config.vocab_size
            config_encoder.rms_norm_eps = getattr(config, "rms_norm_eps", config_encoder.rms_norm_eps)
            self.model = T5GemmaEncoderModelTP(config_encoder).to(device).to(weight_dtype)
            self.is_tp = True
        else:
            self.model = T5GemmaEncoderModel.from_pretrained(
                model_path, is_encoder_decoder=False, dtype=weight_dtype, **hf_kwargs
            ).to(device)
            self.is_tp = False

    @torch.inference_mode()
    def encode(self, prompt: str) -> torch.Tensor:
        inputs = self.tokenizer([prompt], return_tensors="pt").to(self.device)
        outputs = self.model(**inputs)

        if self.is_tp:
            # T5GemmaEncoderModelTP just returns the hidden states tensor
            return outputs.half()
        else:
            # HF model returns BaseModelOutput
            return outputs["last_hidden_state"].half()


def _pad_or_trim(tensor: torch.Tensor, target_size: int, dim: int, pad_value: float = 0.0) -> tuple[torch.Tensor, int]:
    current_size = tensor.size(dim)
    if current_size < target_size:
        padding_amount = target_size - current_size
        padding_tuple = [0] * (2 * tensor.dim())
        padding_dim_index = tensor.dim() - 1 - dim
        padding_tuple[2 * padding_dim_index + 1] = padding_amount
        return F.pad(tensor, tuple(padding_tuple), "constant", pad_value), current_size
    slicing = [slice(None)] * tensor.dim()
    slicing[dim] = slice(0, target_size)
    return tensor[tuple(slicing)], target_size


def _get_padded_t5_gemma_embedding(
    prompt: str,
    encoder: _T5GemmaEncoder,
    target_length: int,
) -> tuple[torch.Tensor, int]:
    txt_feat = encoder.encode(prompt)
    txt_feat, original_len = _pad_or_trim(txt_feat, target_size=target_length, dim=1)
    return txt_feat.to(torch.float32), original_len


def _resizecrop(img: Image.Image, target_height: int, target_width: int) -> Image.Image:
    """Centre-crop resize keeping aspect ratio then letterbox to target."""
    pil_image = img.convert("RGB")
    original_width, original_height = pil_image.size
    scale_x = target_width / original_width
    scale_y = target_height / original_height
    scale = max(scale_x, scale_y)
    new_width = int(round(original_width * scale))
    new_height = int(round(original_height * scale))
    resized_image = pil_image.resize((new_width, new_height), Image.LANCZOS)
    left = (new_width - target_width) // 2
    top = (new_height - target_height) // 2
    return resized_image.crop((left, top, left + target_width, top + target_height))


class ZeroSNRDDPMDiscretization:
    """ZeroSNR DDPM sigma schedule, ported from daVinci-MagiHuman.
    Used to compute sigma values for SR noise injection.
    """

    def __init__(
        self,
        linear_start: float = 0.00085,
        linear_end: float = 0.0120,
        num_timesteps: int = 1000,
        shift_scale: float = 1.0,
        keep_start: bool = False,
        post_shift: bool = False,
    ):
        from functools import partial

        if keep_start and not post_shift:
            linear_start = linear_start / (shift_scale + (1 - shift_scale) * linear_start)
        self.num_timesteps = num_timesteps
        betas = torch.linspace(linear_start**0.5, linear_end**0.5, num_timesteps, dtype=torch.float64) ** 2
        alphas = 1.0 - betas.cpu().numpy()
        self.alphas_cumprod = np.cumprod(alphas, axis=0)
        self.to_torch = partial(torch.tensor, dtype=torch.float32)
        if not post_shift:
            self.alphas_cumprod = self.alphas_cumprod / (shift_scale + (1 - shift_scale) * self.alphas_cumprod)
        self.post_shift = post_shift
        self.shift_scale = shift_scale

    def __call__(
        self,
        n: int,
        do_append_zero: bool = True,
        device: str = "cpu",
        flip: bool = False,
        return_idx: bool = False,
    ):
        from functools import partial

        if n < self.num_timesteps:
            timesteps = np.linspace(self.num_timesteps - 1, 0, n, endpoint=False).astype(int)[::-1]
            alphas_cumprod = self.alphas_cumprod[timesteps]
        elif n == self.num_timesteps:
            alphas_cumprod = self.alphas_cumprod
        else:
            raise ValueError(f"n={n} > num_timesteps={self.num_timesteps}")

        to_torch = partial(torch.tensor, dtype=torch.float32, device=device)
        alphas_cumprod = to_torch(alphas_cumprod)
        alphas_cumprod_sqrt = alphas_cumprod.sqrt()
        alphas_cumprod_sqrt_0 = alphas_cumprod_sqrt[0].clone()
        alphas_cumprod_sqrt_T = alphas_cumprod_sqrt[-1].clone()
        alphas_cumprod_sqrt -= alphas_cumprod_sqrt_T
        alphas_cumprod_sqrt *= alphas_cumprod_sqrt_0 / (alphas_cumprod_sqrt_0 - alphas_cumprod_sqrt_T)

        if self.post_shift:
            alphas_cumprod_sqrt = (
                alphas_cumprod_sqrt**2 / (self.shift_scale + (1 - self.shift_scale) * alphas_cumprod_sqrt**2)
            ) ** 0.5

        sigmas = torch.flip(alphas_cumprod_sqrt, (0,))
        sigmas = torch.cat([sigmas, sigmas.new_zeros([1])]) if do_append_zero else sigmas
        if return_idx:
            return sigmas if not flip else torch.flip(sigmas, (0,)), timesteps
        return sigmas if not flip else torch.flip(sigmas, (0,))


def _schedule_latent_step(
    *,
    video_scheduler: FlowUniPCMultistepScheduler,
    audio_scheduler: FlowUniPCMultistepScheduler,
    latent_video: torch.Tensor,
    latent_audio: torch.Tensor,
    t,
    idx: int,
    steps,
    v_cfg_video: torch.Tensor,
    v_cfg_audio: torch.Tensor,
    is_a2v: bool,
    cfg_number: int,
    using_sde_flag: bool,
    use_sr_model: bool = False,
):
    # Fast DDIM path for cfg_number==1, only used during the BR stage
    if cfg_number == 1 and not use_sr_model:
        latent_video = video_scheduler.step_ddim(v_cfg_video, idx, latent_video)
        latent_audio = audio_scheduler.step_ddim(v_cfg_audio, idx, latent_audio)
        return latent_video, latent_audio

    if using_sde_flag:
        if use_sr_model:
            # SR stage with SDE: only update video, keep audio unchanged
            latent_video = video_scheduler.step(v_cfg_video, t, latent_video, return_dict=False)[0]
            return latent_video, latent_audio
        if idx < int(len(steps) * (3 / 4)):
            noise_theta = 1.0 if (idx + 1) % 2 == 0 else 0.0
        else:
            noise_theta = 1.0 if idx % 3 == 0 else 0.0
        latent_video = video_scheduler.step_sde(v_cfg_video, idx, latent_video, noise_theta=noise_theta)
        if not is_a2v:
            latent_audio = audio_scheduler.step_sde(v_cfg_audio, idx, latent_audio, noise_theta=noise_theta)
        return latent_video, latent_audio

    latent_video = video_scheduler.step(v_cfg_video, t, latent_video, return_dict=False)[0]
    # Do not update audio latent during the SR stage
    if not is_a2v and not use_sr_model:
        latent_audio = audio_scheduler.step(v_cfg_audio, t, latent_audio, return_dict=False)[0]
    return latent_video, latent_audio


_NEGATIVE_PROMPT = (
    "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, "
    "overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, "
    "poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, "
    "still picture, messy background, three legs, many people in the background, walking backwards"
    ", low quality, worst quality, poor quality, noise, background noise, hiss, hum, buzz, crackle, static, "
    "compression artifacts, MP3 artifacts, digital clipping, distortion, muffled, muddy, unclear, echo, "
    "reverb, room echo, over-reverberated, hollow sound, distant, washed out, harsh, shrill, piercing, "
    "grating, tinny, thin sound, boomy, bass-heavy, flat EQ, over-compressed, abrupt cut, jarring transition, "
    "sudden silence, looping artifact, music, instrumental, sirens, alarms, crowd noise, unrelated sound "
    "effects, chaotic, disorganized, messy, cheap sound"
    ", emotionless, flat delivery, deadpan, lifeless, apathetic, robotic, mechanical, monotone, flat "
    "intonation, undynamic, boring, reading from a script, AI voice, synthetic, text-to-speech, TTS, "
    "insincere, fake emotion, exaggerated, overly dramatic, melodramatic, cheesy, cringey, hesitant, "
    "unconfident, tired, weak voice, stuttering, stammering, mumbling, slurred speech, mispronounced, "
    "bad articulation, lisp, vocal fry, creaky voice, mouth clicks, lip smacks, wet mouth sounds, heavy "
    "breathing, audible inhales, plosives, p-pops, coughing, clearing throat, sneezing, speaking too fast, "
    "rushed, speaking too slow, dragged out, unnatural pauses, awkward silence, choppy, disjointed, multiple "
    "speakers, two voices, background talking, out of tune, off-key, autotune artifacts"
)


# ===========================================================================
# Pre/post process funcs (registered in registry)
# ===========================================================================
def get_magi_human_pre_process_func(*args, **kwargs):
    def pre_process(request: OmniDiffusionRequest):
        return request

    return pre_process


def get_magi_human_post_process_func(*args, **kwargs):
    def post_process(output):
        if isinstance(output, tuple) and len(output) == 2:
            video, audio = output
            return {
                "video": video,
                "audio": audio,
                "audio_sample_rate": 44100,
                "fps": 25,
            }
        return output

    return post_process


# ===========================================================================
# HF Hub / local path helpers
# ===========================================================================


def _load_json(model_path: str, filename: str, local_files_only: bool = True) -> dict:
    """Load a JSON config file from a local path or HuggingFace Hub repo."""
    if local_files_only:
        path = os.path.join(model_path, *filename.split("/"))
        with open(path) as f:
            return json.load(f)
    else:
        from huggingface_hub import hf_hub_download

        cached = hf_hub_download(repo_id=model_path, filename=filename)
        with open(cached) as f:
            return json.load(f)


def _resolve_subdir(
    model_path: str,
    subfolder: str,
    local_files_only: bool = True,
    required_files: list[str] | None = None,
) -> str:
    """Resolve a model subfolder to a local directory path.

    For HF Hub repos, downloads all ``required_files`` (default: ``["config.json"]``)
    into the HF cache and returns the parent directory.
    """
    if local_files_only:
        return os.path.join(model_path, subfolder)
    from huggingface_hub import hf_hub_download

    files = required_files or ["config.json"]
    last_cached: str | None = None
    for fname in files:
        last_cached = hf_hub_download(repo_id=model_path, filename=f"{subfolder}/{fname}")
    return os.path.dirname(last_cached)


# ===========================================================================
# Main Pipeline
# ===========================================================================
class MagiHumanPipeline(nn.Module, ProgressBarMixin, SupportsComponentDiscovery, DiffusionPipelineProfilerMixin):
    _dit_modules: ClassVar[list[str]] = ["dit", "sr_dit"]
    _encoder_modules: ClassVar[list[str]] = ["text_encoder"]
    _vae_modules: ClassVar[list[str]] = ["vae", "audio_vae"]

    def __init__(self, od_config: OmniDiffusionConfig, **kwargs):
        super().__init__()
        model_path = od_config.model
        local_files_only = os.path.exists(model_path)
        device = f"cuda:{torch.accelerator.current_device_index()}"
        self.device_str = device
        self.dtype = od_config.dtype or torch.bfloat16

        model_index = _load_json(model_path, "model_index.json", local_files_only)
        eval_cfg = model_index
        dp_cfg = model_index.get("data_proxy", {})

        dit_subfolder = "transformer"

        dit_json = _load_json(model_path, f"{dit_subfolder}/config.json", local_files_only)
        dit_model_config = MagiHumanDiTConfig(**dit_json)

        self.dit = DiTModel(dit_model_config)
        self.dit.eval()

        self.vae = DistributedAutoencoderKLWan.from_pretrained(model_path, subfolder="vae")
        self.vae.to(device)
        self.vae.eval()
        vae_cfg = _load_json(model_path, "vae/config.json", local_files_only)
        self.vae_latent_mean = torch.tensor(vae_cfg["latents_mean"], dtype=torch.float32)
        self.vae_latent_std = torch.tensor(vae_cfg["latents_std"], dtype=torch.float32)

        self.audio_vae = SAAudioFeatureExtractor(
            device=device,
            model_path=_resolve_subdir(
                model_path,
                "audio_vae",
                local_files_only,
                required_files=["config.json", "model_config.json", "model.safetensors"],
            ),
        )

        logger.info("Loading T5Gemma text encoder from %s (subfolder=text_encoder)", model_path)
        if local_files_only:
            txt_enc_path = os.path.join(model_path, "text_encoder")
            txt_enc_subfolder = None
        else:
            txt_enc_path = model_path
            txt_enc_subfolder = "text_encoder"
        self.text_encoder = _T5GemmaEncoder(
            model_path=txt_enc_path,
            device=device,
            weight_dtype=self.dtype,
            subfolder=txt_enc_subfolder,
        )

        self.data_proxy = MagiDataProxy(
            patch_size=dp_cfg.get("patch_size", 2),
            t_patch_size=dp_cfg.get("t_patch_size", 1),
            frame_receptive_field=dp_cfg.get("frame_receptive_field", 11),
            spatial_rope_interpolation=dp_cfg.get("spatial_rope_interpolation", "extra"),
            ref_audio_offset=dp_cfg.get("ref_audio_offset", 1000),
            text_offset=dp_cfg.get("text_offset", 0),
            coords_style=dp_cfg.get("coords_style", "v2"),
        )
        # SR DataProxy forces v1 coordinate style (consistent with the original)
        self.sr_data_proxy = MagiDataProxy(
            patch_size=dp_cfg.get("patch_size", 2),
            t_patch_size=dp_cfg.get("t_patch_size", 1),
            frame_receptive_field=dp_cfg.get("frame_receptive_field", 11),
            spatial_rope_interpolation=dp_cfg.get("spatial_rope_interpolation", "extra"),
            ref_audio_offset=dp_cfg.get("ref_audio_offset", 1000),
            text_offset=dp_cfg.get("text_offset", 0),
            coords_style="v1",
        )

        self.fps = eval_cfg.get("fps", 25)
        self.num_inference_steps_default = eval_cfg.get("num_inference_steps", 32)
        self.video_txt_guidance_scale = eval_cfg.get("video_txt_guidance_scale", 5.0)
        self.audio_txt_guidance_scale = eval_cfg.get("audio_txt_guidance_scale", 5.0)
        self.shift = eval_cfg.get("shift", 5.0)
        self.cfg_number = eval_cfg.get("cfg_number", 2)
        self.use_cfg_trick = eval_cfg.get("use_cfg_trick", True)
        self.cfg_trick_start_frame = eval_cfg.get("cfg_trick_start_frame", 13)
        self.cfg_trick_value = eval_cfg.get("cfg_trick_value", 2.0)
        self.using_sde_flag = eval_cfg.get("using_sde_flag", False)
        self.t5_gemma_target_length = eval_cfg.get("t5_gemma_target_length", 640)
        self.vae_stride = eval_cfg.get("vae_stride", [4, 16, 16])
        self.z_dim = eval_cfg.get("z_dim", 48)
        self.patch_size = eval_cfg.get("patch_size", [1, 2, 2])
        # SR-specific hyperparameters
        self.sr_num_inference_steps_default = eval_cfg.get("sr_num_inference_steps", 5)
        self.sr_cfg_number = eval_cfg.get("sr_cfg_number", 2)
        self.sr_video_txt_guidance_scale = eval_cfg.get("sr_video_txt_guidance_scale", 3.5)
        self.noise_value = eval_cfg.get("noise_value", 220)
        self.sr_audio_noise_scale = eval_cfg.get("sr_audio_noise_scale", 0.7)
        # ZeroSNR sigma schedule for SR noise injection (flip=True, high to low)
        self.zerosnr_sigmas = ZeroSNRDDPMDiscretization()(1000, do_append_zero=False, flip=True)

        self.context_null, self.original_context_null_len = _get_padded_t5_gemma_embedding(
            _NEGATIVE_PROMPT,
            self.text_encoder,
            self.t5_gemma_target_length,
        )
        self.video_processor = VideoProcessor(vae_scale_factor=16)

        # SR DiT model (loaded from the sr/ subdirectory)
        sr_dit_subfolder = "sr"
        sr_dit_json = _load_json(model_path, f"{sr_dit_subfolder}/config.json", local_files_only)
        sr_dit_model_config = MagiHumanDiTConfig(**sr_dit_json)
        self.sr_dit = DiTModel(sr_dit_model_config)
        self.sr_dit.eval()

        self.weights_sources = [
            DiffusersPipelineLoader.ComponentSource(
                model_or_path=model_path,
                subfolder=dit_subfolder,
                revision=None,
                prefix="dit.",
                fall_back_to_pt=True,
            ),
            DiffusersPipelineLoader.ComponentSource(
                model_or_path=model_path,
                subfolder=sr_dit_subfolder,
                revision=None,
                prefix="sr_dit.",
                fall_back_to_pt=True,
            ),
        ]
        if getattr(self.text_encoder, "is_tp", False):
            self.weights_sources.append(
                DiffusersPipelineLoader.ComponentSource(
                    model_or_path=model_path,
                    subfolder="text_encoder",
                    revision=None,
                    prefix="text_encoder.",
                    fall_back_to_pt=True,
                ),
            )

    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
        # Weight loading for MagiHuman DiT with TP support.
        #
        # The checkpoint stores weights with these naming patterns:
        #   - attention.linear_qkv.weight: fused [Q, K, V, G] for shared layers,
        #     or stacked per-expert [expert0_Q|K|V|G, expert1_..., expert2_...] for MoE.
        #   - attention.linear_proj.weight: single for shared, stacked per-expert for MoE.
        #   - mlp.up_gate_proj.weight / mlp.down_proj.weight: similarly stacked for MoE.
        #
        # The model now uses per-expert vLLM parallel layers for MoE blocks:
        #   attention.linear_qkv.experts.{i}.weight  (QKVParallelLinear per expert)
        #   attention.linear_gating.experts.{i}.weight  (ColumnParallelLinear per expert)
        #   attention.linear_proj.experts.{i}.weight  (RowParallelLinear per expert)
        #   mlp.up_gate_proj.experts.{i}.weight  (ColumnParallelLinear per expert)
        #   mlp.down_proj.experts.{i}.weight  (RowParallelLinear per expert)
        #
        # Shared layers keep the same naming (no .experts.).
        params_dict = dict(self.named_parameters())
        modules_dict = dict(self.named_modules())
        loaded_params: set[str] = set()

        for name, loaded_weight in weights:
            # ── Text Encoder weights ──
            if name.startswith("text_encoder."):
                if getattr(self.text_encoder, "is_tp", False):
                    # Strip "text_encoder." prefix for the T5Gemma TP model
                    # The T5GemmaEncoderModelTP load_weights handles the "encoder." prefix itself
                    sub_name = name[len("text_encoder.") :]
                    loaded_params.update(
                        f"text_encoder.{k}" for k in self.text_encoder.model.load_weights([(sub_name, loaded_weight)])
                    )
                else:
                    loaded_params.add(name)
                continue

            # ── Shared attention QKV + Gating split ──
            # Checkpoint: attention.linear_qkv.weight = [Q, K, V, G] fused.
            # Model: attention.linear_qkv.weight (QKVParallelLinear) + attention.linear_gating.weight.
            if "attention.linear_qkv.weight" in name:
                gating_name = name.replace("attention.linear_qkv.weight", "attention.linear_gating.weight")
                # Check if this is a shared layer (direct param exists, no .experts.)
                if name in params_dict and gating_name in params_dict:
                    qkv_param = params_dict[name]
                    gating_param = params_dict[gating_name]

                    mod_path = name[: -len(".weight")]
                    qkv_mod = modules_dict.get(mod_path)
                    if qkv_mod is not None and hasattr(qkv_mod, "total_num_heads"):
                        total_heads_q = qkv_mod.total_num_heads
                        total_heads_kv = qkv_mod.total_num_kv_heads
                        head_dim = qkv_mod.head_size
                    else:
                        head_dim = 128
                        tp_size = get_tensor_model_parallel_world_size()
                        total_heads_q = gating_param.data.shape[0] * tp_size
                        total_heads_kv = (loaded_weight.shape[0] - total_heads_q * head_dim - total_heads_q) // (
                            2 * head_dim
                        )

                    q_size = total_heads_q * head_dim
                    kv_size = total_heads_kv * head_dim

                    q_w = loaded_weight[:q_size]
                    k_w = loaded_weight[q_size : q_size + kv_size]
                    v_w = loaded_weight[q_size + kv_size : q_size + 2 * kv_size]
                    g_w = loaded_weight[q_size + 2 * kv_size :]

                    qkv_loader = getattr(qkv_param, "weight_loader", default_weight_loader)
                    qkv_loader(qkv_param, q_w, "q")
                    qkv_loader(qkv_param, k_w, "k")
                    qkv_loader(qkv_param, v_w, "v")

                    gating_loader = getattr(gating_param, "weight_loader", default_weight_loader)
                    gating_loader(gating_param, g_w)

                    loaded_params.add(name)
                    loaded_params.add(gating_name)
                    continue

                # ── MoE attention QKV + Gating split ──
                # Checkpoint: attention.linear_qkv.weight = stacked [expert0_QKVG, expert1_QKVG, ...].
                # Model: attention.linear_qkv.experts.{i}.weight (QKVParallelLinear per expert)
                #       + attention.linear_gating.experts.{i}.weight (ColumnParallelLinear per expert).
                expert0_name = name.replace("attention.linear_qkv.weight", "attention.linear_qkv.experts.0.weight")
                if expert0_name in params_dict:
                    # Determine num_experts by checking which expert indices exist.
                    moe_qkv_mod_path = name[: -len(".weight")]
                    moe_qkv_mod = modules_dict.get(moe_qkv_mod_path)
                    num_experts = moe_qkv_mod.num_experts if moe_qkv_mod is not None else 3

                    # Get head info from the first expert's QKVParallelLinear.
                    expert0_mod_path = name.replace("attention.linear_qkv.weight", "attention.linear_qkv.experts.0")
                    expert0_mod = modules_dict.get(expert0_mod_path)
                    if expert0_mod is not None and hasattr(expert0_mod, "total_num_heads"):
                        total_heads_q = expert0_mod.total_num_heads
                        total_heads_kv = expert0_mod.total_num_kv_heads
                        head_dim = expert0_mod.head_size
                    else:
                        head_dim = 128
                        # Infer from checkpoint weight shape.
                        # We'll get exact sizes from model config below.
                        total_heads_q = 40  # fallback for default config
                        total_heads_kv = 8

                    q_size = total_heads_q * head_dim
                    kv_size = total_heads_kv * head_dim
                    # Check if gating is present.
                    gating_expert0_name = name.replace(
                        "attention.linear_qkv.weight", "attention.linear_gating.experts.0.weight"
                    )
                    has_gating = gating_expert0_name in params_dict

                    # Split stacked checkpoint weight into per-expert chunks.
                    expert_weights = loaded_weight.chunk(num_experts, dim=0)

                    for i in range(num_experts):
                        expert_w = expert_weights[i]
                        # Each expert chunk: [Q, K, V, G (optional)].
                        q_w = expert_w[:q_size]
                        k_w = expert_w[q_size : q_size + kv_size]
                        v_w = expert_w[q_size + kv_size : q_size + 2 * kv_size]

                        expert_param_name = name.replace(
                            "attention.linear_qkv.weight",
                            f"attention.linear_qkv.experts.{i}.weight",
                        )
                        expert_param = params_dict[expert_param_name]
                        expert_loader = getattr(expert_param, "weight_loader", default_weight_loader)
                        expert_loader(expert_param, q_w, "q")
                        expert_loader(expert_param, k_w, "k")
                        expert_loader(expert_param, v_w, "v")
                        loaded_params.add(expert_param_name)

                        if has_gating:
                            g_w = expert_w[q_size + 2 * kv_size :]
                            gating_param_name = name.replace(
                                "attention.linear_qkv.weight",
                                f"attention.linear_gating.experts.{i}.weight",
                            )
                            gating_param = params_dict[gating_param_name]
                            gating_loader = getattr(gating_param, "weight_loader", default_weight_loader)
                            gating_loader(gating_param, g_w)
                            loaded_params.add(gating_param_name)
                    continue

            # ── MoE stacked weight splitting for proj / MLP layers ──
            # Checkpoint: x.y.weight (stacked [expert0, expert1, ...]).
            # Model: x.y.experts.{i}.weight.
            if name not in params_dict:
                # Check if this is a stacked MoE weight by looking for .experts.0.
                base, _, suffix = name.rpartition(".")
                expert0_name = f"{base}.experts.0.{suffix}" if base else None
                if expert0_name and expert0_name in params_dict:
                    # Determine num_experts.
                    moe_mod = modules_dict.get(base)
                    num_experts = getattr(moe_mod, "num_experts", 3) if moe_mod is not None else 3

                    # Split stacked weight into per-expert chunks.
                    expert_weights = loaded_weight.chunk(num_experts, dim=0)
                    for i in range(num_experts):
                        expert_param_name = f"{base}.experts.{i}.{suffix}"
                        if expert_param_name not in params_dict:
                            continue
                        expert_param = params_dict[expert_param_name]
                        expert_loader = getattr(expert_param, "weight_loader", default_weight_loader)
                        expert_loader(expert_param, expert_weights[i])
                        loaded_params.add(expert_param_name)
                    continue
                # Truly unknown weight — skip.
                continue

            # ── Standard weight loading (shared layers + non-MoE params) ──
            param = params_dict[name]
            weight_loader = getattr(param, "weight_loader", default_weight_loader)
            weight_loader(param, loaded_weight)
            loaded_params.add(name)

        if getattr(self.text_encoder, "is_tp", False):
            self.context_null, self.original_context_null_len = _get_padded_t5_gemma_embedding(
                _NEGATIVE_PROMPT,
                self.text_encoder,
                self.t5_gemma_target_length,
            )

        return loaded_params

    def _dit_forward(self, eval_input: EvalInput) -> tuple[torch.Tensor, torch.Tensor]:
        packed = self.data_proxy.process_input(eval_input)
        noise_pred = self.dit(*packed)
        return self.data_proxy.process_output(noise_pred)

    def _sr_dit_forward(self, eval_input: EvalInput) -> tuple[torch.Tensor, torch.Tensor]:
        """SR stage uses sr_data_proxy (coords_style=v1) and sr_dit model."""
        packed = self.sr_data_proxy.process_input(eval_input)
        noise_pred = self.sr_dit(*packed)
        return self.sr_data_proxy.process_output(noise_pred)

    @torch.inference_mode()
    def _evaluate_with_latent(
        self,
        context: torch.Tensor,
        original_context_len: int,
        latent_image: torch.Tensor | None,
        latent_video: torch.Tensor,
        latent_audio: torch.Tensor,
        num_inference_steps: int,
        is_a2v: bool = False,
        use_sr_model: bool = False,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        # Select cfg_number and guidance_scale based on BR/SR stage
        cfg_number = self.sr_cfg_number if use_sr_model else self.cfg_number
        video_guidance = self.sr_video_txt_guidance_scale if use_sr_model else self.video_txt_guidance_scale
        forward_fn = self._sr_dit_forward if use_sr_model else self._dit_forward

        video_scheduler = FlowUniPCMultistepScheduler()
        audio_scheduler = FlowUniPCMultistepScheduler()
        video_scheduler.set_timesteps(num_inference_steps, device=self.device_str, shift=self.shift)
        audio_scheduler.set_timesteps(num_inference_steps, device=self.device_str, shift=self.shift)
        timesteps = video_scheduler.timesteps

        latent_length = latent_video.shape[2]
        cfg_trick_guidance = (
            torch.tensor(video_guidance, device=self.device_str).expand(1, 1, latent_length, 1, 1).clone()
        )
        if self.use_cfg_trick:
            cfg_trick_guidance[:, :, : self.cfg_trick_start_frame] = min(self.cfg_trick_value, video_guidance)

        with self.progress_bar(total=len(timesteps)) as pbar:
            for idx, t in enumerate(timesteps):
                if latent_image is not None:
                    latent_video[:, :, :1] = latent_image[:, :, :1]

                # Reduce guidance when t<=500 during BR stage (original behavior)
                cur_video_guidance = video_guidance if (use_sr_model or t > 500) else 2.0

                eval_input_cond = EvalInput(
                    x_t=latent_video,
                    audio_x_t=latent_audio,
                    audio_feat_len=[latent_audio.shape[1]],
                    txt_feat=context,
                    txt_feat_len=[original_context_len],
                )

                v_cond_video, v_cond_audio = forward_fn(eval_input_cond)

                if cfg_number == 1:
                    v_cfg_video = v_cond_video
                    v_cfg_audio = v_cond_audio
                elif cfg_number == 2:
                    eval_input_uncond = EvalInput(
                        x_t=latent_video,
                        audio_x_t=latent_audio,
                        audio_feat_len=[latent_audio.shape[1]],
                        txt_feat=self.context_null,
                        txt_feat_len=[self.original_context_null_len],
                    )
                    v_uncond_video, v_uncond_audio = forward_fn(eval_input_uncond)
                    v_cfg_video = v_uncond_video + cur_video_guidance * (v_cond_video - v_uncond_video)
                    v_cfg_audio = v_uncond_audio + self.audio_txt_guidance_scale * (v_cond_audio - v_uncond_audio)
                else:
                    raise ValueError(f"Invalid cfg_number: {cfg_number}")

                latent_video, latent_audio = _schedule_latent_step(
                    video_scheduler=video_scheduler,
                    audio_scheduler=audio_scheduler,
                    latent_video=latent_video,
                    latent_audio=latent_audio,
                    t=t,
                    idx=idx,
                    steps=timesteps,
                    v_cfg_video=v_cfg_video,
                    v_cfg_audio=v_cfg_audio,
                    is_a2v=is_a2v,
                    cfg_number=cfg_number,
                    using_sde_flag=self.using_sde_flag,
                    use_sr_model=use_sr_model,
                )

                pbar.update()

        if latent_image is not None:
            latent_video[:, :, :1] = latent_image[:, :, :1]
        return latent_video, latent_audio

    def _encode_image(self, image: Image.Image, height: int, width: int) -> torch.Tensor:
        image = load_image(image)
        image = _resizecrop(image, height, width)
        image = self.video_processor.preprocess(image, height=height, width=width)
        image = image.to(device=self.device_str, dtype=self.dtype).unsqueeze(2)
        vae_out = self.vae.encode(image)
        if hasattr(vae_out, "latent_dist"):
            return vae_out.latent_dist.mode().to(torch.float32)
        return vae_out.to(torch.float32)

    def _decode_video(self, latent: torch.Tensor) -> list[np.ndarray]:
        mean = self.vae_latent_mean.to(latent.device, dtype=latent.dtype).view(1, -1, 1, 1, 1)
        std = self.vae_latent_std.to(latent.device, dtype=latent.dtype).view(1, -1, 1, 1, 1)
        latent = latent * std + mean

        videos = self.vae.decode(latent.to(self.dtype))
        if hasattr(videos, "sample"):
            videos = videos.sample
        videos.mul_(0.5).add_(0.5).clamp_(0, 1)
        videos = [v.float().cpu().permute(1, 2, 3, 0) * 255 for v in videos]
        return [v.numpy().astype(np.uint8) for v in videos]

    def _decode_audio(self, latent_audio: torch.Tensor) -> np.ndarray:
        latent_audio = latent_audio.squeeze(0).to(self.dtype)
        audio_output = self.audio_vae.decode(latent_audio.T)
        audio_np = audio_output.squeeze(0).T.float().cpu().numpy()
        target_len = int(audio_np.shape[0] * 441 / 512)
        from scipy.signal import resample

        return resample(audio_np, target_len)

    def encode_prompt(
        self,
        prompt: str,
        target_length: int | None = None,
    ) -> tuple[torch.Tensor, int]:
        """Encode *prompt* with the T5-Gemma text encoder and pad to fixed length.

        This is the single text-encoder entrypoint so the runner-level
        prompt-embedding cache (see
        ``vllm_omni/diffusion/cache/prompt_embed_cache.py``) can transparently
        memoize results when the same prompt is submitted repeatedly.

        Returns:
            ``(context, original_context_len)`` matching
            :func:`_get_padded_t5_gemma_embedding`.
        """
        return _get_padded_t5_gemma_embedding(
            prompt,
            self.text_encoder,
            target_length if target_length is not None else self.t5_gemma_target_length,
        )

    @torch.inference_mode()
    def forward(
        self,
        req: DiffusionRequestBatch,
        prompt: str | None = None,
        height: int = 256,
        width: int = 448,
        num_inference_steps: int | None = None,
        seconds: int = 10,
        seed: int | None = None,
        image_path: str | None = None,
        audio_path: str | None = None,
        **kwargs,
    ) -> DiffusionOutput:
        if len(req.prompts) >= 1:
            p = req.prompts[0]
            prompt = p if isinstance(p, str) else p.get("prompt", prompt)
            if not isinstance(p, str):
                image_path = p.get("image_path", image_path)
                audio_path = p.get("audio_path", audio_path)
        if prompt is None:
            raise ValueError("prompt is required")

        height = req.sampling_params.height or height
        width = req.sampling_params.width or width
        seed = req.sampling_params.seed if req.sampling_params.seed is not None else seed
        num_steps = req.sampling_params.num_inference_steps or num_inference_steps or self.num_inference_steps_default
        sr_height: int | None = None
        sr_width: int | None = None
        sr_num_steps: int | None = None
        if hasattr(req.sampling_params, "extra_args") and req.sampling_params.extra_args:
            seconds = req.sampling_params.extra_args.get("seconds", seconds)
            audio_path = req.sampling_params.extra_args.get("audio_path", audio_path)
            image_path = req.sampling_params.extra_args.get("image_path", image_path)
            sr_height = req.sampling_params.extra_args.get("sr_height", None)
            sr_width = req.sampling_params.extra_args.get("sr_width", None)
            sr_num_steps = req.sampling_params.extra_args.get("sr_num_inference_steps", None)

        device = self.device_str

        br_latent_height = height // self.vae_stride[1] // self.patch_size[1] * self.patch_size[1]
        br_latent_width = width // self.vae_stride[2] // self.patch_size[2] * self.patch_size[2]
        br_height = br_latent_height * self.vae_stride[1]
        br_width = br_latent_width * self.vae_stride[2]

        if seed is not None:
            torch.manual_seed(seed)
            torch.cuda.manual_seed_all(seed)

        if audio_path is not None:
            latent_audio = load_audio_and_encode(self.audio_vae, audio_path, seconds)
            latent_audio = latent_audio.permute(0, 2, 1)
            num_frames = latent_audio.shape[1]
            is_a2v = True
        else:
            num_frames = seconds * self.fps + 1
            latent_audio = torch.randn(1, num_frames, 64, dtype=torch.float32, device=device)
            is_a2v = False

        latent_length = (num_frames - 1) // 4 + 1
        latent_video = torch.randn(
            1,
            self.z_dim,
            latent_length,
            br_latent_height,
            br_latent_width,
            dtype=torch.float32,
            device=device,
        )

        context, original_context_len = self.encode_prompt(prompt)

        if image_path is not None:
            br_image = self._encode_image(load_image(image_path), br_height, br_width)
        else:
            br_image = None

        # ── BR stage ─────────────────────────────────────────────────────────
        br_latent_video, br_latent_audio = self._evaluate_with_latent(
            context,
            original_context_len,
            br_image,
            latent_video.clone(),
            latent_audio.clone(),
            num_steps,
            is_a2v,
            use_sr_model=False,
        )

        # ── SR stage (optional, triggered when sr_height/sr_width are provided) ──
        if sr_height is not None and sr_width is not None:
            sr_latent_height = sr_height // self.vae_stride[1] // self.patch_size[1] * self.patch_size[1]
            sr_latent_width = sr_width // self.vae_stride[2] // self.patch_size[2] * self.patch_size[2]
            sr_height = sr_latent_height * self.vae_stride[1]
            sr_width = sr_latent_width * self.vae_stride[2]

            # Image condition (at SR resolution)
            if image_path is not None:
                sr_image = self._encode_image(load_image(image_path), sr_height, sr_width)
            else:
                sr_image = None

            # Trilinear interpolation of BR latent to SR resolution
            sr_latent_video = torch.nn.functional.interpolate(
                br_latent_video,
                size=(latent_length, sr_latent_height, sr_latent_width),
                mode="trilinear",
                align_corners=True,
            )

            # Noise injection: sigma-weighted blend (noise_value indexes the ZeroSNR sigma schedule)
            if self.noise_value != 0:
                noise = torch.randn_like(sr_latent_video)
                sigma = self.zerosnr_sigmas.to(sr_latent_video.device)[self.noise_value]
                sr_latent_video = sr_latent_video * sigma + noise * (1 - sigma**2) ** 0.5

            # Audio: blend with noise (noised version used during SR inference; final audio keeps BR result)
            sr_latent_audio = torch.randn_like(br_latent_audio) * self.sr_audio_noise_scale + br_latent_audio * (
                1 - self.sr_audio_noise_scale
            )

            torch.accelerator.empty_cache()
            sr_steps = sr_num_steps or self.sr_num_inference_steps_default
            final_latent_video, _ = self._evaluate_with_latent(
                context,
                original_context_len,
                sr_image,
                sr_latent_video.clone(),
                sr_latent_audio.clone(),
                sr_steps,
                is_a2v,
                use_sr_model=True,
            )
            # SR stage does not update audio; keep the BR result
            final_latent_video = final_latent_video
            final_latent_audio = br_latent_audio
        else:
            final_latent_video = br_latent_video
            final_latent_audio = br_latent_audio

        torch.accelerator.empty_cache()
        videos_np = self._decode_video(final_latent_video)
        torch.accelerator.empty_cache()
        audio_np = self._decode_audio(final_latent_audio)

        return DiffusionOutput(output=(videos_np, audio_np))