From cacc8cfcb3bbd7b9a09caa02b7c400cd5f46e3e8 Mon Sep 17 00:00:00 2001 From: Kyle Hu <160547796+KyleNeverGivesUp@users.noreply.github.com> Date: Wed, 7 Oct 2026 01:36:35 -0700 Subject: [PATCH] [bugfix]: condition HunyuanVideo 1.5 i2v on the reference image (#1693) Co-authored-by: SolitaryThinker Co-authored-by: Claude Opus 5.5 (1M context) --- .github/scripts/plan_merge_ci.py | 5 + fastvideo/configs/pipelines/hunyuan15.py | 34 +- .../basic/hunyuan15/hunyuan15_i2v_pipeline.py | 10 +- fastvideo/pipelines/stages/image_encoding.py | 124 ++++++- .../tests/contract/test_merge_ci_plan.py | 9 + .../ssim/test_hunyuan15_i2v_similarity.py | 92 +++++ .../stages/test_hy15_image_encoding_stage.py | 319 ++++++++++++++++++ 7 files changed, 579 insertions(+), 14 deletions(-) create mode 100644 fastvideo/tests/ssim/test_hunyuan15_i2v_similarity.py create mode 100644 fastvideo/tests/stages/test_hy15_image_encoding_stage.py diff --git a/.github/scripts/plan_merge_ci.py b/.github/scripts/plan_merge_ci.py index db79723ce..7025d81e7 100644 --- a/.github/scripts/plan_merge_ci.py +++ b/.github/scripts/plan_merge_ci.py @@ -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", ), diff --git a/fastvideo/configs/pipelines/hunyuan15.py b/fastvideo/configs/pipelines/hunyuan15.py index b8a06bae3..266b32452 100644 --- a/fastvideo/configs/pipelines/hunyuan15.py +++ b/fastvideo/configs/pipelines/hunyuan15.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 diff --git a/fastvideo/pipelines/basic/hunyuan15/hunyuan15_i2v_pipeline.py b/fastvideo/pipelines/basic/hunyuan15/hunyuan15_i2v_pipeline.py index d839f955d..1837aafe9 100644 --- a/fastvideo/pipelines/basic/hunyuan15/hunyuan15_i2v_pipeline.py +++ b/fastvideo/pipelines/basic/hunyuan15/hunyuan15_i2v_pipeline.py @@ -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"), diff --git a/fastvideo/pipelines/stages/image_encoding.py b/fastvideo/pipelines/stages/image_encoding.py index 1e2be0443..6ebb6db86 100644 --- a/fastvideo/pipelines/stages/image_encoding.py +++ b/fastvideo/pipelines/stages/image_encoding.py @@ -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 diff --git a/fastvideo/tests/contract/test_merge_ci_plan.py b/fastvideo/tests/contract/test_merge_ci_plan.py index 8cfce78ee..967f1ba4a 100644 --- a/fastvideo/tests/contract/test_merge_ci_plan.py +++ b/fastvideo/tests/contract/test_merge_ci_plan.py @@ -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", diff --git a/fastvideo/tests/ssim/test_hunyuan15_i2v_similarity.py b/fastvideo/tests/ssim/test_hunyuan15_i2v_similarity.py new file mode 100644 index 000000000..eae6318ee --- /dev/null +++ b/fastvideo/tests/ssim/test_hunyuan15_i2v_similarity.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, + ) diff --git a/fastvideo/tests/stages/test_hy15_image_encoding_stage.py b/fastvideo/tests/stages/test_hy15_image_encoding_stage.py new file mode 100644 index 000000000..54cfb4603 --- /dev/null +++ b/fastvideo/tests/stages/test_hy15_image_encoding_stage.py @@ -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())