diff --git a/fastvideo/pipelines/basic/magi_human/magi_human_pipeline.py b/fastvideo/pipelines/basic/magi_human/magi_human_pipeline.py new file mode 100644 index 000000000..0c3a985b6 --- /dev/null +++ b/fastvideo/pipelines/basic/magi_human/magi_human_pipeline.py @@ -0,0 +1,417 @@ +# SPDX-License-Identifier: Apache-2.0 +"""MagiHuman base text-to-AV pipeline. + +Top-level composition for the daVinci-MagiHuman base model. Wires: + + InputValidationStage -> TextEncodingStage (T5-Gemma) + -> MagiHumanLatentPreparationStage + -> MagiHumanDenoisingStage + -> DecodingStage (Wan 2.2 TI2V-5B VAE decode for video) + -> MagiHumanAudioDecodingStage (Stable Audio Open 1.0 VAE decode) + +The base checkpoint is a joint audio-visual generator; both the video +and audio paths run in the denoising loop and both are decoded. + +`load_modules` is overridden so the four cross-variant shared components +(text_encoder, tokenizer, audio_vae, video vae) lazy-load from their +canonical upstream HF repos at first build time instead of being +bundled inside every converted MagiHuman variant. This keeps each +variant's converted repo at ~5-30 GB (transformer + scheduler + +model_index.json) instead of ~30-55 GB, and lets all variants share +the same ~25 GB of cached upstream weights. +""" +from __future__ import annotations + +import json +import os +from pathlib import Path +from typing import Any + +from transformers import AutoTokenizer + +from fastvideo.configs.models.encoders.t5gemma import T5GemmaEncoderConfig +from fastvideo.configs.models.vaes import OobleckVAEConfig +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.logger import init_logger +from fastvideo.models.encoders.t5gemma import T5GemmaEncoderModel +from fastvideo.models.vaes.sa_audio import SAAudioVAEModel +from fastvideo.models.schedulers.scheduling_flow_unipc_multistep import ( + FlowUniPCMultistepScheduler, ) +from fastvideo.pipelines.basic.magi_human.stages import ( + MagiHumanAudioDecodingStage, + MagiHumanDenoisingStage, + MagiHumanLatentPreparationStage, + MagiHumanReferenceImageStage, + MagiHumanSRDenoisingStage, + MagiHumanSRLatentPreparationStage, +) +from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase +from fastvideo.pipelines.stages import ( + DecodingStage, + InputValidationStage, + TextEncodingStage, +) +from fastvideo.utils import maybe_download_model + +logger = init_logger(__name__) + +_T5GEMMA_HF_ID = "google/t5gemma-9b-9b-ul2" +_SA_AUDIO_HF_ID = "stabilityai/stable-audio-open-1.0" +_WAN_VAE_HF_ID = "Wan-AI/Wan2.2-TI2V-5B-Diffusers" + + +def _ensure_hf_token_env() -> str | None: + """Surface any of the three common HF token env vars as `HF_TOKEN`. + + FastVideo workers spawn child processes that inherit env; both + `huggingface_hub` and `transformers.AutoTokenizer.from_pretrained` + look at `HF_TOKEN` / `HUGGINGFACE_HUB_TOKEN` by default but not + `HF_API_KEY`. If only the latter is set, gated downloads fail with + 401. Aliasing at pipeline-load time is the minimum-disruption fix. + """ + for src in ("HF_TOKEN", "HUGGINGFACE_HUB_TOKEN", "HF_API_KEY"): + value = os.environ.get(src) + if value: + os.environ.setdefault("HF_TOKEN", value) + os.environ.setdefault("HUGGINGFACE_HUB_TOKEN", value) + return value + return None + + +class MagiHumanPipeline(ComposedPipelineBase): + """Base MagiHuman text-to-AV pipeline (no LoRA, no distill, no SR).""" + + _required_config_modules = [ + "text_encoder", + "tokenizer", + "vae", + "transformer", + "scheduler", + "audio_vae", + ] + + def load_modules( + self, + fastvideo_args: FastVideoArgs, + loaded_modules: dict[str, Any] | None = None, + ) -> dict[str, Any]: + """Load the variant-specific transformer + scheduler from the + converted MagiHuman repo and lazy-load the four cross-variant + shared components from their canonical upstream HF repos: + + * text_encoder, tokenizer -> ``google/t5gemma-9b-9b-ul2`` + (gated, requires HF token with accepted terms of use) + * audio_vae -> ``stabilityai/stable-audio-open-1.0`` (gated) + * vae -> ``Wan-AI/Wan2.2-TI2V-5B-Diffusers`` + + Backwards-compatible with bundled converted repos: if any of + these subfolders is present locally and listed in + ``model_index.json``, the standard component loader picks it up + via super(). Otherwise the loader is told to skip the entry and + we lazy-load it here. + """ + # T5-Gemma is gated: expose `HF_API_KEY` as `HF_TOKEN` if needed. + _ensure_hf_token_env() + + # Resolve to a local cache path so we can inspect + # model_index.json before invoking super(). `maybe_download_model` + # is idempotent for local paths; super() repeats the call cheaply + # via `_load_config`. + local_path = maybe_download_model(self.model_path) + + # Identify which cross-variant shared keys are bundled in the + # converted repo (declared in model_index.json with a non-null + # spec) versus absent (the umbrella scheme). Bundled keys stay + # in `required_config_modules` and are loaded normally by super() + # from `//`. Absent keys are temporarily + # dropped so super() does not fail the "every required entry + # must appear in model_index.json" check, then lazy-loaded + # below. + model_index: dict[str, Any] = {} + try: + with open(Path(local_path) / "model_index.json") as f: + model_index = json.load(f) + except (FileNotFoundError, json.JSONDecodeError): + pass + + def _is_bundled(key: str) -> bool: + spec = model_index.get(key) + return (isinstance(spec, list | tuple) and len(spec) >= 1 and spec[0] is not None) + + deferred = [] + for key in ("text_encoder", "tokenizer", "audio_vae", "vae"): + if key in self.required_config_modules and not _is_bundled(key): + self.required_config_modules.remove(key) + deferred.append(key) + + try: + modules = super().load_modules(fastvideo_args, loaded_modules) + finally: + for key in deferred: + if key not in self.required_config_modules: + self.required_config_modules.append(key) + + # For each lazy-load key, prefer whatever super() already loaded + # (a bundled subfolder, or a caller-provided override merged in + # via `loaded_modules`). Fall back to the caller-provided + # `loaded_modules` entry for keys absent from model_index.json + # (super() never iterates those). Otherwise lazy-load from the + # canonical upstream HF repo. + def _resolve(key: str) -> bool: + """Return True if `modules[key]` is already populated.""" + if modules.get(key) is not None: + return True + if loaded_modules and key in loaded_modules: + modules[key] = loaded_modules[key] + return True + return False + + if not _resolve("text_encoder"): + logger.info("Building T5-Gemma text encoder (lazy-load from %s)", _T5GEMMA_HF_ID) + enc_config = T5GemmaEncoderConfig() + enc_config.arch_config.t5gemma_model_path = _T5GEMMA_HF_ID + modules["text_encoder"] = T5GemmaEncoderModel(enc_config) + + if not _resolve("tokenizer"): + logger.info("Loading T5-Gemma tokenizer from %s", _T5GEMMA_HF_ID) + modules["tokenizer"] = AutoTokenizer.from_pretrained(_T5GEMMA_HF_ID) + + if not _resolve("audio_vae"): + logger.info( + "Building Stable Audio Open 1.0 VAE (lazy-load from %s) — " + "requires HF terms accepted for gated repo", + _SA_AUDIO_HF_ID, + ) + audio_config = OobleckVAEConfig() + audio_config.pretrained_path = _SA_AUDIO_HF_ID + modules["audio_vae"] = SAAudioVAEModel(audio_config) + + if not _resolve("vae"): + modules["vae"] = self._load_video_vae(fastvideo_args) + + return modules + + def _load_video_vae(self, fastvideo_args: FastVideoArgs) -> Any: + """Resolve the video VAE: prefer a bundled ``vae/`` subfolder in + the converted repo (legacy), fall back to lazy-downloading the + Wan 2.2 TI2V-5B VAE shards from upstream. + + Either way the load goes through FastVideo's standard + ``VAELoader`` so the result is the same FV ``AutoencoderKLWan`` + nn.Module that the bundled path produces. + """ + from fastvideo.models.loader.component_loader import VAELoader + + bundled = Path(self.model_path) / "vae" + if bundled.is_dir() and (bundled / "config.json").is_file(): + logger.info("Loading bundled video VAE from %s", bundled) + return VAELoader().load(str(bundled), fastvideo_args) + + from huggingface_hub import snapshot_download + + logger.info( + "Bundled vae/ not found at %s; lazy-loading Wan 2.2 TI2V-5B VAE from %s", + self.model_path, + _WAN_VAE_HF_ID, + ) + snapshot = snapshot_download( + repo_id=_WAN_VAE_HF_ID, + allow_patterns=["vae/*"], + ) + vae_dir = os.path.join(snapshot, "vae") + if not os.path.isdir(vae_dir): + raise RuntimeError( + f"snapshot_download returned {snapshot} but no vae/ " + f"subfolder was found inside it. Check that {_WAN_VAE_HF_ID} " + "still exposes a Diffusers-format vae/ folder.", ) + return VAELoader().load(vae_dir, fastvideo_args) + + def initialize_pipeline(self, fastvideo_args: FastVideoArgs) -> None: + # MagiHuman applies `flow_shift` during timestep setup; keep the + # scheduler constructor at its default no-op shift. + self.modules["scheduler"] = FlowUniPCMultistepScheduler() + + def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None: + self._add_input_and_conditioning_stages(fastvideo_args) + self._add_base_latent_and_denoising_stages(fastvideo_args) + self._add_decode_stages() + + def _add_input_and_conditioning_stages(self, fastvideo_args: FastVideoArgs) -> None: + self.add_stage( + stage_name="input_validation_stage", + stage=InputValidationStage(), + ) + + self.add_stage( + stage_name="prompt_encoding_stage", + stage=TextEncodingStage( + text_encoders=[self.get_module("text_encoder")], + tokenizers=[self.get_module("tokenizer")], + ), + ) + + self._add_reference_image_stage(fastvideo_args) + + def _add_base_latent_and_denoising_stages(self, fastvideo_args: FastVideoArgs) -> None: + pc = fastvideo_args.pipeline_config + dit_arch = pc.dit_config.arch_config + + # Data-proxy + eval knobs come from the PipelineConfig (`pc`). + # Only DiT-architecture fields live on `dit_arch` now. + self.add_stage( + stage_name="latent_preparation_stage", + stage=MagiHumanLatentPreparationStage( + vae_stride=tuple(pc.vae_stride), + z_dim=pc.z_dim, + patch_size=tuple(dit_arch.patch_size), + fps=pc.fps, + t5_gemma_target_length=pc.t5_gemma_target_length, + coords_style=pc.coords_style, + text_offset=pc.text_offset, + audio_in_channels=dit_arch.audio_in_channels, + ), + ) + + self.add_stage( + stage_name="denoising_stage", + stage=MagiHumanDenoisingStage( + transformer=self.get_module("transformer"), + scheduler=self.get_module("scheduler"), + patch_size=tuple(dit_arch.patch_size), + video_in_channels=dit_arch.video_in_channels, + audio_in_channels=dit_arch.audio_in_channels, + video_txt_guidance_scale=pc.video_txt_guidance_scale, + audio_txt_guidance_scale=pc.audio_txt_guidance_scale, + cfg_number=pc.cfg_number, + coords_style=pc.coords_style, + video_guidance_high_t_threshold=pc.video_guidance_high_t_threshold, + video_guidance_low_t_value=pc.video_guidance_low_t_value, + ), + ) + + def _add_decode_stages(self) -> None: + self.add_stage( + stage_name="decoding_stage", + stage=DecodingStage(vae=self.get_module("vae"), pipeline=self), + ) + + self.add_stage( + stage_name="audio_decoding_stage", + stage=MagiHumanAudioDecodingStage(audio_vae=self.get_module("audio_vae"), ), + ) + + def _add_reference_image_stage(self, fastvideo_args: FastVideoArgs) -> None: + return + + +class MagiHumanI2VPipeline(MagiHumanPipeline): + """MagiHuman text+image-to-AV pipeline using the T2V DiT weights.""" + + def _add_reference_image_stage(self, fastvideo_args: FastVideoArgs) -> None: + pc = fastvideo_args.pipeline_config + self.add_stage( + stage_name="reference_image_stage", + stage=MagiHumanReferenceImageStage( + vae=self.get_module("vae"), + vae_scale_factor=pc.vae_stride[1], + ), + ) + + +class MagiHumanSRPipeline(MagiHumanPipeline): + """Two-stage MagiHuman base + SR-540p text-to-AV pipeline.""" + + _required_config_modules = [ + "text_encoder", + "tokenizer", + "vae", + "transformer", + "sr_transformer", + "scheduler", + "audio_vae", + ] + + def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None: + self._add_input_and_conditioning_stages(fastvideo_args) + self._add_base_latent_and_denoising_stages(fastvideo_args) + self._add_sr_latent_and_denoising_stages(fastvideo_args) + self._add_decode_stages() + + def _add_sr_latent_and_denoising_stages(self, fastvideo_args: FastVideoArgs) -> None: + pc = fastvideo_args.pipeline_config + dit_arch = pc.dit_config.arch_config + sr_transformer = self.get_module("sr_transformer") + sr_local_attn_layers = tuple(getattr(pc, "sr_local_attn_layers", ())) + if sr_local_attn_layers and hasattr(sr_transformer, "configure_local_attention"): + sr_transformer.configure_local_attention( + sr_local_attn_layers, + frame_receptive_field=pc.frame_receptive_field, + ) + + self.add_stage( + stage_name="sr_latent_preparation_stage", + stage=MagiHumanSRLatentPreparationStage( + vae=self.get_module("vae"), + vae_stride=tuple(pc.vae_stride), + patch_size=tuple(dit_arch.patch_size), + noise_value=pc.noise_value, + sr_audio_noise_scale=pc.sr_audio_noise_scale, + sr_height=pc.sr_height, + sr_width=pc.sr_width, + vae_scale_factor=pc.vae_stride[1], + ), + ) + self.add_stage( + stage_name="sr_denoising_stage", + stage=MagiHumanSRDenoisingStage( + transformer=sr_transformer, + scheduler=self.get_module("scheduler"), + patch_size=tuple(dit_arch.patch_size), + video_in_channels=dit_arch.video_in_channels, + audio_in_channels=dit_arch.audio_in_channels, + sr_num_inference_steps=pc.sr_num_inference_steps, + sr_video_txt_guidance_scale=pc.sr_video_txt_guidance_scale, + use_cfg_trick=pc.use_cfg_trick, + cfg_trick_start_frame=pc.cfg_trick_start_frame, + cfg_trick_value=pc.cfg_trick_value, + cfg_number=pc.cfg_number, + coords_style="v1", + ), + ) + + +class MagiHumanSRI2VPipeline(MagiHumanSRPipeline): + """Two-stage MagiHuman base + SR-540p text+image-to-AV pipeline.""" + + def _add_reference_image_stage(self, fastvideo_args: FastVideoArgs) -> None: + pc = fastvideo_args.pipeline_config + self.add_stage( + stage_name="reference_image_stage", + stage=MagiHumanReferenceImageStage( + vae=self.get_module("vae"), + vae_scale_factor=pc.vae_stride[1], + ), + ) + + +class MagiHumanSR1080pPipeline(MagiHumanSRPipeline): + """Two-stage MagiHuman base + SR-1080p text-to-AV pipeline. + + The stage chain is identical to SR-540p. The paired pipeline config enables + block-sparse local-window attention on 32 SR-DiT layers and requests the + 1080p latent target. + """ + + +class MagiHumanSR1080pI2VPipeline(MagiHumanSRI2VPipeline): + """Two-stage MagiHuman base + SR-1080p text+image-to-AV pipeline.""" + + +EntryClass = [ + MagiHumanPipeline, + MagiHumanI2VPipeline, + MagiHumanSRPipeline, + MagiHumanSRI2VPipeline, + MagiHumanSR1080pPipeline, + MagiHumanSR1080pI2VPipeline, +] diff --git a/fastvideo/pipelines/basic/magi_human/pipeline_configs.py b/fastvideo/pipelines/basic/magi_human/pipeline_configs.py new file mode 100644 index 000000000..03feb79e8 --- /dev/null +++ b/fastvideo/pipelines/basic/magi_human/pipeline_configs.py @@ -0,0 +1,236 @@ +# SPDX-License-Identifier: Apache-2.0 +"""PipelineConfig for the daVinci-MagiHuman base text-to-AV pipeline.""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass, field + +import torch + +from fastvideo.configs.models import DiTConfig, EncoderConfig, VAEConfig +from fastvideo.configs.models.dits import MagiHumanVideoConfig +from fastvideo.configs.models.encoders import ( + BaseEncoderOutput, + T5GemmaEncoderConfig, +) +from fastvideo.configs.models.vaes import OobleckVAEConfig, WanVAEConfig +from fastvideo.configs.pipelines.base import PipelineConfig + + +def t5gemma_postprocess_text(outputs: BaseEncoderOutput) -> torch.Tensor: + """Return per-prompt last_hidden_state as a batched [B, L, D] tensor. + + MagiHuman pads/trims the embedding to a fixed length in its own + `pad_or_trim` helper at pipeline time. Here we simply hand through + whatever the tokenizer produced — the latent-prep stage is responsible + for pad/trim so that the original context length can be preserved. + """ + hidden = outputs.last_hidden_state + assert torch.isnan(hidden).sum() == 0 + # Keep the shape the tokenizer emitted; the pipeline stage handles + # pad-or-trim to t5_gemma_target_length=640. + return hidden + + +@dataclass +class MagiHumanBaseConfig(PipelineConfig): + """Base MagiHuman text-to-AV pipeline config (prompt → video + audio). + + MagiHuman's base model is a joint audio-visual generator. This config + wires up both the video VAE (Wan 2.2 TI2V-5B) and the audio VAE + (Stable Audio Open 1.0); the pipeline produces an mp4 with a muxed + audio track. The framework's `WorkloadType` enum has no `T2AV` + variant yet, so the registry entry uses `WorkloadType.T2V` as a + placeholder. + """ + + # DiT + dit_config: DiTConfig = field(default_factory=MagiHumanVideoConfig) + # VAE — Wan 2.2 TI2V-5B. Diffusers `vae/config.json` drives arch_config + # at load time, including z_dim=48 and scale_factor_temporal=4 / + # scale_factor_spatial=16. + vae_config: VAEConfig = field(default_factory=WanVAEConfig) + vae_tiling: bool = False + vae_sp: bool = False + + # Audio VAE — Stable Audio Open 1.0 (Oobleck), shared with the + # standalone Stable Audio pipeline. Lazy-loaded from + # `stabilityai/stable-audio-open-1.0` (HF gated, Apache 2.0). + audio_vae_config: VAEConfig = field(default_factory=OobleckVAEConfig) + + # Denoising (flow-matching UniPC). + flow_shift: float | None = 5.0 + + # Text encoding + text_encoder_configs: tuple[EncoderConfig, ...] = field(default_factory=lambda: (T5GemmaEncoderConfig(), )) + postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.Tensor], + ...] = field(default_factory=lambda: (t5gemma_postprocess_text, )) + + # Precisions — the DiT runs bf16 internally, the text encoder is + # bf16-native, and the VAE decode path benefits from fp32 for long + # sequences. + precision: str = "bf16" + vae_precision: str = "fp32" + text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("bf16", )) + + # MagiHuman-specific defaults surfaced for the pipeline stages. These + # are pipeline-level knobs sourced from the upstream + # `EvaluationConfig` / `DataProxyConfig` (not `ModelConfig`), so they + # belong here, NOT on `MagiHumanArchConfig`. + t5_gemma_target_length: int = 640 + fps: int = 25 + num_inference_steps: int = 32 + video_txt_guidance_scale: float = 5.0 + audio_txt_guidance_scale: float = 5.0 + cfg_number: int = 2 + + # VAE / data-proxy knobs (were on ArchConfig before; moved here). + vae_stride: tuple[int, int, int] = (4, 16, 16) + z_dim: int = 48 + frame_receptive_field: int = 11 + coords_style: str = "v2" + ref_audio_offset: int = 1000 + text_offset: int = 0 + + # Video CFG step-dependent guidance: low-t steps use a relaxed scale. + # Upstream daVinci-MagiHuman/inference/pipeline/video_generate.py:426 + # uses 5.0 for high-t and 2.0 for low-t with cutoff at t=500. + video_guidance_high_t_threshold: int = 500 + video_guidance_low_t_value: float = 2.0 + + def __post_init__(self) -> None: + # Base text-to-AV does not need the VAE encoder (no reference-image + # conditioning). Keep decoder only to save memory. + self.vae_config.load_encoder = False + self.vae_config.load_decoder = True + + +@dataclass +class MagiHumanBaseI2VConfig(MagiHumanBaseConfig): + """Base MagiHuman text+image-to-AV pipeline config. + + TI2V reuses the T2V DiT weights; the only pipeline-side difference is + that a reference image is encoded with the Wan VAE and reinserted into + the first video-latent frame before every denoise step. + """ + + image_conditioning: bool = True + + def __post_init__(self) -> None: + self.vae_config.load_encoder = True + self.vae_config.load_decoder = True + + +@dataclass +class MagiHumanDistillConfig(MagiHumanBaseConfig): + """DMD-2 distilled MagiHuman text-to-AV pipeline config. + + Same arch as base (identical 331 keys, same shapes, same module tree), + but trained via DMD-2 for 8-step inference without classifier-free + guidance. Weights are stored in fp32 upstream; the conversion script's + `--cast-bf16` flag reduces the checkpoint to ~30 GB on disk. + """ + + num_inference_steps: int = 8 + cfg_number: int = 1 # DMD distilled models skip CFG. + # Lower flow_shift matches the distilled DMD schedule; if parity later + # shows drift, measure against `scheduler_config.json` generated by the + # conversion script for the distill subfolder. + flow_shift: float | None = 5.0 + + +@dataclass +class MagiHumanDistillI2VConfig(MagiHumanDistillConfig): + """DMD-2 distilled MagiHuman text+image-to-AV pipeline config.""" + + image_conditioning: bool = True + + def __post_init__(self) -> None: + self.vae_config.load_encoder = True + self.vae_config.load_decoder = True + + +@dataclass +class MagiHumanSR540pConfig(MagiHumanBaseConfig): + """Two-stage MagiHuman base + SR-540p text-to-AV pipeline config.""" + + noise_value: int = 220 + sr_audio_noise_scale: float = 0.7 + sr_num_inference_steps: int = 5 + sr_video_txt_guidance_scale: float = 3.5 + use_cfg_trick: bool = True + cfg_trick_start_frame: int = 13 + cfg_trick_value: float = 2.0 + # Upstream example/sr_540p uses sr_height=512, sr_width=896. Despite the + # marketing name, these are the VAE/patch-aligned dimensions actually run. + sr_height: int = 512 + sr_width: int = 896 + sr_local_attn_layers: tuple[int, ...] = () + + +@dataclass +class MagiHumanSR540pI2VConfig(MagiHumanSR540pConfig): + """Two-stage MagiHuman base + SR-540p text+image-to-AV config.""" + + image_conditioning: bool = True + + def __post_init__(self) -> None: + self.vae_config.load_encoder = True + self.vae_config.load_decoder = True + + +_SR_1080P_LOCAL_ATTN_LAYERS: tuple[int, ...] = ( + 0, + 1, + 2, + 4, + 5, + 6, + 8, + 9, + 10, + 12, + 13, + 14, + 16, + 17, + 18, + 20, + 21, + 22, + 24, + 25, + 26, + 28, + 29, + 30, + 32, + 33, + 34, + 35, + 36, + 37, + 38, + 39, +) + + +@dataclass +class MagiHumanSR1080pConfig(MagiHumanSR540pConfig): + """Two-stage MagiHuman base + SR-1080p text-to-AV pipeline config.""" + + sr_height: int = 1080 + sr_width: int = 1920 + sr_local_attn_layers: tuple[int, ...] = _SR_1080P_LOCAL_ATTN_LAYERS + + +@dataclass +class MagiHumanSR1080pI2VConfig(MagiHumanSR1080pConfig): + """Two-stage MagiHuman base + SR-1080p text+image-to-AV config.""" + + image_conditioning: bool = True + + def __post_init__(self) -> None: + self.vae_config.load_encoder = True + self.vae_config.load_decoder = True diff --git a/fastvideo/pipelines/basic/magi_human/presets.py b/fastvideo/pipelines/basic/magi_human/presets.py new file mode 100644 index 000000000..f28b51fc7 --- /dev/null +++ b/fastvideo/pipelines/basic/magi_human/presets.py @@ -0,0 +1,225 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Presets for the daVinci-MagiHuman pipelines.""" +from fastvideo.api.presets import InferencePreset, PresetStageSpec + +# Keep this in sync with upstream MagiEvaluator.negative_prompt +# (daVinci-MagiHuman/inference/pipeline/video_generate.py:222-224): the +# video, audio-quality, and speech-delivery blocks all condition CFG. +_MAGI_HUMAN_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") + +_DENOISE_STAGE = PresetStageSpec( + name="denoise", + kind="denoising", + description="Joint video+audio UniPC flow-matching denoise pass.", + allowed_overrides=frozenset({ + "num_inference_steps", + "guidance_scale", + }), +) + +MAGI_HUMAN_BASE = InferencePreset( + name="magi_human_base", + version=1, + model_family="magi_human", + description=("daVinci-MagiHuman base text-to-AV at 256x480, 4s @ 25 fps. " + "Produces an mp4 with muxed audio + video. workload_type " + "is `t2v` because the framework enum has no `t2av` variant yet."), + workload_type="t2v", + stage_schemas=(_DENOISE_STAGE, ), + defaults={ + "seed": 42, + "height": 256, + # Upstream pipeline.py:61-64 defaults br_width=480, br_height=272, + # and video_generate.py:254-261 snaps height to 256 while width stays + # 480, so the rendered default is 256x480. + "width": 480, + # num_frames is derived by the pipeline as `seconds*fps + 1`; we + # surface it here for APIs that expect a concrete default. + "num_frames": 101, + "fps": 25, + "guidance_scale": 5.0, # used as video_txt_guidance_scale + "num_inference_steps": 32, + "negative_prompt": _MAGI_HUMAN_NEGATIVE_PROMPT, + }, +) + +MAGI_HUMAN_DISTILL = InferencePreset( + name="magi_human_distill", + version=1, + model_family="magi_human", + description=("daVinci-MagiHuman DMD-2 distilled text-to-AV at 256x480, 4s @ " + "25 fps. 8-step inference, no classifier-free guidance. Produces " + "an mp4 with muxed audio + video."), + workload_type="t2v", + stage_schemas=(_DENOISE_STAGE, ), + defaults={ + "seed": 42, + "height": 256, + "width": 480, + "num_frames": 101, + "fps": 25, + # DMD: cfg=1 at the pipeline level. guidance_scale is kept at 1.0 + # for interop; the DenoisingStage ignores it when cfg_number=1. + "guidance_scale": 1.0, + "num_inference_steps": 8, + "negative_prompt": _MAGI_HUMAN_NEGATIVE_PROMPT, + }, +) + +MAGI_HUMAN_BASE_TI2V = InferencePreset( + name="magi_human_base_ti2v", + version=1, + model_family="magi_human", + description=("daVinci-MagiHuman base text+image-to-AV at 256x480, 4s @ 25 fps. " + "The reference image is VAE-encoded and pinned to the first " + "video latent frame at each denoise step."), + workload_type="i2v", + stage_schemas=(_DENOISE_STAGE, ), + defaults={ + "seed": 42, + "height": 256, + "width": 480, + "num_frames": 101, + "fps": 25, + "guidance_scale": 5.0, + "num_inference_steps": 32, + "negative_prompt": _MAGI_HUMAN_NEGATIVE_PROMPT, + }, +) + +MAGI_HUMAN_DISTILL_TI2V = InferencePreset( + name="magi_human_distill_ti2v", + version=1, + model_family="magi_human", + description=("daVinci-MagiHuman DMD-2 distilled text+image-to-AV at 256x480, " + "4s @ 25 fps. 8-step inference, no CFG."), + workload_type="i2v", + stage_schemas=(_DENOISE_STAGE, ), + defaults={ + "seed": 42, + "height": 256, + "width": 480, + "num_frames": 101, + "fps": 25, + "guidance_scale": 1.0, + "num_inference_steps": 8, + "negative_prompt": _MAGI_HUMAN_NEGATIVE_PROMPT, + }, +) + +MAGI_HUMAN_SR_540P = InferencePreset( + name="magi_human_sr_540p", + version=1, + model_family="magi_human", + description=("daVinci-MagiHuman two-stage base + SR-540p text-to-AV. " + "Base pass runs at 256x480; SR pass refines to upstream's " + "aligned 512x896 output with muxed audio."), + workload_type="t2v", + stage_schemas=(_DENOISE_STAGE, ), + defaults={ + "seed": 42, + "height": 256, + "width": 480, + "num_frames": 101, + "fps": 25, + "guidance_scale": 5.0, + "num_inference_steps": 32, + "negative_prompt": _MAGI_HUMAN_NEGATIVE_PROMPT, + }, +) + +MAGI_HUMAN_SR_540P_TI2V = InferencePreset( + name="magi_human_sr_540p_ti2v", + version=1, + model_family="magi_human", + description=("daVinci-MagiHuman two-stage base + SR-540p text+image-to-AV. " + "The reference image is encoded at base resolution and then " + "re-encoded at SR resolution before the SR denoise pass."), + workload_type="i2v", + stage_schemas=(_DENOISE_STAGE, ), + defaults={ + "seed": 42, + "height": 256, + "width": 480, + "num_frames": 101, + "fps": 25, + "guidance_scale": 5.0, + "num_inference_steps": 32, + "negative_prompt": _MAGI_HUMAN_NEGATIVE_PROMPT, + }, +) + +MAGI_HUMAN_SR_1080P = InferencePreset( + name="magi_human_sr_1080p", + version=1, + model_family="magi_human", + description=("daVinci-MagiHuman two-stage base + SR-1080p text-to-AV. " + "The SR DiT uses upstream local-window attention in 32 of " + "40 layers and refines to 1080p-class output."), + workload_type="t2v", + stage_schemas=(_DENOISE_STAGE, ), + defaults={ + "seed": 42, + "height": 256, + "width": 480, + "num_frames": 101, + "fps": 25, + "guidance_scale": 5.0, + "num_inference_steps": 32, + "negative_prompt": _MAGI_HUMAN_NEGATIVE_PROMPT, + }, +) + +MAGI_HUMAN_SR_1080P_TI2V = InferencePreset( + name="magi_human_sr_1080p_ti2v", + version=1, + model_family="magi_human", + description=("daVinci-MagiHuman two-stage base + SR-1080p text+image-to-AV. " + "The SR DiT uses upstream local-window attention in 32 of " + "40 layers; the reference image is re-encoded at SR resolution."), + workload_type="i2v", + stage_schemas=(_DENOISE_STAGE, ), + defaults={ + "seed": 42, + "height": 256, + "width": 480, + "num_frames": 101, + "fps": 25, + "guidance_scale": 5.0, + "num_inference_steps": 32, + "negative_prompt": _MAGI_HUMAN_NEGATIVE_PROMPT, + }, +) + +ALL_PRESETS = ( + MAGI_HUMAN_BASE, + MAGI_HUMAN_DISTILL, + MAGI_HUMAN_BASE_TI2V, + MAGI_HUMAN_DISTILL_TI2V, + MAGI_HUMAN_SR_540P, + MAGI_HUMAN_SR_540P_TI2V, + MAGI_HUMAN_SR_1080P, + MAGI_HUMAN_SR_1080P_TI2V, +) diff --git a/tests/local_tests/magi_human/test_magi_human_pipeline_parity.py b/tests/local_tests/magi_human/test_magi_human_pipeline_parity.py new file mode 100644 index 000000000..4b47e6dc4 --- /dev/null +++ b/tests/local_tests/magi_human/test_magi_human_pipeline_parity.py @@ -0,0 +1,532 @@ +# SPDX-License-Identifier: Apache-2.0 +"""End-to-end latent parity test for the daVinci-MagiHuman base text-to-AV pipeline. + +Runs the joint video+audio FlowUniPC denoise loop with CFG=2 on both: + - FastVideo MagiHumanDiT loaded from `converted_weights/magi_human_base/transformer/`. + - Upstream daVinci-MagiHuman DiTModel loaded from the HF `base/` shards, + via the `magi_compiler` / distributed stubs in + `tests/local_tests/helpers/magi_human_upstream.py`. + +Both sides use the **same** `FlowUniPCMultistepScheduler` (FastVideo's +implementation), identical latent / text inputs, identical scheduler +state, and SDPA-routed attention — so drift here is purely the +compound of per-call DiT parity drift through the denoise loop + CFG +mixing amplification. + +What this catches (that the component-level DiT parity does NOT): + - Scheduler integration mistakes (state leaks between video/audio + schedulers, wrong shift, wrong `step()` args). + - CFG math errors (guidance scale switchover at t=500, per-modality + guidance scale wiring, unconditional-path text padding). + - Latent-preparation / token-unpacking drift between my + `build_packed_inputs` / `unpack_tokens` and the upstream + `MagiDataProxy` equivalents. + - Compounding behavior: 1% per-call DiT drift compounding through + `num_steps * cfg_number` calls. + +Skips when: + - `daVinci-MagiHuman/` clone or GAIR/daVinci-MagiHuman base shards + are not available locally. + - Converted transformer weights are missing (run the conversion + script first). + - CUDA is unavailable. + +Tolerance: `atol=0.35, rtol=0.05` on bf16 denoise-loop latents. The +atol absorbs the observed worst-element drift (~0.31 on a signal of +abs_mean ~2.4 — bf16 + CFG amplification + UniPC accumulation). The +tight rtol still flags gross structural bugs (sign flip, scheduler +state leak, modality branch drop). If tighter parity is wanted, +chase the per-call drift first (see the DiT component parity test). +""" +from __future__ import annotations + +import gc +import glob +import os +from pathlib import Path + +import pytest +import torch +import torch.nn.functional as F +from torch.testing import assert_close + + +# Force SDPA on both sides so the attention kernel is shared. +os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA") +os.environ.setdefault("MASTER_ADDR", "localhost") +os.environ.setdefault("MASTER_PORT", "29519") + +_T5GEMMA_ID = os.getenv("MAGI_HUMAN_T5GEMMA_ID", "google/t5gemma-9b-9b-ul2") +_T5_GEMMA_TARGET_LENGTH = 640 +_SAMPLE_PROMPT = ( + "A warm afternoon scene: a person sits on a park bench reading a book, " + "surrounded by softly swaying trees." +) + + +def _hf_token() -> str | None: + for key in ("HF_TOKEN", "HUGGINGFACE_HUB_TOKEN", "HF_API_KEY"): + token = os.environ.get(key) + if token: + return token + return None + + +def _can_access_t5gemma() -> bool: + token = _hf_token() + if token is None: + return False + try: + from huggingface_hub import hf_hub_download + hf_hub_download( + repo_id=_T5GEMMA_ID, + filename="config.json", + token=token, + ) + return True + except Exception: + return False + + +def _pad_or_trim_dim1(t: torch.Tensor, target: int) -> tuple[torch.Tensor, int]: + """Mirror MagiHumanLatentPreparationStage's text pad-or-trim.""" + current = t.size(1) + if current < target: + pad = [0, 0, 0, target - current] + return F.pad(t, pad, "constant", 0.0), current + return t[:, :target], target + + +def _encode_magi_human_prompt_pair(device: torch.device): + """Encode the production preset prompt pair once via T5-Gemma.""" + if not _can_access_t5gemma(): + pytest.skip( + f"{_T5GEMMA_ID} not accessible — gated Google repo; set " + "HF_TOKEN / HF_API_KEY and accept the terms of use." + ) + + for src in ("HF_TOKEN", "HUGGINGFACE_HUB_TOKEN", "HF_API_KEY"): + token = os.environ.get(src) + if token: + os.environ.setdefault("HF_TOKEN", token) + os.environ.setdefault("HUGGINGFACE_HUB_TOKEN", token) + break + + try: + from transformers import AutoTokenizer + + from fastvideo.configs.models.encoders.t5gemma import ( + T5GemmaEncoderConfig, + ) + from fastvideo.models.encoders.t5gemma import ( + T5GemmaEncoderModel, + ) + from fastvideo.pipelines.basic.magi_human.presets import ( + _MAGI_HUMAN_NEGATIVE_PROMPT, + ) + except Exception as exc: + pytest.skip(f"T5-Gemma prompt encoding dependencies unavailable: {exc}") + + tokenizer = AutoTokenizer.from_pretrained(_T5GEMMA_ID) + enc_config = T5GemmaEncoderConfig() + enc_config.arch_config.t5gemma_model_path = _T5GEMMA_ID + encoder = T5GemmaEncoderModel(enc_config) + + def encode(text: str, text_encoder=encoder) -> tuple[torch.Tensor, int]: + inputs = tokenizer( + [text], + return_tensors="pt", + padding=True, + truncation=False, + ).to(device) + with torch.inference_mode(): + hidden = text_encoder( + input_ids=inputs["input_ids"], + attention_mask=inputs.get("attention_mask"), + ).last_hidden_state + return _pad_or_trim_dim1(hidden.to(torch.float32), _T5_GEMMA_TARGET_LENGTH) + + txt_feat, txt_feat_len = encode(_SAMPLE_PROMPT) + neg_txt_feat, neg_txt_feat_len = encode(_MAGI_HUMAN_NEGATIVE_PROMPT) + del encoder + _cleanup_gpu() + return txt_feat, txt_feat_len, neg_txt_feat, neg_txt_feat_len + + +def _find_base_shard_dir() -> Path | None: + """Return the local path to GAIR/daVinci-MagiHuman/base/ with shards present, or None.""" + override = os.getenv("MAGI_HUMAN_BASE_SHARD_DIR") + if override: + p = Path(override) + return p if p.is_dir() else None + try: + from huggingface_hub import snapshot_download + snap = snapshot_download( + repo_id="GAIR/daVinci-MagiHuman", + allow_patterns=[ + "base/*.safetensors", + "base/model.safetensors.index.json", + ], + ) + candidate = Path(snap) / "base" + if candidate.is_dir() and any(candidate.glob("*.safetensors")): + return candidate + return None + except Exception: + return None + + +def _cleanup_gpu() -> None: + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + +def _dit_forward_fv( + dit, video_latent, audio_latent, audio_feat_len, + txt_feat, txt_feat_len, patch_size, coords_style, + video_in_channels, audio_in_channels, +): + """One FastVideo DiT call — same as MagiHumanDenoisingStage._dit_forward.""" + from fastvideo.pipelines.basic.magi_human.stages.latent_preparation import ( + build_packed_inputs, unpack_tokens, + ) + x, coords, mm = build_packed_inputs( + video_latent=video_latent, audio_latent=audio_latent, + audio_feat_len=audio_feat_len, txt_feat=txt_feat, + txt_feat_len=txt_feat_len, patch_size=patch_size, + coords_style=coords_style, + ) + video_token_num = x.shape[0] - audio_feat_len - txt_feat_len + out = dit(x, coords, mm) + return unpack_tokens( + out, video_token_num=video_token_num, + audio_feat_len=audio_feat_len, + video_in_channels=video_in_channels, + audio_in_channels=audio_in_channels, + latent_shape=tuple(video_latent.shape), + patch_size=patch_size, + ) + + +def _dit_forward_upstream( + dit, video_latent, audio_latent, audio_feat_len, + txt_feat, txt_feat_len, patch_size, coords_style, + video_in_channels, audio_in_channels, +): + """One upstream DiT call — identical input construction + output + unpacking to the FastVideo path. The only thing that differs is + the DiT module and the extra `varlen_handler` / `local_attn_handler` + kwargs the upstream expects. + """ + from fastvideo.pipelines.basic.magi_human.stages.latent_preparation import ( + build_packed_inputs, unpack_tokens, + ) + from inference.common import VarlenHandler + x, coords, mm = build_packed_inputs( + video_latent=video_latent, audio_latent=audio_latent, + audio_feat_len=audio_feat_len, txt_feat=txt_feat, + txt_feat_len=txt_feat_len, patch_size=patch_size, + coords_style=coords_style, + ) + video_token_num = x.shape[0] - audio_feat_len - txt_feat_len + total = x.shape[0] + cu = torch.tensor([0, total], dtype=torch.int32, device=x.device) + varlen = VarlenHandler( + cu_seqlens_q=cu, cu_seqlens_k=cu, + max_seqlen_q=total, max_seqlen_k=total, + ) + out = dit( + x=x, coords_mapping=coords, modality_mapping=mm, + varlen_handler=varlen, local_attn_handler=None, + ) + return unpack_tokens( + out, video_token_num=video_token_num, + audio_feat_len=audio_feat_len, + video_in_channels=video_in_channels, + audio_in_channels=audio_in_channels, + latent_shape=tuple(video_latent.shape), + patch_size=patch_size, + ) + + +def _build_fastvideo_schedulers(shift: float, num_inference_steps: int, device): + """Mirror current FastVideo production at `magi_human_pipeline.py:146-149` + and `denoising.py:105-116`: default scheduler constructor (`shift=1`, + no-op) followed by `set_timesteps(..., shift=shift)` so the temporal + shift is applied exactly once. The earlier double-shift pattern was + reverted with the Wave 11 single-shift fix; if both __init__ and + set_timesteps applied non-trivial shift, the schedule would diverge + from upstream. + """ + from fastvideo.models.schedulers.scheduling_flow_unipc_multistep import ( + FlowUniPCMultistepScheduler, + ) + video_sched = FlowUniPCMultistepScheduler() + audio_sched = FlowUniPCMultistepScheduler() + video_sched.set_timesteps(num_inference_steps, device=device, shift=shift) + audio_sched.set_timesteps(num_inference_steps, device=device, shift=shift) + return video_sched, audio_sched + + +def _build_upstream_schedulers(shift: float, num_inference_steps: int, device): + """Construct schedulers the way the official `MagiEvaluator.eval_with_text` + does (`daVinci-MagiHuman/inference/pipeline/video_generate.py:404-407`): + `FlowUniPCMultistepScheduler()` with default shift=1.0 in __init__ + (no-op), then `set_timesteps(num_inference_steps, device, shift=self.shift)` + applies shift exactly once. Uses FastVideo's scheduler class for + the orchestration (algorithmically identical to the upstream copy + of the same Diffusers-derived class) but matches the upstream's + *call pattern*. + """ + from fastvideo.models.schedulers.scheduling_flow_unipc_multistep import ( + FlowUniPCMultistepScheduler, + ) + video_sched = FlowUniPCMultistepScheduler() + audio_sched = FlowUniPCMultistepScheduler() + video_sched.set_timesteps(num_inference_steps, device=device, shift=shift) + audio_sched.set_timesteps(num_inference_steps, device=device, shift=shift) + return video_sched, audio_sched + + +def _run_denoise_loop( + dit, dit_forward_fn, video_latent, audio_latent, + txt_feat, txt_feat_len, neg_txt_feat, neg_txt_feat_len, + *, video_sched, audio_sched, cfg_number, + video_txt_guidance_scale, audio_txt_guidance_scale, + patch_size, coords_style, video_in_channels, audio_in_channels, + image_latent=None, +): + """Joint video+audio FlowUniPC denoise. The schedulers are passed + in pre-constructed so each side can mirror its production scheduler + init pattern (see `_build_*_schedulers`). + """ + from fastvideo.forward_context import set_forward_context + audio_feat_len = int(audio_latent.shape[1]) + + with torch.inference_mode(): + for idx, t in enumerate(video_sched.timesteps): + if image_latent is not None: + video_latent[:, :, :1] = image_latent.to( + device=video_latent.device, + dtype=video_latent.dtype, + )[:, :, :1] + t_int = int(t.item()) if torch.is_tensor(t) else int(t) + with set_forward_context(current_timestep=t_int, attn_metadata=None): + v_cond_video, v_cond_audio = dit_forward_fn( + dit, video_latent, audio_latent, audio_feat_len, + txt_feat, txt_feat_len, patch_size, coords_style, + video_in_channels, audio_in_channels, + ) + if cfg_number == 2: + v_uncond_video, v_uncond_audio = dit_forward_fn( + dit, video_latent, audio_latent, audio_feat_len, + neg_txt_feat, neg_txt_feat_len, patch_size, coords_style, + video_in_channels, audio_in_channels, + ) + # Upstream's video-guidance drop-at-t<=500 trick. + video_guidance = ( + video_txt_guidance_scale if t > 500 else 2.0 + ) + v_video = v_uncond_video + video_guidance * ( + v_cond_video - v_uncond_video + ) + v_audio = v_uncond_audio + audio_txt_guidance_scale * ( + v_cond_audio - v_uncond_audio + ) + else: + v_video = v_cond_video + v_audio = v_cond_audio + + video_latent = video_sched.step( + v_video, t, video_latent, return_dict=False, + )[0] + audio_latent = audio_sched.step( + v_audio, t, audio_latent, return_dict=False, + )[0] + if image_latent is not None: + video_latent[:, :, :1] = image_latent.to( + device=video_latent.device, + dtype=video_latent.dtype, + )[:, :, :1] + return video_latent, audio_latent + + +@pytest.mark.skipif( + not torch.cuda.is_available(), + reason="MagiHuman pipeline parity requires CUDA.", +) +def test_magi_human_pipeline_latent_parity(): + repo_root = Path(__file__).resolve().parents[3] + upstream_src = repo_root / "daVinci-MagiHuman" + if not upstream_src.exists(): + pytest.skip( + "Upstream daVinci-MagiHuman/ clone missing. Run " + "`git clone --depth 1 https://github.com/GAIR-NLP/daVinci-MagiHuman.git`" + ) + + base_shard_dir = _find_base_shard_dir() + if base_shard_dir is None or not base_shard_dir.is_dir(): + pytest.skip( + "GAIR/daVinci-MagiHuman base/ shards not available locally." + ) + + converted_dir = Path(os.getenv( + "MAGI_HUMAN_DIFFUSERS_PATH", + repo_root / "converted_weights" / "magi_human_base", + )) + transformer_dir = converted_dir / "transformer" + if not transformer_dir.is_dir(): + pytest.skip(f"Converted transformer dir missing at {transformer_dir}") + + from tests.local_tests.helpers.magi_human_upstream import ( + install_stubs, load_upstream_dit, + ) + install_stubs() + + # --- Shared pipeline inputs --- + device = torch.device("cuda:0") + torch.manual_seed(0) + + # Deliberately tiny so 2 * CFG=2 = 4 DiT calls per side fit in + # CI/dev runtime budget. + z_dim = 48 + patch_size = (1, 2, 2) + lat_T, lat_H, lat_W = 2, 6, 6 + video_latent = torch.randn( + (1, z_dim, lat_T, lat_H, lat_W), + dtype=torch.float32, device=device, + ) + audio_latent = torch.randn( + (1, 4, 64), dtype=torch.float32, device=device, + ) + # Production-facing text embeddings: encode the example prompt and the + # preset negative prompt via T5-Gemma once, then feed the identical cached + # tensors to upstream and FastVideo. This keeps the DiT comparison focused + # while still validating prompt/preset content such as the full + # three-block MagiHuman negative prompt. + txt_feat, txt_feat_len, neg_txt_feat, neg_txt_feat_len = ( + _encode_magi_human_prompt_pair(device) + ) + + num_inference_steps = 4 # 4 steps × CFG=2 = 8 DiT calls / side; surfaces compounding drift that 1-step hides + shift = 5.0 + common_kwargs = dict( + cfg_number=2, + video_txt_guidance_scale=5.0, + audio_txt_guidance_scale=5.0, + patch_size=patch_size, + coords_style="v2", + video_in_channels=192, + audio_in_channels=64, + ) + + # --- Upstream side first (so we can free it before loading FastVideo). --- + # Upstream uses single-shift scheduler init (matches MagiEvaluator). + up_video_sched, up_audio_sched = _build_upstream_schedulers( + shift=shift, num_inference_steps=num_inference_steps, device=device, + ) + print("Loading upstream DiTModel from base shards...") + upstream_dit = load_upstream_dit(base_shard_dir, device=device, dtype=None) + print("Running upstream denoise loop...") + ref_video, ref_audio = _run_denoise_loop( + upstream_dit, _dit_forward_upstream, + video_latent.clone(), audio_latent.clone(), + txt_feat.clone(), txt_feat_len, + neg_txt_feat.clone(), neg_txt_feat_len, + video_sched=up_video_sched, audio_sched=up_audio_sched, + **common_kwargs, + ) + ref_video = ref_video.detach().float().cpu() + ref_audio = ref_audio.detach().float().cpu() + del upstream_dit + _cleanup_gpu() + + # --- FastVideo side --- + # FastVideo uses double-shift scheduler init (matches + # `MagiHumanDenoisingStage` in production: shift in __init__ via + # `magi_human_pipeline.initialize_pipeline` AND in set_timesteps). + fv_video_sched, fv_audio_sched = _build_fastvideo_schedulers( + shift=shift, num_inference_steps=num_inference_steps, device=device, + ) + from fastvideo.configs.models.dits.magi_human import MagiHumanVideoConfig + from fastvideo.models.dits.magi_human import MagiHumanDiT + from safetensors.torch import load_file + + print("Loading FastVideo MagiHumanDiT from converted transformer/...") + fv_cfg = MagiHumanVideoConfig() + fv_dit = MagiHumanDiT(fv_cfg) + fv_state = {} + for shard in sorted(glob.glob(str(transformer_dir / "*.safetensors"))): + fv_state.update(load_file(shard)) + missing, unexpected = fv_dit.load_state_dict(fv_state, strict=False) + assert not missing, f"FastVideo DiT missing {len(missing)} keys: {missing[:5]}" + assert not unexpected, f"FastVideo DiT unexpected {len(unexpected)} keys: {unexpected[:5]}" + fv_dit = fv_dit.to(device=device) + fv_dit.eval() + + print("Running FastVideo denoise loop...") + fv_video, fv_audio = _run_denoise_loop( + fv_dit, _dit_forward_fv, + video_latent.clone(), audio_latent.clone(), + txt_feat.clone(), txt_feat_len, + neg_txt_feat.clone(), neg_txt_feat_len, + video_sched=fv_video_sched, audio_sched=fv_audio_sched, + **common_kwargs, + ) + fv_video = fv_video.detach().float().cpu() + fv_audio = fv_audio.detach().float().cpu() + + # --- Report + assertions --- + v_diff = (ref_video - fv_video).abs() + a_diff = (ref_audio - fv_audio).abs() + print( + f"video ref_abs={ref_video.abs().mean().item():.4f} " + f"fv_abs={fv_video.abs().mean().item():.4f} " + f"diff_max={v_diff.max().item():.4f} " + f"diff_mean={v_diff.mean().item():.4f} " + f"diff_median={v_diff.median().item():.4f}" + ) + print( + f"audio ref_abs={ref_audio.abs().mean().item():.4f} " + f"fv_abs={fv_audio.abs().mean().item():.4f} " + f"diff_max={a_diff.max().item():.4f} " + f"diff_mean={a_diff.mean().item():.4f} " + f"diff_median={a_diff.median().item():.4f}" + ) + + assert ref_video.shape == fv_video.shape + assert ref_audio.shape == fv_audio.shape + + # Tolerance budget for 1-step / CFG=2 (bf16 DiT + bf16 CFG mix): + # * Single-DiT bf16 drift: diff_mean ~0.008 on `abs ~ 1.0` + # (see DiT component parity, `test_magi_human_dit_parity`). + # * CFG mixes `v = v_uncond + guidance * (v_cond - v_uncond)` + # with guidance=5; cond and uncond drift independently in bf16, + # so the post-CFG `diff_mean` scales by ~guidance (~5x). + # * One FlowUniPC scheduler step passes that through unchanged. + # `diff_max` is the noisiest statistic for bf16 transformer parity + # (a single fma quantization can blow it up). Use it only as a loose + # guard. The two ratio guards below catch real structural bugs: + # `abs_mean` drift signals scale errors / dropped branches, and + # `diff_mean / ref_abs` signals systematic per-element bias far + # beyond what bf16+CFG noise can produce. + assert_close(fv_video, ref_video, atol=0.40, rtol=0.05) + assert_close(fv_audio, ref_audio, atol=0.40, rtol=0.05) + + # Global-magnitude guard — tightest single assertion. A gross bug + # (scheduler state leak, dropped modality branch, CFG sign flip) + # would shift `abs_mean` far beyond the bf16+CFG noise floor. + ref_v_abs = ref_video.abs().mean().item() + ref_a_abs = ref_audio.abs().mean().item() + rel_v = abs(ref_v_abs - fv_video.abs().mean().item()) / max(ref_v_abs, 1e-6) + rel_a = abs(ref_a_abs - fv_audio.abs().mean().item()) / max(ref_a_abs, 1e-6) + assert rel_v < 0.01, f"video abs_mean drift {rel_v:.2%} > 1%" + assert rel_a < 0.01, f"audio abs_mean drift {rel_a:.2%} > 1%" + + # Per-element mean-bias guard — catches systematic shift that + # `abs_mean` misses (e.g. equal-magnitude flip across many elements). + mean_rel_v = v_diff.mean().item() / max(ref_v_abs, 1e-6) + mean_rel_a = a_diff.mean().item() / max(ref_a_abs, 1e-6) + assert mean_rel_v < 0.04, f"video mean_diff/ref_abs {mean_rel_v:.2%} > 4%" + assert mean_rel_a < 0.04, f"audio mean_diff/ref_abs {mean_rel_a:.2%} > 4%" diff --git a/tests/local_tests/magi_human/test_magi_human_pipeline_smoke.py b/tests/local_tests/magi_human/test_magi_human_pipeline_smoke.py new file mode 100644 index 000000000..a3027fa9f --- /dev/null +++ b/tests/local_tests/magi_human/test_magi_human_pipeline_smoke.py @@ -0,0 +1,246 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Smoke / preflight tests for the daVinci-MagiHuman base text-to-AV pipeline. + +Two tests: + + * `test_magi_human_typed_surface_preflight` — pure-Python, no GPU, no + weights. Verifies that the scaffold is importable, that the preset + registers cleanly, and that the DiT module tree matches the upstream + HuggingFace checkpoint shape-for-shape on `meta` device. This is what + CI should run on every PR. + + * `test_magi_human_pipeline_smoke` — end-to-end pipeline construction + + a tiny generate_video call, gated on local converted-weights paths. + Skips cleanly when weights are missing. +""" +from __future__ import annotations + +import json +import os +from pathlib import Path + +import pytest +import torch + + +def test_magi_human_typed_surface_preflight() -> None: + """Import + registry + module-tree surface check. + + Covers regressions that would otherwise only surface on a GPU host: + preset drop from ALL_PRESETS, renamed modules, registry mis-wiring, + or DiT module-tree drift from the upstream checkpoint. + """ + import fastvideo.registry # noqa: F401 — triggers preset registration + from fastvideo.api.presets import get_preset, get_presets_for_family + from fastvideo.configs.models.dits.magi_human import ( + MagiHumanArchConfig, + MagiHumanVideoConfig, + ) + from fastvideo.configs.models.encoders.t5gemma import ( + T5GemmaEncoderArchConfig, + T5GemmaEncoderConfig, + ) + from fastvideo.models.dits.magi_human import MagiHumanDiT + from fastvideo.pipelines.basic.magi_human.magi_human_pipeline import ( # noqa: F401 + MagiHumanI2VPipeline, + MagiHumanPipeline, + MagiHumanSRI2VPipeline, + MagiHumanSRPipeline, + ) + from fastvideo.pipelines.basic.magi_human.pipeline_configs import ( + MagiHumanBaseConfig, + MagiHumanBaseI2VConfig, + MagiHumanDistillI2VConfig, + MagiHumanSR540pConfig, + MagiHumanSR540pI2VConfig, + ) + from fastvideo.pipelines.basic.magi_human.stages import ( # noqa: F401 + MagiHumanDenoisingStage, + MagiHumanLatentPreparationStage, + MagiHumanReferenceImageStage, + MagiHumanSRDenoisingStage, + MagiHumanSRLatentPreparationStage, + ) + + # Presets are registered under the expected family. + names = {p.name for p in get_presets_for_family("magi_human")} + assert names == { + "magi_human_base", + "magi_human_distill", + "magi_human_base_ti2v", + "magi_human_distill_ti2v", + "magi_human_sr_540p", + "magi_human_sr_540p_ti2v", + "magi_human_sr_1080p", + "magi_human_sr_1080p_ti2v", + } + + base_preset = get_preset("magi_human_base", "magi_human") + assert base_preset.workload_type == "t2v" + assert base_preset.defaults["num_inference_steps"] == 32 + assert base_preset.defaults["fps"] == 25 + + distill_preset = get_preset("magi_human_distill", "magi_human") + assert distill_preset.workload_type == "t2v" + assert distill_preset.defaults["num_inference_steps"] == 8 + assert distill_preset.defaults["guidance_scale"] == 1.0 + + base_ti2v_preset = get_preset("magi_human_base_ti2v", "magi_human") + assert base_ti2v_preset.workload_type == "i2v" + assert base_ti2v_preset.defaults["num_inference_steps"] == 32 + + distill_ti2v_preset = get_preset("magi_human_distill_ti2v", "magi_human") + assert distill_ti2v_preset.workload_type == "i2v" + assert distill_ti2v_preset.defaults["num_inference_steps"] == 8 + + sr_preset = get_preset("magi_human_sr_540p", "magi_human") + assert sr_preset.workload_type == "t2v" + assert sr_preset.defaults["num_inference_steps"] == 32 + + sr_ti2v_preset = get_preset("magi_human_sr_540p_ti2v", "magi_human") + assert sr_ti2v_preset.workload_type == "i2v" + assert sr_ti2v_preset.defaults["num_inference_steps"] == 32 + + # Distill pipeline config: same arch as base, CFG=1, 8 steps. + from fastvideo.pipelines.basic.magi_human.pipeline_configs import MagiHumanDistillConfig + distill_pc = MagiHumanDistillConfig() + assert distill_pc.num_inference_steps == 8 + assert distill_pc.cfg_number == 1 + assert distill_pc.dit_config.arch_config.num_layers == 40 + + base_i2v_pc = MagiHumanBaseI2VConfig() + assert base_i2v_pc.image_conditioning is True + assert base_i2v_pc.vae_config.load_encoder is True + assert base_i2v_pc.vae_config.load_decoder is True + + distill_i2v_pc = MagiHumanDistillI2VConfig() + assert distill_i2v_pc.num_inference_steps == 8 + assert distill_i2v_pc.cfg_number == 1 + assert distill_i2v_pc.image_conditioning is True + assert distill_i2v_pc.vae_config.load_encoder is True + + sr_pc = MagiHumanSR540pConfig() + assert sr_pc.num_inference_steps == 32 + assert sr_pc.sr_num_inference_steps == 5 + assert sr_pc.noise_value == 220 + assert sr_pc.sr_audio_noise_scale == 0.7 + assert sr_pc.sr_video_txt_guidance_scale == 3.5 + assert sr_pc.sr_height == 512 + assert sr_pc.sr_width == 896 + + sr_i2v_pc = MagiHumanSR540pI2VConfig() + assert sr_i2v_pc.image_conditioning is True + assert sr_i2v_pc.vae_config.load_encoder is True + + # Config constructs with the documented defaults. + pc = MagiHumanBaseConfig() + assert pc.flow_shift == 5.0 + assert pc.cfg_number == 2 + assert pc.num_inference_steps == 32 + assert pc.dit_config.arch_config.num_layers == 40 + assert pc.dit_config.arch_config.hidden_size == 5120 + assert pc.dit_config.arch_config.num_attention_heads == 40 + assert pc.dit_config.arch_config.num_heads_kv == 8 + assert pc.dit_config.arch_config.mm_layers == (0, 1, 2, 3, 36, 37, 38, 39) + assert pc.text_encoder_configs[0].arch_config.hidden_size == 3584 + + # The DiT module tree matches the upstream HF base/ checkpoint + # shape-for-shape. This is checkpoint-loading parity, not numerical + # parity — but a regression here means loaded weights won't align. + dit_cfg = MagiHumanVideoConfig() + with torch.device("meta"): + dit = MagiHumanDiT(dit_cfg) + fv_shapes = {n: tuple(p.shape) for n, p in dit.state_dict().items()} + + index_path = _hf_index_path_or_none() + if index_path is None: + pytest.skip("HF repo unavailable (no network / no token) — " + "skipping cross-check against GAIR/daVinci-MagiHuman.") + with open(index_path) as f: + wmap = json.load(f)["weight_map"] + hf_keys = set(wmap.keys()) + fv_keys = set(fv_shapes.keys()) + + missing_in_fv = sorted(hf_keys - fv_keys) + extra_in_fv = sorted(fv_keys - hf_keys) + assert not missing_in_fv, f"fastvideo missing keys: {missing_in_fv[:5]}" + assert not extra_in_fv, f"fastvideo extra keys: {extra_in_fv[:5]}" + assert len(fv_keys) == 331 + + +def _hf_index_path_or_none() -> str | None: + """Return a local path to the base/ index.json, or None if unavailable.""" + try: + from huggingface_hub import hf_hub_download + except ImportError: + return None + try: + return hf_hub_download( + repo_id="GAIR/daVinci-MagiHuman", + filename="base/model.safetensors.index.json", + ) + except Exception: + return None + + +@pytest.mark.skipif( + not torch.cuda.is_available(), + reason="MagiHuman pipeline smoke test requires CUDA.", +) +def test_magi_human_pipeline_smoke() -> None: + """End-to-end smoke: build the pipeline and run a tiny denoise. + + Skips cleanly when the converted-weights directory is not present. + """ + diffusers_path = os.getenv( + "MAGI_HUMAN_DIFFUSERS_PATH", + "converted_weights/magi_human_base", + ) + if not os.path.isdir(diffusers_path): + pytest.skip( + f"Missing converted MagiHuman repo at {diffusers_path}. " + f"Run scripts/checkpoint_conversion/convert_magi_human_to_diffusers.py " + f"first." + ) + if not os.path.isfile(os.path.join(diffusers_path, "model_index.json")): + pytest.skip(f"Missing model_index.json in {diffusers_path}") + + os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA") + + from fastvideo import VideoGenerator + + # Small shapes to keep the smoke test cheap. + prompt = "A cheerful person waving at the camera in a well-lit room." + seed = 42 + height = 256 + width = 448 + num_frames = 13 # seconds=1, 12fps for smoke; the pipeline derives it + fps = 12.0 + steps = 2 + + generator = VideoGenerator.from_pretrained( + diffusers_path, + num_gpus=1, + use_fsdp_inference=False, + dit_cpu_offload=False, + vae_cpu_offload=False, + text_encoder_cpu_offload=False, + ) + try: + result = generator.generate_video( + prompt=prompt, + output_path="outputs_video/magi_human_smoke", + save_video=False, + height=height, + width=width, + num_frames=num_frames, + fps=fps, + num_inference_steps=steps, + seed=seed, + ) + finally: + generator.shutdown() + + samples = result["samples"] + assert samples.ndim == 5, f"expected [B,C,T,H,W], got {samples.shape}" + assert samples.shape[0] == 1 diff --git a/tests/local_tests/magi_human/test_magi_human_sa_audio_official_parity.py b/tests/local_tests/magi_human/test_magi_human_sa_audio_official_parity.py new file mode 100644 index 000000000..88f9f599b --- /dev/null +++ b/tests/local_tests/magi_human/test_magi_human_sa_audio_official_parity.py @@ -0,0 +1,162 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Parity test: FastVideo MagiHuman Stable-Audio wrapper vs the official +daVinci-MagiHuman Stable-Audio usage path. + +The sibling `test_magi_human_sa_audio_parity.py` compares FastVideo's +`SAAudioVAEModel` against Diffusers `AutoencoderOobleck`, which validates the +low-level VAE weights/API. This test instead follows the official +daVinci-MagiHuman integration layer: + +* `inference.model.sa_audio.SAAudioFeatureExtractor` is constructed from the + full Stable Audio checkpoint (`model_config.json` + `model.safetensors`). +* The official loader rebuilds its local `AudioAutoencoder` from + `model.pretransform.config` and filters `pretransform.model.*` weights. +* The official decode entry point is `SAAudioFeatureExtractor.decode(latents)`, + which calls `vae_model.decode(latents)` directly. There is no latent + mean/std normalization or reference-audio injection inside this decode layer. +* Pipeline post-processing is outside the SA module: `MagiEvaluator` transposes + `[B, L, C] -> [C, L]` before decode, then transposes waveform samples and + applies `resample_audio_sinc(..., 441 / 512)`. + +This catches drift between FastVideo's full SA wrapper path and the official +repo's custom Stable-Audio wrapper/module, not just the bare Diffusers VAE. + +Skips when: + * CUDA is unavailable. + * `daVinci-MagiHuman/` is not checked out under the repo root. + * `stabilityai/stable-audio-open-1.0` is inaccessible (gated; user must have + accepted terms and set HF_TOKEN / HUGGINGFACE_HUB_TOKEN / HF_API_KEY). +""" +from __future__ import annotations + +import os +from pathlib import Path + +import pytest +import torch +from torch.testing import assert_close + + +_SA_AUDIO_ID = "stabilityai/stable-audio-open-1.0" + + +def _repo_root() -> Path: + return Path(__file__).resolve().parents[3] + + +def _upstream_root() -> Path: + return _repo_root() / "daVinci-MagiHuman" + + +def _hf_token(): + for key in ("HF_TOKEN", "HUGGINGFACE_HUB_TOKEN", "HF_API_KEY"): + value = os.environ.get(key) + if value: + return value + return None + + +def _can_access() -> bool: + token = _hf_token() + if token is None: + return False + try: + from huggingface_hub import hf_hub_download + + hf_hub_download( + repo_id=_SA_AUDIO_ID, + filename="model_config.json", + token=token, + ) + return True + except Exception: + return False + + +def _stable_audio_snapshot() -> str: + from huggingface_hub import snapshot_download + + return snapshot_download( + repo_id=_SA_AUDIO_ID, + token=_hf_token(), + allow_patterns=["model_config.json", "model.safetensors"], + ) + + +@pytest.mark.skipif( + not torch.cuda.is_available(), + reason="MagiHuman official Stable-Audio parity requires CUDA.", +) +@pytest.mark.skipif( + not _upstream_root().exists(), + reason="daVinci-MagiHuman checkout is required under the repo root.", +) +@pytest.mark.skipif( + not _can_access(), + reason=( + f"{_SA_AUDIO_ID} not accessible — gated Stability AI repo; set " + "HF_TOKEN / HF_API_KEY and accept the terms on " + f"https://huggingface.co/{_SA_AUDIO_ID}." + ), +) +def test_magi_human_sa_audio_official_decode_parity(): + # Make sure both HF helpers and FastVideo's loader see the same token alias. + for src in ("HF_TOKEN", "HUGGINGFACE_HUB_TOKEN", "HF_API_KEY"): + value = os.environ.get(src) + if value: + os.environ.setdefault("HF_TOKEN", value) + os.environ.setdefault("HUGGINGFACE_HUB_TOKEN", value) + break + + device = torch.device("cuda:0") + + # --- Official daVinci-MagiHuman path: custom SAAudioFeatureExtractor + # rebuilds AudioAutoencoder and filters `pretransform.model.*` from + # the full Stable Audio checkpoint. + from tests.local_tests.helpers.magi_human_upstream import install_stubs + + install_stubs() + from inference.model.sa_audio import SAAudioFeatureExtractor + + upstream_vae = SAAudioFeatureExtractor( + device=device, + model_path=_stable_audio_snapshot(), + ) + + # --- FastVideo MagiHuman wrapper path: lazy loader around the native + # OobleckVAE port, exactly what the MagiHuman pipeline constructs. + from fastvideo.configs.models.vaes import OobleckVAEConfig + from fastvideo.models.vaes.sa_audio import SAAudioVAEModel + + fv_config = OobleckVAEConfig() + fv_config.pretrained_path = _SA_AUDIO_ID + fv_config.pretrained_dtype = "float32" + fv_vae = SAAudioVAEModel(fv_config) + + torch.manual_seed(0) + latent = torch.randn( + (1, fv_config.arch_config.decoder_input_channels, 8), + dtype=torch.float32, + device=device, + ) + + with torch.inference_mode(): + upstream_out = upstream_vae.decode(latent).detach().float().cpu() + fv_out = fv_vae.decode(latent).detach().float().cpu() + + print( + f"upstream shape={tuple(upstream_out.shape)} " + f"abs_mean={upstream_out.abs().mean().item():.6f} " + f"range=[{upstream_out.min().item():.4f}, " + f"{upstream_out.max().item():.4f}]" + ) + print( + f"fv shape={tuple(fv_out.shape)} " + f"abs_mean={fv_out.abs().mean().item():.6f} " + f"range=[{fv_out.min().item():.4f}, {fv_out.max().item():.4f}]" + ) + diff = (upstream_out - fv_out).abs() + print(f"diff max={diff.max().item():.6e} mean={diff.mean().item():.6e}") + + assert upstream_out.shape == fv_out.shape + assert_close(fv_out, upstream_out, atol=1e-5, rtol=1e-5) diff --git a/tests/local_tests/magi_human/test_magi_human_sa_audio_parity.py b/tests/local_tests/magi_human/test_magi_human_sa_audio_parity.py new file mode 100644 index 000000000..05c06294c --- /dev/null +++ b/tests/local_tests/magi_human/test_magi_human_sa_audio_parity.py @@ -0,0 +1,123 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Parity test: MagiHuman's audio-VAE path (FastVideo `SAAudioVAEModel` +lazy-loader around the native `OobleckVAE` port, shared with the +standalone Stable Audio pipeline) vs `diffusers.AutoencoderOobleck. +from_pretrained(...)` on the Stable Audio Open 1.0 VAE. + +Companion to `tests/local_tests/vaes/test_oobleck_vae_parity.py`, which +already validates `OobleckVAE` itself; this test exercises the wrapper +layer that MagiHuman uses (lazy load, device migration, decode output +unwrap) so wrapper-level regressions don't slip past the underlying-VAE +parity test. + +Skips when: + * CUDA is unavailable (VAE is 156M params, small enough for CPU but + we keep the test GPU-only to match the pipeline's runtime). + * `stabilityai/stable-audio-open-1.0` is inaccessible (gated; user + must have accepted terms on the HF repo page). +""" +from __future__ import annotations + +import os + +import pytest +import torch +from torch.testing import assert_close + + +_SA_AUDIO_ID = "stabilityai/stable-audio-open-1.0" + + +def _hf_token(): + for k in ("HF_TOKEN", "HUGGINGFACE_HUB_TOKEN", "HF_API_KEY"): + v = os.environ.get(k) + if v: + return v + return None + + +def _can_access() -> bool: + token = _hf_token() + if token is None: + return False + try: + from huggingface_hub import hf_hub_download + hf_hub_download( + repo_id=_SA_AUDIO_ID, filename="vae/config.json", token=token, + ) + return True + except Exception: + return False + + +@pytest.mark.skipif( + not torch.cuda.is_available(), + reason="MagiHuman Stable-Audio VAE parity requires CUDA.", +) +@pytest.mark.skipif( + not _can_access(), + reason=(f"{_SA_AUDIO_ID} not accessible — gated Stability AI repo; " + "set HF_TOKEN / HF_API_KEY and accept the terms on " + f"https://huggingface.co/{_SA_AUDIO_ID}."), +) +def test_magi_human_sa_audio_vae_decode_parity(): + # Make sure HF_TOKEN is the alias the Diffusers loader actually reads. + for src in ("HF_TOKEN", "HUGGINGFACE_HUB_TOKEN", "HF_API_KEY"): + v = os.environ.get(src) + if v: + os.environ.setdefault("HF_TOKEN", v) + os.environ.setdefault("HUGGINGFACE_HUB_TOKEN", v) + break + + device = torch.device("cuda:0") + + # --- Reference: direct HF Diffusers call, same as upstream's + # SAAudioFeatureExtractor but via the Diffusers Oobleck port. --- + from diffusers import AutoencoderOobleck + ref_vae = AutoencoderOobleck.from_pretrained( + _SA_AUDIO_ID, subfolder="vae", torch_dtype=torch.float32, + ).to(device).eval() + + # --- FastVideo wrapper path (shared with the standalone Stable + # Audio pipeline that landed in main: `OobleckVAEConfig` + + # `SAAudioVAEModel` lazy-loader around the first-class + # `OobleckVAE` port). --- + from fastvideo.configs.models.vaes import OobleckVAEConfig + from fastvideo.models.vaes.sa_audio import SAAudioVAEModel + fv_config = OobleckVAEConfig() + fv_config.pretrained_path = _SA_AUDIO_ID + # The default `pretrained_dtype="float16"` matches official stable- + # audio-tools, but this parity test runs the reference path in fp32 + # — so override here. + fv_config.pretrained_dtype = "float32" + fv_vae = SAAudioVAEModel(fv_config) + + # --- Tiny shared latent --- + torch.manual_seed(0) + # decoder_input_channels=64, latent length ~8 frames for a quick test. + latent = torch.randn( + (1, fv_config.arch_config.decoder_input_channels, 8), + dtype=torch.float32, device=device, + ) + + with torch.inference_mode(): + ref_out = ref_vae.decode(latent).sample.detach().float().cpu() + fv_out = fv_vae.decode(latent).detach().float().cpu() + + print( + f"ref shape={tuple(ref_out.shape)} " + f"abs_mean={ref_out.abs().mean().item():.6f} " + f"range=[{ref_out.min().item():.4f}, {ref_out.max().item():.4f}]" + ) + print( + f"fv shape={tuple(fv_out.shape)} " + f"abs_mean={fv_out.abs().mean().item():.6f} " + f"range=[{fv_out.min().item():.4f}, {fv_out.max().item():.4f}]" + ) + diff = (ref_out - fv_out).abs() + print(f"diff max={diff.max().item():.6e} mean={diff.mean().item():.6e}") + + assert ref_out.shape == fv_out.shape + # Both sides call the same HF class on the same weights in fp32 — + # should agree to machine epsilon. + assert_close(fv_out, ref_out, atol=1e-5, rtol=1e-5) diff --git a/tests/local_tests/magi_human/test_magi_human_sr1080p_pipeline_parity.py b/tests/local_tests/magi_human/test_magi_human_sr1080p_pipeline_parity.py new file mode 100644 index 000000000..785d492f2 --- /dev/null +++ b/tests/local_tests/magi_human/test_magi_human_sr1080p_pipeline_parity.py @@ -0,0 +1,340 @@ +# SPDX-License-Identifier: Apache-2.0 +"""SR-1080p local-window latent-loop parity for daVinci-MagiHuman. + +This mirrors the SR-540p two-stage parity test but enables upstream's +SR2_1080 local-attention layer set on the SR DiT. The reference side uses the +test helper's SDPA implementation of FFAHandler's segmented accumulator, so the +assertion is a kernel-noise tolerance rather than bit-exact. +""" +from __future__ import annotations + +import glob +import os +from pathlib import Path + +import pytest +import torch +import torch.nn.functional as F +from torch.testing import assert_close + +from fastvideo.pipelines.basic.magi_human.pipeline_configs import ( + _SR_1080P_LOCAL_ATTN_LAYERS, +) +from tests.local_tests.magi_human.test_magi_human_pipeline_parity import ( + _build_fastvideo_schedulers, + _build_upstream_schedulers, + _cleanup_gpu, + _dit_forward_fv, + _dit_forward_upstream, + _find_base_shard_dir, + _run_denoise_loop, +) +from tests.local_tests.magi_human.test_magi_human_sr540p_pipeline_parity import ( + _load_fv_dit, + _prepare_sr_latents, + _run_sr_denoise_loop, +) + + +os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA") +os.environ.setdefault("MASTER_ADDR", "localhost") +os.environ.setdefault("MASTER_PORT", "29522") + + +def _find_sr1080p_shard_dir() -> Path | None: + override = os.getenv("MAGI_HUMAN_SR1080P_SHARD_DIR") + if override: + path = Path(override) + return path if path.is_dir() else None + try: + from huggingface_hub import snapshot_download + snap = snapshot_download( + repo_id="GAIR/daVinci-MagiHuman", + allow_patterns=[ + "1080p_sr/*.safetensors", + "1080p_sr/model.safetensors.index.json", + ], + ) + candidate = Path(snap) / "1080p_sr" + if candidate.is_dir() and any(candidate.glob("*.safetensors")): + return candidate + return None + except Exception: + return None + + +def _dit_forward_upstream_local( + dit, + video_latent, + audio_latent, + audio_feat_len, + txt_feat, + txt_feat_len, + patch_size, + coords_style, + video_in_channels, + audio_in_channels, +): + from fastvideo.pipelines.basic.magi_human.stages.latent_preparation import ( + build_packed_inputs, + unpack_tokens, + ) + from inference.common import VarlenHandler + from inference.pipeline.data_proxy import calc_local_attn_ffa_handler + + x, coords, mm = build_packed_inputs( + video_latent=video_latent, + audio_latent=audio_latent, + audio_feat_len=audio_feat_len, + txt_feat=txt_feat, + txt_feat_len=txt_feat_len, + patch_size=patch_size, + coords_style=coords_style, + ) + video_token_num = x.shape[0] - audio_feat_len - txt_feat_len + total = x.shape[0] + cu = torch.tensor([0, total], dtype=torch.int32, device=x.device) + varlen = VarlenHandler( + cu_seqlens_q=cu, + cu_seqlens_k=cu, + max_seqlen_q=total, + max_seqlen_k=total, + ) + local_attn = calc_local_attn_ffa_handler( + video_token_num, + audio_feat_len + txt_feat_len, + video_latent.shape[2] // patch_size[0], + 11, + ) + out = dit( + x=x, + coords_mapping=coords, + modality_mapping=mm, + varlen_handler=varlen, + local_attn_handler=local_attn, + ) + return unpack_tokens( + out, + video_token_num=video_token_num, + audio_feat_len=audio_feat_len, + video_in_channels=video_in_channels, + audio_in_channels=audio_in_channels, + latent_shape=tuple(video_latent.shape), + patch_size=patch_size, + ) + + +@pytest.mark.skipif( + not torch.cuda.is_available(), + reason="MagiHuman SR-1080p pipeline parity requires CUDA.", +) +@pytest.mark.parametrize("use_image", [False, True], ids=["t2v", "ti2v"]) +def test_magi_human_sr1080p_pipeline_latent_parity(use_image: bool): + repo_root = Path(__file__).resolve().parents[3] + if not (repo_root / "daVinci-MagiHuman").exists(): + pytest.skip("Upstream daVinci-MagiHuman/ clone missing.") + + base_shard_dir = _find_base_shard_dir() + sr_shard_dir = _find_sr1080p_shard_dir() + if base_shard_dir is None or not base_shard_dir.is_dir(): + pytest.skip("GAIR/daVinci-MagiHuman base/ shards not available locally.") + if sr_shard_dir is None or not sr_shard_dir.is_dir(): + pytest.skip("GAIR/daVinci-MagiHuman 1080p_sr/ shards not available locally.") + + converted_dir = Path(os.getenv( + "MAGI_HUMAN_SR1080P_DIFFUSERS_PATH", + repo_root / "converted_weights" / "magi_human_sr_1080p", + )) + transformer_dir = converted_dir / "transformer" + sr_transformer_dir = converted_dir / "sr_transformer" + if not transformer_dir.is_dir(): + pytest.skip(f"Converted base transformer dir missing at {transformer_dir}") + if not sr_transformer_dir.is_dir(): + pytest.skip(f"Converted SR transformer dir missing at {sr_transformer_dir}") + + from tests.local_tests.helpers.magi_human_upstream import ( + install_stubs, + load_upstream_dit, + ) + install_stubs() + + device = torch.device("cuda:0") + torch.manual_seed(1080) + z_dim = 48 + patch_size = (1, 2, 2) + base_lat_T, base_lat_H, base_lat_W = 24, 4, 4 + sr_lat_H, sr_lat_W = 6, 8 + video_latent = torch.randn( + (1, z_dim, base_lat_T, base_lat_H, base_lat_W), + dtype=torch.float32, + device=device, + ) + audio_latent = torch.randn((1, 32, 64), dtype=torch.float32, device=device) + base_image_latent = None + sr_image_latent = None + if use_image: + base_image_latent = torch.randn( + (1, z_dim, 1, base_lat_H, base_lat_W), + dtype=torch.float32, + device=device, + ) + sr_image_latent = F.interpolate( + base_image_latent, + size=(1, sr_lat_H, sr_lat_W), + mode="trilinear", + align_corners=True, + ) + + txt_feat_len = 7 + neg_txt_feat_len = 11 + txt_feat = torch.randn((1, 640, 3584), dtype=torch.float32, device=device) + neg_txt_feat = torch.randn((1, 640, 3584), dtype=torch.float32, device=device) + base_steps = 4 + sr_steps = 2 + shift = 5.0 + base_kwargs = dict( + cfg_number=2, + video_txt_guidance_scale=5.0, + audio_txt_guidance_scale=5.0, + patch_size=patch_size, + coords_style="v2", + video_in_channels=192, + audio_in_channels=64, + image_latent=base_image_latent, + ) + sr_kwargs = dict( + patch_size=patch_size, + coords_style="v1", + video_in_channels=192, + audio_in_channels=64, + image_latent=sr_image_latent, + ) + + up_video_sched, up_audio_sched = _build_upstream_schedulers( + shift=shift, + num_inference_steps=base_steps, + device=device, + ) + upstream_base = load_upstream_dit(base_shard_dir, device=device, dtype=None) + ref_base_video, ref_base_audio = _run_denoise_loop( + upstream_base, + _dit_forward_upstream, + video_latent.clone(), + audio_latent.clone(), + txt_feat.clone(), + txt_feat_len, + neg_txt_feat.clone(), + neg_txt_feat_len, + video_sched=up_video_sched, + audio_sched=up_audio_sched, + **base_kwargs, + ) + del upstream_base + _cleanup_gpu() + + torch.manual_seed(1081) + ref_sr_video_in, ref_sr_audio_in = _prepare_sr_latents( + ref_base_video, + ref_base_audio, + latent_h=sr_lat_H, + latent_w=sr_lat_W, + noise_value=220, + ) + up_sr_video_sched, _ = _build_upstream_schedulers( + shift=shift, + num_inference_steps=sr_steps, + device=device, + ) + upstream_sr = load_upstream_dit( + sr_shard_dir, + device=device, + dtype=None, + local_attn_layers=_SR_1080P_LOCAL_ATTN_LAYERS, + ) + ref_video, ref_audio = _run_sr_denoise_loop( + upstream_sr, + _dit_forward_upstream_local, + ref_sr_video_in.clone(), + ref_sr_audio_in.clone(), + txt_feat.clone(), + txt_feat_len, + neg_txt_feat.clone(), + neg_txt_feat_len, + video_sched=up_sr_video_sched, + **sr_kwargs, + ) + ref_video = ref_video.detach().float().cpu() + ref_audio = ref_audio.detach().float().cpu() + del upstream_sr + _cleanup_gpu() + + fv_video_sched, fv_audio_sched = _build_fastvideo_schedulers( + shift=shift, + num_inference_steps=base_steps, + device=device, + ) + fv_base = _load_fv_dit(transformer_dir, device) + fv_base_video, fv_base_audio = _run_denoise_loop( + fv_base, + _dit_forward_fv, + video_latent.clone(), + audio_latent.clone(), + txt_feat.clone(), + txt_feat_len, + neg_txt_feat.clone(), + neg_txt_feat_len, + video_sched=fv_video_sched, + audio_sched=fv_audio_sched, + **base_kwargs, + ) + del fv_base + _cleanup_gpu() + + torch.manual_seed(1081) + fv_sr_video_in, fv_sr_audio_in = _prepare_sr_latents( + fv_base_video, + fv_base_audio, + latent_h=sr_lat_H, + latent_w=sr_lat_W, + noise_value=220, + ) + fv_sr_video_sched, _ = _build_fastvideo_schedulers( + shift=shift, + num_inference_steps=sr_steps, + device=device, + ) + fv_sr = _load_fv_dit(sr_transformer_dir, device) + fv_sr.configure_local_attention(_SR_1080P_LOCAL_ATTN_LAYERS, frame_receptive_field=11) + fv_video, fv_audio = _run_sr_denoise_loop( + fv_sr, + _dit_forward_fv, + fv_sr_video_in.clone(), + fv_sr_audio_in.clone(), + txt_feat.clone(), + txt_feat_len, + neg_txt_feat.clone(), + neg_txt_feat_len, + video_sched=fv_sr_video_sched, + **sr_kwargs, + ) + fv_video = fv_video.detach().float().cpu() + fv_audio = fv_audio.detach().float().cpu() + + v_diff = (ref_video - fv_video).abs() + a_diff = (ref_audio - fv_audio).abs() + print( + f"sr1080p {('ti2v' if use_image else 't2v')} " + f"video diff_max={v_diff.max().item():.4f} diff_mean={v_diff.mean().item():.4f}" + ) + print( + f"sr1080p {('ti2v' if use_image else 't2v')} " + f"audio diff_max={a_diff.max().item():.4f} diff_mean={a_diff.mean().item():.4f}" + ) + + assert ref_video.shape == fv_video.shape + assert ref_audio.shape == fv_audio.shape + assert v_diff.max().item() < 0.05 + assert_close(fv_audio, ref_audio, atol=0.0, rtol=0.0) + if use_image: + assert_close(fv_video[:, :, :1], sr_image_latent.detach().cpu(), atol=0.0, rtol=0.0) + assert_close(ref_video[:, :, :1], sr_image_latent.detach().cpu(), atol=0.0, rtol=0.0) diff --git a/tests/local_tests/magi_human/test_magi_human_sr540p_pipeline_parity.py b/tests/local_tests/magi_human/test_magi_human_sr540p_pipeline_parity.py new file mode 100644 index 000000000..a23eb3c7b --- /dev/null +++ b/tests/local_tests/magi_human/test_magi_human_sr540p_pipeline_parity.py @@ -0,0 +1,400 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Two-stage SR-540p latent-loop parity for daVinci-MagiHuman.""" +from __future__ import annotations + +import glob +import os +from pathlib import Path + +import pytest +import torch +import torch.nn.functional as F +from torch.testing import assert_close + +from fastvideo.pipelines.basic.magi_human.stages.sr_latent_preparation import ( + ZeroSNRDDPMDiscretization, +) +from tests.local_tests.magi_human.test_magi_human_pipeline_parity import ( + _build_fastvideo_schedulers, + _build_upstream_schedulers, + _cleanup_gpu, + _dit_forward_fv, + _dit_forward_upstream, + _find_base_shard_dir, + _run_denoise_loop, +) + + +os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA") +os.environ.setdefault("MASTER_ADDR", "localhost") +os.environ.setdefault("MASTER_PORT", "29521") + + +def _find_sr540p_shard_dir() -> Path | None: + override = os.getenv("MAGI_HUMAN_SR540P_SHARD_DIR") + if override: + path = Path(override) + return path if path.is_dir() else None + try: + from huggingface_hub import snapshot_download + snap = snapshot_download( + repo_id="GAIR/daVinci-MagiHuman", + allow_patterns=[ + "540p_sr/*.safetensors", + "540p_sr/model.safetensors.index.json", + ], + ) + candidate = Path(snap) / "540p_sr" + if candidate.is_dir() and any(candidate.glob("*.safetensors")): + return candidate + return None + except Exception: + return None + + +def _prepare_sr_latents( + br_video: torch.Tensor, + br_audio: torch.Tensor, + *, + latent_h: int, + latent_w: int, + noise_value: int, +) -> tuple[torch.Tensor, torch.Tensor]: + latent_video = F.interpolate( + br_video, + size=(br_video.shape[2], latent_h, latent_w), + mode="trilinear", + align_corners=True, + ) + if noise_value != 0: + noise = torch.randn_like(latent_video, device=latent_video.device) + sigmas = ZeroSNRDDPMDiscretization()( + 1000, + do_append_zero=False, + flip=True, + device=latent_video.device, + ) + sigma = sigmas[noise_value] + latent_video = latent_video * sigma + noise * (1 - sigma**2)**0.5 + sr_audio = torch.randn_like(br_audio, device=br_audio.device) * 0.7 + br_audio * 0.3 + return latent_video, sr_audio + + +def _run_sr_denoise_loop( + dit, + dit_forward_fn, + video_latent, + audio_latent, + txt_feat, + txt_feat_len, + neg_txt_feat, + neg_txt_feat_len, + *, + video_sched, + patch_size, + coords_style, + video_in_channels, + audio_in_channels, + image_latent=None, +): + from fastvideo.forward_context import set_forward_context + + audio_feat_len = int(audio_latent.shape[1]) + latent_length = video_latent.shape[2] + guidance = torch.tensor(3.5, device=video_latent.device).expand( + 1, + 1, + latent_length, + 1, + 1, + ).clone() + guidance[:, :, :13] = 2.0 + + with torch.inference_mode(): + for t in video_sched.timesteps: + if image_latent is not None: + video_latent[:, :, :1] = image_latent.to( + device=video_latent.device, + dtype=video_latent.dtype, + )[:, :, :1] + with set_forward_context( + current_timestep=int(t.item()) if torch.is_tensor(t) else int(t), + attn_metadata=None, + ): + v_cond_video, _ = dit_forward_fn( + dit, + video_latent, + audio_latent, + audio_feat_len, + txt_feat, + txt_feat_len, + patch_size, + coords_style, + video_in_channels, + audio_in_channels, + ) + v_uncond_video, _ = dit_forward_fn( + dit, + video_latent, + audio_latent, + audio_feat_len, + neg_txt_feat, + neg_txt_feat_len, + patch_size, + coords_style, + video_in_channels, + audio_in_channels, + ) + v_video = v_uncond_video + guidance * (v_cond_video - v_uncond_video) + video_latent = video_sched.step( + v_video, + t, + video_latent, + return_dict=False, + )[0] + if image_latent is not None: + video_latent[:, :, :1] = image_latent.to( + device=video_latent.device, + dtype=video_latent.dtype, + )[:, :, :1] + return video_latent, audio_latent + + +def _load_fv_dit(transformer_dir: Path, device: torch.device): + from fastvideo.configs.models.dits.magi_human import MagiHumanVideoConfig + from fastvideo.models.dits.magi_human import MagiHumanDiT + from safetensors.torch import load_file + + dit = MagiHumanDiT(MagiHumanVideoConfig()) + state = {} + for shard in sorted(glob.glob(str(transformer_dir / "*.safetensors"))): + state.update(load_file(shard)) + missing, unexpected = dit.load_state_dict(state, strict=False) + assert not missing, f"FastVideo DiT missing {len(missing)} keys: {missing[:5]}" + assert not unexpected, f"FastVideo DiT unexpected {len(unexpected)} keys: {unexpected[:5]}" + return dit.to(device=device).eval() + + +@pytest.mark.skipif( + not torch.cuda.is_available(), + reason="MagiHuman SR-540p pipeline parity requires CUDA.", +) +@pytest.mark.parametrize("use_image", [False, True], ids=["t2v", "ti2v"]) +def test_magi_human_sr540p_pipeline_latent_parity(use_image: bool): + repo_root = Path(__file__).resolve().parents[3] + if not (repo_root / "daVinci-MagiHuman").exists(): + pytest.skip("Upstream daVinci-MagiHuman/ clone missing.") + + base_shard_dir = _find_base_shard_dir() + sr_shard_dir = _find_sr540p_shard_dir() + if base_shard_dir is None or not base_shard_dir.is_dir(): + pytest.skip("GAIR/daVinci-MagiHuman base/ shards not available locally.") + if sr_shard_dir is None or not sr_shard_dir.is_dir(): + pytest.skip("GAIR/daVinci-MagiHuman 540p_sr/ shards not available locally.") + + converted_dir = Path(os.getenv( + "MAGI_HUMAN_SR540P_DIFFUSERS_PATH", + repo_root / "converted_weights" / "magi_human_sr_540p", + )) + transformer_dir = converted_dir / "transformer" + sr_transformer_dir = converted_dir / "sr_transformer" + if not transformer_dir.is_dir(): + pytest.skip(f"Converted base transformer dir missing at {transformer_dir}") + if not sr_transformer_dir.is_dir(): + pytest.skip(f"Converted SR transformer dir missing at {sr_transformer_dir}") + + from tests.local_tests.helpers.magi_human_upstream import ( + install_stubs, + load_upstream_dit, + ) + install_stubs() + + device = torch.device("cuda:0") + torch.manual_seed(540) + z_dim = 48 + patch_size = (1, 2, 2) + base_lat_T, base_lat_H, base_lat_W = 2, 6, 6 + sr_lat_H, sr_lat_W = 8, 10 + video_latent = torch.randn( + (1, z_dim, base_lat_T, base_lat_H, base_lat_W), + dtype=torch.float32, + device=device, + ) + audio_latent = torch.randn((1, 4, 64), dtype=torch.float32, device=device) + base_image_latent = None + sr_image_latent = None + if use_image: + base_image_latent = torch.randn( + (1, z_dim, 1, base_lat_H, base_lat_W), + dtype=torch.float32, + device=device, + ) + sr_image_latent = F.interpolate( + base_image_latent, + size=(1, sr_lat_H, sr_lat_W), + mode="trilinear", + align_corners=True, + ) + + txt_feat_len = 7 + neg_txt_feat_len = 11 + txt_feat = torch.randn( + (1, 640, 3584), + dtype=torch.float32, + device=device, + ) + neg_txt_feat = torch.randn( + (1, 640, 3584), + dtype=torch.float32, + device=device, + ) + base_steps = 4 + sr_steps = 2 + shift = 5.0 + base_kwargs = dict( + cfg_number=2, + video_txt_guidance_scale=5.0, + audio_txt_guidance_scale=5.0, + patch_size=patch_size, + coords_style="v2", + video_in_channels=192, + audio_in_channels=64, + image_latent=base_image_latent, + ) + sr_kwargs = dict( + patch_size=patch_size, + coords_style="v1", + video_in_channels=192, + audio_in_channels=64, + image_latent=sr_image_latent, + ) + + up_video_sched, up_audio_sched = _build_upstream_schedulers( + shift=shift, + num_inference_steps=base_steps, + device=device, + ) + upstream_base = load_upstream_dit(base_shard_dir, device=device, dtype=None) + ref_base_video, ref_base_audio = _run_denoise_loop( + upstream_base, + _dit_forward_upstream, + video_latent.clone(), + audio_latent.clone(), + txt_feat.clone(), + txt_feat_len, + neg_txt_feat.clone(), + neg_txt_feat_len, + video_sched=up_video_sched, + audio_sched=up_audio_sched, + **base_kwargs, + ) + del upstream_base + _cleanup_gpu() + + torch.manual_seed(541) + ref_sr_video_in, ref_sr_audio_in = _prepare_sr_latents( + ref_base_video, + ref_base_audio, + latent_h=sr_lat_H, + latent_w=sr_lat_W, + noise_value=220, + ) + up_sr_video_sched, _ = _build_upstream_schedulers( + shift=shift, + num_inference_steps=sr_steps, + device=device, + ) + upstream_sr = load_upstream_dit(sr_shard_dir, device=device, dtype=None) + ref_video, ref_audio = _run_sr_denoise_loop( + upstream_sr, + _dit_forward_upstream, + ref_sr_video_in.clone(), + ref_sr_audio_in.clone(), + txt_feat.clone(), + txt_feat_len, + neg_txt_feat.clone(), + neg_txt_feat_len, + video_sched=up_sr_video_sched, + **sr_kwargs, + ) + ref_video = ref_video.detach().float().cpu() + ref_audio = ref_audio.detach().float().cpu() + del upstream_sr + _cleanup_gpu() + + fv_video_sched, fv_audio_sched = _build_fastvideo_schedulers( + shift=shift, + num_inference_steps=base_steps, + device=device, + ) + fv_base = _load_fv_dit(transformer_dir, device) + fv_base_video, fv_base_audio = _run_denoise_loop( + fv_base, + _dit_forward_fv, + video_latent.clone(), + audio_latent.clone(), + txt_feat.clone(), + txt_feat_len, + neg_txt_feat.clone(), + neg_txt_feat_len, + video_sched=fv_video_sched, + audio_sched=fv_audio_sched, + **base_kwargs, + ) + del fv_base + _cleanup_gpu() + + torch.manual_seed(541) + fv_sr_video_in, fv_sr_audio_in = _prepare_sr_latents( + fv_base_video, + fv_base_audio, + latent_h=sr_lat_H, + latent_w=sr_lat_W, + noise_value=220, + ) + fv_sr_video_sched, _ = _build_fastvideo_schedulers( + shift=shift, + num_inference_steps=sr_steps, + device=device, + ) + fv_sr = _load_fv_dit(sr_transformer_dir, device) + fv_video, fv_audio = _run_sr_denoise_loop( + fv_sr, + _dit_forward_fv, + fv_sr_video_in.clone(), + fv_sr_audio_in.clone(), + txt_feat.clone(), + txt_feat_len, + neg_txt_feat.clone(), + neg_txt_feat_len, + video_sched=fv_sr_video_sched, + **sr_kwargs, + ) + fv_video = fv_video.detach().float().cpu() + fv_audio = fv_audio.detach().float().cpu() + + v_diff = (ref_video - fv_video).abs() + a_diff = (ref_audio - fv_audio).abs() + print( + f"sr540p {('ti2v' if use_image else 't2v')} " + f"video diff_max={v_diff.max().item():.4f} diff_mean={v_diff.mean().item():.4f}" + ) + print( + f"sr540p {('ti2v' if use_image else 't2v')} " + f"audio diff_max={a_diff.max().item():.4f} diff_mean={a_diff.mean().item():.4f}" + ) + + assert ref_video.shape == fv_video.shape + assert ref_audio.shape == fv_audio.shape + assert_close(fv_audio, ref_audio, atol=0.0, rtol=0.0) + assert_close(fv_video, ref_video, atol=0.0, rtol=0.0) + if use_image: + assert_close(fv_video[:, :, :1], sr_image_latent.detach().cpu(), atol=0.0, rtol=0.0) + assert_close(ref_video[:, :, :1], sr_image_latent.detach().cpu(), atol=0.0, rtol=0.0) + + ref_v_abs = ref_video.abs().mean().item() + ref_a_abs = ref_audio.abs().mean().item() + assert abs(ref_v_abs - fv_video.abs().mean().item()) / max(ref_v_abs, 1e-6) < 0.02 + assert abs(ref_a_abs - fv_audio.abs().mean().item()) / max(ref_a_abs, 1e-6) < 0.02 + assert v_diff.mean().item() / max(ref_v_abs, 1e-6) < 0.06 + assert a_diff.mean().item() / max(ref_a_abs, 1e-6) < 0.04 diff --git a/tests/local_tests/magi_human/test_magi_human_ti2v_pipeline_parity.py b/tests/local_tests/magi_human/test_magi_human_ti2v_pipeline_parity.py new file mode 100644 index 000000000..51232b4b9 --- /dev/null +++ b/tests/local_tests/magi_human/test_magi_human_ti2v_pipeline_parity.py @@ -0,0 +1,186 @@ +# SPDX-License-Identifier: Apache-2.0 +"""TI2V latent-loop parity for daVinci-MagiHuman. + +This mirrors the base MagiHuman pipeline parity test but enables the upstream +`latent_image is not None` branch: the clean image latent is copied into +`latent_video[:, :, :1]` before every DiT call and once more after denoising. +""" +from __future__ import annotations + +import glob +import os +from pathlib import Path + +import pytest +import torch +from torch.testing import assert_close + +from tests.local_tests.magi_human.test_magi_human_pipeline_parity import ( + _build_fastvideo_schedulers, + _build_upstream_schedulers, + _cleanup_gpu, + _dit_forward_fv, + _dit_forward_upstream, + _encode_magi_human_prompt_pair, + _find_base_shard_dir, + _run_denoise_loop, +) + + +os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA") +os.environ.setdefault("MASTER_ADDR", "localhost") +os.environ.setdefault("MASTER_PORT", "29520") + + +@pytest.mark.skipif( + not torch.cuda.is_available(), + reason="MagiHuman TI2V pipeline parity requires CUDA.", +) +def test_magi_human_ti2v_pipeline_latent_parity(): + repo_root = Path(__file__).resolve().parents[3] + upstream_src = repo_root / "daVinci-MagiHuman" + if not upstream_src.exists(): + pytest.skip("Upstream daVinci-MagiHuman/ clone missing.") + + base_shard_dir = _find_base_shard_dir() + if base_shard_dir is None or not base_shard_dir.is_dir(): + pytest.skip("GAIR/daVinci-MagiHuman base/ shards not available locally.") + + converted_dir = Path(os.getenv( + "MAGI_HUMAN_DIFFUSERS_PATH", + repo_root / "converted_weights" / "magi_human_base", + )) + transformer_dir = converted_dir / "transformer" + if not transformer_dir.is_dir(): + pytest.skip(f"Converted transformer dir missing at {transformer_dir}") + + from tests.local_tests.helpers.magi_human_upstream import ( + install_stubs, + load_upstream_dit, + ) + install_stubs() + + device = torch.device("cuda:0") + torch.manual_seed(123) + + z_dim = 48 + patch_size = (1, 2, 2) + lat_T, lat_H, lat_W = 2, 6, 6 + video_latent = torch.randn( + (1, z_dim, lat_T, lat_H, lat_W), + dtype=torch.float32, + device=device, + ) + audio_latent = torch.randn((1, 4, 64), dtype=torch.float32, device=device) + image_latent = torch.randn( + (1, z_dim, 1, lat_H, lat_W), + dtype=torch.float32, + device=device, + ) + txt_feat, txt_feat_len, neg_txt_feat, neg_txt_feat_len = ( + _encode_magi_human_prompt_pair(device) + ) + + num_inference_steps = 4 + shift = 5.0 + common_kwargs = dict( + cfg_number=2, + video_txt_guidance_scale=5.0, + audio_txt_guidance_scale=5.0, + patch_size=patch_size, + coords_style="v2", + video_in_channels=192, + audio_in_channels=64, + image_latent=image_latent, + ) + + up_video_sched, up_audio_sched = _build_upstream_schedulers( + shift=shift, + num_inference_steps=num_inference_steps, + device=device, + ) + print("Loading upstream DiTModel from base shards...") + upstream_dit = load_upstream_dit(base_shard_dir, device=device, dtype=None) + print("Running upstream TI2V denoise loop...") + ref_video, ref_audio = _run_denoise_loop( + upstream_dit, + _dit_forward_upstream, + video_latent.clone(), + audio_latent.clone(), + txt_feat.clone(), + txt_feat_len, + neg_txt_feat.clone(), + neg_txt_feat_len, + video_sched=up_video_sched, + audio_sched=up_audio_sched, + **common_kwargs, + ) + ref_video = ref_video.detach().float().cpu() + ref_audio = ref_audio.detach().float().cpu() + del upstream_dit + _cleanup_gpu() + + fv_video_sched, fv_audio_sched = _build_fastvideo_schedulers( + shift=shift, + num_inference_steps=num_inference_steps, + device=device, + ) + from fastvideo.configs.models.dits.magi_human import MagiHumanVideoConfig + from fastvideo.models.dits.magi_human import MagiHumanDiT + from safetensors.torch import load_file + + print("Loading FastVideo MagiHumanDiT from converted transformer/...") + fv_cfg = MagiHumanVideoConfig() + fv_dit = MagiHumanDiT(fv_cfg) + fv_state = {} + for shard in sorted(glob.glob(str(transformer_dir / "*.safetensors"))): + fv_state.update(load_file(shard)) + missing, unexpected = fv_dit.load_state_dict(fv_state, strict=False) + assert not missing, f"FastVideo DiT missing {len(missing)} keys: {missing[:5]}" + assert not unexpected, f"FastVideo DiT unexpected {len(unexpected)} keys: {unexpected[:5]}" + fv_dit = fv_dit.to(device=device) + fv_dit.eval() + + print("Running FastVideo TI2V denoise loop...") + fv_video, fv_audio = _run_denoise_loop( + fv_dit, + _dit_forward_fv, + video_latent.clone(), + audio_latent.clone(), + txt_feat.clone(), + txt_feat_len, + neg_txt_feat.clone(), + neg_txt_feat_len, + video_sched=fv_video_sched, + audio_sched=fv_audio_sched, + **common_kwargs, + ) + fv_video = fv_video.detach().float().cpu() + fv_audio = fv_audio.detach().float().cpu() + + v_diff = (ref_video - fv_video).abs() + a_diff = (ref_audio - fv_audio).abs() + print( + f"ti2v video diff_max={v_diff.max().item():.4f} " + f"diff_mean={v_diff.mean().item():.4f}" + ) + print( + f"ti2v audio diff_max={a_diff.max().item():.4f} " + f"diff_mean={a_diff.mean().item():.4f}" + ) + + assert ref_video.shape == fv_video.shape + assert ref_audio.shape == fv_audio.shape + assert_close(fv_video, ref_video, atol=0.40, rtol=0.05) + assert_close(fv_audio, ref_audio, atol=0.40, rtol=0.05) + assert_close(fv_video[:, :, :1], image_latent.detach().cpu(), atol=0.0, rtol=0.0) + assert_close(ref_video[:, :, :1], image_latent.detach().cpu(), atol=0.0, rtol=0.0) + + ref_v_abs = ref_video.abs().mean().item() + ref_a_abs = ref_audio.abs().mean().item() + rel_v = abs(ref_v_abs - fv_video.abs().mean().item()) / max(ref_v_abs, 1e-6) + rel_a = abs(ref_a_abs - fv_audio.abs().mean().item()) / max(ref_a_abs, 1e-6) + assert rel_v < 0.01, f"video abs_mean drift {rel_v:.2%} > 1%" + assert rel_a < 0.01, f"audio abs_mean drift {rel_a:.2%} > 1%" + assert v_diff.mean().item() / max(ref_v_abs, 1e-6) < 0.04 + assert a_diff.mean().item() / max(ref_a_abs, 1e-6) < 0.04 diff --git a/tests/local_tests/magi_human/test_magi_human_vae_parity.py b/tests/local_tests/magi_human/test_magi_human_vae_parity.py new file mode 100644 index 000000000..d57bad089 --- /dev/null +++ b/tests/local_tests/magi_human/test_magi_human_vae_parity.py @@ -0,0 +1,180 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Parity test: FastVideo `AutoencoderKLWan` vs upstream `Wan2_2_VAE`. + +MagiHuman uses the Wan 2.2 TI2V-5B VAE. The two implementations +compared here are: + + * Upstream (SandAI port) — `inference/model/vae2_2/vae2_2_module.py::Wan2_2_VAE` + loaded from `Wan-AI/Wan2.2-TI2V-5B/Wan2.2_VAE.pth` (the official .pth + inside the daVinci-MagiHuman repo). This is the reference. + * FastVideo — `fastvideo.models.vaes.wanvae.AutoencoderKLWan` (the + class registered as `EntryClass` and resolved by the VAE component + loader at runtime; this is what `MagiHumanBaseConfig.vae_config` + materializes when the magi pipeline runs). Weights are loaded from + a Diffusers-format `vae/` subdir (`config.json` + + `diffusion_pytorch_model.safetensors`). + +This test decodes the same random latent through both and asserts the +decoded videos are close. Catches regressions in: + - FastVideo's `AutoencoderKLWan` weight load / scale / shift handling. + - Any deviation in `latents_mean` / `latents_std` baked into the + Diffusers-format config vs the upstream constants. + +Skips when: + - CUDA is unavailable. + - The .pth is not locally available (requires ~2.8 GB download). + - The converted MagiHuman Diffusers repo (or any `Wan-AI/*-Diffusers` + repo with a `vae/` subdir) is not available locally. +""" +from __future__ import annotations + +import os +from pathlib import Path + +import pytest +import torch +from torch.testing import assert_close + + +@pytest.mark.skipif( + not torch.cuda.is_available(), + reason="VAE parity test requires CUDA.", +) +def test_magi_human_vae_decode_parity(): + repo_root = Path(__file__).resolve().parents[3] + upstream_src = repo_root / "daVinci-MagiHuman" + if not upstream_src.exists(): + pytest.skip( + "Upstream daVinci-MagiHuman/ clone missing — no Wan2_2_VAE source." + ) + + fv_vae_dir = Path(os.getenv( + "MAGI_HUMAN_VAE_DIR", + repo_root / "converted_weights" / "magi_human_base" / "vae", + )) + if not (fv_vae_dir / "config.json").is_file(): + pytest.skip(f"FastVideo VAE dir missing at {fv_vae_dir}") + + # Upstream Wan2_2_VAE needs the raw .pth shipped by Wan-AI/Wan2.2-TI2V-5B + # (NOT the -Diffusers variant; that one has safetensors, not .pth). + try: + from huggingface_hub import hf_hub_download + pth_path = hf_hub_download( + repo_id="Wan-AI/Wan2.2-TI2V-5B", filename="Wan2.2_VAE.pth", + ) + except Exception as exc: + pytest.skip(f"Wan2.2_VAE.pth not available: {exc}") + + # Push upstream + install compiler stubs (the VAE module itself doesn't + # need magi_compiler, but `inference.*` imports pull in siblings that do). + from tests.local_tests.helpers.magi_human_upstream import install_stubs + install_stubs() + + device = torch.device("cuda:0") + torch.manual_seed(0) + + # Tiny latent so the test stays well inside GPU memory budget. + # z_dim=48, T=1, H=4, W=4 -> VAE decodes to [1, 3, 1 (or 1+4*0), 64, 64] + z = torch.randn((1, 48, 1, 4, 4), dtype=torch.float32, device=device) + + # --- Upstream decode --- + from inference.model.vae2_2 import Wan2_2_VAE + up_vae = Wan2_2_VAE( + vae_pth=pth_path, + device=device, + dtype=torch.float32, + ) + with torch.inference_mode(): + # Wan2_2_VAE.decode expects a (C, T, H, W) latent (no batch dim); + # see inference/pipeline/video_generate.py:494 — `self.vae.decode(latent.squeeze(0).to(self.dtype), ...)`. + up_out = up_vae.decode(z[0]).detach().float().cpu() + + del up_vae + import gc; gc.collect(); torch.cuda.empty_cache() + + # --- FastVideo decode --- + # Upstream `Wan2_2_VAE.decode(z)` internally normalizes via + # `(z - latents_mean) / latents_std` before feeding the decoder + # (see `scale = [mean, 1.0/std]` and the _video_vae.decode call). + # FastVideo's `AutoencoderKLWan.decode(z)` expects the input to + # ALREADY be in "decoder-input space" (the normalization is the + # caller's job — `DecodingStage` applies it). So we mirror the + # upstream transform here before calling decode. + import glob + + from safetensors.torch import load_file as safetensors_load_file + + from fastvideo.configs.models.vaes import WanVAEConfig + from fastvideo.models.loader.component_loader import get_diffusers_config + from fastvideo.models.vaes.wanvae import AutoencoderKLWan + + diffusers_cfg = get_diffusers_config(model=str(fv_vae_dir)) + diffusers_cfg.pop("_class_name", None) + diffusers_cfg.pop("_name_or_path", None) + fv_config = WanVAEConfig() + fv_config.load_encoder = False + fv_config.load_decoder = True + fv_config.update_model_arch(diffusers_cfg) + fv_vae = AutoencoderKLWan(fv_config).to(device=device, dtype=torch.float32) + + # Mirror the VAE component loader: glob `*.safetensors`, merge, load + # non-strictly so any unused buffers (per_channel_statistics, etc.) + # don't fail the load. + sf_files = glob.glob(os.path.join(str(fv_vae_dir), "*.safetensors")) + assert sf_files, f"No safetensors files in {fv_vae_dir}" + state = {} + for sf in sf_files: + state.update(safetensors_load_file(sf)) + fv_vae.load_state_dict(state, strict=False) + fv_vae.eval() + + # Upstream's inner `_video_vae.decode(z, scale)` (line 874-877 of + # inference/model/vae2_2/vae2_2_module.py) does: + # z = z / scale[1] + scale[0] # where scale = [mean, 1/std] + # = z * std + mean + # FastVideo's `AutoencoderKLWan.decode` expects the pre-denormalized + # latent — apply the same transform externally to feed both paths + # equivalently. + latents_mean = torch.tensor( + fv_config.arch_config.latents_mean, dtype=torch.float32, device=device, + ) + latents_std = torch.tensor( + fv_config.arch_config.latents_std, dtype=torch.float32, device=device, + ) + z_denormalized = z * latents_std.view(1, -1, 1, 1, 1) + latents_mean.view(1, -1, 1, 1, 1) + with torch.inference_mode(): + fv_out_tensor = fv_vae.decode(z_denormalized) + fv_out = fv_out_tensor.detach().float().cpu() + + # Both sides should return a video tensor of shape [..., C, T_dec, H_dec, W_dec]. + # Normalize shapes for comparison — upstream returns a list per-video or a + # single tensor depending on CP group; we just squeeze batch dims. + def _squeeze(t): + while t.ndim > 4 and t.shape[0] == 1: + t = t[0] + return t + + up_s = _squeeze(up_out) + fv_s = _squeeze(fv_out) + print( + f"up shape={tuple(up_s.shape)} abs_mean={up_s.abs().mean().item():.4f} " + f"range=[{up_s.min().item():.4f}, {up_s.max().item():.4f}]" + ) + print( + f"fv shape={tuple(fv_s.shape)} abs_mean={fv_s.abs().mean().item():.4f} " + f"range=[{fv_s.min().item():.4f}, {fv_s.max().item():.4f}]" + ) + + # Wan VAE has a known fp32 op-ordering drift of ~8e-4 caused by + # `z * std + mean` (FV) vs `z / (1/std) + mean` (upstream) at decode + # normalization. This is a SHARED Wan-family bug, not magi-specific. + # Tracked as OQ-7 in tests/local_tests/magi-human.md; tighten to + # atol=1e-4 once the Wan VAE op-order fix lands. + assert up_s.shape == fv_s.shape, ( + f"shape mismatch: up={up_s.shape} fv={fv_s.shape}" + ) + diff = (up_s - fv_s).abs() + print( + f"diff max={diff.max().item():.6f} mean={diff.mean().item():.6f}" + ) + assert_close(fv_s, up_s, atol=1e-3, rtol=1e-3)