models/magi_human/magi_human_model.py

import math
import os
import sys
from dataclasses import dataclass
from pathlib import Path
import numpy as np
import torch
import torch.nn.functional as F
import torchaudio.functional as taF
from accelerate import init_empty_weights
from diffusers.video_processor import VideoProcessor
from mmgp import offload
from PIL import Image
from tqdm import tqdm
from transformers import AutoTokenizer
from transformers.models.t5gemma import T5GemmaEncoderModel

from shared.utils import files_locator as fl
from shared.utils.loras_mutipliers import update_loras_slists
from shared.utils.text_encoder_cache import TextEncoderCache
from shared.utils.utils import calculate_new_dimensions
from models.wan.modules.vae2_2 import Wan2_2_VAE

_UPSTREAM_ROOT = Path(__file__).resolve().parent / "upstream"
if str(_UPSTREAM_ROOT) not in sys.path:
    sys.path.insert(0, str(_UPSTREAM_ROOT))

from inference.model.dit.dit_module import DiTModel, MagiAbortRequested  # noqa: E402
from inference.model.sa_audio import SAAudioFeatureExtractor  # noqa: E402
from inference.model.turbo_vaed import get_turbo_vaed  # noqa: E402
from inference.pipeline.data_proxy import MagiDataProxy  # noqa: E402
from inference.pipeline.scheduler_unipc import FlowUniPCMultistepScheduler  # noqa: E402
from inference.pipeline.video_process import load_audio_and_encode, resample_audio_sinc, resizecrop  # noqa: E402


MODEL_CONFIG = {
    "num_layers": 40,
    "hidden_size": 5120,
    "head_dim": 128,
    "num_query_groups": 8,
    "video_in_channels": 48 * 4,
    "audio_in_channels": 64,
    "text_in_channels": 3584,
    "checkpoint_qk_layernorm_rope": False,
    "params_dtype": torch.bfloat16,
    "tread_config": {"selection_rate": 0.5, "start_layer_idx": 2, "end_layer_idx": 25},
    "mm_layers": [0, 1, 2, 3, 36, 37, 38, 39],
    "local_attn_layers": [],
    "enable_attn_gating": True,
    "activation_type": "swiglu7",
    "gelu7_layers": [0, 1, 2, 3],
    "post_norm_layers": [],
}

DEFAULT_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"
)


@dataclass
class _EvalInput:
    x_t: torch.Tensor
    audio_x_t: torch.Tensor
    audio_feat_len: list[int]
    txt_feat: torch.Tensor
    txt_feat_len: list[int]


class _ConfigObject:
    def __init__(self, **kwargs):
        for key, value in kwargs.items():
            setattr(self, key, value)


class _ZeroSNRDDPMDiscretization:
    def __init__(
        self,
        linear_start=0.00085,
        linear_end=0.0120,
        num_timesteps=1000,
        shift_scale=1.0,
        keep_start=False,
        post_shift=False,
    ):
        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).square().cpu().numpy()
        alphas_cumprod = np.cumprod(1.0 - betas, axis=0)
        if not post_shift:
            alphas_cumprod = alphas_cumprod / (shift_scale + (1 - shift_scale) * alphas_cumprod)
        self.alphas_cumprod = alphas_cumprod
        self.post_shift = post_shift
        self.shift_scale = shift_scale

    def get_sigmas(self, n=None, device="cpu"):
        n = self.num_timesteps if n is None else int(n)
        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
        alphas_cumprod = torch.tensor(alphas_cumprod, dtype=torch.float32, device=device)
        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
        return torch.flip(alphas_cumprod_sqrt, dims=(0,))

    def __call__(self, n=None, do_append_zero=True, device="cpu", flip=False):
        sigmas = self.get_sigmas(n=n, device=device)
        if do_append_zero:
            sigmas = torch.cat([sigmas, sigmas.new_zeros((1,))])
        return sigmas if not flip else torch.flip(sigmas, dims=(0,))


