[misc] Add Helios-Distilled inference example

Author: FlappyBob

examples/inference/basic/basic_helios_distilled_t2v.py

@@ -0,0 +1,85 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Generate one Helios-Distilled T2V chunk through FastVideo's typed API.
+
+Set ``HELIOS_MODEL_PATH`` to a local snapshot to avoid downloading the public
+checkpoint again. The 33-frame example is a short integration run; increase
+``num_frames`` to 240 for the official eight-chunk default.
+"""
+
+from __future__ import annotations
+
+import os
+
+from fastvideo import VideoGenerator
+from fastvideo.api import (
+    EngineConfig,
+    GenerationRequest,
+    GeneratorConfig,
+    OffloadConfig,
+    OutputConfig,
+    SamplingConfig,
+)
+
+MODEL_PATH = os.getenv("HELIOS_MODEL_PATH", "BestWishYsh/Helios-Distilled")
+OUTPUT_PATH = os.getenv(
+    "HELIOS_OUTPUT_PATH",
+    "outputs_video/helios/helios_distilled_t2v.mp4",
+)
+PROMPT = ("A vibrant tropical fish swims gracefully through a colorful coral reef "
+          "in clear turquoise water, cinematic close-up, fluid motion, vivid detail.")
+NEGATIVE_PROMPT = ("Bright tones, overexposed, static, blurred details, subtitles, paintings, "
+                   "images, overall gray, worst quality, low quality, JPEG artifacts, ugly, "
+                   "deformed, disfigured, still picture, messy background.")
+
+
+def main() -> None:
+    generator = VideoGenerator.from_config(
+        GeneratorConfig(
+            model_path=MODEL_PATH,
+            engine=EngineConfig(
+                num_gpus=1,
+                use_fsdp_inference=False,
+                offload=OffloadConfig(
+                    dit=False,
+                    dit_layerwise=True,
+                    text_encoder=True,
+                    vae=True,
+                    pin_cpu_memory=False,
+                ),
+            ),
+        ))
+    request = GenerationRequest(
+        prompt=PROMPT,
+        negative_prompt=NEGATIVE_PROMPT,
+        sampling=SamplingConfig(
+            seed=42,
+            height=384,
+            width=640,
+            num_frames=33,
+            fps=24,
+            num_inference_steps=2,
+            guidance_scale=1.0,
+            pyramid_num_inference_steps_list=[2, 2, 2],
+            history_sizes=[16, 2, 1],
+            num_latent_frames_per_chunk=9,
+            keep_first_frame=True,
+            is_skip_first_chunk=False,
+            use_zero_init=True,
+            zero_steps=1,
+            is_amplify_first_chunk=True,
+        ),
+        output=OutputConfig(
+            output_path=OUTPUT_PATH,
+            save_video=True,
+            return_frames=False,
+        ),
+    )
+
+    try:
+        generator.generate(request=request)
+    finally:
+        generator.shutdown()
+
+
+if __name__ == "__main__":
+    main()

fastvideo/pipelines/basic/helios/stages.py

@@ -363,7 +363,9 @@ class HeliosChunkDecodingStage(DecodingStage):
     def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
         if fastvideo_args.output_type == "latent":
             assert batch.latents is not None
-            batch.output = batch.latents.to(torch.float32)
+            batch.output = batch.latents.detach().to(dtype=torch.float32, device="cpu")
+            batch.latents = None
+            batch.helios_latent_chunks = None
             return batch
 
         pipeline = self.pipeline() if self.pipeline else None
@@ -381,7 +383,9 @@ def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> Forward
 
         temporal_scale = fastvideo_args.pipeline_config.vae_config.arch_config.scale_factor_temporal
         generated_frames = get_generated_pixel_frames(frames.shape[2], temporal_scale)
-        batch.output = frames[:, :, :generated_frames].to(torch.float32)
+        batch.output = frames[:, :, :generated_frames].detach().to(dtype=torch.float32, device="cpu")
+        batch.latents = None
+        batch.helios_latent_chunks = None
 
         if fastvideo_args.vae_cpu_offload:
             self.vae.to("cpu")

tests/local_tests/helios/test_helios_quality_regression.py

@@ -0,0 +1,74 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Asset-optional video-level quality and container regression checks."""
+
+from __future__ import annotations
+
+import json
+import os
+from pathlib import Path
+import subprocess
+
+import numpy as np
+import pytest
+
+
+def _video_summary(path: Path) -> dict:
+    probe = subprocess.run(
+        [
+            "ffprobe", "-v", "error", "-select_streams", "v:0", "-count_frames",
+            "-show_entries", "stream=codec_name,width,height,avg_frame_rate,nb_read_frames",
+            "-of", "json", str(path),
+        ],
+        check=True,
+        capture_output=True,
+        text=True,
+    )
+    stream = json.loads(probe.stdout)["streams"][0]
+    width, height = int(stream["width"]), int(stream["height"])
+    raw = subprocess.run(
+        ["ffmpeg", "-v", "error", "-i", str(path), "-f", "rawvideo", "-pix_fmt", "gray", "-"],
+        check=True,
+        capture_output=True,
+    ).stdout
+    frames = np.frombuffer(raw, dtype=np.uint8).reshape(-1, height, width)
+    means = frames.mean(axis=(1, 2))
+    return {
+        "codec": stream["codec_name"],
+        "width": width,
+        "height": height,
+        "fps": stream["avg_frame_rate"],
+        "frames": int(stream["nb_read_frames"]),
+        "decoded_frames": len(frames),
+        "mean": float(frames.mean()),
+        "std": float(frames.std()),
+        "black_frame_count": int((means < 1.0).sum()),
+        "frame_mean_std": float(means.std()),
+    }
+
+
+def test_helios_quality_candidate_has_valid_384x640_video():
+    candidate = os.environ.get("HELIOS_QUALITY_CANDIDATE")
+    if not candidate:
+        pytest.skip("Set HELIOS_QUALITY_CANDIDATE to run the real-video quality gate")
+    path = Path(candidate)
+    if not path.is_file():
+        pytest.skip(f"Quality candidate does not exist: {path}")
+    summary = _video_summary(path)
+    assert summary["codec"] in {"h264", "hevc", "av1"}
+    assert (summary["width"], summary["height"]) == (640, 384)
+    assert summary["fps"] == "24/1"
+    assert summary["frames"] == summary["decoded_frames"] == 33
+    assert summary["std"] > 5.0
+    assert summary["black_frame_count"] == 0
+
+
+def test_helios_quality_candidate_is_stable_against_reference():
+    candidate = os.environ.get("HELIOS_QUALITY_CANDIDATE")
+    reference = os.environ.get("HELIOS_QUALITY_REFERENCE")
+    if not candidate or not reference:
+        pytest.skip("Set HELIOS_QUALITY_CANDIDATE and HELIOS_QUALITY_REFERENCE for comparison")
+    candidate_summary = _video_summary(Path(candidate))
+    reference_summary = _video_summary(Path(reference))
+    assert abs(candidate_summary["mean"] - reference_summary["mean"]) < 35.0
+    assert abs(candidate_summary["std"] - reference_summary["std"]) < 35.0
+    assert abs(candidate_summary["frame_mean_std"] - reference_summary["frame_mean_std"]) < 25.0

