[bugfix]: load the VAE encoder for HunyuanVideo 1.5 i2v and honor the latent batch

Both I2V configs inherited Hunyuan15T2V480PConfig.__post_init__, which sets
vae_config.load_encoder = False. The i2v image-encoding stage calls
self.vae.encode, and AutoencoderKLHunyuanVideo15 builds its encoder only when
load_encoder is set, so every real checkpoint raised AttributeError before
denoising. The shipped test hid this behind a fake VAE. Both I2V configs now
enable the encoder after super().__post_init__ and check_pipeline_config
rejects a config that turns it back off.

The stage also assumed a batch of one. SigLIP embeddings, the VAE conditioning
latents, and the mask now expand to the latent batch, and an incompatible
batch raises instead of broadcasting wrong. The generic denoising stage still
asserts batch size one, so this fixes the stage contract, not end-to-end
batched generation.

The encode path gained two guards it could not reach before this fix made it
live. Tiling now follows pipeline_config.vae_tiling like every other VAE
encoding stage in the file, since the decoding stage enables tiling on the
shared VAE and nothing turns it off, which made the first and later
generations encode the same reference differently. The reference image is
aligned to the VAE's spatial compression ratio of 16 rather than a hardcoded
8, which let a height like 552 crash inside the encoder reshape.

The test builds the real AutoencoderKLHunyuanVideo15 from the real I2V config
with a lightweight encoder and decoder, so the load_encoder wiring is what is
under test. Reverting the production half fails ten test ids, five of them
with the original AttributeError.

Author: KyleNeverGivesUp

fastvideo/configs/pipelines/hunyuan15.py

@@ -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:
@@ -128,6 +128,15 @@ class Hunyuan15I2V480PStepDistilledConfig(Hunyuan15T2V480PConfig):
     # makes the loader build one.
     image_encoder_config: EncoderConfig = field(default_factory=SiglipVisionConfig)
 
+    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 Hunyuan15T2V720PConfig(Hunyuan15T2V480PConfig):
@@ -146,6 +155,15 @@ class Hunyuan15I2V720PConfig(Hunyuan15T2V720PConfig):
 
     image_encoder_config: EncoderConfig = field(default_factory=SiglipVisionConfig)
 
+    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 Hunyuan15SR1080PConfig(Hunyuan15T2V720PConfig):

fastvideo/pipelines/stages/image_encoding.py