class MagiHumanTextEncoder:
    def __init__(self, checkpoint_path: str, tokenizer_path: str, dtype: torch.dtype):
        self.checkpoint_path = checkpoint_path
        self.tokenizer_path = tokenizer_path
        self.dtype = dtype
        self.device = torch.device("cpu")
        self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, local_files_only=True)
        model_prefix = "encoder" if "quanto" in os.path.basename(checkpoint_path).lower() else "model.encoder"
        self.model = offload.fast_load_transformers_model(
            checkpoint_path,
            writable_tensors=False,
            modelClass=T5GemmaEncoderModel,
            defaultConfigPath=os.path.join(tokenizer_path, "config.json"),
            modelPrefix=model_prefix,
            configKwargs={"is_encoder_decoder": False},
            default_dtype=dtype,
        )
        self.model.eval().requires_grad_(False)

    @torch.inference_mode()
    def encode(self, prompts):
        if isinstance(prompts, str):
            prompts = [prompts]
        inputs = self.tokenizer(prompts, return_tensors="pt", padding=True).to(self.device)
        outputs = self.model(**inputs)
        return outputs.last_hidden_state.to(self.dtype)


class MagiHumanModel:
    def __init__(
        self,
        model_filename,
        model_type,
        base_model_type,
        model_def,
        text_encoder_filename=None,
        quantizeTransformer=False,
        dtype=torch.bfloat16,
        VAE_dtype=torch.float32,
        mixed_precision_transformer=False,
        save_quantized=False,
        **kwargs,
    ):
        self.device = torch.device("cuda")
        self.dtype = dtype
        self.VAE_dtype = VAE_dtype
        self.model_def = model_def or {}
        self.model_type = model_type
        self.base_model_type = base_model_type
        self._interrupt = False

        text_encoder_folder = self.model_def.get("text_encoder_folder", "t5gemma-9b-9b-ul2")
        text_encoder_path = text_encoder_filename or fl.locate_file(
            os.path.join(text_encoder_folder, "t5gemma-9b-9b-ul2_bf16.safetensors")
        )
        tokenizer_path = fl.locate_folder(text_encoder_folder)
        self.text_encoder = MagiHumanTextEncoder(text_encoder_path, tokenizer_path, dtype)

        transformer_paths = list(model_filename) if isinstance(model_filename, (list, tuple)) else [model_filename]
        transformer_path = transformer_paths[0]
        self.transformer_config_path = fl.locate_file(self.model_def.get("config_file", "models/magi_human/configs/magi_human_distill.json"))
        model_cfg = self._build_model_config()
        self.transformer = self._load_transformer(transformer_path, model_cfg, quantizeTransformer, save_quantized, dtype)
        self.transformer2 = None
        if len(transformer_paths) > 1:
            self.transformer2 = self._load_transformer(transformer_paths[1], model_cfg, quantizeTransformer, save_quantized, dtype)

        if save_quantized:
            from wgp import save_quantized_model

            save_quantized_model(self.transformer, model_type, transformer_path, dtype, self.transformer_config_path)
            if self.transformer2 is not None:
                save_quantized_model(self.transformer2, model_type, transformer_paths[1], dtype, self.transformer_config_path, submodel_no=2)

        self.data_proxy = self._build_data_proxy("v2")
        self.sr_data_proxy = self._build_data_proxy("v1")
        self.fps = int(self.model_def.get("fps", 25))
        self.vae_stride = (4, 16, 16)
        self.latent_channels = 48
        self.audio_channels = 64
        self.text_target_length = 640
        self.text_encoder_cache = TextEncoderCache()
        self._default_negative_prompt_embeds = None
        self._default_negative_prompt_len = None
        self._turbo_vae_decode_max_pixels = 832 * 480
        self.sr_num_inference_steps = int(self.model_def.get("sr_num_inference_steps", 5))
        self.sr_cfg_number = int(self.model_def.get("sr_cfg_number", 1))
        self.sr_noise_value = int(self.model_def.get("sr_noise_value", 220))
        self.sr_video_txt_guidance_scale = float(self.model_def.get("sr_video_txt_guidance_scale", 3.5))
        self.use_cfg_trick = bool(self.model_def.get("use_cfg_trick", True))
        self.cfg_trick_start_frame = int(self.model_def.get("cfg_trick_start_frame", 13))
        self.cfg_trick_value = float(self.model_def.get("cfg_trick_value", 2.0))
        self.using_sde_flag = bool(self.model_def.get("using_sde_flag", False))
        self.sr_audio_noise_scale = float(self.model_def.get("sr_audio_noise_scale", 0.7))
        self._sr_sigmas = _ZeroSNRDDPMDiscretization()(1000, do_append_zero=False, flip=True)
        self.video_processor = VideoProcessor(vae_scale_factor=16)

        self.vae = Wan2_2_VAE(vae_pth=fl.locate_file("Wan2.2_VAE.safetensors"), dtype=VAE_dtype, device="cpu")
        self.vae.device = self.device
        self.audio_vae = SAAudioFeatureExtractor(device="cpu", model_path=fl.locate_folder("stable-audio-open-1.0"))
        self.turbo_vae = get_turbo_vaed(
            fl.locate_file("turbo_vae/TurboV3-Wan22-TinyShallow_7_7.json"),
            fl.locate_file("turbo_vae/TurboV3-Wan22-TinyShallow_7_7.safetensors"),
            device="cpu",
            weight_dtype=VAE_dtype,
        )

    def _build_model_config(self):
        cfg = dict(MODEL_CONFIG)
        cfg["num_heads_q"] = cfg["hidden_size"] // cfg["head_dim"]
        cfg["num_heads_kv"] = cfg["num_query_groups"]
        return _ConfigObject(**cfg)

    def _build_data_proxy(self, coords_style: str):
        return MagiDataProxy(
            _ConfigObject(
                t_patch_size=1,
                patch_size=2,
                frame_receptive_field=11,
                spatial_rope_interpolation="extra",
                ref_audio_offset=1000,
                text_offset=0,
                coords_style=coords_style,
            )
        )

    def _load_transformer(self, checkpoint_path: str, model_cfg, quantizeTransformer: bool, save_quantized: bool, dtype: torch.dtype):
        with init_empty_weights():
            transformer = DiTModel(model_cfg)
        offload.load_model_data(
            transformer,
            checkpoint_path,
            do_quantize=quantizeTransformer and not save_quantized,
            writable_tensors=False,
            default_dtype=dtype,
        )
        transformer.eval().requires_grad_(False)
        transformer._interrupt_check = self._interrupt_requested
        transformer.block._interrupt_check = self._interrupt_requested
        return transformer

    def _interrupt_requested(self):
        return bool(self._interrupt)

    def request_early_stop(self):
        self._interrupt = True

    def _run_transformer(self, transformer, *inputs):
        try:
            return transformer(*inputs)
        except MagiAbortRequested:
            self._interrupt = True
            return None

    def prepare_preview_payload(self, latents, preview_meta=None):
        if not torch.is_tensor(latents):
            return None
        return {"latents": latents.float()}

    def _pad_or_trim(self, tensor: torch.Tensor, target_size: int):
        current_size = tensor.shape[1]
        if current_size < target_size:
            pad = torch.zeros((tensor.shape[0], target_size - current_size, tensor.shape[2]), device=tensor.device, dtype=tensor.dtype)
            return torch.cat([tensor, pad], dim=1), current_size
        return tensor[:, :target_size], target_size

    def _pad_or_trim_prompt(self, tensor: torch.Tensor, target_size: int):
        current_size = tensor.shape[0]
        if current_size < target_size:
            pad = torch.zeros((target_size - current_size, tensor.shape[1]), device=tensor.device, dtype=tensor.dtype)
            return torch.cat([tensor, pad], dim=0), current_size
        return tensor[:target_size], target_size

    def _normalize_prompt_batch(self, prompt, batch_size: int):
        if isinstance(prompt, str):
            return [prompt] * batch_size
        prompt = list(prompt)
        if len(prompt) == 1 and batch_size > 1:
            return prompt * batch_size
        if len(prompt) != batch_size:
            raise ValueError("Prompt batch size does not match Magi Human batch size.")
        return prompt

    def _encode_prompt(self, prompt, batch_size: int):
        prompts = self._normalize_prompt_batch(prompt, batch_size)
        encode_fn = lambda prompt_batch: list(self.text_encoder.encode(prompt_batch))
        encoded_prompts = self.text_encoder_cache.encode(encode_fn, prompts, device=self.device, parallel=True)
        embeddings = []
        prompt_lens = []
        for encoded_prompt in encoded_prompts:
            padded_prompt, prompt_len = self._pad_or_trim_prompt(encoded_prompt, self.text_target_length)
            embeddings.append(padded_prompt)
            prompt_lens.append(int(prompt_len))
        stacked = torch.stack(embeddings, dim=0).to(self.device)
        unique_lens = set(prompt_lens)
        return stacked, prompt_lens[0] if len(unique_lens) == 1 else prompt_lens

    def _get_uncond_prompt(self, negative_prompt: str):
        if negative_prompt:
            embeds, prompt_len = self._encode_prompt(negative_prompt, 1)
            return embeds, [int(prompt_len)]
        if self._default_negative_prompt_embeds is None:
            embeds, prompt_len = self._encode_prompt(DEFAULT_NEGATIVE_PROMPT, 1)
            self._default_negative_prompt_embeds = embeds
            self._default_negative_prompt_len = [int(prompt_len)]
        return self._default_negative_prompt_embeds, self._default_negative_prompt_len

    def _coerce_image_batch(self, image, batch_size: int):
        if image is None:
            return None
        if isinstance(image, list):
            image = image[0] if image else None
        if image is None:
            return None
        if not torch.is_tensor(image):
            raise ValueError("Magi Human expects image tensors prepared by WanGP.")
        if image.dim() == 3:
            image = image.unsqueeze(0)
        if image.dim() == 5 and image.shape[2] == 1:
            image = image.squeeze(2)
        if image.dim() != 4:
            raise ValueError(f"Unexpected Magi Human image shape: {tuple(image.shape)}")
        image = image.to(device=self.device, dtype=self.VAE_dtype)
        if image.shape[0] == 1 and batch_size > 1:
            image = image.repeat(batch_size, 1, 1, 1)
        return image

    def _coerce_last_video_frame_batch(self, video, batch_size: int):
        if video is None:
            return None
        if not torch.is_tensor(video):
            raise ValueError("Magi Human expects video tensors prepared by WanGP.")
        if video.dim() == 4:
            return self._coerce_image_batch(video[:, -1], batch_size)
        if video.dim() == 5:
            return self._coerce_image_batch(video[:, :, -1], batch_size)
        raise ValueError(f"Unexpected Magi Human video shape: {tuple(video.shape)}")

    def _encode_image_latent(self, image_batch: torch.Tensor, height: int, width: int, tile_size: int):
        latents = []
        for image in image_batch:
            image_uint8 = image.float().clamp(-1, 1).add(1).mul(127.5).to(torch.uint8).permute(1, 2, 0).cpu().numpy()
            image_pil = resizecrop(Image.fromarray(image_uint8), height, width)
            image_tensor = self.video_processor.preprocess(image_pil, height=height, width=width)[0].to(device=self.device, dtype=self.VAE_dtype)
            latents.append(self.vae.encode([image_tensor.unsqueeze(1)], tile_size=tile_size)[0])
        return torch.stack(latents, dim=0).to(device=self.device, dtype=torch.float32)

    def _resolve_vae_tile_size(self, VAE_tile_size, height: int, width: int):
        if isinstance(VAE_tile_size, dict):
            tile_size = int(VAE_tile_size.get("tile_sample_min_size", VAE_tile_size.get("tile_latent_min_size", 0)) or 0)
        elif isinstance(VAE_tile_size, (list, tuple)):
            if not VAE_tile_size:
                tile_size = 0
            elif len(VAE_tile_size) >= 2 and isinstance(VAE_tile_size[0], bool):
                tile_size = int(VAE_tile_size[1] if VAE_tile_size[0] else 0)
            else:
                tile_size = int(VAE_tile_size[-1] or 0)
        else:
            tile_size = int(VAE_tile_size or 0)

        if height * width >= 1280 * 704:
            print(
                f"[Magi][VAE] resolved tile_size={tile_size} for {width}x{height} "
                f"(input={VAE_tile_size})",
                flush=True,
            )
        return tile_size

    def _load_audio_latent(self, audio_path: str, frame_num: int, fps: int):
        seconds = max(1, math.ceil((frame_num - 1) / max(fps, 1)))
        audio_latent = load_audio_and_encode(self.audio_vae, audio_path, seconds=seconds).permute(0, 2, 1).to(torch.float32)
        return self._finalize_audio_latent(audio_latent, frame_num)

    def _finalize_audio_latent(self, audio_latent: torch.Tensor, frame_num: int):
        if audio_latent.shape[1] < frame_num:
            pad = torch.zeros((audio_latent.shape[0], frame_num - audio_latent.shape[1], audio_latent.shape[2]), device=audio_latent.device, dtype=audio_latent.dtype)
            audio_latent = torch.cat([audio_latent, pad], dim=1)
        return audio_latent[:, :frame_num].to(self.device)

    def _encode_audio_waveform_latent(self, input_waveform, input_waveform_sample_rate: int, frame_num: int, fps: int):
        if input_waveform_sample_rate is None or int(input_waveform_sample_rate) <= 0:
            raise ValueError("Magi Human requires a valid input_waveform_sample_rate when input_waveform is provided.")
        waveform = torch.from_numpy(input_waveform) if isinstance(input_waveform, np.ndarray) else input_waveform
        if not torch.is_tensor(waveform):
            raise ValueError("Magi Human expects input_waveform as a numpy array or torch tensor.")
        waveform = waveform.to(dtype=torch.float32)
        if waveform.ndim == 1:
            waveform = waveform.unsqueeze(0)
        elif waveform.ndim == 2:
            waveform = waveform.T
        elif waveform.ndim == 3 and waveform.shape[0] == 1:
            waveform = waveform.squeeze(0).T
        if waveform.ndim != 2:
            raise ValueError(f"Unexpected Magi Human waveform shape: {tuple(waveform.shape)}")
        target_sample_rate = 51200
        if int(input_waveform_sample_rate) != target_sample_rate:
            waveform = taF.resample(waveform, int(input_waveform_sample_rate), target_sample_rate)
        seconds = max(1, math.ceil((frame_num - 1) / max(fps, 1)))
        waveform = waveform[:, : min(waveform.shape[-1], int(seconds * target_sample_rate))]
        if waveform.shape[0] == 1:
            waveform = waveform.expand(2, -1)
        elif waveform.shape[0] > 2:
            waveform = waveform[:2]
        encode_device = next(self.audio_vae.vae_model.parameters()).device
        waveform = waveform.to(device=encode_device)
        audio_latent = self.audio_vae.encode(waveform).permute(0, 2, 1).to(torch.float32)
        return self._finalize_audio_latent(audio_latent, frame_num)

    def _decode_video(self, latents: torch.Tensor, tile_size: int):
        turbo_dtype = next(self.turbo_vae.parameters()).dtype
        sample_h = int(latents.shape[-2]) * self.vae_stride[1]
        sample_w = int(latents.shape[-1]) * self.vae_stride[2]
        output_offload = int(tile_size) > 0 or sample_h * sample_w > self._turbo_vae_decode_max_pixels
        temporal_chunk_size = 0
        # if sample_h * sample_w >= 1920 * 1088:
        #     temporal_chunk_size = 3
        # elif sample_h * sample_w >= 1280 * 704 and latents.shape[2] > self.turbo_vae.step_size:
        #     temporal_chunk_size = 5
        # print(
        #     f"[Magi][VAE] decode {sample_w}x{sample_h} latent_shape={tuple(latents[:1].shape)} "
        #     f"tile_size={tile_size} temporal_chunk_size={temporal_chunk_size} output_offload={output_offload}",
        #     flush=True,
        # )
        decoded = self.turbo_vae.decode(
            latents[:1].to(device=self.device, dtype=turbo_dtype),
            output_offload=output_offload,
            tile_size=tile_size,
            temporal_chunk_size=temporal_chunk_size,
        )
        with torch.inference_mode(False):
            return decoded[0].float().clamp(-1, 1).clone()

    def _decode_audio(self, latent_audio: torch.Tensor):
        audio_dtype = next(self.audio_vae.vae_model.parameters()).dtype
        audio_output = self.audio_vae.decode(latent_audio.squeeze(0).T.to(audio_dtype))
        audio_output = audio_output.float().squeeze(0).T.detach().cpu().numpy()
        return resample_audio_sinc(audio_output, 441 / 512)

    def _prepare_proxy_input(self, eval_input: _EvalInput, data_proxy: MagiDataProxy):
        packed = data_proxy.process_input(eval_input)
        return packed, dict(data_proxy._saved_data)

    def _process_proxy_output(self, pred: torch.Tensor, saved_state: dict, data_proxy: MagiDataProxy):
        data_proxy._saved_data = dict(saved_state)
        return data_proxy.process_output(pred)

    def _resolve_base_phase_dimensions(self, height: int, width: int):
        # return 256, 448
        phase_height, phase_width = calculate_new_dimensions(256, 448, height, width, 0, block_size=self.model_def.get("vae_block_size", 32))
        return int(phase_height), int(phase_width)

    def _run_diffusion_phase(
        self,
        transformer,
        data_proxy,
        latent_video,
        latent_audio,
        image_latent,
        prompt_embeds,
        prompt_lens,
        sampling_steps,
        shift,
        guide_scale,
        audio_cfg_scale,
        use_audio_guide,
        callback,
        pass_no,
        total_passes,
        joint_pass,
        uncond_prompt_embeds=None,
        uncond_prompt_lens=None,
        update_audio=True,
        use_sr_model=False,
        sr_cfg_scale=None,
    ):
        video_scheduler = FlowUniPCMultistepScheduler()
        audio_scheduler = FlowUniPCMultistepScheduler()
        video_scheduler.set_timesteps(int(sampling_steps), device=self.device, shift=float(shift))
        audio_scheduler.set_timesteps(int(sampling_steps), device=self.device, shift=float(shift))
        timesteps = video_scheduler.timesteps
        cfg_number = int(self.sr_cfg_number if use_sr_model else 2)
        if not use_sr_model:
            cfg_number = 1 if float(guide_scale) == 1.0 and (not update_audio or float(audio_cfg_scale) == 1.0) else 2
        sr_guidance = None
        if use_sr_model and cfg_number == 2:
            latent_frames = latent_video.shape[2]
            sr_guidance = torch.full((latent_video.shape[0], 1, latent_frames, 1, 1), float(self.sr_video_txt_guidance_scale), device=self.device, dtype=latent_video.dtype)
            if self.use_cfg_trick:
                sr_guidance[:, :, : self.cfg_trick_start_frame] = min(self.cfg_trick_value, self.sr_video_txt_guidance_scale)
            if sr_cfg_scale is not None:
                sr_guidance.mul_(float(sr_cfg_scale) / float(self.sr_video_txt_guidance_scale))
        if callback is not None:
            callback(-1, None, True, override_num_inference_steps=len(timesteps), pass_no=pass_no)
        progress_bar = tqdm(timesteps, desc=f"Phase {pass_no}/{total_passes}" if total_passes > 1 else None)
        for step_idx, t in enumerate(progress_bar):
            if self._interrupt:
                return None, None
            latent_video[:, :, :1] = image_latent[:, :, :1]
            eval_input = _EvalInput(
                x_t=latent_video,
                audio_x_t=latent_audio,
                audio_feat_len=[latent_audio.shape[1]] * latent_audio.shape[0],
                txt_feat=prompt_embeds,
                txt_feat_len=prompt_lens,
            )
            packed_inputs, packed_state = self._prepare_proxy_input(eval_input, data_proxy)
            if cfg_number == 2:
                uncond_eval_input = _EvalInput(
                    x_t=latent_video,
                    audio_x_t=latent_audio,
                    audio_feat_len=[latent_audio.shape[1]] * latent_audio.shape[0],
                    txt_feat=uncond_prompt_embeds,
                    txt_feat_len=uncond_prompt_lens,
                )
                uncond_inputs, uncond_state = self._prepare_proxy_input(uncond_eval_input, data_proxy)
                if joint_pass:
                    pred_pair = self._run_transformer(transformer, *[[packed_inputs[i], uncond_inputs[i]] for i in range(len(packed_inputs))])
                    if pred_pair is None:
                        return None, None
                    pred_cond, pred_uncond = pred_pair
                else:
                    pred_cond = self._run_transformer(transformer, *packed_inputs)
                    if pred_cond is None:
                        return None, None
                    pred_uncond = self._run_transformer(transformer, *uncond_inputs)
                    if pred_uncond is None:
                        return None, None
                pred_video, pred_audio = self._process_proxy_output(pred_cond, packed_state, data_proxy)
                pred_video_uncond, pred_audio_uncond = self._process_proxy_output(pred_uncond, uncond_state, data_proxy)
                current_video_guidance = sr_guidance if use_sr_model else float(guide_scale if t > 500 else 2.0)
                pred_video = pred_video_uncond + current_video_guidance * (pred_video - pred_video_uncond)
                pred_audio = pred_audio_uncond + float(audio_cfg_scale) * (pred_audio - pred_audio_uncond)
                latent_video = video_scheduler.step(pred_video, t, latent_video, return_dict=False)[0]
                if update_audio and not use_audio_guide:
                    latent_audio = audio_scheduler.step(pred_audio, t, latent_audio, return_dict=False)[0]
            else:
                pred = self._run_transformer(transformer, *packed_inputs)
                if pred is None:
                    return None, None
                pred_video, pred_audio = self._process_proxy_output(pred, packed_state, data_proxy)
                if use_sr_model:
                    latent_video = video_scheduler.step(pred_video, t, latent_video, return_dict=False)[0]
                else:
                    latent_video = video_scheduler.step_ddim(pred_video, step_idx, latent_video)
                if update_audio and not use_audio_guide:
                    latent_audio = audio_scheduler.step_ddim(pred_audio, step_idx, latent_audio)
            if callback is not None:
                callback(step_idx, latent_video[0].detach(), pass_no=pass_no)
        latent_video[:, :, :1] = image_latent[:, :, :1]
        return latent_video, latent_audio

    @torch.inference_mode()
    def generate(
        self,
        seed=None,
        input_prompt="",
        alt_prompt="",
        n_prompt="",
        sampling_steps=8,
        input_ref_images=None,
        input_frames=None,
        input_frames2=None,
        input_masks=None,
        input_masks2=None,
        input_video=None,
        image_start=None,
        image_end=None,
        frame_num=101,
        batch_size=1,
        height=256,
        width=448,
        guide_scale=1.0,
        guide2_scale=1.0,
        guide3_scale=1.0,
        switch_threshold=0.0,
        switch2_threshold=0.0,
        guide_phases=1,
        model_switch_phase=0,
        embedded_guidance_scale=0.0,
        shift=5.0,
        sample_solver="unipc",
        denoising_strength=1.0,
        masking_strength=1.0,
        callback=None,
        VAE_tile_size=None,
        joint_pass=True,
        audio_cfg_scale=1.0,
        prefix_video=None,
        prefix_frames_count=0,
        input_video_strength=1.0,
        input_waveform=None,
        input_waveform_sample_rate=None,
        audio_guide=None,
        audio_guide2=None,
        audio_prompt_type="",
        fps=None,
        offloadobj=None,
        set_header_text=None,
        loras_slists=None,
        **kwargs,
    ):
        self._interrupt = False
        if seed is None or seed == -1:
            seed = torch.seed() % (2**32 - 1)
        if fps is None or fps <= 0:
            fps = self.fps
        if batch_size < 1:
            raise ValueError("Magi Human batch_size must be positive.")
        if batch_size != 1:
            raise ValueError("Magi Human currently supports batch_size=1 only.")

        generator = torch.Generator(device=self.device)
        generator.manual_seed(seed)

        prompt = alt_prompt or input_prompt
        prompt_embeds, prompt_len = self._encode_prompt(prompt, batch_size)
        if isinstance(prompt_len, int):
            prompt_lens = [prompt_len] * batch_size
        else:
            prompt_lens = [int(prompt_len)] * batch_size
        if self.base_model_type == "magi_human_distill":
            guide_scale = 1.0
            audio_cfg_scale = 1.0
        need_cfg = float(guide_scale) != 1.0 or float(audio_cfg_scale) != 1.0 or (self.transformer2 is not None and self.sr_cfg_number == 2)
        if need_cfg:
            uncond_prompt_embeds, uncond_prompt_lens = self._get_uncond_prompt(n_prompt)
        else:
            uncond_prompt_embeds, uncond_prompt_lens = None, None
        image_batch = input_video[:, -1].unsqueeze(0)
        # image_batch = self._coerce_last_video_frame_batch(input_video, batch_size)
        # if image_batch is None:
        #     image_batch = self._coerce_last_video_frame_batch(prefix_video, batch_size)
        # if image_batch is None:
        #     source_image = image_start
        #     if source_image is None and input_ref_images is not None:
        #         source_image = input_ref_images[0] if isinstance(input_ref_images, list) and input_ref_images else input_ref_images
        #     image_batch = self._coerce_image_batch(source_image, batch_size)
        # if image_batch is None:
        #     raise ValueError("Magi Human requires a start image.")

        total_passes = 2 if self.transformer2 is not None else 1
        phase1_height, phase1_width = self._resolve_base_phase_dimensions(height, width) if self.transformer2 is not None else (height, width)
        tile_size = self._resolve_vae_tile_size(VAE_tile_size, phase1_height, phase1_width)
        image_latent = self._encode_image_latent(image_batch, phase1_height, phase1_width, tile_size)

        latent_frames = (frame_num - 1) // self.vae_stride[0] + 1
        latent_h = phase1_height // self.vae_stride[1]
        latent_w = phase1_width // self.vae_stride[2]
        latent_video = torch.randn(
            (batch_size, self.latent_channels, latent_frames, latent_h, latent_w),
            generator=generator,
            device=self.device,
            dtype=torch.float32,
        )
        latent_video[:, :, :1] = image_latent[:, :, :1]

        use_audio_guide = bool("A" in (audio_prompt_type or "") and (input_waveform is not None or audio_guide))
        if "A" in (audio_prompt_type or "") and input_waveform is not None:
            latent_audio = self._encode_audio_waveform_latent(input_waveform, input_waveform_sample_rate, frame_num, fps)
        elif use_audio_guide:
            latent_audio = self._load_audio_latent(audio_guide, frame_num, fps)
        else:
            latent_audio = torch.randn((batch_size, frame_num, self.audio_channels), generator=generator, device=self.device, dtype=torch.float32)

        if loras_slists is not None:
            update_loras_slists(self.transformer, loras_slists, int(sampling_steps), phase_switch_step=int(sampling_steps), phase_switch_step2=int(sampling_steps))
        latent_video, latent_audio = self._run_diffusion_phase(
            self.transformer,
            self.data_proxy,
            latent_video,
            latent_audio,
            image_latent,
            prompt_embeds,
            prompt_lens,
            sampling_steps,
            shift,
            guide_scale,
            audio_cfg_scale,
            use_audio_guide,
            callback,
            1,
            total_passes,
            joint_pass,
            uncond_prompt_embeds=uncond_prompt_embeds,
            uncond_prompt_lens=uncond_prompt_lens,
            update_audio=True,
            use_sr_model=False,
        )
        if latent_video is None or latent_audio is None:
            return None

        if self._interrupt:
            return None

        final_video_latent = latent_video
        final_audio_latent = latent_audio
        final_tile_size = tile_size
        if self.transformer2 is not None:
            final_tile_size = self._resolve_vae_tile_size(VAE_tile_size, height, width)
            sr_image_latent = self._encode_image_latent(image_batch, height, width, final_tile_size)
            sr_latent_h = height // self.vae_stride[1]
            sr_latent_w = width // self.vae_stride[2]
            sr_video_latent = F.interpolate(latent_video, size=(latent_frames, sr_latent_h, sr_latent_w), mode="trilinear", align_corners=True)
            if self.sr_noise_value > 0:
                sigma = self._sr_sigmas.to(sr_video_latent.device)[self.sr_noise_value]
                sr_video_latent = sr_video_latent * sigma + torch.randn_like(sr_video_latent) * (1 - sigma**2).sqrt()
            sr_audio_latent = torch.randn_like(latent_audio) * self.sr_audio_noise_scale + latent_audio * (1 - self.sr_audio_noise_scale)
            if loras_slists is not None:
                update_loras_slists(self.transformer2, loras_slists, self.sr_num_inference_steps, phase_switch_step=0, phase_switch_step2=self.sr_num_inference_steps)
            final_video_latent, _ = self._run_diffusion_phase(
                self.transformer2,
                self.sr_data_proxy,
                sr_video_latent,
                sr_audio_latent,
                sr_image_latent,
                prompt_embeds,
                prompt_lens,
                self.sr_num_inference_steps,
                shift,
                1.0,
                1.0,
                True,
                callback,
                2,
                total_passes,
                False,
                uncond_prompt_embeds=uncond_prompt_embeds,
                uncond_prompt_lens=uncond_prompt_lens,
                update_audio=False,
                use_sr_model=True,
            )
            if final_video_latent is None:
                return None

        video = self._decode_video(final_video_latent, final_tile_size)
        audio = self._decode_audio(final_audio_latent[:1])
        return {"x": video, "audio": audio, "audio_sampling_rate": 44100}