diff --git a/examples/inference/basic/basic_cosmos3_t2vs_new_api.py b/examples/inference/basic/basic_cosmos3_t2vs_new_api.py new file mode 100644 index 000000000..60eef1a2d --- /dev/null +++ b/examples/inference/basic/basic_cosmos3_t2vs_new_api.py @@ -0,0 +1,67 @@ +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() diff --git a/fastvideo/pipelines/basic/cosmos3/cosmos3_pipeline.py b/fastvideo/pipelines/basic/cosmos3/cosmos3_pipeline.py index dee2837ae..8348531f0 100644 --- a/fastvideo/pipelines/basic/cosmos3/cosmos3_pipeline.py +++ b/fastvideo/pipelines/basic/cosmos3/cosmos3_pipeline.py @@ -38,7 +38,7 @@ third-party, both explicitly allowed. from __future__ import annotations import math -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import Any import torch @@ -49,6 +49,7 @@ from fastvideo.models.schedulers.scheduling_unipc_multistep import ( UniPCMultistepScheduler, ) from fastvideo.pipelines.basic.cosmos3.sequence_packing import ( Cosmos3SampleInputs, + Cosmos3SoundItem, Cosmos3VisionItem, pack_cosmos3_video_sequence, ) @@ -191,6 +192,35 @@ def _split_flat_latent(flat: torch.Tensor, specs: list[Cosmos3VisionSpec]) -> li return out +@dataclass +class Cosmos3SoundSpec: + """Geometry + conditioning for one sound item in a denoise run. + + Args: + shape: ``(C, T)`` of the sound latent (channels, temporal frames). + condition_frame_indexes: Latent-frame indices kept clean (``[]`` for t2vs). + fps: Sound latent FPS (``sound_latent_fps``); used iff fps modulation is on. + """ + + shape: tuple[int, int] + condition_frame_indexes: list[int] = field(default_factory=list) + fps: float | None = None + + @property + def numel(self) -> int: + return int(math.prod(self.shape)) + + +def _split_sound_latent(flat: torch.Tensor, specs: list[Cosmos3SoundSpec]) -> list[torch.Tensor]: + """Split a flat sound vector into per-item ``[C, T]`` tensors.""" + out: list[torch.Tensor] = [] + offset = 0 + for spec in specs: + out.append(flat[offset:offset + spec.numel].reshape(spec.shape)) + offset += spec.numel + return out + + def cosmos3_get_cfg_velocity( *, transformer: Any, @@ -210,6 +240,8 @@ def cosmos3_get_cfg_velocity( include_end_of_generation_token: bool = False, fps_per_item: list[float] | None = None, normalize_cfg: bool = False, + sound_specs: list[Cosmos3SoundSpec] | None = None, + sound_fps_per_item: list[float] | None = None, ) -> torch.Tensor: """Sequential-CFG velocity for one denoise step (framework math). @@ -231,10 +263,23 @@ def cosmos3_get_cfg_velocity( assert timestep.numel() == 1, "timestep must be a scalar" timestep_value = float(timestep.reshape(()).item()) - noise_x_vision = _split_flat_latent(flat_latent, specs) + # Combined flat layout (t2vs): [all vision | all sound], matching the + # framework per-sample concat order ([vision_i | sound_i]); single sample here. + vision_total = sum(spec.numel for spec in specs) + noise_x_vision = _split_flat_latent(flat_latent[:vision_total], specs) + noise_x_sound = (_split_sound_latent(flat_latent[vision_total:], sound_specs) if sound_specs else None) device = next(transformer.parameters()).device def _run(token_ids: list[int]) -> torch.Tensor: + sound_items: list[Cosmos3SoundItem] = [] + if sound_specs is not None and noise_x_sound is not None: + sound_items = [ + Cosmos3SoundItem( + latent=noise_x_sound[i], + condition_frame_indexes=list(ss.condition_frame_indexes), + fps=(sound_fps_per_item[i] if sound_fps_per_item is not None else None), + ) for i, ss in enumerate(sound_specs) + ] samples = [ Cosmos3SampleInputs( text_ids=list(token_ids), @@ -243,6 +288,7 @@ def cosmos3_get_cfg_velocity( condition_frame_indexes=list(spec.condition_frame_indexes), fps=(fps_per_item[i] if fps_per_item is not None else None), ), + sound=(sound_items[i] if i < len(sound_items) else None), timestep=timestep_value, ) for i, (latent, spec) in enumerate(zip(noise_x_vision, specs, strict=False)) ] @@ -258,20 +304,35 @@ def cosmos3_get_cfg_velocity( temporal_compression_factor=temporal_compression_factor, ) out = transformer(**packed.to_dit_kwargs(device=device)) + + # Vision velocity: zero on conditioning frames, per item, flattened. + vision_vel = torch.zeros(vision_total, device=flat_latent.device, dtype=flat_latent.dtype) preds = out.get("preds_vision") - if preds is None: - # Fully-conditioned clip: no noisy frames, zero velocity. - return torch.zeros_like(flat_latent) - # Zero the prediction on conditioning frames, per item. - velocity_items: list[torch.Tensor] = [] - for pred, cond_mask in zip(preds, packed.vision_condition_mask, strict=False): - pred = pred.squeeze(0) if pred.dim() == 5 else pred # [C, T, H, W] - keep = (1.0 - cond_mask).to(dtype=pred.dtype, device=pred.device) # [T,1,1] - if keep.sum() > 0: - velocity_items.append(pred * keep) - else: - velocity_items.append(torch.zeros_like(pred)) - return torch.cat([v.reshape(-1) for v in velocity_items]) + if preds is not None: + items: list[torch.Tensor] = [] + for pred, cond_mask in zip(preds, packed.vision_condition_mask, strict=False): + pred = pred.squeeze(0) if pred.dim() == 5 else pred # [C, T, H, W] + keep = (1.0 - cond_mask).to(dtype=pred.dtype, device=pred.device) # [T,1,1] + items.append(pred * keep if keep.sum() > 0 else torch.zeros_like(pred)) + vision_vel = torch.cat([v.reshape(-1) for v in items]).to(flat_latent.dtype) + + if not sound_specs: + return vision_vel + + # Sound velocity: preds_sound are per-item [C, T], already zero on clean + # frames (unpack fills only noisy frames); zero on cond frames defensively. + sound_total = sum(spec.numel for spec in sound_specs) + sound_vel = torch.zeros(sound_total, device=flat_latent.device, dtype=flat_latent.dtype) + preds_s = out.get("preds_sound") + if preds_s is not None: + s_items: list[torch.Tensor] = [] + for pred, cond_mask in zip(preds_s, packed.sound_condition_mask, strict=False): + pred = pred.squeeze(0) if pred.dim() == 3 else pred # [C, T] + keep = (1.0 - cond_mask).reshape(1, -1).to(dtype=pred.dtype, device=pred.device) # [1, T] + s_items.append(pred * keep) + sound_vel = torch.cat([v.reshape(-1) for v in s_items]).to(flat_latent.dtype) + + return torch.cat([vision_vel, sound_vel]) cond_v = _run(cond_token_ids) uncond_v = _run(uncond_token_ids) @@ -325,6 +386,8 @@ class Cosmos3DenoiseEngine: cond_token_ids: list[int], uncond_token_ids: list[int], fps_per_item: list[float] | None = None, + sound_specs: list[Cosmos3SoundSpec] | None = None, + sound_fps_per_item: list[float] | None = None, ) -> torch.Tensor: return cosmos3_get_cfg_velocity( transformer=self.transformer, @@ -343,6 +406,8 @@ class Cosmos3DenoiseEngine: temporal_compression_factor=self.temporal_compression_factor, include_end_of_generation_token=self.include_end_of_generation_token, fps_per_item=fps_per_item, + sound_specs=sound_specs, + sound_fps_per_item=sound_fps_per_item, ) def denoise( @@ -356,12 +421,16 @@ class Cosmos3DenoiseEngine: uncond_token_ids: list[int], fps_per_item: list[float] | None = None, progress_bar: Any | None = None, + sound_specs: list[Cosmos3SoundSpec] | None = None, + sound_fps_per_item: list[float] | None = None, ) -> torch.Tensor: """Run the full UniPC denoise loop, returning the final flat latent. For each timestep: compute the sequential-CFG velocity, then ``scheduler.step(model_output=v, timestep, sample=latent.unsqueeze(0))`` (the framework steps with a leading batch axis), squeezing back to flat. + For t2vs the flat latent is ``[vision | sound]`` and the velocity covers + both; the scheduler steps the combined vector jointly. """ latent = flat_latent iterator = progress_bar(timesteps) if progress_bar is not None else timesteps @@ -374,6 +443,8 @@ class Cosmos3DenoiseEngine: cond_token_ids=cond_token_ids, uncond_token_ids=uncond_token_ids, fps_per_item=fps_per_item, + sound_specs=sound_specs, + sound_fps_per_item=sound_fps_per_item, ) stepped = self.scheduler.step( model_output=v_pred, diff --git a/fastvideo/pipelines/stages/cosmos3_stages.py b/fastvideo/pipelines/stages/cosmos3_stages.py index 6dd02784e..628b0008f 100644 --- a/fastvideo/pipelines/stages/cosmos3_stages.py +++ b/fastvideo/pipelines/stages/cosmos3_stages.py @@ -20,6 +20,7 @@ This mirrors the framework ``Cosmos3OmniDiffusersPipeline.__call__``. """ from __future__ import annotations +import os import weakref from typing import Any @@ -31,6 +32,7 @@ from fastvideo.fastvideo_args import FastVideoArgs from fastvideo.logger import init_logger from fastvideo.pipelines.basic.cosmos3.cosmos3_pipeline import ( Cosmos3DenoiseEngine, + Cosmos3SoundSpec, Cosmos3VisionSpec, _VaeNorm, cosmos3_special_tokens, @@ -186,6 +188,29 @@ class Cosmos3DenoisingStage(PipelineStage): flat_latent = init_latent.reshape(-1) fps_per_item = [fps] if bool(arch.enable_fps_modulation) else None + + # ---- t2vs: jointly generate sound (combined [vision | sound] latent) ---- + # Mirrors the framework: a placeholder audio sized to the video duration + # sets the sound latent length; sound shares the denoise/CFG with vision. + with_audio = is_video and os.environ.get("COSMOS3_T2VS", "") not in ("", "0") + sound_specs = None + sound_fps_per_item = None + sound_vae = None + sound_shape: tuple[int, int] | None = None + if with_audio: + sound_vae = self._get_sound_vae(pipe, device, dtype) + sound_dim = int(arch.sound_dim) + sound_latent_fps = float(arch.sound_latent_fps) + # Framework ``create_placeholder_audio`` + ``get_latent_num_samples``. + num_audio_samples = int(num_frames / fps * sound_vae.sample_rate) + sound_latent_t = max(1, sound_vae.get_latent_num_samples(num_audio_samples)) + sound_shape = (sound_dim, sound_latent_t) + sound_noise = randn_tensor((sound_dim, sound_latent_t), generator=generator, device=device, + dtype=dtype).float() + flat_latent = torch.cat([flat_latent, sound_noise.reshape(-1)]) + sound_specs = [Cosmos3SoundSpec(shape=sound_shape, condition_frame_indexes=[], fps=sound_latent_fps)] + sound_fps_per_item = [sound_latent_fps] if bool(arch.enable_fps_modulation) else None + final_flat = engine.denoise( flat_latent=flat_latent, timesteps=timesteps, @@ -195,15 +220,25 @@ class Cosmos3DenoisingStage(PipelineStage): uncond_token_ids=uncond_ids, fps_per_item=fps_per_item, progress_bar=lambda it: tqdm(it, desc="Cosmos3 denoising"), + sound_specs=sound_specs, + sound_fps_per_item=sound_fps_per_item, ) - # ---- Decode: [C, T, H, W] -> pixels [B, 3, T, H, W] in [0, 1] ---- - result_latent = final_flat.reshape(latent_shape).unsqueeze(0).to(device=device, dtype=dtype) + # ---- Decode vision: [C, T, H, W] -> pixels [B, 3, T, H, W] in [0, 1] ---- + vision_flat = final_flat[:spec.numel] + result_latent = vision_flat.reshape(latent_shape).unsqueeze(0).to(device=device, dtype=dtype) decoded = cosmos3_vae_decode(self.vae, result_latent, norm) # [B, 3, T, H, W] in [-1, 1] video = ((1.0 + decoded) / 2.0).clamp(0.0, 1.0) batch.latents = result_latent batch.output = video + + # ---- Decode sound: AVAE latent [C, T] -> waveform [C, N] in [-1, 1] ---- + if with_audio and sound_vae is not None and sound_shape is not None: + sound_latent = final_flat[spec.numel:].reshape(sound_shape).unsqueeze(0).to(device=device, dtype=dtype) + waveform = sound_vae.decode(sound_latent) # [1, C_audio, N] + batch.extra["audio"] = waveform[0].detach().float().cpu() # [C_audio, N] + batch.extra["audio_sample_rate"] = int(sound_vae.sample_rate) return batch # ------------------------------------------------------------------ @@ -223,6 +258,24 @@ class Cosmos3DenoisingStage(PipelineStage): img = TF.resize(img, [resize_h, resize_w]) return TF.center_crop(img, [target_h, target_w]) + @staticmethod + def _get_sound_vae(pipe: Any, device: torch.device, dtype: torch.dtype) -> Any: + """Lazily load + cache the Cosmos3 sound AVAE decoder from the checkpoint. + + The video path does not load ``sound_tokenizer``; t2vs needs only its + decoder, so we load it on first use from ``/sound_tokenizer``. + """ + cached = getattr(pipe, "_sound_vae", None) if pipe is not None else None + if cached is not None: + return cached + from fastvideo.models.audio.cosmos3_avae import Cosmos3SoundVAE + model_path = pipe.model_path + sound_dir = os.path.join(model_path, "sound_tokenizer") + sound_vae = Cosmos3SoundVAE.from_pretrained(sound_dir, torch_dtype=dtype).to(device) + if pipe is not None: + pipe._sound_vae = sound_vae + return sound_vae + @classmethod def _image_to_video_tensor( cls, diff --git a/tests/local_tests/cosmos3/test_cosmos3_sound_parity.py b/tests/local_tests/cosmos3/test_cosmos3_sound_parity.py index 98ae95bec..9a655327a 100644 --- a/tests/local_tests/cosmos3/test_cosmos3_sound_parity.py +++ b/tests/local_tests/cosmos3/test_cosmos3_sound_parity.py @@ -233,3 +233,56 @@ class TestCosmos3SoundParity: torch.testing.assert_close(fv_out["preds_vision"][0], fw_out["preds_vision"][0], atol=1e-4, rtol=1e-3) torch.testing.assert_close(fv_out["preds_sound"][0], fw_out["preds_sound"][0], atol=1e-4, rtol=1e-3) torch.testing.assert_close(fv_on_fw["preds_sound"][0], fw_out["preds_sound"][0], atol=1e-4, rtol=1e-3) + + @pytest.mark.parametrize(("grid_t", "lh", "lw", "snd_t", "n_text", "cond_v", "cond_s"), _CASES) + def test_t2vs_cfg_velocity_matches_framework(self, grid_t, lh, lw, snd_t, n_text, cond_v, cond_s): + """The combined [vision|sound] sequential-CFG velocity (the t2vs denoise + step's pipeline glue) matches a framework-DiT oracle.""" + from fastvideo.pipelines.basic.cosmos3.cosmos3_pipeline import ( + Cosmos3SoundSpec, + Cosmos3VisionSpec, + cosmos3_get_cfg_velocity, + ) + + vfm, dit = self._build() + vision_shape = (_LATENT_CHANNEL, grid_t, lh // _LATENT_PATCH_SIZE, lw // _LATENT_PATCH_SIZE) + # _make_inputs builds the vision LATENT [1,C,T,H,W]; here we drive the + # combined flat latent directly, so use the un-patchified latent shape. + vlat_shape = (_LATENT_CHANNEL, grid_t, lh, lw) + sound_shape = (_SOUND_DIM, snd_t) + torch.manual_seed(3) + cond_ids = torch.randint(0, 60, (n_text,)).tolist() + uncond_ids = torch.randint(0, 60, (max(1, n_text - 1),)).tolist() + vis_numel = int(torch.tensor(vlat_shape).prod()) + snd_numel = int(torch.tensor(sound_shape).prod()) + flat = torch.randn(vis_numel + snd_numel) + guidance, ts = 6.0, 500.0 + + def _fw_velocity(ids): + vision = flat[:vis_numel].reshape(vlat_shape).unsqueeze(0) # [1,C,T,H,W] + sound = flat[vis_numel:].reshape(sound_shape) # [C,T] + ps = _framework_pack_sound(text_ids=ids, vision=vision, sound=sound, + cond_vision=cond_v, cond_sound=cond_s, + timestep=ts, is_image_batch=(grid_t == 1)) + with torch.no_grad(): + out = vfm(packed_seq=ps) + pv = out["preds_vision"][0].squeeze(0) # [C,T,H,W] (zero on clean) + psd = out["preds_sound"][0] # [C,T] (zero on clean) + return torch.cat([pv.reshape(-1), psd.reshape(-1)]) + + fw_cond, fw_uncond = _fw_velocity(cond_ids), _fw_velocity(uncond_ids) + fw_v = fw_uncond + guidance * (fw_cond - fw_uncond) + + fv_v = cosmos3_get_cfg_velocity( + transformer=dit, flat_latent=flat, timestep=torch.tensor([ts]), guidance=guidance, + specs=[Cosmos3VisionSpec(shape=vlat_shape, condition_frame_indexes=list(cond_v))], + sound_specs=[Cosmos3SoundSpec(shape=sound_shape, condition_frame_indexes=list(cond_s))], + cond_token_ids=cond_ids, uncond_token_ids=uncond_ids, + special_tokens=_SPECIAL_TOKENS, latent_patch_size=_LATENT_PATCH_SIZE, + temporal_modality_margin=_TEMPORAL_MODALITY_MARGIN, reset_spatial_ids=_RESET_SPATIAL_IDS, + enable_fps_modulation=False, base_fps=24.0, temporal_compression_factor=_TCF, + ) + assert fv_v.shape == fw_v.shape, f"shape fv={fv_v.shape} fw={fw_v.shape}" + mx, mn = _diffs(fv_v, fw_v) + print(f"\n[t2vs_cfg_velocity {grid_t}x{lh}x{lw} snd={snd_t}] max={mx:.3e} mean={mn:.3e}") + torch.testing.assert_close(fv_v, fw_v, atol=1e-4, rtol=1e-3)