tests/local_tests/pipelines/test_helios_pipeline_stages.py

@@ -296,6 +296,10 @@ def decode(self, latent):
 )
 decoded = HeliosChunkDecodingStage(vae).forward(decode_batch, decode_args).output
 
+latent_batch = ForwardBatch(data_type="video", latents=expected)
+latent_args = SimpleNamespace(output_type="latent")
+latent_output = HeliosChunkDecodingStage(vae).forward(latent_batch, latent_args).output
+
 print(json.dumps({
     "cuda_available": True,
     "latent_max_diff": (actual_batch.latents - expected).abs().max().item(),
@@ -308,8 +312,10 @@ def decode(self, latent):
     "autoregressive_second_short_prefix_mean": autoregressive_model.calls[6]["short_prefix_mean"],
     "vae_calls": vae.calls,
     "decoded_shape": list(decoded.shape),
+    "decoded_device": decoded.device.type,
     "decoded_first_mean": decoded[:, :, :33].mean().item(),
     "decoded_second_mean": decoded[:, :, 33:].mean().item(),
+    "latent_output_device": latent_output.device.type,
 }))
 """
     result = subprocess.run(
@@ -368,5 +374,7 @@ def test_chunk_decoder_calls_vae_per_chunk_and_matches_frame_rounding():
     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_device"] == "cpu"
     assert result["decoded_first_mean"] == 0
     assert result["decoded_second_mean"] == 1
+    assert result["latent_output_device"] == "cpu"