From a5be058cf1ae37d20f095dced9a85fca8a1f2c08 Mon Sep 17 00:00:00 2001 From: SolitaryThinker Date: Sun, 7 Jun 2026 04:45:32 -0700 Subject: [PATCH] [feat] cosmos3: I2V image conditioning (static-repeat) + parity + example MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Wire up and verify the Cosmos3 image-to-video path. The denoise stage already handled the I2V conditioning math (encode -> condition mask -> keep frame 0 clean -> velocity zeroed on condition frames, all matching the framework), but the conditioning pixel clip was built wrong. Bug: `_image_to_video_tensor` zero-filled every frame after frame 0. The framework (`cosmos_framework.inference.vision.build_conditioned_video_batch`) instead fills frame 0 with the image and **repeats the last conditioning frame** for the rest of the clip (a static video) before VAE-encoding. Because the Wan VAE is temporal (4x), latent frame 0 (the kept-clean condition frame) depends on several pixel frames, so zero-filling produces a wrong conditioning latent. Fix `_image_to_video_tensor` to be framework-faithful: - aspect-preserving resize + center crop + uint8 quantization, then `/127.5 - 1` (mirrors `load_conditioning_image` / `_resize_and_center_crop`); - static-repeat fill across all frames (mirrors `build_conditioned_video_batch`). Parity: add `test_cosmos3_i2v_conditioning_parity.py` comparing the FastVideo conditioning clip against the framework's `load_conditioning_image` + repeat-fill across aspect/size/frame-count cases — bit-exact (max abs diff 0.0). Example: `examples/inference/basic/basic_cosmos3_i2v_new_api.py` (new VideoGenerator API, image via `InputConfig(image_path=...)`, `assets/images/cyclist.jpg` default). Verified on B200 (1280x704, 29f, 35 steps, real weights): frame 0 reproduces the conditioning image and the clip generates coherent forward motion following the prompt. Full cosmos3 suite: 98 passed, 0 skipped. --- .../basic/basic_cosmos3_i2v_new_api.py | 81 +++++++++++++++++ fastvideo/pipelines/stages/cosmos3_stages.py | 72 ++++++++++----- .../test_cosmos3_i2v_conditioning_parity.py | 91 +++++++++++++++++++ 3 files changed, 219 insertions(+), 25 deletions(-) create mode 100644 examples/inference/basic/basic_cosmos3_i2v_new_api.py create mode 100644 tests/local_tests/cosmos3/test_cosmos3_i2v_conditioning_parity.py diff --git a/examples/inference/basic/basic_cosmos3_i2v_new_api.py b/examples/inference/basic/basic_cosmos3_i2v_new_api.py new file mode 100644 index 000000000..84de60377 --- /dev/null +++ b/examples/inference/basic/basic_cosmos3_i2v_new_api.py @@ -0,0 +1,81 @@ +import os +import time + +from fastvideo import VideoGenerator +from fastvideo.api import ( + EngineConfig, + GenerationRequest, + GeneratorConfig, + InputConfig, + OffloadConfig, + OutputConfig, + SamplingConfig, +) + +# NVIDIA Cosmos3-Nano omni world model — image-to-video (I2V) path through +# FastVideo's native Cosmos3 pipeline. The input image conditions latent frame 0 +# (kept clean during denoising); the rest of the clip is generated to follow it. +# Point COSMOS3_MODEL_PATH at a local diffusers checkpoint (e.g. +# ``official_weights/cosmos3``) to skip the Hugging Face download. +OUTPUT_PATH = "video_samples_cosmos3_i2v" + + +def main(): + model_name = os.environ.get("COSMOS3_MODEL_PATH", "nvidia/Cosmos3-Nano") + image_path = os.environ.get("COSMOS3_IMAGE_PATH", "assets/images/cyclist.jpg") + + generator_config = GeneratorConfig( + model_path=model_name, + engine=EngineConfig( + num_gpus=1, + use_fsdp_inference=False, + offload=OffloadConfig( + text_encoder=True, + pin_cpu_memory=True, + dit=False, + vae=False, + ), + ), + ) + + load_start_time = time.perf_counter() + generator = VideoGenerator.from_config(generator_config) + load_time = time.perf_counter() - load_start_time + + prompt = ( + "A mountain biker rides forward along the sunlit forest trail, wheels " + "kicking up dust as trees and dappled light sweep past, smooth cinematic " + "tracking shot from behind." + ) + request = GenerationRequest( + prompt=prompt, + inputs=InputConfig(image_path=image_path), + sampling=SamplingConfig( + # cosmos3_nano native defaults are 704x1280, 189 frames, 35 steps; + # overridable via env for quick smoke runs. + num_frames=int(os.environ.get("COSMOS3_NUM_FRAMES", "189")), + height=int(os.environ.get("COSMOS3_HEIGHT", "704")), + width=int(os.environ.get("COSMOS3_WIDTH", "1280")), + num_inference_steps=int(os.environ.get("COSMOS3_STEPS", "35")), + guidance_scale=6.0, + fps=24, + seed=1024, + ), + output=OutputConfig( + output_path=OUTPUT_PATH, + save_video=True, + return_frames=False, + ), + ) + + start_time = time.perf_counter() + result = generator.generate(request) + gen_time = time.perf_counter() - start_time + + print(f"Time taken to load model: {load_time} seconds") + print(f"Time taken to generate video: {gen_time} seconds") + print(f"Output written to: {result.video_path}") + + +if __name__ == "__main__": + main() diff --git a/fastvideo/pipelines/stages/cosmos3_stages.py b/fastvideo/pipelines/stages/cosmos3_stages.py index b0bf596b2..3c5b5d864 100644 --- a/fastvideo/pipelines/stages/cosmos3_stages.py +++ b/fastvideo/pipelines/stages/cosmos3_stages.py @@ -187,7 +187,22 @@ class Cosmos3DenoisingStage(PipelineStage): # Image preprocessing # ------------------------------------------------------------------ @staticmethod + def _resize_and_center_crop(img: torch.Tensor, target_h: int, target_w: int) -> torch.Tensor: + """Aspect-ratio-preserving resize + center crop, matching the framework + (``cosmos_framework.inference.vision._resize_and_center_crop``).""" + import math + + import torchvision.transforms.functional as TF + orig_h, orig_w = img.shape[-2], img.shape[-1] + scaling_ratio = max(target_w / orig_w, target_h / orig_h) + resize_h = int(math.ceil(scaling_ratio * orig_h)) + resize_w = int(math.ceil(scaling_ratio * orig_w)) + img = TF.resize(img, [resize_h, resize_w]) + return TF.center_crop(img, [target_h, target_w]) + + @classmethod def _image_to_video_tensor( + cls, image: Any, num_frames: int, height: int, @@ -195,36 +210,43 @@ class Cosmos3DenoisingStage(PipelineStage): device: torch.device, dtype: torch.dtype, ) -> torch.Tensor: - """Coerce a conditioning image into a ``[1, 3, T, H, W]`` pixel tensor in [-1, 1]. + """Build the I2V conditioning pixel video ``[1, 3, T, H, W]`` in [-1, 1]. - The first frame holds the (already preprocessed) conditioning image and - the remaining frames are zero-filled; only frame 0 is kept clean by the - condition mask, so the rest of the clip is denoised from noise. + Faithful to the framework (``cosmos_framework.inference.vision``): + ``load_conditioning_image`` (aspect-preserving resize + center crop + + uint8 quantization, then ``/127.5 - 1``) followed by + ``build_conditioned_video_batch``, which fills frame 0 with the image and + **repeats the last conditioning frame** for the rest of the clip (a static + video), NOT zeros. The whole clip is VAE-encoded by the caller; only the + latent condition frame(s) are kept clean by the condition mask, but the + VAE is temporal, so the repeated (not zeroed) frames change the condition + latent — zero-filling here produces a wrong conditioning latent. """ - import torchvision.transforms.functional as TF + import numpy as np - if isinstance(image, torch.Tensor): - img = image.float() - # Accept [B,3,T,H,W], [3,T,H,W], [3,H,W], or [B,3,H,W]. - if img.dim() == 5: - img = img[0] - if img.dim() == 4 and img.shape[1] == 1: - img = img[:, 0] # [3, H, W] - elif img.dim() == 4: - img = img[:, 0] # take first frame -> [3, H, W] - # Normalize to [-1, 1] if it looks like [0, 1] or [0, 255]. - if img.max() > 1.5: - img = img / 127.5 - 1.0 - elif hasattr(image, "convert"): # PIL.Image - import numpy as np + # 1. Coerce to a [3, H, W] tensor in the [0, 255] pixel domain. + if hasattr(image, "convert"): # PIL.Image arr = np.array(image.convert("RGB")) - img = torch.from_numpy(arr).permute(2, 0, 1).float() / 127.5 - 1.0 + img = torch.from_numpy(arr).permute(2, 0, 1).float() # [3, H, W] + elif isinstance(image, torch.Tensor): + img = image.float() + if img.dim() == 5: # [B,3,T,H,W] + img = img[0] + if img.dim() == 4: # [3,T,H,W] or [B,3,H,W] -> first frame + img = img[:, 0] + if img.min() < -0.01: # [-1, 1] -> [0, 255] + img = (img + 1.0) * 127.5 + elif img.max() <= 1.5: # [0, 1] -> [0, 255] + img = img * 255.0 else: raise TypeError(f"Unsupported conditioning image type: {type(image)}") - if img.shape[-2:] != (height, width): - img = TF.resize(img.unsqueeze(0), [height, width]).squeeze(0) + # 2. Framework resize + center crop + uint8 quantization, then -> [-1, 1]. + img = cls._resize_and_center_crop(img.unsqueeze(0), height, width).squeeze(0) + img = img.round().clamp(0, 255) / 127.5 - 1.0 # [3, H, W] in [-1, 1] - video = torch.zeros((1, 3, num_frames, height, width), device=device, dtype=dtype) - video[0, :, 0] = img.to(device=device, dtype=dtype) - return video + # 3. Static-repeat video (build_conditioned_video_batch: frame 0 = image, + # remaining frames repeat the last conditioning frame). + img = img.to(device=device, dtype=dtype) + video = img.unsqueeze(0).unsqueeze(2).expand(1, 3, num_frames, height, width) + return video.contiguous() diff --git a/tests/local_tests/cosmos3/test_cosmos3_i2v_conditioning_parity.py b/tests/local_tests/cosmos3/test_cosmos3_i2v_conditioning_parity.py new file mode 100644 index 000000000..9b4c5c8be --- /dev/null +++ b/tests/local_tests/cosmos3/test_cosmos3_i2v_conditioning_parity.py @@ -0,0 +1,91 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Numerical-parity test: FastVideo I2V conditioning pixel video vs the framework. + +The Cosmos3 I2V path conditions on a *static repeat* of the input image. The +framework (``cosmos_framework.inference.vision``): + + * ``load_conditioning_image``: aspect-preserving resize + center crop + uint8 + quantization, then ``/127.5 - 1`` -> ``[3, 1, h, w]`` in [-1, 1]; + * ``build_conditioned_video_batch``: frame 0 = the image, and every remaining + frame **repeats the last conditioning frame** (a static video) -> the clip + is then VAE-encoded and only the latent condition frame(s) are kept clean. + +Because the VAE is temporal, zero-filling the non-condition frames (the earlier +FastVideo behavior) changes the encoded condition latent, so the repeat-fill is +correctness-critical. This pins FastVideo's +``Cosmos3DenoisingStage._image_to_video_tensor`` against the framework's +image preprocessing + repeat-fill. + +CPU / float32. The framework is the parity ORACLE. + +Run: + cd && -m pytest \ + tests/local_tests/cosmos3/test_cosmos3_i2v_conditioning_parity.py -q +""" +from __future__ import annotations + +import numpy as np +import pytest +import torch +from PIL import Image + +# The official framework provides the parity oracle. +vision = pytest.importorskip( + "cosmos_framework.inference.vision", + reason="cosmos_framework not installed; run in fv-cosmos3 env.", +) + +from fastvideo.pipelines.stages.cosmos3_stages import ( # noqa: E402 + Cosmos3DenoisingStage, +) + +pytestmark = [pytest.mark.local] + + +def _make_image(path, h_in: int, w_in: int, seed: int = 0) -> None: + rng = np.random.default_rng(seed) + arr = rng.integers(0, 256, size=(h_in, w_in, 3), dtype=np.uint8) + Image.fromarray(arr, "RGB").save(path) + + +# (input H, input W, target H, target W, num_frames) +_CASES = [ + pytest.param(120, 200, 256, 256, 9, id="square_from_landscape"), + pytest.param(200, 120, 704, 1280, 13, id="wide_from_portrait"), + pytest.param(256, 256, 256, 256, 5, id="same_size"), +] + + +class TestCosmos3I2VConditioningParity: + + @pytest.mark.parametrize(("h_in", "w_in", "h", "w", "num_frames"), _CASES) + def test_conditioning_video_matches_framework(self, tmp_path, h_in, w_in, h, w, num_frames): + img_path = tmp_path / "cond.png" + _make_image(img_path, h_in, w_in) + + # ---- Framework oracle ---- + # load_conditioning_image -> [3, 1, h, w] in [-1, 1]. + cond = vision.load_conditioning_image(img_path, target_h=h, target_w=w).float() + # Mirror build_conditioned_video_batch (vision.py lines 117-123) in fp32/CPU: + # frame 0 = image; remaining frames repeat the last conditioning frame. + t_cond = cond.shape[1] + expected = torch.zeros(1, 3, num_frames, h, w, dtype=torch.float32) + t_fill = min(t_cond, num_frames) + expected[0, :, :t_fill] = cond[:, :t_fill] + if t_fill < num_frames: + expected[0, :, t_fill:] = expected[0, :, t_fill - 1:t_fill].expand(-1, num_frames - t_fill, -1, -1) + + # ---- FastVideo: same PIL image through the stage helper ---- + pil = Image.open(img_path).convert("RGB") + got = Cosmos3DenoisingStage._image_to_video_tensor( + pil, num_frames, h, w, torch.device("cpu"), torch.float32) + + assert got.shape == expected.shape, f"shape: got={got.shape} expected={expected.shape}" + max_abs = (got - expected).abs().max().item() + print(f"\n[i2v_cond {h}x{w} nf={num_frames}] max abs diff = {max_abs:.3e}") + torch.testing.assert_close(got, expected) + + # Static repeat (not zero-fill): every frame equals frame 0, and the + # frames past frame 0 are non-zero. + assert torch.equal(got[0, :, 0], got[0, :, -1]), "non-condition frames must repeat the image" + assert got[0, :, 1:].abs().sum() > 0, "non-condition frames must not be zero-filled"