nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[feat] cosmos3: I2V image conditioning (static-repeat) + parity + example

Browse Source 
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.
This commit is contained in:
SolitaryThinker
2026-07-21 15:45:37 -07:00
parent 6ca684539c
commit bf054e5b79
3 changed files with 219 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/inference/basic/basic_cosmos3_i2v_new_api.py
+81
View File
@@ -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()
fastvideo/pipelines/stages/cosmos3_stages.py
+47 -25
View File
@@ -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()
tests/local_tests/cosmos3/test_cosmos3_i2v_conditioning_parity.py
+91
View File
@@ -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 <worktree> && <fv-cosmos3 python> -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"