[bugfix]: condition HunyuanVideo 1.5 i2v on the reference image (#1693)
Co-authored-by: SolitaryThinker <wlsaidhi@gmail.com> Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Kyle Hu
co-authored by
SolitaryThinker
Claude Opus 5.5
SolitaryThinker
Claude Opus 5.5
parent
ef62ac48d1
commit
cacc8cfcb3
@@ -160,6 +160,11 @@ FAMILY_COVERAGE = (
|
||||
("test_glm_image.py", ),
|
||||
("test_glm_image_similarity.py", ),
|
||||
),
|
||||
FamilyCoverage(
|
||||
re.compile(r"(^|[/_.-])hunyuan(video)?15([a-z0-9_-]*)([/_.-]|$)"),
|
||||
(),
|
||||
("test_hunyuan15_i2v_similarity.py", ),
|
||||
),
|
||||
FamilyCoverage(
|
||||
re.compile(r"(^|[/_.-])kandinsky[_-]?5([/_.-]|$)"),
|
||||
("test_kandinsky5.py", ),
|
||||
|
||||
@@ -8,7 +8,7 @@ import torch
|
||||
|
||||
from fastvideo.configs.models import DiTConfig, EncoderConfig, VAEConfig
|
||||
from fastvideo.configs.models.dits import HunyuanVideo15Config
|
||||
from fastvideo.configs.models.encoders import (BaseEncoderOutput, Qwen2_5_VLConfig, T5Config)
|
||||
from fastvideo.configs.models.encoders import (BaseEncoderOutput, Qwen2_5_VLConfig, SiglipVisionConfig, T5Config)
|
||||
from fastvideo.configs.models.vaes import Hunyuan15VAEConfig
|
||||
from fastvideo.configs.models.upsamplers import SRTo720pUpsamplerConfig, SRTo1080pUpsamplerConfig
|
||||
from fastvideo.configs.pipelines.base import PipelineConfig, UpsamplerConfig
|
||||
@@ -112,7 +112,7 @@ class Hunyuan15T2V480PConfig(PipelineConfig):
|
||||
|
||||
vae_tiling: bool = True
|
||||
|
||||
def __post_init__(self):
|
||||
def __post_init__(self) -> None:
|
||||
self.vae_config.load_encoder = False
|
||||
self.vae_config.load_decoder = True
|
||||
if self.text_encoder_configs:
|
||||
@@ -121,7 +121,30 @@ class Hunyuan15T2V480PConfig(PipelineConfig):
|
||||
|
||||
|
||||
@dataclass
|
||||
class Hunyuan15I2V480PStepDistilledConfig(Hunyuan15T2V480PConfig):
|
||||
class Hunyuan15I2VConfig(Hunyuan15T2V480PConfig):
|
||||
"""Shared configuration for the HunyuanVideo 1.5 image-to-video checkpoints."""
|
||||
|
||||
# The i2v checkpoints ship a SigLIP vision tower; declaring it here is what
|
||||
# makes the loader build one.
|
||||
image_encoder_config: EncoderConfig = field(default_factory=SiglipVisionConfig)
|
||||
# HYWorldConfig runs the same tower in fp16; fp32 only doubles the
|
||||
# footprint and slows the matmuls.
|
||||
image_encoder_precision: str = "fp16"
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
super().__post_init__()
|
||||
self.vae_config.load_encoder = True
|
||||
|
||||
def check_pipeline_config(self) -> None:
|
||||
super().check_pipeline_config()
|
||||
if not self.vae_config.load_encoder:
|
||||
raise ValueError("HunyuanVideo 1.5 I2V requires the VAE encoder.")
|
||||
|
||||
|
||||
@dataclass
|
||||
class Hunyuan15I2V480PStepDistilledConfig(Hunyuan15I2VConfig):
|
||||
"""480p step-distilled image-to-video checkpoint."""
|
||||
|
||||
flow_shift: int = 7
|
||||
|
||||
|
||||
@@ -134,10 +157,9 @@ class Hunyuan15T2V720PConfig(Hunyuan15T2V480PConfig):
|
||||
|
||||
|
||||
@dataclass
|
||||
class Hunyuan15I2V720PConfig(Hunyuan15T2V720PConfig):
|
||||
"""Base configuration for HunYuan pipeline architecture."""
|
||||
class Hunyuan15I2V720PConfig(Hunyuan15I2VConfig):
|
||||
"""720p distilled image-to-video checkpoint."""
|
||||
|
||||
# HunyuanConfig-specific parameters with defaults
|
||||
flow_shift: int = 7
|
||||
|
||||
|
||||
|
||||
@@ -20,8 +20,12 @@ logger = init_logger(__name__)
|
||||
|
||||
class HunyuanVideo15ImageToVideoPipeline(ComposedPipelineBase):
|
||||
|
||||
# image_encoder and feature_extractor are the SigLIP pair the i2v
|
||||
# checkpoints ship; without them the reference image has no route into the
|
||||
# transformer and the run is text-to-video wearing an i2v label.
|
||||
_required_config_modules = [
|
||||
"text_encoder", "text_encoder_2", "tokenizer", "tokenizer_2", "vae", "transformer", "scheduler"
|
||||
"text_encoder", "text_encoder_2", "tokenizer", "tokenizer_2", "vae", "transformer", "scheduler",
|
||||
"image_encoder", "feature_extractor"
|
||||
]
|
||||
|
||||
def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
|
||||
@@ -47,7 +51,9 @@ class HunyuanVideo15ImageToVideoPipeline(ComposedPipelineBase):
|
||||
transformer=self.get_module("transformer")))
|
||||
|
||||
self.add_stage(stage_name="image_encoding_stage",
|
||||
stage=Hy15ImageEncodingStage(image_encoder=None, image_processor=None))
|
||||
stage=Hy15ImageEncodingStage(image_encoder=self.get_module("image_encoder"),
|
||||
image_processor=self.get_module("feature_extractor"),
|
||||
vae=self.get_module("vae")))
|
||||
|
||||
self.add_stage(stage_name="denoising_stage",
|
||||
stage=DenoisingStage(transformer=self.get_module("transformer"),
|
||||
|
||||
@@ -15,6 +15,7 @@ import torch
|
||||
from fastvideo.distributed import get_local_torch_device
|
||||
from fastvideo.fastvideo_args import ExecutionMode, FastVideoArgs
|
||||
from fastvideo.forward_context import set_forward_context
|
||||
from fastvideo.image_processor import ImageProcessor
|
||||
from fastvideo.logger import init_logger
|
||||
from fastvideo.models.vaes.common import ParallelTiledVAE
|
||||
from fastvideo.models.vision_utils import (get_default_height_width, normalize, numpy_to_pt, pil_to_numpy, resize,
|
||||
@@ -97,22 +98,133 @@ class ImageEncodingStage(PipelineStage):
|
||||
class Hy15ImageEncodingStage(ImageEncodingStage):
|
||||
"""
|
||||
Stage for encoding image prompts into embeddings for HunyuanVideo1.5 models.
|
||||
|
||||
Without a reference image (T2V and the super-resolution pipelines) the
|
||||
conditioning slots are zero-filled and the transformer takes its T2V branch.
|
||||
With one, SigLIP supplies ``image_embeds`` and the VAE supplies the first
|
||||
frame of the conditioning latent, which is what turns that branch off.
|
||||
"""
|
||||
|
||||
def __init__(self, image_encoder=None, image_processor=None, vae=None):
|
||||
super().__init__(image_encoder=image_encoder, image_processor=image_processor)
|
||||
self.vae = vae
|
||||
|
||||
def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
|
||||
"""Verify image encoding stage inputs."""
|
||||
return VerificationResult()
|
||||
|
||||
@torch.no_grad()
|
||||
def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
|
||||
"""
|
||||
Encode the prompt into image encoder hidden states.
|
||||
"""
|
||||
if batch.pil_image is None:
|
||||
batch.image_embeds = [torch.zeros(1, 729, 1152, device=get_local_torch_device())]
|
||||
|
||||
raw_latent_shape = list(batch.raw_latent_shape)
|
||||
raw_latent_shape[1] = 1
|
||||
batch.video_latent = torch.zeros(tuple(raw_latent_shape), device=get_local_torch_device())
|
||||
The reference image is encoded once and reused for the whole latent
|
||||
batch. ``DenoisingStage`` still asserts a batch of one, so the repeat is
|
||||
a no-op today; the stage no longer assumes a batch of one itself.
|
||||
"""
|
||||
device = get_local_torch_device()
|
||||
batch_size = batch.raw_latent_shape[0]
|
||||
|
||||
if batch.pil_image is None:
|
||||
batch.image_embeds = [torch.zeros(batch_size, 729, 1152, device=device)]
|
||||
raw_latent_shape = list(batch.raw_latent_shape)
|
||||
raw_latent_shape[1] = 1
|
||||
batch.video_latent = torch.zeros(tuple(raw_latent_shape), device=device)
|
||||
return batch
|
||||
|
||||
if self.image_encoder is None or self.image_processor is None or self.vae is None:
|
||||
raise ValueError("HunyuanVideo 1.5 needs an image encoder, its processor and the VAE to condition on a "
|
||||
"reference image, but this pipeline was built without them. Image-to-video requires "
|
||||
"HunyuanVideo15ImageToVideoPipeline and an i2v checkpoint.")
|
||||
|
||||
_, latent_channels, latent_temporal, latent_height, latent_width = batch.raw_latent_shape
|
||||
|
||||
# 1. SigLIP hidden states. The transformer reads all-zero image
|
||||
# embeddings as "this is T2V", so this is what selects the I2V branch.
|
||||
self.image_encoder = self.image_encoder.to(device)
|
||||
encoder_dtype = next(self.image_encoder.parameters()).dtype
|
||||
image_inputs = self.image_processor.preprocess(
|
||||
images=batch.pil_image,
|
||||
do_resize=True,
|
||||
do_convert_rgb=True,
|
||||
return_tensors="pt",
|
||||
).to(device=device, dtype=encoder_dtype)
|
||||
with set_forward_context(current_timestep=0, attn_metadata=None):
|
||||
image_embeds = self.image_encoder(**image_inputs).last_hidden_state
|
||||
# The public input is one reference image. Encode it once, then reuse
|
||||
# the conditioning for every latent requested from that image.
|
||||
if image_embeds.shape[0] == 1:
|
||||
image_embeds = image_embeds.repeat(batch_size, 1, 1)
|
||||
elif image_embeds.shape[0] != batch_size:
|
||||
raise ValueError(f"HunyuanVideo 1.5 image embeddings have batch size {image_embeds.shape[0]}, "
|
||||
f"but the latent batch size is {batch_size}.")
|
||||
batch.image_embeds = [image_embeds]
|
||||
if fastvideo_args.image_encoder_cpu_offload:
|
||||
self.image_encoder.to("cpu")
|
||||
|
||||
# 2. The reference image occupies the first latent frame of the
|
||||
# conditioning stream; every later frame is zero and the mask channel
|
||||
# marks which one is real.
|
||||
vae_dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.vae_precision]
|
||||
vae_autocast_enabled = (vae_dtype != torch.float32) and not fastvideo_args.disable_autocast
|
||||
|
||||
# The HunyuanVideo 1.5 VAE compresses space by 16 and latent
|
||||
# preparation divides by the same ratio, so aligning the reference
|
||||
# image to 8 lets a height like 552 reach the encoder and fail
|
||||
# inside its reshape.
|
||||
image_processor = ImageProcessor(
|
||||
vae_scale_factor=fastvideo_args.pipeline_config.vae_config.arch_config.spatial_compression_ratio)
|
||||
pixels = image_processor.preprocess(batch.pil_image, batch.height, batch.width)
|
||||
pixels = pixels.unsqueeze(2).to(device=device, dtype=vae_dtype)
|
||||
|
||||
self.vae = self.vae.to(device)
|
||||
with torch.autocast(device_type="cuda", dtype=vae_dtype, enabled=vae_autocast_enabled):
|
||||
# The decoding stage turns tiling on and nothing turns it back off,
|
||||
# and both stages hold the same VAE. Without this the first
|
||||
# generation in a process encodes the reference image untiled and
|
||||
# every later one encodes it tiled, so the same image and seed give
|
||||
# different conditioning. Decide from the config, as the other
|
||||
# VAE-encoding stages in this file do.
|
||||
if fastvideo_args.pipeline_config.vae_tiling:
|
||||
self.vae.enable_tiling()
|
||||
cond_latents = self.vae.encode(pixels).mode()
|
||||
cond_latents = cond_latents * self.vae.config.scaling_factor
|
||||
if fastvideo_args.vae_cpu_offload:
|
||||
self.vae.to("cpu")
|
||||
|
||||
if cond_latents.shape[0] == 1:
|
||||
cond_latents = cond_latents.repeat(batch_size, 1, 1, 1, 1)
|
||||
elif cond_latents.shape[0] != batch_size:
|
||||
raise ValueError(f"HunyuanVideo 1.5 conditioning latents have batch size {cond_latents.shape[0]}, "
|
||||
f"but the requested latent batch size is {batch_size}.")
|
||||
|
||||
first_frame_latents = cond_latents[:, :, :1]
|
||||
cond_latents = cond_latents.new_zeros(
|
||||
batch_size,
|
||||
latent_channels,
|
||||
latent_temporal,
|
||||
latent_height,
|
||||
latent_width,
|
||||
)
|
||||
cond_latents[:, :, :1] = first_frame_latents
|
||||
|
||||
mask = torch.zeros(
|
||||
batch_size,
|
||||
1,
|
||||
latent_temporal,
|
||||
latent_height,
|
||||
latent_width,
|
||||
device=device,
|
||||
dtype=cond_latents.dtype,
|
||||
)
|
||||
mask[:, :, 0] = 1.0
|
||||
|
||||
# Latents, then conditioning, then mask: the denoising stage appends
|
||||
# ``image_latent`` as one block, so the channel order has to be built
|
||||
# here. Leaving ``video_latent`` unset keeps it out of the earlier
|
||||
# branch, which would append the mask before the conditioning.
|
||||
batch.image_latent = torch.cat([cond_latents, mask], dim=1)
|
||||
assert batch.image_latent.shape[1] == latent_channels + 1
|
||||
return batch
|
||||
|
||||
|
||||
|
||||
@@ -53,6 +53,15 @@ def test_model_family_change_selects_focused_golden_and_ssim(path, goldens):
|
||||
"test_causal_similarity.py,test_wan_i2v_similarity.py,test_wan_t2v_similarity.py")
|
||||
|
||||
|
||||
def test_hunyuan15_change_selects_its_i2v_ssim_test():
|
||||
plan = PLAN_MERGE_CI.classify_paths(["fastvideo/pipelines/basic/hunyuan15/hunyuan15_i2v_pipeline.py"])
|
||||
|
||||
assert plan.encoded_lanes() == ",golden-gate,ssim,"
|
||||
# No HunyuanVideo 1.5 golden test exists yet, so golden coverage stays "all".
|
||||
assert plan.encoded_golden_tests() == "all"
|
||||
assert plan.encoded_ssim_tests() == "test_hunyuan15_i2v_similarity.py"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("path", [
|
||||
"fastvideo/configs/models/dits/wanvideo.py",
|
||||
"fastvideo/models/wan/config.py",
|
||||
|
||||
@@ -0,0 +1,92 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""End-to-end SSIM regression test for HunyuanVideo 1.5 image-to-video.
|
||||
|
||||
The CPU stage tests fake the SigLIP tower and the VAE, so the two things that
|
||||
can only fail against the real checkpoint are unverified there: the loader
|
||||
wiring (the i2v checkpoints ship ``image_encoder``/``feature_extractor``) and
|
||||
the numerics (scaling factor, reference-image resize, mask placement). This
|
||||
test runs the real checkpoint with an image and a prompt and compares the clip
|
||||
against a device-specific reference, which is what catches a regression that
|
||||
silently drops the reference image.
|
||||
"""
|
||||
import os
|
||||
|
||||
import pytest
|
||||
|
||||
from fastvideo.api.sampling_param import SamplingParam
|
||||
from fastvideo.logger import init_logger
|
||||
from fastvideo.tests.ssim.inference_similarity_utils import (
|
||||
resolve_inference_device_reference_folder,
|
||||
run_image_to_video_similarity_test,
|
||||
)
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
REQUIRED_GPUS = 2
|
||||
|
||||
device_reference_folder = resolve_inference_device_reference_folder(logger)
|
||||
|
||||
HUNYUAN15_I2V_PARAMS = {
|
||||
"num_gpus": 2,
|
||||
"model_path": "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v_step_distilled",
|
||||
"height": 480,
|
||||
"width": 848,
|
||||
"num_frames": 45,
|
||||
"num_inference_steps": 6,
|
||||
"guidance_scale": 1.0,
|
||||
"seed": 1024,
|
||||
"sp_size": 2,
|
||||
"tp_size": 1,
|
||||
"fps": 24,
|
||||
}
|
||||
_HUNYUAN15_I2V_FULL_QUALITY_DEFAULTS = SamplingParam.from_pretrained(HUNYUAN15_I2V_PARAMS["model_path"])
|
||||
HUNYUAN15_I2V_FULL_QUALITY_PARAMS = {
|
||||
"num_gpus": HUNYUAN15_I2V_PARAMS["num_gpus"],
|
||||
"model_path": HUNYUAN15_I2V_PARAMS["model_path"],
|
||||
"height": _HUNYUAN15_I2V_FULL_QUALITY_DEFAULTS.height,
|
||||
"width": _HUNYUAN15_I2V_FULL_QUALITY_DEFAULTS.width,
|
||||
"num_frames": _HUNYUAN15_I2V_FULL_QUALITY_DEFAULTS.num_frames,
|
||||
"num_inference_steps": _HUNYUAN15_I2V_FULL_QUALITY_DEFAULTS.num_inference_steps,
|
||||
"guidance_scale": _HUNYUAN15_I2V_FULL_QUALITY_DEFAULTS.guidance_scale,
|
||||
"seed": _HUNYUAN15_I2V_FULL_QUALITY_DEFAULTS.seed,
|
||||
"sp_size": HUNYUAN15_I2V_PARAMS["sp_size"],
|
||||
"tp_size": HUNYUAN15_I2V_PARAMS["tp_size"],
|
||||
"fps": _HUNYUAN15_I2V_FULL_QUALITY_DEFAULTS.fps,
|
||||
}
|
||||
|
||||
HUNYUAN15_I2V_MODEL_TO_PARAMS = {
|
||||
"HunyuanVideo-1.5-Diffusers-480p_i2v_step_distilled": HUNYUAN15_I2V_PARAMS,
|
||||
}
|
||||
FULL_QUALITY_HUNYUAN15_I2V_MODEL_TO_PARAMS = {
|
||||
"HunyuanVideo-1.5-Diffusers-480p_i2v_step_distilled": HUNYUAN15_I2V_FULL_QUALITY_PARAMS,
|
||||
}
|
||||
|
||||
HUNYUAN15_I2V_TEST_CASES = [
|
||||
(
|
||||
"An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space realised in the background. High quality, ultrarealistic detail and breath-taking movie-like camera shot.",
|
||||
"https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg",
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("prompt", "image_path"), HUNYUAN15_I2V_TEST_CASES)
|
||||
@pytest.mark.parametrize("attention_backend_name", ["FLASH_ATTN"])
|
||||
@pytest.mark.parametrize("model_id", list(HUNYUAN15_I2V_MODEL_TO_PARAMS.keys()))
|
||||
def test_hunyuan15_i2v_inference_similarity(
|
||||
prompt: str,
|
||||
image_path: str,
|
||||
attention_backend_name: str,
|
||||
model_id: str,
|
||||
) -> None:
|
||||
run_image_to_video_similarity_test(
|
||||
logger=logger,
|
||||
script_dir=os.path.dirname(os.path.abspath(__file__)),
|
||||
device_reference_folder=device_reference_folder,
|
||||
prompt=prompt,
|
||||
image_path=image_path,
|
||||
attention_backend_name=attention_backend_name,
|
||||
model_id=model_id,
|
||||
default_params_map=HUNYUAN15_I2V_MODEL_TO_PARAMS,
|
||||
full_quality_params_map=FULL_QUALITY_HUNYUAN15_I2V_MODEL_TO_PARAMS,
|
||||
min_acceptable_ssim=0.95,
|
||||
)
|
||||
@@ -0,0 +1,319 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""CPU tests for ``Hy15ImageEncodingStage``.
|
||||
|
||||
The failure this guards against is quiet: with no reference image conditioning,
|
||||
HunyuanVideo 1.5 still trains and still samples, it just ignores the image. So
|
||||
these assert the channel layout rather than only that a tensor came back.
|
||||
|
||||
Ordering matters and is invisible in T2V. ``DenoisingStage`` appends
|
||||
``image_latent`` as one block, so the 65 channel input is
|
||||
``[latents 32][conditioning 32][mask 1]``. For T2V the last 33 are all zero and
|
||||
either order produces the same tensor, which is why a swapped layout would go
|
||||
unnoticed until an image is actually supplied.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import PIL.Image
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
import fastvideo.models.vaes.hunyuan15vae as hunyuan15vae
|
||||
from fastvideo.configs.pipelines.hunyuan15 import (Hunyuan15I2V480PStepDistilledConfig, Hunyuan15I2V720PConfig,
|
||||
Hunyuan15SR1080PConfig, Hunyuan15T2V480PConfig,
|
||||
Hunyuan15T2V720PConfig)
|
||||
from fastvideo.models.vaes.hunyuan15vae import AutoencoderKLHunyuanVideo15
|
||||
from fastvideo.pipelines.stages.image_encoding import Hy15ImageEncodingStage
|
||||
|
||||
VAE_SPATIAL_COMPRESSION_RATIO = 16
|
||||
LATENT_SHAPE = (1, 32, 5, 4, 6) # B, C, T, H, W
|
||||
REFERENCE_IMAGE = PIL.Image.new(
|
||||
"RGB",
|
||||
(LATENT_SHAPE[4] * VAE_SPATIAL_COMPRESSION_RATIO, LATENT_SHAPE[3] * VAE_SPATIAL_COMPRESSION_RATIO),
|
||||
(120, 60, 30),
|
||||
)
|
||||
VISION_TOKENS = 729
|
||||
VISION_DIM = 1152
|
||||
|
||||
|
||||
class _FakeSiglip(torch.nn.Module):
|
||||
"""Returns a constant non-zero hidden state, which is all the stage reads."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.marker = torch.nn.Parameter(torch.ones(1))
|
||||
self.forward_calls = 0
|
||||
|
||||
def forward(self, pixel_values: torch.Tensor):
|
||||
self.forward_calls += 1
|
||||
batch = pixel_values.shape[0]
|
||||
return SimpleNamespace(last_hidden_state=torch.full((batch, VISION_TOKENS, VISION_DIM), 0.5))
|
||||
|
||||
|
||||
class _FakeProcessor:
|
||||
|
||||
def preprocess(self, images, **kwargs):
|
||||
return SimpleNamespace(to=lambda **_: {"pixel_values": torch.zeros(1, 3, 384, 384)})
|
||||
|
||||
|
||||
class _ConstantEncoder(torch.nn.Module):
|
||||
"""Small encoder used inside the real HunyuanVideo 1.5 VAE wrapper."""
|
||||
|
||||
def __init__(self, out_channels: int, spatial_compression_ratio: int, **kwargs) -> None:
|
||||
super().__init__()
|
||||
self.out_channels = out_channels
|
||||
self.spatial_compression_ratio = spatial_compression_ratio
|
||||
self.forward_calls = 0
|
||||
self.pixel_shapes: list[tuple[int, ...]] = []
|
||||
|
||||
def forward(self, pixels: torch.Tensor) -> torch.Tensor:
|
||||
self.forward_calls += 1
|
||||
self.pixel_shapes.append(tuple(pixels.shape))
|
||||
batch_size, _, _, height, width = pixels.shape
|
||||
# The real encoder packs 2x2 spatial blocks at each of its four stages,
|
||||
# so a side that is a multiple of 8 but not 16 blows up in its view().
|
||||
assert height % 16 == 0, height
|
||||
assert width % 16 == 0, width
|
||||
return torch.full(
|
||||
(
|
||||
batch_size,
|
||||
self.out_channels,
|
||||
1,
|
||||
height // self.spatial_compression_ratio,
|
||||
width // self.spatial_compression_ratio,
|
||||
),
|
||||
3.0,
|
||||
device=pixels.device,
|
||||
dtype=pixels.dtype,
|
||||
)
|
||||
|
||||
|
||||
class _UnusedDecoder(torch.nn.Module):
|
||||
|
||||
def __init__(self, **kwargs) -> None:
|
||||
super().__init__()
|
||||
|
||||
def forward(self, latents: torch.Tensor) -> torch.Tensor:
|
||||
return latents
|
||||
|
||||
|
||||
def _batch(pil_image, batch_size: int = 1, height: int | None = None, width: int | None = None):
|
||||
height = height or LATENT_SHAPE[3] * VAE_SPATIAL_COMPRESSION_RATIO
|
||||
width = width or LATENT_SHAPE[4] * VAE_SPATIAL_COMPRESSION_RATIO
|
||||
raw_latent_shape = (batch_size, LATENT_SHAPE[1], LATENT_SHAPE[2], height // VAE_SPATIAL_COMPRESSION_RATIO,
|
||||
width // VAE_SPATIAL_COMPRESSION_RATIO)
|
||||
return SimpleNamespace(
|
||||
pil_image=pil_image,
|
||||
raw_latent_shape=raw_latent_shape,
|
||||
image_embeds=[],
|
||||
image_latent=None,
|
||||
video_latent=None,
|
||||
height=height,
|
||||
width=width,
|
||||
)
|
||||
|
||||
|
||||
def _args(vae_tiling: bool = True):
|
||||
return SimpleNamespace(
|
||||
pipeline_config=SimpleNamespace(
|
||||
vae_precision="fp32",
|
||||
vae_tiling=vae_tiling,
|
||||
vae_config=SimpleNamespace(
|
||||
arch_config=SimpleNamespace(spatial_compression_ratio=VAE_SPATIAL_COMPRESSION_RATIO)),
|
||||
),
|
||||
disable_autocast=True,
|
||||
image_encoder_cpu_offload=False,
|
||||
vae_cpu_offload=False,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _cpu_device(monkeypatch):
|
||||
"""Pin the stage to CPU so the fakes and the real code agree on device.
|
||||
|
||||
``get_local_torch_device`` resolves to MPS on macOS and CUDA elsewhere, and
|
||||
nothing here needs an accelerator.
|
||||
"""
|
||||
monkeypatch.setattr(
|
||||
"fastvideo.pipelines.stages.image_encoding.get_local_torch_device",
|
||||
lambda: torch.device("cpu"),
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def lightweight_vae_factory(monkeypatch):
|
||||
monkeypatch.setattr(hunyuan15vae, "HunyuanVideo15Encoder3D", _ConstantEncoder)
|
||||
monkeypatch.setattr(hunyuan15vae, "HunyuanVideo15Decoder3D", _UnusedDecoder)
|
||||
|
||||
def build(pipeline_config=None):
|
||||
if pipeline_config is None:
|
||||
pipeline_config = Hunyuan15I2V480PStepDistilledConfig()
|
||||
return AutoencoderKLHunyuanVideo15(pipeline_config.vae_config)
|
||||
|
||||
return build
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def i2v_stage(lightweight_vae_factory):
|
||||
return Hy15ImageEncodingStage(
|
||||
image_encoder=_FakeSiglip(),
|
||||
image_processor=_FakeProcessor(),
|
||||
vae=lightweight_vae_factory(),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"config_cls",
|
||||
[Hunyuan15I2V480PStepDistilledConfig, Hunyuan15I2V720PConfig],
|
||||
)
|
||||
def test_i2v_configs_build_the_vae_encoder_and_decoder(config_cls, lightweight_vae_factory) -> None:
|
||||
pipeline_config = config_cls()
|
||||
vae = lightweight_vae_factory(pipeline_config)
|
||||
|
||||
assert pipeline_config.vae_config.load_encoder is True
|
||||
assert pipeline_config.vae_config.load_decoder is True
|
||||
assert pipeline_config.text_encoder_configs[0].arch_config.output_hidden_states is True
|
||||
assert isinstance(vae.encoder, _ConstantEncoder)
|
||||
assert isinstance(vae.decoder, _UnusedDecoder)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"config_cls",
|
||||
[Hunyuan15I2V480PStepDistilledConfig, Hunyuan15I2V720PConfig],
|
||||
)
|
||||
def test_i2v_configs_run_the_image_encoder_in_fp16(config_cls) -> None:
|
||||
# fp32 is the PipelineConfig default; the SigLIP tower only doubles its
|
||||
# footprint for it. HYWorldConfig runs the same tower in fp16.
|
||||
assert config_cls().image_encoder_precision == "fp16"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"config_cls",
|
||||
[Hunyuan15I2V480PStepDistilledConfig, Hunyuan15I2V720PConfig],
|
||||
)
|
||||
def test_i2v_configs_reject_disabling_the_vae_encoder(config_cls) -> None:
|
||||
pipeline_config = config_cls()
|
||||
pipeline_config.update_config_from_dict({"vae_config.load_encoder": False})
|
||||
|
||||
with pytest.raises(ValueError, match="requires the VAE encoder"):
|
||||
pipeline_config.check_pipeline_config()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"config_cls",
|
||||
[Hunyuan15T2V480PConfig, Hunyuan15T2V720PConfig, Hunyuan15SR1080PConfig],
|
||||
)
|
||||
def test_non_i2v_configs_keep_the_vae_encoder_disabled(config_cls) -> None:
|
||||
assert config_cls().vae_config.load_encoder is False
|
||||
|
||||
|
||||
class TestTextToVideoPath:
|
||||
"""No reference image: the transformer must still see the T2V layout."""
|
||||
|
||||
def test_image_embeds_are_zero(self) -> None:
|
||||
batch = Hy15ImageEncodingStage().forward(_batch(None), _args())
|
||||
|
||||
assert len(batch.image_embeds) == 1
|
||||
assert batch.image_embeds[0].shape == (1, VISION_TOKENS, VISION_DIM)
|
||||
# The transformer keys its T2V branch off this being all zero.
|
||||
assert torch.all(batch.image_embeds[0] == 0)
|
||||
|
||||
def test_conditioning_stays_on_the_video_latent_slot(self) -> None:
|
||||
batch = Hy15ImageEncodingStage().forward(_batch(None), _args())
|
||||
|
||||
# The super-resolution stages read video_latent directly, so T2V keeps
|
||||
# populating it rather than switching to image_latent.
|
||||
assert batch.video_latent is not None
|
||||
assert batch.video_latent.shape == (1, 1) + LATENT_SHAPE[2:]
|
||||
assert torch.all(batch.video_latent == 0)
|
||||
assert batch.image_latent is None
|
||||
|
||||
def test_placeholders_follow_the_latent_batch(self) -> None:
|
||||
batch_size = 2
|
||||
batch = Hy15ImageEncodingStage().forward(_batch(None, batch_size=batch_size), _args())
|
||||
|
||||
assert batch.image_embeds[0].shape == (batch_size, VISION_TOKENS, VISION_DIM)
|
||||
assert batch.video_latent.shape == (batch_size, 1) + LATENT_SHAPE[2:]
|
||||
|
||||
|
||||
class TestImageToVideoPath:
|
||||
|
||||
def test_image_embeds_come_from_siglip(self, i2v_stage) -> None:
|
||||
batch = i2v_stage.forward(_batch(REFERENCE_IMAGE), _args())
|
||||
|
||||
assert batch.image_embeds[0].shape == (1, VISION_TOKENS, VISION_DIM)
|
||||
# Non-zero is the whole point: all-zero would re-select the T2V branch.
|
||||
assert not torch.all(batch.image_embeds[0] == 0)
|
||||
|
||||
def test_channel_layout_is_conditioning_then_mask(self, i2v_stage) -> None:
|
||||
batch = i2v_stage.forward(_batch(REFERENCE_IMAGE), _args())
|
||||
|
||||
assert batch.image_latent.shape == (1, LATENT_SHAPE[1] + 1) + LATENT_SHAPE[2:]
|
||||
|
||||
conditioning = batch.image_latent[:, :LATENT_SHAPE[1]]
|
||||
mask = batch.image_latent[:, LATENT_SHAPE[1]:]
|
||||
|
||||
# A swapped layout would put a 1-wide block first, so check that the
|
||||
# trailing channel is the mask and not part of the conditioning.
|
||||
assert torch.all(mask[:, :, 0] == 1.0)
|
||||
assert torch.all(mask[:, :, 1:] == 0.0)
|
||||
assert torch.all(conditioning[:, :, 0] != 0)
|
||||
assert torch.all(conditioning[:, :, 1:] == 0)
|
||||
|
||||
def test_scaling_factor_is_applied(self, i2v_stage) -> None:
|
||||
batch = i2v_stage.forward(_batch(REFERENCE_IMAGE), _args())
|
||||
|
||||
expected = 3.0 * i2v_stage.vae.config.scaling_factor
|
||||
assert torch.allclose(batch.image_latent[:, :LATENT_SHAPE[1], 0],
|
||||
torch.full((1, LATENT_SHAPE[1]) + LATENT_SHAPE[3:], expected))
|
||||
|
||||
def test_video_latent_left_unset(self, i2v_stage) -> None:
|
||||
batch = i2v_stage.forward(_batch(REFERENCE_IMAGE), _args())
|
||||
|
||||
# DenoisingStage checks video_latent first and would append the mask
|
||||
# ahead of the conditioning, so this slot has to stay empty.
|
||||
assert batch.video_latent is None
|
||||
|
||||
def test_reference_image_is_aligned_to_the_vae_compression_ratio(self, i2v_stage) -> None:
|
||||
# 552 is a multiple of 8, which is all InputValidationStage checks, but
|
||||
# not of 16, which the encoder's 2x2 rearrange needs, so the stage has
|
||||
# to hand the VAE 544 -- the same ratio the latent preparation divides
|
||||
# by. Tiling is off so the encoder sees the aligned image and not tiles.
|
||||
height, width = 552, 96
|
||||
batch = i2v_stage.forward(_batch(REFERENCE_IMAGE, height=height, width=width), _args(vae_tiling=False))
|
||||
|
||||
assert i2v_stage.vae.encoder.pixel_shapes == [(1, 3, 1, 544, width)]
|
||||
assert batch.image_latent.shape == (1, LATENT_SHAPE[1] + 1, LATENT_SHAPE[2],
|
||||
height // VAE_SPATIAL_COMPRESSION_RATIO,
|
||||
width // VAE_SPATIAL_COMPRESSION_RATIO)
|
||||
|
||||
def test_single_reference_conditioning_repeats_to_the_latent_batch(self, i2v_stage) -> None:
|
||||
batch_size = 2
|
||||
batch = i2v_stage.forward(_batch(REFERENCE_IMAGE, batch_size=batch_size), _args())
|
||||
|
||||
expected_shape = (batch_size, LATENT_SHAPE[1] + 1) + LATENT_SHAPE[2:]
|
||||
assert batch.image_latent.shape == expected_shape
|
||||
assert batch.image_embeds[0].shape == (batch_size, VISION_TOKENS, VISION_DIM)
|
||||
assert torch.equal(batch.image_latent[0], batch.image_latent[1])
|
||||
assert torch.equal(batch.image_embeds[0][0], batch.image_embeds[0][1])
|
||||
|
||||
conditioning = batch.image_latent[:, :LATENT_SHAPE[1]]
|
||||
mask = batch.image_latent[:, LATENT_SHAPE[1]:]
|
||||
assert torch.all(conditioning[:, :, 0] != 0)
|
||||
assert torch.all(conditioning[:, :, 1:] == 0)
|
||||
assert torch.all(mask[:, :, 0] == 1)
|
||||
assert torch.all(mask[:, :, 1:] == 0)
|
||||
|
||||
assert i2v_stage.vae.encoder.forward_calls == 1
|
||||
assert i2v_stage.image_encoder.forward_calls == 1
|
||||
|
||||
noise_latents = torch.zeros((batch_size, ) + LATENT_SHAPE[1:])
|
||||
assert torch.cat([noise_latents, batch.image_latent], dim=1).shape == (
|
||||
batch_size,
|
||||
LATENT_SHAPE[1] * 2 + 1,
|
||||
) + LATENT_SHAPE[2:]
|
||||
|
||||
def test_missing_encoder_is_an_explicit_error(self) -> None:
|
||||
with pytest.raises(ValueError, match="image encoder"):
|
||||
Hy15ImageEncodingStage().forward(_batch(REFERENCE_IMAGE), _args())
|
||||
Reference in New Issue
Block a user