SolitaryThinker
2114d19ced
Add and verify the text-to-image path, and fix the UniPC flow_shift selection
it exposed.
Bug: the stage chose flow_shift by task (`3.0 if is_t2i else 10.0`). The
framework selects it purely by the named resolution bucket
(`OmniSampleArgs._RESOLUTION_SHIFT_DEFAULTS`, 8B backbone: 256->3.0, 480->5.0,
720/768->10.0); T2V/I2V/T2I share a shift at a given resolution. The task-based
rule only coincidentally matched (T2V@720, T2I@256). The canonical Cosmos3 T2I
is 960x960 (the "720" bucket -> 10.0), so `is_t2i->3.0` was wrong for real T2I.
- Replace with `_flow_shift_for_resolution(h, w)` mapping the longest side to the
framework bucket (<=320:3.0, 640-832:5.0, 960-1360:10.0), applied to all tasks.
- Parity: `test_cosmos3_flow_shift_parity.py` checks the mapping against the
framework's `{VIDEO,IMAGE}_RES_SIZE_INFO` x `_RESOLUTION_SHIFT_DEFAULTS` for
every (resolution, aspect) in the 8B rows (20 cases).
- Update the call-graph mode-dispatch tests for resolution-based shift.
Also harden `_image_to_video_tensor`'s tensor branch to respect the FastVideo
[-1,1] convention for already-preprocessed conditioning tensors (the PIL path,
used by the real pipeline, stays framework-exact and bit-parity-tested).
Example: `examples/inference/basic/basic_cosmos3_t2i_new_api.py` (num_frames=1,
960x960). Verified on B200 (real weights, 35 steps): a coherent red-panda image
matching the prompt, flow_shift=10.0. Full cosmos3 suite: 118 passed, 0 skipped.
79 lines
2.4 KiB
Python
79 lines
2.4 KiB
Python
import os
|
|
import time
|
|
|
|
from fastvideo import VideoGenerator
|
|
from fastvideo.api import (
|
|
EngineConfig,
|
|
GenerationRequest,
|
|
GeneratorConfig,
|
|
OffloadConfig,
|
|
OutputConfig,
|
|
SamplingConfig,
|
|
)
|
|
|
|
# NVIDIA Cosmos3-Nano omni world model — text-to-image (T2I) path through
|
|
# FastVideo's native Cosmos3 pipeline. T2I is the single-frame case
|
|
# (num_frames=1); the canonical Cosmos3 T2I resolution is 960x960 (the model's
|
|
# "720" bucket, UniPC flow_shift=10.0). Point COSMOS3_MODEL_PATH at a local
|
|
# diffusers checkpoint (e.g. ``official_weights/cosmos3``) to skip the download.
|
|
OUTPUT_PATH = "video_samples_cosmos3_t2i"
|
|
|
|
|
|
def main():
|
|
model_name = os.environ.get("COSMOS3_MODEL_PATH", "nvidia/Cosmos3-Nano")
|
|
|
|
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 photograph of a red panda sitting on a mossy log in a misty bamboo "
|
|
"forest, soft golden morning light filtering through the leaves, shallow "
|
|
"depth of field, crisp fur detail, serene atmosphere."
|
|
)
|
|
request = GenerationRequest(
|
|
prompt=prompt,
|
|
sampling=SamplingConfig(
|
|
# T2I is single-frame; canonical Cosmos3 T2I is 960x960. Overridable
|
|
# via env for quick smoke runs.
|
|
num_frames=1,
|
|
height=int(os.environ.get("COSMOS3_HEIGHT", "960")),
|
|
width=int(os.environ.get("COSMOS3_WIDTH", "960")),
|
|
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 image: {gen_time} seconds")
|
|
print(f"Output written to: {result.video_path}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|