[feat]: MagiHuman pipeline orchestrator + 10-test parity battery (5/8) (#1299)
This commit is contained in:
William Lin
@@ -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 `<model_path>/<key>/`. 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,
|
||||
]
|
||||
@@ -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
|
||||
@@ -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,
|
||||
)
|
||||
@@ -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%"
|
||||
@@ -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
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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)
|
||||
Reference in New Issue
Block a user