SolitaryThinker
1ec268cf80
Wire text-to-video+sound end-to-end, completing the audio PR. - CFG velocity (`cosmos3_get_cfg_velocity`) + `Cosmos3DenoiseEngine` take an optional `sound_specs`: the denoise latent becomes a combined [vision | sound] flat vector; each step packs vision+sound, forwards the DiT, zeros the prediction on conditioning frames for both, and returns the concatenated velocity. The UniPC scheduler steps the joint vector. `Cosmos3SoundSpec` + `_split_sound_latent` added. - `Cosmos3DenoisingStage` t2vs path (gated on `COSMOS3_T2VS`): size the sound latent from the video duration (framework `create_placeholder_audio` + `get_latent_num_samples`), append sound noise, joint-denoise, split, AVAE-decode the sound latent, and set `batch.extra["audio"]` / `["audio_sample_rate"]` so the generator muxes a stereo 48 kHz AAC track into the mp4. Sound AVAE is lazy-loaded from `<model_path>/sound_tokenizer` (`_get_sound_vae`). - Parity (`test_cosmos3_sound_parity::test_t2vs_cfg_velocity_matches_framework`): the combined [vision|sound] sequential-CFG velocity vs a framework-DiT oracle — max abs diff = 0.0, mean abs diff = 0.0 across t2vs/i2vs cases. - Example `basic_cosmos3_t2vs_new_api.py`. Verified on B200 (real weights, 1280x704, 29f, 35 steps): coherent ocean/sunset video WITH a real stereo 48 kHz audio track (mean -10.2 dB, peak -0.4 dB; not silence). Full cosmos3 suite: 130 passed, 0 skipped.
68 lines
2.2 KiB
Python
68 lines
2.2 KiB
Python
import os
|
|
import time
|
|
|
|
# t2vs (text -> video + sound). The Cosmos3 denoise stage generates a joint
|
|
# [vision | sound] latent and AVAE-decodes the sound to a waveform muxed into the
|
|
# mp4. The joint-sound path is gated on COSMOS3_T2VS (set here for the example).
|
|
os.environ.setdefault("COSMOS3_T2VS", "1")
|
|
|
|
from fastvideo import VideoGenerator # noqa: E402
|
|
from fastvideo.api import ( # noqa: E402
|
|
EngineConfig,
|
|
GenerationRequest,
|
|
GeneratorConfig,
|
|
OffloadConfig,
|
|
OutputConfig,
|
|
SamplingConfig,
|
|
)
|
|
|
|
OUTPUT_PATH = "video_samples_cosmos3_t2vs"
|
|
|
|
|
|
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.perf_counter()
|
|
generator = VideoGenerator.from_config(generator_config)
|
|
load_time = time.perf_counter() - load_start
|
|
|
|
prompt = (
|
|
"Ocean waves crash against a rocky shore at sunset, white foam spraying "
|
|
"into the air as seagulls wheel overhead. Golden light, cinematic wide "
|
|
"shot, the rhythmic roar of the surf."
|
|
)
|
|
request = GenerationRequest(
|
|
prompt=prompt,
|
|
sampling=SamplingConfig(
|
|
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.perf_counter()
|
|
result = generator.generate(request)
|
|
gen_time = time.perf_counter() - start
|
|
|
|
print(f"Time taken to load model: {load_time} seconds")
|
|
print(f"Time taken to generate video+sound: {gen_time} seconds")
|
|
print(f"Output written to: {result.video_path}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|