The first real-weights Cosmos3-Nano T2V run produced an all-black video.
Root cause: the checkpoint's diffusers-style scheduler_config.json sets
use_karras_sigmas=true, and FastVideo's vendored UniPC checks karras
*before* use_flow_sigmas, so it built diffusion (beta) sigmas instead of
flow-matching sigmas -> scheduler.step diverged to NaN latents -> black
frames. The DiT/CFG velocity itself was clean.
The framework samples with FlowUniPCMultistepScheduler (pure flow:
shift + num_train_timesteps only). FastVideo's vendored UniPC is the same
algorithm, so the fix is config, not a port:
- Coerce the loaded scheduler to the flow setup in
Cosmos3OmniDiffusersPipeline.initialize_pipeline (kill karras/exp/beta
sigmas, force use_flow_sigmas + flow_prediction + final_sigmas_type=zero)
and rebuild the runtime scheduler from it.
- Use FastVideo's native UniPC (not diffusers) in the pipeline and tests,
honoring the "no diffusers at runtime" constraint.
Parity:
- Add test_cosmos3_scheduler_parity.py: FastVideo vendored UniPC (flow
config) vs the framework FlowUniPCMultistepScheduler. timesteps bit-exact,
sigmas within fp32 epsilon (~1e-8), full multi-step trajectory < ~1e-6
across shift in {10 (video), 3 (t2i)} and steps in {4,10,35}.
- Fix test_cosmos3_denoise_cfg_parity to use the framework scheduler as the
oracle (it previously compared diffusers-vs-diffusers, so the scheduler
was never actually checked against the framework).
Integration fixes to run end-to-end on real weights:
- registry: map checkpoint DiT class Cosmos3OmniTransformer ->
Cosmos3VFMTransformer.
- component_loader: route the "text_tokenizer" module to TokenizerLoader;
filter scheduler config to the class's __init__ params (robust to
diffusers schema drift like shift_terminal/sigma_min/sigma_max).
- DiT: materialize_non_persistent_buffers (recompute rotary_emb.inv_freq
after meta-device FSDP load) + cast vision latents / timestep embeds to
the compute dtype (no-ops in the fp32 parity tests).
- sequence_packing: move packed tensors to the model device in to_dit_kwargs.
- config: drop the unused identity text-preprocess (Cosmos3 has no text
encoder; the DiT tokenizes in the denoise stage).
Adds examples/inference/basic/basic_cosmos3_new_api.py (new VideoGenerator
API). Verified: 1280x704, 29 frames, 35 steps on B200 -> coherent video
matching the prompt (no NaNs, real pixel variance). Full cosmos3 parity
suite green.
78 lines
2.4 KiB
Python
78 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 — this example exercises the
|
|
# text-to-video (T2V) path through FastVideo's native Cosmos3 pipeline.
|
|
# 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"
|
|
|
|
|
|
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 golden retriever puppy runs across a sunlit meadow toward the camera, "
|
|
"ears flopping and wildflowers swaying in the breeze. Shallow depth of "
|
|
"field, warm afternoon light, smooth cinematic tracking shot."
|
|
)
|
|
request = GenerationRequest(
|
|
prompt=prompt,
|
|
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()
|