fix: handle Helios output boundary cases

Author: FlappyBob

fastvideo/configs/pipelines/helios.py

@@ -46,7 +46,7 @@ def helios_postprocess_text(output: BaseEncoderOutput) -> torch.Tensor:
     sequence_lengths = output.attention_mask.gt(0).sum(dim=1).long()
     if torch.isnan(hidden_states).any():
         raise ValueError("Helios UMT5 produced NaN hidden states")
-    trimmed = [hidden[:length] for hidden, length in zip(hidden_states, sequence_lengths, strict=True)]
+    trimmed = [hidden[:min(int(length), 512)] for hidden, length in zip(hidden_states, sequence_lengths, strict=True)]
     return torch.stack(
         [torch.cat([hidden, hidden.new_zeros(512 - hidden.shape[0], hidden.shape[1])]) for hidden in trimmed],
         dim=0,

fastvideo/pipelines/basic/helios/stages.py

@@ -382,7 +382,11 @@ def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> Forward
         frames = torch.cat(decoded_chunks, dim=2)
 
         temporal_scale = fastvideo_args.pipeline_config.vae_config.arch_config.scale_factor_temporal
-        generated_frames = get_generated_pixel_frames(frames.shape[2], temporal_scale)
+        assert isinstance(batch.num_frames, int)
+        generated_frames = min(
+            batch.num_frames,
+            get_generated_pixel_frames(frames.shape[2], temporal_scale),
+        )
         batch.output = frames[:, :, :generated_frames].detach().to(dtype=torch.float32, device="cpu")
         batch.latents = None
         batch.helios_latent_chunks = None

tests/local_tests/pipelines/test_helios_pipeline_smoke.py

@@ -196,6 +196,12 @@ def test_helios_pipeline_config_and_text_contract() -> None:
     assert torch.count_nonzero(output[0, 3:]) == 0
     assert torch.equal(output[1, :5], hidden[1, :5])
 
+    long_hidden = torch.arange(520 * 4, dtype=torch.float32).reshape(1, 520, 4)
+    long_mask = torch.ones(1, 520, dtype=torch.long)
+    long_output = postprocess(BaseEncoderOutput(last_hidden_state=long_hidden, attention_mask=long_mask))
+    assert long_output.shape == (1, 512, 4)
+    assert torch.equal(long_output[0], long_hidden[0, :512])
+
 
 def test_helios_preset_matches_official_distilled_defaults() -> None:
     _, _, _, _, preset = _pipeline_symbols()

tests/local_tests/pipelines/test_helios_pipeline_stages.py

@@ -279,7 +279,7 @@ def decode(self, latent):
 
 
 vae = TinyVAE().to(device)
-decode_batch = ForwardBatch(data_type="video")
+decode_batch = ForwardBatch(data_type="video", num_frames=34)
 decode_batch.latents = torch.cat([expected, expected], dim=2)
 decode_batch.helios_latent_chunks = [expected, expected]
 decode_args = SimpleNamespace(
@@ -370,10 +370,10 @@ def test_tiny_pyramid_stage_uses_history_on_second_chunk():
     assert abs(result["autoregressive_second_short_prefix_mean"]) > 1e-5
 
 
-def test_chunk_decoder_calls_vae_per_chunk_and_matches_frame_rounding():
+def test_chunk_decoder_calls_vae_per_chunk_and_trims_to_requested_frames():
     result = _results()
     assert result["vae_calls"] == [[1, 2, 9, 8, 8], [1, 2, 9, 8, 8]]
-    assert result["decoded_shape"] == [1, 3, 65, 64, 64]
+    assert result["decoded_shape"] == [1, 3, 34, 64, 64]
     assert result["decoded_device"] == "cpu"
     assert result["decoded_first_mean"] == 0
     assert result["decoded_second_mean"] == 1