@@ -119,9 +119,10 @@ def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> Forward
         Encode the prompt into image encoder hidden states.
         """
         device = get_local_torch_device()
+        batch_size = batch.raw_latent_shape[0]
 
         if batch.pil_image is None:
-            batch.image_embeds = [torch.zeros(1, 729, 1152, device=device)]
+            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)
@@ -145,7 +146,15 @@ def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> Forward
             return_tensors="pt",
         ).to(device=device, dtype=encoder_dtype)
         with set_forward_context(current_timestep=0, attn_metadata=None):
-            batch.image_embeds = [self.image_encoder(**image_inputs).last_hidden_state]
+            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")
 
@@ -155,22 +164,48 @@ def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> Forward
         vae_dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.vae_precision]
         vae_autocast_enabled = (vae_dtype != torch.float32) and not fastvideo_args.disable_autocast
 
-        image_processor = ImageProcessor(vae_scale_factor=8)
+        # 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")
 
-        cond_latents = cond_latents.repeat(1, 1, latent_temporal, 1, 1)
-        cond_latents[:, :, 1:] = 0.0
+        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(
-            1,
+            batch_size,
             1,
             latent_temporal,
             latent_height,

fastvideo/tests/stages/test_hy15_image_encoding_stage.py

@@ -19,10 +19,20 @@
 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
 
-LATENT_SHAPE = (1, 32, 5, 8, 12)  # B, C, T, H, W
-REFERENCE_IMAGE = PIL.Image.new("RGB", (LATENT_SHAPE[4] * 8, LATENT_SHAPE[3] * 8), (120, 60, 30))
+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
 
@@ -33,8 +43,10 @@ class _FakeSiglip(torch.nn.Module):
     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))
 
@@ -45,37 +57,62 @@ def preprocess(self, images, **kwargs):
         return SimpleNamespace(to=lambda **_: {"pixel_values": torch.zeros(1, 3, 384, 384)})
 
 
-class _FakeVAE(torch.nn.Module):
-    """Encodes to a constant so the first frame is distinguishable from zero."""
+class _ConstantEncoder(torch.nn.Module):
+    """Small encoder used inside the real HunyuanVideo 1.5 VAE wrapper."""
 
-    def __init__(self) -> None:
+    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
+
+    def forward(self, pixels: torch.Tensor) -> torch.Tensor:
+        self.forward_calls += 1
+        batch_size, _, _, height, width = pixels.shape
+        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__()
-        self.config = SimpleNamespace(scaling_factor=2.0)
-
-    def encode(self, pixels: torch.Tensor):
-        _, _, _, height, width = pixels.shape
-        latent = torch.full((1, LATENT_SHAPE[1], 1, height // 8, width // 8), 3.0)
-        return SimpleNamespace(mode=lambda: latent)
 
-    def to(self, *args, **kwargs):
-        return self
+    def forward(self, latents: torch.Tensor) -> torch.Tensor:
+        return latents
 
 
-def _batch(pil_image):
+def _batch(pil_image, batch_size: int = 1):
+    raw_latent_shape = (batch_size, ) + LATENT_SHAPE[1:]
     return SimpleNamespace(
         pil_image=pil_image,
-        raw_latent_shape=LATENT_SHAPE,
+        raw_latent_shape=raw_latent_shape,
         image_embeds=[],
         image_latent=None,
         video_latent=None,
-        height=LATENT_SHAPE[3] * 8,
-        width=LATENT_SHAPE[4] * 8,
+        height=LATENT_SHAPE[3] * VAE_SPATIAL_COMPRESSION_RATIO,
+        width=LATENT_SHAPE[4] * VAE_SPATIAL_COMPRESSION_RATIO,
     )
 
 
 def _args():
     return SimpleNamespace(
-        pipeline_config=SimpleNamespace(vae_precision="fp32"),
+        pipeline_config=SimpleNamespace(
+            vae_precision="fp32",
+            vae_tiling=True,
+            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,
@@ -95,14 +132,63 @@ def _cpu_device(monkeypatch):
     )
 
 
-def _i2v_stage():
+@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=_FakeVAE(),
+        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_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."""
 
@@ -124,18 +210,25 @@ def test_conditioning_stays_on_the_video_latent_slot(self) -> None:
         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) -> None:
-        batch = _i2v_stage().forward(_batch(REFERENCE_IMAGE), _args())
+    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) -> None:
-        batch = _i2v_stage().forward(_batch(REFERENCE_IMAGE), _args())
+    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:]
 
@@ -149,20 +242,46 @@ def test_channel_layout_is_conditioning_then_mask(self) -> None:
         assert torch.all(conditioning[:, :, 0] != 0)
         assert torch.all(conditioning[:, :, 1:] == 0)
 
-    def test_scaling_factor_is_applied(self) -> None:
-        batch = _i2v_stage().forward(_batch(REFERENCE_IMAGE), _args())
+    def test_scaling_factor_is_applied(self, i2v_stage) -> None:
+        batch = i2v_stage.forward(_batch(REFERENCE_IMAGE), _args())
 
-        # Fake VAE emits 3.0 and declares scaling_factor 2.0.
+        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:], 6.0))
+                              torch.full((1, LATENT_SHAPE[1]) + LATENT_SHAPE[3:], expected))
 
-    def test_video_latent_left_unset(self) -> None:
-        batch = _i2v_stage().forward(_batch(REFERENCE_IMAGE), _args())
+    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_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())