diff --git a/docs/source/en/api/pipelines/cosmos3.md b/docs/source/en/api/pipelines/cosmos3.md index 1af70fc44f8c..93efbbb5e51e 100644 --- a/docs/source/en/api/pipelines/cosmos3.md +++ b/docs/source/en/api/pipelines/cosmos3.md @@ -740,7 +740,7 @@ pipe = Cosmos3OmniPipeline.from_pretrained( ## Cosmos3OmniModularPipeline -Cosmos 3 is also available as a Modular Diffusers pipeline. The task-based [`Cosmos3OmniPipeline`] remains available; the modular pipeline coexists with it and covers the same modes (`text2image`, `text2video`, `image2video`, `video2video`, and action-conditioned generation, with optional sound when supported by the checkpoint). +Cosmos 3 is also available as a Modular Diffusers pipeline. The task-based [`Cosmos3OmniPipeline`] remains available; the modular pipeline coexists with it and covers the same modes (`text2image`, `text2video`, `image2video`, `video2video`, action-conditioned generation, and `transfer` (structural control), with optional sound when supported by the checkpoint). ```python import torch @@ -788,7 +788,7 @@ image2video_blocks = pipe.blocks.get_workflow("image2video") ### Modular examples for all existing workflows -The modular pipeline supports the same call signatures as the task pipeline. The snippets below mirror every generation example shown above (`text2video`, `text2image`, `image2video`, `video2video`, `video2video_sound`, `text2video_sound`, and `action_policy`). +The modular pipeline supports the same call signatures as the task pipeline. The snippets below mirror every generation example shown above (`text2video`, `text2image`, `image2video`, `video2video`, `video2video_sound`, `text2video_sound`, and `action_policy`). Transfer (structural control) has its own inputs and is shown separately in [Modular transfer](#modular-transfer-structural-control) below. ```python import json @@ -934,6 +934,61 @@ if outputs["action"] is not None: json.dump(outputs["action"][0].tolist(), f) ``` +### Modular transfer (structural control) + +Transfer follows a **precomputed control video** (edge, blur, depth, segmentation, or a world-scenario map) passed through `control_videos=` as a `{hint: video}` mapping. It is video-only (no `image` / `video` / `action` / `enable_sound`), the prompt is a pre-upsampled JSON caption (see [Prompt upsampling](#prompt-upsampling)), and long clips are generated autoregressively in chunks of `num_video_frames_per_chunk` and stitched automatically. `guidance_scale` is the usual text CFG; `control_guidance` (`!= 1.0`) additionally amplifies the control signal. Recommended starting values per hint: + +| Hint | `guidance_scale` | `control_guidance` | `flow_shift` | Geometry | +| --- | --- | --- | --- | --- | +| Edge / Blur / Depth | 3.0 | 1.5 | 10.0 | 121 frames @ 30 FPS | +| Segmentation | 3.0 | 2.0 | 10.0 | 121 frames @ 30 FPS | +| World scenario (WSM) | 1.0 | 3.0 | 10.0 | 101 frames @ 10 FPS | + +Diffusers does not ship the control assets. Ready-made ones (a control video + matching `prompt.json` per hint, plus a shared `negative_prompt.json`) live in the [Cosmos cookbook](https://github.com/NVIDIA/cosmos/tree/main/cookbooks/cosmos3/generator/transfer/assets). For the edge example below, download them into a local `assets/` folder: + +```bash +base=https://github.com/NVIDIA/cosmos/raw/refs/heads/main/cookbooks/cosmos3/generator/transfer/assets +mkdir -p assets/edge +curl -sL "$base/edge/control_edge.mp4" -o assets/edge/control_edge.mp4 +curl -sL "$base/edge/prompt.json" -o assets/edge/prompt.json +curl -sL "$base/negative_prompt.json" -o assets/negative_prompt.json +``` + +```python +import json +import torch +from diffusers import Cosmos3OmniModularPipeline +from diffusers.schedulers.scheduling_unipc_multistep import UniPCMultistepScheduler +from diffusers.utils import export_to_video, load_video + +pipe = Cosmos3OmniModularPipeline.from_pretrained("nvidia/Cosmos3-Nano", torch_dtype=torch.bfloat16) +pipe.load_components(torch_dtype=torch.bfloat16) +pipe.to("cuda") +pipe.scheduler = UniPCMultistepScheduler.from_config( + pipe.scheduler.config, flow_shift=10.0, use_karras_sigmas=False +) + +# Downloaded into assets/ from the Cosmos cookbook (see the curl snippet above). +json_prompt = json.load(open("assets/edge/prompt.json")) +negative_prompt = json.load(open("assets/negative_prompt.json")) +control_edge = load_video("assets/edge/control_edge.mp4") + +videos = pipe( + prompt=json.dumps(json_prompt), + negative_prompt=json.dumps(negative_prompt), + control_videos={"edge": control_edge}, + num_frames=121, + height=720, + width=1280, + fps=30.0, + num_inference_steps=35, + guidance_scale=3.0, + control_guidance=1.5, + output="videos", +) +export_to_video(videos, "cosmos3_modular_transfer_edge.mp4", fps=30, macro_block_size=1) +``` + [[autodoc]] Cosmos3OmniModularPipeline - all diff --git a/src/diffusers/modular_pipelines/cosmos/before_denoise.py b/src/diffusers/modular_pipelines/cosmos/before_denoise.py index 6c1b852419c7..aa9fd94d32b3 100644 --- a/src/diffusers/modular_pipelines/cosmos/before_denoise.py +++ b/src/diffusers/modular_pipelines/cosmos/before_denoise.py @@ -960,3 +960,267 @@ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) ) self.set_block_state(state, block_state) return components, state + + +class Cosmos3TransferPrepareLatentsStep(ModularPipelineBlocks): + model_name = "cosmos3-omni" + + @property + def description(self) -> str: + return ( + "Per-chunk transfer latent prep: takes the clean target latents encoded by " + "Cosmos3TransferChunkVaeEncoderStep and builds the noisy target latents, velocity mask, condition latents " + "and conditioned-frame indexes for this chunk." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("transformer", Cosmos3OmniTransformer)] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + name="x0_tokens_vision", + type_hint=torch.Tensor, + required=True, + description="Clean target vision latents encoded from the seeded target frames.", + ), + InputParam( + name="current_conditional_frames", + type_hint=int, + required=True, + description="Number of pixel frames used to seed this chunk's target.", + ), + InputParam.template("generator"), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam("latents", type_hint=torch.Tensor, description="Noisy target latents for this chunk."), + OutputParam( + "velocity_mask", + type_hint=torch.Tensor, + description="Mask that zeroes the velocity on conditioned (clean) latent frames.", + ), + OutputParam( + "condition_latents", + type_hint=torch.Tensor, + description="Clean target latents on the conditioned frames (the autoregressive seed).", + ), + OutputParam( + "target_condition_indexes", + type_hint=list[int], + description="Latent-frame indexes fixed by the chunk's conditioning.", + ), + ] + + @torch.no_grad() + def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + device = components._execution_device + dtype = components.transformer.dtype + tcf = components.vae_scale_factor_temporal + + target_x0 = block_state.x0_tokens_vision.to(device=device) + current_conditional_frames = block_state.current_conditional_frames + + # Build the noisy target latents + conditioning mask from the clean target latents. + latent_t = target_x0.shape[2] + condition_mask = torch.zeros((latent_t, 1, 1), device=device, dtype=dtype) + latent_condition_frames = 0 + if current_conditional_frames > 0: + latent_condition_frames = (current_conditional_frames - 1) // tcf + 1 + condition_mask[:latent_condition_frames] = 1.0 + noise = randn_tensor(tuple(target_x0.shape), generator=block_state.generator, device=device, dtype=dtype) + block_state.latents = condition_mask * target_x0 + (1.0 - condition_mask) * noise + block_state.velocity_mask = 1.0 - condition_mask + block_state.condition_latents = condition_mask * target_x0 + block_state.target_condition_indexes = list(range(latent_condition_frames)) + + self.set_block_state(state, block_state) + return components, state + + +class Cosmos3TransferPackSequenceStep(ModularPipelineBlocks): + model_name = "cosmos3-omni" + + @property + def description(self) -> str: + return ( + "Pre-packs the three transfer CFG sequence variants: cond_full / uncond_full carry every control item, " + "the no-control branch drops them (only [text, target]) so the control axis can be amplified." + ) + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + name="cond_text_segment", type_hint=dict, required=True, description="Conditional text segment." + ), + InputParam( + name="uncond_text_segment", type_hint=dict, required=True, description="Unconditional text segment." + ), + InputParam( + name="control_latents", + type_hint=list[torch.Tensor], + required=True, + description="Clean control latents for this chunk, one per hint in canonical order.", + ), + InputParam( + name="latents", + type_hint=torch.Tensor, + required=True, + description="Noisy target latents for this chunk.", + ), + InputParam( + name="target_condition_indexes", + type_hint=list[int], + required=True, + description="Latent-frame indexes fixed by the chunk's conditioning.", + ), + InputParam(name="fps", type_hint=float, default=24.0, description="Frame rate of the generated video."), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "cond_full_static", + type_hint=dict, + kwargs_type="denoiser_input_fields", + description="Conditional [control..., target] transfer sequence carrying every control item.", + ), + OutputParam( + "cond_no_control_static", + type_hint=dict, + kwargs_type="denoiser_input_fields", + description="Conditional [target] transfer sequence with the control items dropped.", + ), + OutputParam( + "uncond_full_static", + type_hint=dict, + kwargs_type="denoiser_input_fields", + description="Unconditional [control..., target] transfer sequence for text CFG.", + ), + OutputParam( + "num_noisy_vision_tokens", + type_hint=int, + description="Number of noisy target vision tokens denoised each step.", + ), + ] + + @torch.no_grad() + def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + device = components._execution_device + num_hints = len(block_state.control_latents) + + def _vision_pack(text_segment: dict, include_controls: bool) -> dict: + if include_controls: + vision_items = [*block_state.control_latents, block_state.latents] + condition_indexes = [None] * num_hints + [block_state.target_condition_indexes] + clean_flags = [True] * num_hints + [False] + else: + vision_items = [block_state.latents] + condition_indexes = [block_state.target_condition_indexes] + clean_flags = [False] + + # Transfer packs [ctrl_1, ..., ctrl_N, target] into one vision segment + mrope_offset = text_segment["vision_start_temporal_offset"] + item_curr = text_segment["und_len"] + token_shapes = [] + sequence_index_parts = [] + mse_loss_index_parts = [] + noisy_frame_indexes_per_item = [] + mrope_id_parts = [] + num_vision_tokens = 0 + num_noisy_vision_tokens = 0 + for item, item_condition, is_clean in zip(vision_items, condition_indexes, clean_flags): + latent_t = item.shape[2] + if is_clean: + frame_condition = list(range(latent_t)) + else: + frame_condition = item_condition if item_condition is not None else [] + item_segment = components._prepare_vision_segment( + input_vision_tokens=item, + has_image_condition=False, + mrope_offset=mrope_offset, + vision_fps=block_state.fps, + curr=item_curr, + device=device, + condition_frame_indexes=frame_condition, + ) + token_shapes.extend(item_segment["vision_token_shapes"]) + sequence_index_parts.append(item_segment["vision_sequence_indexes"]) + mse_loss_index_parts.append(item_segment["vision_mse_loss_indexes"]) + noisy_frame_indexes_per_item.extend(item_segment["vision_noisy_frame_indexes"]) + mrope_id_parts.append(item_segment["vision_mrope_ids"]) + num_vision_tokens += item_segment["num_vision_tokens"] + num_noisy_vision_tokens += item_segment["num_noisy_vision_tokens"] + item_curr += item_segment["num_vision_tokens"] + + vision_segment = { + "vision_token_shapes": token_shapes, + "vision_sequence_indexes": torch.cat(sequence_index_parts, dim=0), + "vision_mse_loss_indexes": torch.cat(mse_loss_index_parts, dim=0), + "vision_noisy_frame_indexes": noisy_frame_indexes_per_item, + "vision_mrope_ids": torch.cat(mrope_id_parts, dim=1), + "num_vision_tokens": num_vision_tokens, + "num_noisy_vision_tokens": num_noisy_vision_tokens, + } + return { + **text_segment, + **vision_segment, + "position_ids": torch.cat([text_segment["text_mrope_ids"], vision_segment["vision_mrope_ids"]], dim=1), + "sequence_length": text_segment["und_len"] + vision_segment["num_vision_tokens"], + } + + block_state.cond_full_static = _vision_pack(block_state.cond_text_segment, include_controls=True) + block_state.cond_no_control_static = _vision_pack(block_state.cond_text_segment, include_controls=False) + block_state.uncond_full_static = _vision_pack(block_state.uncond_text_segment, include_controls=True) + block_state.num_noisy_vision_tokens = block_state.cond_full_static["num_noisy_vision_tokens"] + + self.set_block_state(state, block_state) + return components, state + + +class Cosmos3TransferSetTimestepsStep(ModularPipelineBlocks): + model_name = "cosmos3-omni" + + @property + def description(self) -> str: + return ( + "Resets the scheduler and computes timesteps for a single transfer chunk. UniPCMultistepScheduler keeps " + "per-step state on the instance, so it is reset per chunk (each autoregressive chunk is a full denoise)." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("scheduler", UniPCMultistepScheduler)] + + @property + def inputs(self) -> list[InputParam]: + return [InputParam.template("num_inference_steps", required=True)] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam("timesteps", type_hint=torch.Tensor, description="Scheduler timesteps for this chunk."), + OutputParam( + "num_warmup_steps", type_hint=int, description="Number of scheduler warmup steps for this chunk." + ), + ] + + @torch.no_grad() + def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + device = components._execution_device + components.scheduler.set_timesteps(block_state.num_inference_steps, device=device) + block_state.timesteps = components.scheduler.timesteps + block_state.num_warmup_steps = ( + len(block_state.timesteps) - block_state.num_inference_steps * components.scheduler.order + ) + self.set_block_state(state, block_state) + return components, state diff --git a/src/diffusers/modular_pipelines/cosmos/before_encoder.py b/src/diffusers/modular_pipelines/cosmos/before_encoder.py new file mode 100644 index 000000000000..2cdf68712cdf --- /dev/null +++ b/src/diffusers/modular_pipelines/cosmos/before_encoder.py @@ -0,0 +1,166 @@ +import math + +import torch + +from ...configuration_utils import FrozenDict +from ...video_processor import VideoProcessor +from ..modular_pipeline import ModularPipelineBlocks, PipelineState +from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam +from .modular_pipeline import Cosmos3OmniModularPipeline + + +class Cosmos3TransferSetupStep(ModularPipelineBlocks): + model_name = "cosmos3-omni" + + @property + def description(self) -> str: + return ( + "Preprocesses the transfer control videos and resolves the autoregressive chunk geometry " + "(total_frames / chunk_frames / num_chunks / stride). Chunk-invariant, so it runs once before the loop." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec( + "video_processor", + VideoProcessor, + config=FrozenDict({"vae_scale_factor": 16, "resample": "bilinear"}), + default_creation_method="from_config", + ), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + name="control_videos", + type_hint=dict, + required=True, + description="Mapping of hint name (edge/blur/depth/seg/wsm) to the control video for that modality.", + ), + InputParam( + name="height", type_hint=int, default=None, description="Height of the generated video in pixels." + ), + InputParam( + name="width", type_hint=int, default=None, description="Width of the generated video in pixels." + ), + InputParam( + name="num_frames", + type_hint=int, + default=None, + description="Optional cap on the number of output frames (defaults to the control video length).", + ), + InputParam( + name="num_video_frames_per_chunk", + type_hint=int, + default=None, + description="Number of pixel frames generated per autoregressive chunk.", + ), + InputParam( + name="num_conditional_frames", + type_hint=int, + default=1, + description="Number of frames each chunk reuses from the previous chunk's tail.", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam("height", type_hint=int, description="Resolved output height in pixels."), + OutputParam("width", type_hint=int, description="Resolved output width in pixels."), + OutputParam( + "control_frames", + type_hint=dict, + description="Preprocessed, time-padded control maps in canonical hint order.", + ), + OutputParam("total_frames", type_hint=int, description="Total number of output frames to generate."), + OutputParam("chunk_frames", type_hint=int, description="Number of pixel frames per autoregressive chunk."), + OutputParam("num_chunks", type_hint=int, description="Number of autoregressive chunks."), + OutputParam("stride", type_hint=int, description="Frame stride between consecutive chunks."), + ] + + @torch.no_grad() + def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + device = components._execution_device + dtype = components.transformer.dtype + + if block_state.height is None: + block_state.height = 720 + if block_state.width is None: + block_state.width = 1280 + + # Canonical hint order used both to validate and to order the preprocessed control maps. + hint_order = ["edge", "blur", "depth", "seg", "wsm"] + control_videos = block_state.control_videos + if not isinstance(control_videos, dict) or not control_videos: + raise ValueError("`control_videos` must be a non-empty dict mapping hint name -> control video.") + unknown = [k for k in control_videos if k not in hint_order] + if unknown: + raise ValueError(f"`control_videos` has unknown hint(s) {unknown}; expected keys from {hint_order}.") + if any(v is None for v in control_videos.values()): + raise ValueError("`control_videos` entries must be loaded videos, not None.") + + tcf = components.vae_scale_factor_temporal + sf = components.vae_scale_factor_spatial + if block_state.height % sf != 0 or block_state.width % sf != 0: + raise ValueError( + f"`height` and `width` must be multiples of {sf}, got ({block_state.height}, {block_state.width})." + ) + + # Preprocess every control map to [1, 3, T, H, W] in [-1, 1] at target geometry, in canonical hint order. + # The dict preserves this order, so downstream blocks just iterate control_frames (no separate hint_keys). + hint_keys = [k for k in hint_order if k in control_videos] + control_frames = { + key: components.video_processor.preprocess_video( + control_videos[key], height=block_state.height, width=block_state.width + ).to(device=device, dtype=dtype) + for key in hint_keys + } + + # Output frame count / chunking come from the (first) control video, optionally capped by num_frames. + total_frames = next(iter(control_frames.values())).shape[2] + if block_state.num_frames is not None: + total_frames = min(total_frames, block_state.num_frames) + total_frames = max(1, total_frames) + + per_chunk = ( + block_state.num_video_frames_per_chunk + if block_state.num_video_frames_per_chunk is not None + else total_frames + ) + chunk_frames = 1 if total_frames == 1 else per_chunk + chunk_frames = math.ceil((chunk_frames - 1) / tcf) * tcf + 1 + + if total_frames <= chunk_frames: + num_chunks, stride = 1, chunk_frames + else: + stride = chunk_frames - block_state.num_conditional_frames + if stride <= 0: + raise ValueError("`num_conditional_frames` must be smaller than `num_video_frames_per_chunk`.") + remaining = total_frames - chunk_frames + num_chunks = 1 + (remaining // stride + (1 if remaining % stride else 0)) + + # Reflect-pad each control map along time up to `padded` (repeat the last frame once the clip is too short to + # keep reflecting). No truncation here; per-chunk slicing happens later. + padded = max(total_frames, chunk_frames) + control_frames_padded = {} + for key, frames in control_frames.items(): + while frames.shape[2] < padded: + pad_len = min(frames.shape[2] - 1, padded - frames.shape[2]) + if pad_len <= 0: + pad_frame = frames[:, :, -1:].repeat(1, 1, padded - frames.shape[2], 1, 1) + frames = torch.cat([frames, pad_frame], dim=2) + break + frames = torch.cat([frames, frames.flip(dims=[2])[:, :, :pad_len]], dim=2) + control_frames_padded[key] = frames + block_state.control_frames = control_frames_padded + block_state.total_frames = total_frames + block_state.chunk_frames = chunk_frames + block_state.num_chunks = num_chunks + block_state.stride = stride + + self.set_block_state(state, block_state) + return components, state diff --git a/src/diffusers/modular_pipelines/cosmos/decoders.py b/src/diffusers/modular_pipelines/cosmos/decoders.py index c557e98ab9cc..a76e48501d85 100644 --- a/src/diffusers/modular_pipelines/cosmos/decoders.py +++ b/src/diffusers/modular_pipelines/cosmos/decoders.py @@ -114,3 +114,146 @@ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) block_state.sampling_rate = int(components.sound_tokenizer.config.sampling_rate) self.set_block_state(state, block_state) return components, state + + +class Cosmos3TransferDecodeChunkStep(ModularPipelineBlocks): + model_name = "cosmos3-omni" + + @property + def description(self) -> str: + return ( + "Decodes one transfer chunk's latents to pixels (float32, clamped to [-1, 1]), records it as the " + "autoregressive seed for the next chunk, and appends it to output_chunks (dropping the overlap that " + "later chunks share with the previous chunk's conditioning frames)." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("vae", AutoencoderKLWan)] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + name="latents", + type_hint=torch.Tensor, + required=True, + description="Denoised target latents for this chunk.", + ), + InputParam(name="chunk_id", type_hint=int, default=0, description="Index of the current chunk."), + InputParam( + name="current_conditional_frames", + type_hint=int, + required=True, + description="Number of pixel frames this chunk reused from the previous chunk.", + ), + InputParam( + name="output_chunks", + type_hint=list[torch.Tensor], + required=True, + description="Decoded pixel chunks accumulated so far.", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "previous_output", + type_hint=torch.Tensor, + description="Decoded pixels of this chunk, used to seed the next chunk.", + ), + OutputParam( + "output_chunks", + type_hint=list[torch.Tensor], + description="Decoded pixel chunks accumulated so far (with this chunk appended).", + ), + ] + + @torch.no_grad() + def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + latents = block_state.latents + vae_dtype = components.vae.dtype + mean = components._vae_latents_mean.to(device=latents.device, dtype=vae_dtype) + inv_std = components._vae_latents_inv_std.to(device=latents.device, dtype=vae_dtype) + z_raw = latents.to(vae_dtype) / inv_std.view(1, -1, 1, 1, 1) + mean.view(1, -1, 1, 1, 1) + output_video = components.vae.decode(z_raw).sample.to(torch.float32).clamp(-1, 1) + block_state.previous_output = output_video + chunk = ( + output_video if block_state.chunk_id == 0 else output_video[:, :, block_state.current_conditional_frames :] + ) + block_state.output_chunks = [*block_state.output_chunks, chunk] + self.set_block_state(state, block_state) + return components, state + + +class Cosmos3TransferStitchStep(ModularPipelineBlocks): + model_name = "cosmos3-omni" + + @property + def description(self) -> str: + return ( + "Concatenates the decoded transfer chunks along time, truncates to total_frames, and post-processes to " + "the requested output type. Transfer produces no audio, so sound / sampling_rate are None." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("vae", AutoencoderKLWan), + ComponentSpec( + "video_processor", + VideoProcessor, + config=FrozenDict({"vae_scale_factor": 16, "resample": "bilinear"}), + default_creation_method="from_config", + ), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + name="output_chunks", + type_hint=list[torch.Tensor], + required=True, + description="Decoded pixel chunks to stitch together.", + ), + InputParam( + name="total_frames", type_hint=int, required=True, description="Total number of output frames to keep." + ), + InputParam.template("output_type", default="pil"), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam("videos", description="The generated transfer video."), + OutputParam("sound", description="Always None for transfer (no audio)."), + OutputParam("sampling_rate", description="Always None for transfer (no audio)."), + ] + + @torch.no_grad() + def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + device = components._execution_device + + decoded = torch.cat(block_state.output_chunks, dim=2)[:, :, : block_state.total_frames] + block_state.videos = components.video_processor.postprocess_video( + decoded, output_type=block_state.output_type + )[0] + + if components.requires_safety_checker and block_state.output_type != "latent": + if getattr(components, "safety_checker", None) is None: + raise ValueError( + "Cosmos3 requires a safety checker by default. Call `pipe.enable_safety_checker()` to load it " + "(or pass your own), or opt out explicitly with `pipe.disable_safety_checker()`." + ) + block_state.videos = components._apply_video_safety_check( + block_state.videos, output_type=block_state.output_type, device=device + ) + + block_state.sound = None + block_state.sampling_rate = None + self.set_block_state(state, block_state) + return components, state diff --git a/src/diffusers/modular_pipelines/cosmos/denoise.py b/src/diffusers/modular_pipelines/cosmos/denoise.py index 7442ac77d2ad..7d2cdfe16fff 100644 --- a/src/diffusers/modular_pipelines/cosmos/denoise.py +++ b/src/diffusers/modular_pipelines/cosmos/denoise.py @@ -4,7 +4,12 @@ from ...models.transformers.transformer_cosmos3 import Cosmos3OmniTransformer from ...schedulers import UniPCMultistepScheduler -from ..modular_pipeline import BlockState, LoopSequentialPipelineBlocks, ModularPipelineBlocks, PipelineState +from ..modular_pipeline import ( + BlockState, + LoopSequentialPipelineBlocks, + ModularPipelineBlocks, + PipelineState, +) from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam from .modular_pipeline import Cosmos3OmniModularPipeline @@ -475,3 +480,324 @@ class Cosmos3VisionSoundActionDenoiseStep(Cosmos3DenoiseLoopWrapper): @property def description(self) -> str: return "Runs the vision, sound, and action Cosmos3 denoising loop." + + +class Cosmos3TransferLoopPrepareStep(ModularPipelineBlocks): + model_name = "cosmos3-omni" + + @property + def description(self) -> str: + return "Prepares the full [control..., target] and target-only vision token lists plus timesteps for one transfer iteration." + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("transformer", Cosmos3OmniTransformer)] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + name="control_latents", + type_hint=list[torch.Tensor], + required=True, + description="Clean control latents for this chunk, one per hint in canonical order.", + ), + InputParam( + name="latents", type_hint=torch.Tensor, required=True, description="Noisy target latents to denoise." + ), + InputParam( + name="num_noisy_vision_tokens", + type_hint=int, + required=True, + description="Number of noisy target vision tokens denoised each step.", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "vision_tokens_full", + type_hint=list[torch.Tensor], + description="Token list for the [control..., target] forward passes.", + ), + OutputParam( + "vision_tokens_target", + type_hint=list[torch.Tensor], + description="Token list for the target-only (no-control) forward pass.", + ), + OutputParam( + "vision_timesteps", type_hint=torch.Tensor, description="Timesteps for the noisy target tokens." + ), + ] + + @torch.no_grad() + def __call__(self, components: Cosmos3OmniModularPipeline, block_state: BlockState, i: int, t: torch.Tensor): + device = components._execution_device + dtype = components.transformer.dtype + block_state.vision_tokens_full = [c.to(device=device, dtype=dtype) for c in block_state.control_latents] + [ + block_state.latents.to(device=device, dtype=dtype) + ] + block_state.vision_tokens_target = [block_state.latents.to(device=device, dtype=dtype)] + block_state.vision_timesteps = torch.full((block_state.num_noisy_vision_tokens,), t.item(), device=device) + return components, block_state + + +class Cosmos3TransferLoopDenoiser(ModularPipelineBlocks): + # Dedicated (not Cosmos3LoopDenoiser): transfer runs up to 3 passes over different token sequences with nested + # control/text CFG and interval gating, which the generic cond/uncond denoiser cannot express. + model_name = "cosmos3-omni" + + @property + def description(self) -> str: + return ( + "Predicts the transfer velocity with nested control/text CFG over [control..., target]. Each branch is " + "gated by its guidance interval, and the result is masked so conditioned frames get zero velocity." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("transformer", Cosmos3OmniTransformer)] + + @property + def inputs(self) -> list[InputParam]: + return [ + # The three pre-packed CFG sequence variants (cond_full / cond_no_control / uncond_full) flow in as + # denoiser_input_fields, gathered generically like the other Cosmos3 denoisers. + InputParam.template("denoiser_input_fields"), + InputParam( + name="vision_tokens_full", + type_hint=list[torch.Tensor], + required=True, + description="Token list for the [control..., target] forward passes.", + ), + InputParam( + name="vision_tokens_target", + type_hint=list[torch.Tensor], + required=True, + description="Token list for the target-only (no-control) forward pass.", + ), + InputParam( + name="vision_timesteps", + type_hint=torch.Tensor, + required=True, + description="Timesteps for the noisy target tokens.", + ), + InputParam( + name="velocity_mask", + type_hint=torch.Tensor, + required=True, + description="Mask that zeroes the velocity on conditioned (clean) latent frames.", + ), + InputParam( + name="guidance_scale", + type_hint=float, + default=6.0, + description="Scale for text classifier-free guidance.", + ), + InputParam( + name="control_guidance", + type_hint=float, + default=1.0, + description="Scale for the control (structural) guidance axis.", + ), + InputParam( + name="guidance_interval", + type_hint=tuple, + default=None, + description="Timestep interval [lo, hi] over which text guidance is active (None = always).", + ), + InputParam( + name="control_guidance_interval", + type_hint=tuple, + default=None, + description="Timestep interval [lo, hi] over which control guidance is active (None = always).", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [OutputParam("velocity", type_hint=torch.Tensor, description="Predicted (masked) transfer velocity.")] + + @staticmethod + def _forward(components, static, vision_tokens, vision_timesteps): + preds_vision, _, _ = components.transformer( + input_ids=static["input_ids"], + text_indexes=static["text_indexes"], + position_ids=static["position_ids"], + und_len=static["und_len"], + sequence_length=static["sequence_length"], + vision_tokens=vision_tokens, + vision_token_shapes=static["vision_token_shapes"], + vision_sequence_indexes=static["vision_sequence_indexes"], + vision_mse_loss_indexes=static["vision_mse_loss_indexes"], + vision_timesteps=vision_timesteps, + vision_noisy_frame_indexes=static["vision_noisy_frame_indexes"], + ) + return preds_vision[-1] + + @torch.no_grad() + def __call__(self, components: Cosmos3OmniModularPipeline, block_state: BlockState, i: int, t: torch.Tensor): + # active-at: a None interval is always active; otherwise the timestep must fall within [lo, hi]. + guidance_interval = block_state.guidance_interval + guidance_active = guidance_interval is None or ( + float(guidance_interval[0]) <= float(t.item()) <= float(guidance_interval[1]) + ) + control_interval = block_state.control_guidance_interval + control_active = control_interval is None or ( + float(control_interval[0]) <= float(t.item()) <= float(control_interval[1]) + ) + step_guidance = block_state.guidance_scale if guidance_active else 1.0 + step_control = block_state.control_guidance if control_active else 1.0 + needs_text_cfg = step_guidance > 1.0 + needs_control_cfg = step_control != 1.0 + + denoiser_input_fields = block_state.denoiser_input_fields + cond_full_static = denoiser_input_fields["cond_full_static"] + cond_no_control_static = denoiser_input_fields["cond_no_control_static"] + uncond_full_static = denoiser_input_fields["uncond_full_static"] + + cond_full = self._forward( + components, cond_full_static, block_state.vision_tokens_full, block_state.vision_timesteps + ) + + cond_no_control = None + if needs_control_cfg: + cond_no_control = self._forward( + components, + cond_no_control_static, + block_state.vision_tokens_target, + block_state.vision_timesteps, + ) + + uncond_full = None + if needs_text_cfg: + uncond_full = self._forward( + components, + uncond_full_static, + block_state.vision_tokens_full, + block_state.vision_timesteps, + ) + + if needs_control_cfg and needs_text_cfg: + control_cond = cond_no_control + step_control * (cond_full - cond_no_control) + velocity = uncond_full + step_guidance * (control_cond - uncond_full) + elif needs_control_cfg: + velocity = cond_no_control + step_control * (cond_full - cond_no_control) + elif needs_text_cfg: + velocity = uncond_full + step_guidance * (cond_full - uncond_full) + else: + velocity = cond_full + + block_state.velocity = velocity * block_state.velocity_mask + return components, block_state + + +class Cosmos3TransferLoopSchedulerStep(ModularPipelineBlocks): + model_name = "cosmos3-omni" + + @property + def description(self) -> str: + return "Steps the scheduler and re-pins the conditioned frames exactly for one transfer iteration." + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("scheduler", UniPCMultistepScheduler)] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + name="latents", type_hint=torch.Tensor, required=True, description="Noisy target latents to update." + ), + InputParam( + name="velocity", + type_hint=torch.Tensor, + required=True, + description="Predicted (masked) transfer velocity.", + ), + InputParam( + name="velocity_mask", + type_hint=torch.Tensor, + required=True, + description="Mask that zeroes the velocity on conditioned (clean) latent frames.", + ), + InputParam( + name="condition_latents", + type_hint=torch.Tensor, + required=True, + description="Clean target latents on the conditioned frames (the autoregressive seed).", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [OutputParam("latents", type_hint=torch.Tensor, description="Updated target latents for this chunk.")] + + @torch.no_grad() + def __call__(self, components: Cosmos3OmniModularPipeline, block_state: BlockState, i: int, t: torch.Tensor): + block_state.latents = components.scheduler.step( + block_state.velocity.unsqueeze(0), t, block_state.latents.unsqueeze(0), return_dict=False + )[0].squeeze(0) + # Re-pin conditioned frames exactly (the autoregressive seed), guarding multistep drift. + block_state.latents = ( + block_state.velocity_mask * block_state.latents + + (1.0 - block_state.velocity_mask) * block_state.condition_latents + ) + return components, block_state + + +# auto_docstring +class Cosmos3TransferDenoiseStep(Cosmos3DenoiseLoopWrapper): + """ + Runs the per-chunk transfer denoising loop over scheduler timesteps. + + Components: + transformer (`Cosmos3OmniTransformer`) + scheduler (`UniPCMultistepScheduler`) + + Inputs: + timesteps (`Tensor`): + Timesteps for the denoising process. + num_inference_steps (`int`): + The number of denoising steps. + num_warmup_steps (`int`): + Number of scheduler warmup steps. + control_latents (`list`): + Clean control latents for this chunk, one per hint in canonical order. + latents (`Tensor`): + Noisy target latents to denoise. + num_noisy_vision_tokens (`int`): + Number of noisy target vision tokens denoised each step. + **denoiser_input_fields (`None`, *optional*): + conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + velocity_mask (`Tensor`): + Mask that zeroes the velocity on conditioned (clean) latent frames. + guidance_scale (`float`, *optional*, defaults to 6.0): + Scale for text classifier-free guidance. + control_guidance (`float`, *optional*, defaults to 1.0): + Scale for the control (structural) guidance axis. + guidance_interval (`tuple`, *optional*): + Timestep interval [lo, hi] over which text guidance is active (None = always). + control_guidance_interval (`tuple`, *optional*): + Timestep interval [lo, hi] over which control guidance is active (None = always). + latents (`Tensor`): + Noisy target latents to update. + condition_latents (`Tensor`): + Clean target latents on the conditioned frames (the autoregressive seed). + + Outputs: + latents (`Tensor`): + Updated target latents for this chunk. + """ + + block_classes = [ + Cosmos3TransferLoopPrepareStep, + Cosmos3TransferLoopDenoiser, + Cosmos3TransferLoopSchedulerStep, + ] + block_names = ["prepare_transfer", "denoiser", "update_transfer"] + + @property + def description(self) -> str: + return "Runs the per-chunk transfer denoising loop over scheduler timesteps." diff --git a/src/diffusers/modular_pipelines/cosmos/encoders.py b/src/diffusers/modular_pipelines/cosmos/encoders.py index 5d999b0121ba..001bb35aa581 100644 --- a/src/diffusers/modular_pipelines/cosmos/encoders.py +++ b/src/diffusers/modular_pipelines/cosmos/encoders.py @@ -17,6 +17,14 @@ logger = logging.get_logger(__name__) +# Transfer conditions on control signals (edge/blur/depth/seg/wsm), so it uses its own system prompt instead of the +# plain image/video ones. Defined here (not on the task pipeline) so the transfer text block is self-contained. +_SYSTEM_PROMPT_TRANSFER = ( + "You are a helpful assistant that generates images or videos following the user's instructions" + " and control signals (edge maps, blur, depth, or segmentation)." +) + + class Cosmos3TextEncoderStep(ModularPipelineBlocks): model_name = "cosmos3-omni" @@ -144,6 +152,125 @@ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) return components, state +class Cosmos3TransferTextStep(ModularPipelineBlocks): + model_name = "cosmos3-omni" + + @property + def description(self) -> str: + return ( + "Tokenizes the transfer prompt with the transfer system prompt. Transfer prompts are pre-upsampled JSON " + "captions passed through verbatim (no resolution/duration templates), so this is self-contained and does " + "not reuse the standard text step." + ) + + @staticmethod + def _check_inputs(block_state) -> None: + prompt = block_state.prompt + negative_prompt = block_state.negative_prompt + + if not isinstance(prompt, (str, list)) or ( + isinstance(prompt, list) and not all(isinstance(p, str) for p in prompt) + ): + raise ValueError(f"`prompt` must be a str or list of str, got {type(prompt).__name__}.") + if negative_prompt is not None and not isinstance(negative_prompt, (str, list)): + raise ValueError( + f"`negative_prompt` must be a str, list of str, or None, got {type(negative_prompt).__name__}." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("text_tokenizer", AutoTokenizer), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + name="prompt", + type_hint=str, + required=True, + description="The text prompt that guides Cosmos3 generation.", + ), + InputParam( + name="negative_prompt", + type_hint=str, + default=None, + description="The negative text prompt used for classifier-free guidance.", + ), + InputParam( + name="use_system_prompt", + type_hint=bool, + default=True, + description="Whether to prepend the Cosmos3 transfer system prompt.", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam("cond_input_ids", type_hint=torch.Tensor, description="Token IDs for the conditional prompt."), + OutputParam( + "uncond_input_ids", type_hint=torch.Tensor, description="Token IDs for the unconditional prompt." + ), + ] + + @torch.no_grad() + def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + self._check_inputs(block_state) + + if isinstance(block_state.prompt, list): + block_state.prompt = block_state.prompt[0] + if isinstance(block_state.negative_prompt, list): + block_state.negative_prompt = block_state.negative_prompt[0] + + if components.requires_safety_checker: + if getattr(components, "safety_checker", None) is None: + raise ValueError( + "Cosmos3 requires a safety checker by default. Call `pipe.enable_safety_checker()` to load it " + "(or pass your own), or opt out explicitly with `pipe.disable_safety_checker()`." + ) + device = components._execution_device + components.safety_checker.to(device) + try: + if not components.safety_checker.check_text_safety(block_state.prompt): + raise ValueError( + f"Cosmos Guardrail detected unsafe text in the prompt: {block_state.prompt}. " + "Please ensure that the prompt abides by the NVIDIA Open Model License Agreement." + ) + finally: + components.safety_checker.to("cpu") + + # Transfer prompts are pre-upsampled JSON captions: tokenize them verbatim (no resolution/duration templates) + # under the transfer system prompt. Kept self-contained here rather than adding a flag to the standard step. + negative_prompt = block_state.negative_prompt if block_state.negative_prompt is not None else "" + special_tokens = components.llm_special_tokens + + def _tokenize(text: str) -> list[int]: + conversations = [] + if block_state.use_system_prompt: + conversations.append({"role": "system", "content": _SYSTEM_PROMPT_TRANSFER}) + conversations.append({"role": "user", "content": text}) + encoding = components.text_tokenizer.apply_chat_template( + conversations, + tokenize=True, + add_generation_prompt=True, + add_vision_id=False, + return_dict=True, + ) + return list(encoding.input_ids) + [ + special_tokens["eos_token_id"], + special_tokens["start_of_generation"], + ] + + block_state.cond_input_ids = _tokenize(block_state.prompt) + block_state.uncond_input_ids = _tokenize(negative_prompt) + + self.set_block_state(state, block_state) + return components, state + + class Cosmos3ActionTextStep(ModularPipelineBlocks): model_name = "cosmos3-omni" @@ -531,6 +658,157 @@ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) return components, state +class Cosmos3TransferChunkVaeEncoderStep(ModularPipelineBlocks): + model_name = "cosmos3-omni" + + @property + def description(self) -> str: + return ( + "Per-chunk transfer VAE encode: slices + pads this chunk's control maps, seeds the target's conditioning " + "frames (first chunk from the input video, later chunks from the previous chunk's tail), and encodes both " + "the controls and the seeded target into clean Cosmos3 vision latents. Runs inside the autoregressive " + "chunk loop." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("vae", AutoencoderKLWan), + ComponentSpec( + "video_processor", + VideoProcessor, + config=FrozenDict({"vae_scale_factor": 16, "resample": "bilinear"}), + default_creation_method="from_config", + ), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam(name="chunk_id", type_hint=int, default=0, description="Index of the current chunk."), + InputParam( + name="previous_output", + default=None, + description="Decoded pixels of the previous chunk, used to seed later chunks.", + ), + InputParam( + name="control_frames", + type_hint=dict, + required=True, + description="Preprocessed, time-padded control maps in canonical hint order.", + ), + InputParam(name="chunk_frames", type_hint=int, required=True, description="Pixel frames per chunk."), + InputParam( + name="total_frames", type_hint=int, required=True, description="Total number of output frames." + ), + InputParam(name="stride", type_hint=int, required=True, description="Frame stride between chunks."), + InputParam( + name="height", type_hint=int, required=True, description="Height of the generated video in pixels." + ), + InputParam( + name="width", type_hint=int, required=True, description="Width of the generated video in pixels." + ), + InputParam( + name="video", + default=None, + description="Optional input video that seeds the first chunk's conditioning.", + ), + InputParam( + name="num_first_chunk_conditional_frames", + type_hint=int, + default=0, + description="Number of frames the first chunk reuses from the input video.", + ), + InputParam( + name="num_conditional_frames", + type_hint=int, + default=1, + description="Number of frames each later chunk reuses from the previous chunk's tail.", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "control_latents", + type_hint=list[torch.Tensor], + description="Clean control latents for this chunk, one per hint in canonical order.", + ), + OutputParam( + "x0_tokens_vision", + type_hint=torch.Tensor, + description="Clean target vision latents encoded from the seeded target frames.", + ), + OutputParam( + "current_conditional_frames", + type_hint=int, + description="Number of pixel frames actually used to seed this chunk's target.", + ), + ] + + @torch.no_grad() + def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + device = components._execution_device + dtype = components.transformer.dtype + + chunk_id = block_state.chunk_id + chunk_frames = block_state.chunk_frames + height = block_state.height + width = block_state.width + + # Slice this chunk's window out of the (padded) control maps and reflect-pad it up to a full chunk (repeat the + # last frame once too short to keep reflecting). control_frames is already in canonical hint order. + start_frame = chunk_id * block_state.stride + end_frame = min(start_frame + chunk_frames, block_state.total_frames) + chunk_controls = [] + for frames in block_state.control_frames.values(): + frames = frames[:, :, start_frame:end_frame] + while frames.shape[2] < chunk_frames: + pad_len = min(frames.shape[2] - 1, chunk_frames - frames.shape[2]) + if pad_len <= 0: + pad_frame = frames[:, :, -1:].repeat(1, 1, chunk_frames - frames.shape[2], 1, 1) + frames = torch.cat([frames, pad_frame], dim=2) + break + frames = torch.cat([frames, frames.flip(dims=[2])[:, :, :pad_len]], dim=2) + chunk_controls.append(frames) + + # Seed the target with conditioning frames (first chunk from the input video, later chunks from the + # previous chunk's tail), repeat-padding the remaining frames so the whole clip is well-defined. + target = torch.zeros(1, 3, chunk_frames, height, width, device=device, dtype=dtype) + current_conditional_frames = 0 + if chunk_id == 0 and block_state.num_first_chunk_conditional_frames > 0 and block_state.video is not None: + input_frames = components.video_processor.preprocess_video( + block_state.video, height=height, width=width + ).to(device=device, dtype=dtype) + current_conditional_frames = min( + block_state.num_first_chunk_conditional_frames, input_frames.shape[2], chunk_frames + ) + if current_conditional_frames > 0: + target[:, :, :current_conditional_frames] = input_frames[:, :, :current_conditional_frames] + elif chunk_id > 0 and block_state.previous_output is not None: + current_conditional_frames = min( + block_state.num_conditional_frames, block_state.previous_output.shape[2], chunk_frames + ) + if current_conditional_frames > 0: + target[:, :, :current_conditional_frames] = block_state.previous_output[ + :, :, -current_conditional_frames: + ].to(device=device, dtype=dtype) + if 0 < current_conditional_frames < chunk_frames: + fill = target[:, :, current_conditional_frames - 1 : current_conditional_frames] + target[:, :, current_conditional_frames:] = fill.expand( + -1, -1, chunk_frames - current_conditional_frames, -1, -1 + ) + + block_state.control_latents = [components._encode_video(ctrl).contiguous().float() for ctrl in chunk_controls] + block_state.x0_tokens_vision = components._encode_video(target).contiguous().float() + block_state.current_conditional_frames = current_conditional_frames + + self.set_block_state(state, block_state) + return components, state + + class Cosmos3ActionVisionVaeEncoderStep(ModularPipelineBlocks): model_name = "cosmos3-omni" diff --git a/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3.py b/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3.py index 191b940768d9..f04226c79150 100644 --- a/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3.py +++ b/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3.py @@ -1,6 +1,11 @@ import torch -from ..modular_pipeline import AutoPipelineBlocks, ConditionalPipelineBlocks, SequentialPipelineBlocks +from ..modular_pipeline import ( + AutoPipelineBlocks, + ConditionalPipelineBlocks, + PipelineState, + SequentialPipelineBlocks, +) from ..modular_pipeline_utils import InputParam, OutputParam from .after_decode import Cosmos3ActionOutputStep from .before_denoise import ( @@ -12,12 +17,22 @@ Cosmos3SoundDenoiseInputStep, Cosmos3SoundPackSequenceStep, Cosmos3SoundPrepareLatentsStep, + Cosmos3TransferPackSequenceStep, + Cosmos3TransferPrepareLatentsStep, + Cosmos3TransferSetTimestepsStep, Cosmos3VisionDenoiseInputStep, Cosmos3VisionPackSequenceStep, Cosmos3VisionPrepareLatentsStep, ) -from .decoders import Cosmos3SoundDecodeStep, Cosmos3VideoDecodeStep +from .before_encoder import Cosmos3TransferSetupStep +from .decoders import ( + Cosmos3SoundDecodeStep, + Cosmos3TransferDecodeChunkStep, + Cosmos3TransferStitchStep, + Cosmos3VideoDecodeStep, +) from .denoise import ( + Cosmos3TransferDenoiseStep, Cosmos3VisionActionDenoiseStep, Cosmos3VisionDenoiseStep, Cosmos3VisionSoundActionDenoiseStep, @@ -28,66 +43,150 @@ Cosmos3ActionVisionVaeEncoderStep, Cosmos3ImageVaeEncoderStep, Cosmos3TextEncoderStep, + Cosmos3TransferChunkVaeEncoderStep, + Cosmos3TransferTextStep, Cosmos3VideoVaeEncoderStep, ) +from .modular_pipeline import Cosmos3OmniModularPipeline + + +# auto_docstring +class Cosmos3TransferTextBlocks(SequentialPipelineBlocks): + """ + Transfer text branch: resolves the control-video chunk geometry, then tokenizes the (pre-upsampled) prompt in transfer mode using the per-chunk frame count. + + Components: + video_processor (`VideoProcessor`) + text_tokenizer (`AutoTokenizer`) + + Inputs: + control_videos (`dict`): + Mapping of hint name (edge/blur/depth/seg/wsm) to the control video for that modality. + height (`int`, *optional*): + Height of the generated video in pixels. + width (`int`, *optional*): + Width of the generated video in pixels. + num_frames (`int`, *optional*): + Optional cap on the number of output frames (defaults to the control video length). + num_video_frames_per_chunk (`int`, *optional*): + Number of pixel frames generated per autoregressive chunk. + num_conditional_frames (`int`, *optional*, defaults to 1): + Number of frames each chunk reuses from the previous chunk's tail. + prompt (`str`): + The text prompt that guides Cosmos3 generation. + negative_prompt (`str`, *optional*): + The negative text prompt used for classifier-free guidance. + use_system_prompt (`bool`, *optional*, defaults to True): + Whether to prepend the Cosmos3 transfer system prompt. + + Outputs: + height (`int`): + Resolved output height in pixels. + width (`int`): + Resolved output width in pixels. + control_frames (`dict`): + Preprocessed, time-padded control maps in canonical hint order. + total_frames (`int`): + Total number of output frames to generate. + chunk_frames (`int`): + Number of pixel frames per autoregressive chunk. + num_chunks (`int`): + Number of autoregressive chunks. + stride (`int`): + Frame stride between consecutive chunks. + cond_input_ids (`Tensor`): + Token IDs for the conditional prompt. + uncond_input_ids (`Tensor`): + Token IDs for the unconditional prompt. + """ + + model_name = "cosmos3-omni" + block_classes = [Cosmos3TransferSetupStep, Cosmos3TransferTextStep] + block_names = ["setup", "transfer_text"] + + @property + def description(self): + return ( + "Transfer text branch: resolves the control-video chunk geometry, then tokenizes the (pre-upsampled) " + "prompt in transfer mode using the per-chunk frame count." + ) # auto_docstring class Cosmos3AutoTextEncoderStep(AutoPipelineBlocks): """ Auto text encoder block for Cosmos3. + - Cosmos3TransferTextBlocks runs when control_videos are provided. - Cosmos3ActionTextStep runs when action is provided. - Cosmos3TextEncoderStep runs otherwise. Components: - text_tokenizer (`AutoTokenizer`) video_processor (`VideoProcessor`) + video_processor (`VideoProcessor`) + text_tokenizer (`AutoTokenizer`) Inputs: + control_videos (`dict`, *optional*): + Mapping of hint name (edge/blur/depth/seg/wsm) to the control video for that modality. + height (`int`, *optional*): + Height of the generated video in pixels. + width (`int`, *optional*): + Width of the generated video in pixels. + num_frames (`int`, *optional*): + Optional cap on the number of output frames (defaults to the control video length). + num_video_frames_per_chunk (`int`, *optional*): + Number of pixel frames generated per autoregressive chunk. + num_conditional_frames (`int`, *optional*, defaults to 1): + Number of frames each chunk reuses from the previous chunk's tail. prompt (`str`): The text prompt that guides Cosmos3 generation. negative_prompt (`str`, *optional*): The negative text prompt used for classifier-free guidance. + use_system_prompt (`bool`, *optional*, defaults to True): + Whether to prepend the Cosmos3 transfer system prompt. action (`CosmosActionCondition`, *optional*): Action-conditioning metadata and its reference visual input. - num_frames (`int`, *optional*): - Number of frames to generate. - height (`int`, *optional*): - Height of the generated video or image in pixels. - width (`int`, *optional*): - Width of the generated video or image in pixels. fps (`float`, *optional*, defaults to 24.0): Frame rate of the generated video. - use_system_prompt (`bool`, *optional*, defaults to True): - Whether to prepend the Cosmos3 system prompt. add_resolution_template (`bool`, *optional*, defaults to True): Whether to add resolution metadata to the prompt. add_duration_template (`bool`, *optional*, defaults to True): Whether to add duration metadata to the prompt. Outputs: - action_mode (`str`): - Requested action-generation mode. - num_frames (`int`): - Number of frames to generate. height (`int`): - Height of the generated video or image in pixels. + Resolved output height in pixels. width (`int`): - Width of the generated video or image in pixels. + Resolved output width in pixels. + control_frames (`dict`): + Preprocessed, time-padded control maps in canonical hint order. + total_frames (`int`): + Total number of output frames to generate. + chunk_frames (`int`): + Number of pixel frames per autoregressive chunk. + num_chunks (`int`): + Number of autoregressive chunks. + stride (`int`): + Frame stride between consecutive chunks. cond_input_ids (`Tensor`): Token IDs for the conditional prompt. uncond_input_ids (`Tensor`): Token IDs for the unconditional prompt. + action_mode (`str`): + Requested action-generation mode. + num_frames (`int`): + Number of frames to generate. """ model_name = "cosmos3-omni" - block_classes = [Cosmos3ActionTextStep, Cosmos3TextEncoderStep] - block_names = ["action_text", "text"] - block_trigger_inputs = ["action", None] + block_classes = [Cosmos3TransferTextBlocks, Cosmos3ActionTextStep, Cosmos3TextEncoderStep] + block_names = ["transfer_text", "action_text", "text"] + block_trigger_inputs = ["control_videos", "action", None] @property def description(self): return ( "Auto text encoder block for Cosmos3.\n" + + " - Cosmos3TransferTextBlocks runs when control_videos are provided.\n" + " - Cosmos3ActionTextStep runs when action is provided.\n" + " - Cosmos3TextEncoderStep runs otherwise." ) @@ -103,7 +202,8 @@ class Cosmos3AutoVaeEncoderStep(ConditionalPipelineBlocks): - when no action, image, or video conditioning is provided, this block is skipped. Components: - vae (`AutoencoderKLWan`) video_processor (`VideoProcessor`) + vae (`AutoencoderKLWan`) + video_processor (`VideoProcessor`) Inputs: action (`CosmosActionCondition`, *optional*): @@ -135,13 +235,17 @@ class Cosmos3AutoVaeEncoderStep(ConditionalPipelineBlocks): model_name = "cosmos3-omni" block_classes = [Cosmos3ActionVisionVaeEncoderStep, Cosmos3VideoVaeEncoderStep, Cosmos3ImageVaeEncoderStep] block_names = ["action_conditioning", "video_conditioning", "image_conditioning"] - block_trigger_inputs = ["action", "video", "image"] + block_trigger_inputs = ["action", "video", "image", "control_videos"] default_block_name = None def select_block(self, **kwargs) -> str | None: action = kwargs.get("action") image = kwargs.get("image") video = kwargs.get("video") + # Transfer preprocesses/encodes its control maps inside the denoise chunk loop, so the standard VAE + # conditioning stage is skipped when control_videos drive the workflow. + if kwargs.get("control_videos") is not None: + return None if action is not None: if image is not None or video is not None: raise ValueError( @@ -208,7 +312,9 @@ class Cosmos3DecodeStep(SequentialPipelineBlocks): Decodes denoised latents into modality outputs. Components: - vae (`AutoencoderKLWan`) video_processor (`VideoProcessor`) sound_tokenizer (`Cosmos3AVAEAudioTokenizer`) + vae (`AutoencoderKLWan`) + video_processor (`VideoProcessor`) + sound_tokenizer (`Cosmos3AVAEAudioTokenizer`) Inputs: latents (`Tensor`): @@ -236,13 +342,35 @@ def description(self) -> str: return "Decodes denoised latents into modality outputs." +class Cosmos3AutoDecodeStep(ConditionalPipelineBlocks): + model_name = "cosmos3-omni" + block_classes = [Cosmos3TransferStitchStep, Cosmos3DecodeStep] + block_names = ["transfer", "standard"] + block_trigger_inputs = ["control_videos"] + default_block_name = "standard" + + def select_block(self, **kwargs) -> str | None: + if kwargs.get("control_videos") is not None: + return "transfer" + return "standard" + + @property + def description(self) -> str: + return ( + "Selects the Cosmos3 decode workflow.\n" + + " - Cosmos3TransferStitchStep stitches the decoded transfer chunks when control_videos are provided.\n" + + " - Cosmos3DecodeStep decodes the denoised latents otherwise." + ) + + # auto_docstring class Cosmos3VisionCoreDenoiseStep(SequentialPipelineBlocks): """ Runs the text-and-vision Cosmos3 denoising workflow. Components: - transformer (`Cosmos3OmniTransformer`) scheduler (`UniPCMultistepScheduler`) + transformer (`Cosmos3OmniTransformer`) + scheduler (`UniPCMultistepScheduler`) Inputs: cond_input_ids (`None`): @@ -310,7 +438,8 @@ class Cosmos3VisionSoundCoreDenoiseStep(SequentialPipelineBlocks): Runs the text, vision, and sound Cosmos3 denoising workflow. Components: - transformer (`Cosmos3OmniTransformer`) scheduler (`UniPCMultistepScheduler`) + transformer (`Cosmos3OmniTransformer`) + scheduler (`UniPCMultistepScheduler`) Inputs: cond_input_ids (`None`): @@ -391,7 +520,8 @@ class Cosmos3VisionActionCoreDenoiseStep(SequentialPipelineBlocks): Runs the text, vision, and action Cosmos3 denoising workflow. Components: - transformer (`Cosmos3OmniTransformer`) scheduler (`UniPCMultistepScheduler`) + transformer (`Cosmos3OmniTransformer`) + scheduler (`UniPCMultistepScheduler`) Inputs: cond_input_ids (`None`): @@ -476,7 +606,8 @@ class Cosmos3VisionSoundActionCoreDenoiseStep(SequentialPipelineBlocks): Runs the text, vision, sound, and action Cosmos3 denoising workflow. Components: - transformer (`Cosmos3OmniTransformer`) scheduler (`UniPCMultistepScheduler`) + transformer (`Cosmos3OmniTransformer`) + scheduler (`UniPCMultistepScheduler`) Inputs: cond_input_ids (`None`): @@ -566,41 +697,337 @@ def outputs(self): ] +# auto_docstring +class Cosmos3TransferChunkDenoiseStep(SequentialPipelineBlocks): + """ + Autoregressive transfer chunk loop. Overrides __call__ to iterate chunks (the inner timestep loop is a non-leaf LoopSequentialPipelineBlocks, so this outer loop cannot itself be a LoopSequentialPipelineBlocks). Per-chunk cross-carry (previous_output, output_chunks) lives on PipelineState. + + Components: + vae (`AutoencoderKLWan`) + video_processor (`VideoProcessor`) + transformer (`Cosmos3OmniTransformer`) + scheduler (`UniPCMultistepScheduler`) + + Inputs: + chunk_id (`int`, *optional*, defaults to 0): + Index of the current chunk. + previous_output (`None`, *optional*): + Decoded pixels of the previous chunk, used to seed later chunks. + control_frames (`dict`): + Preprocessed, time-padded control maps in canonical hint order. + chunk_frames (`int`): + Pixel frames per chunk. + total_frames (`int`): + Total number of output frames. + stride (`int`): + Frame stride between chunks. + height (`int`): + Height of the generated video in pixels. + width (`int`): + Width of the generated video in pixels. + video (`None`, *optional*): + Optional input video that seeds the first chunk's conditioning. + num_first_chunk_conditional_frames (`int`, *optional*, defaults to 0): + Number of frames the first chunk reuses from the input video. + num_conditional_frames (`int`, *optional*, defaults to 1): + Number of frames each later chunk reuses from the previous chunk's tail. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + cond_text_segment (`dict`): + Conditional text segment. + uncond_text_segment (`dict`): + Unconditional text segment. + fps (`float`, *optional*, defaults to 24.0): + Frame rate of the generated video. + num_inference_steps (`int`): + The number of denoising steps. + **denoiser_input_fields (`None`, *optional*): + conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + guidance_scale (`float`, *optional*, defaults to 6.0): + Scale for text classifier-free guidance. + control_guidance (`float`, *optional*, defaults to 1.0): + Scale for the control (structural) guidance axis. + guidance_interval (`tuple`, *optional*): + Timestep interval [lo, hi] over which text guidance is active (None = always). + control_guidance_interval (`tuple`, *optional*): + Timestep interval [lo, hi] over which control guidance is active (None = always). + output_chunks (`list`): + Decoded pixel chunks accumulated so far. + num_chunks (`int`): + Number of autoregressive chunks. + + Outputs: + control_latents (`list`): + Clean control latents for this chunk, one per hint in canonical order. + x0_tokens_vision (`Tensor`): + Clean target vision latents encoded from the seeded target frames. + current_conditional_frames (`int`): + Number of pixel frames actually used to seed this chunk's target. + latents (`Tensor`): + Noisy target latents for this chunk. + velocity_mask (`Tensor`): + Mask that zeroes the velocity on conditioned (clean) latent frames. + condition_latents (`Tensor`): + Clean target latents on the conditioned frames (the autoregressive seed). + target_condition_indexes (`list`): + Latent-frame indexes fixed by the chunk's conditioning. + cond_full_static (`dict`): + Conditional [control..., target] transfer sequence carrying every control item. + cond_no_control_static (`dict`): + Conditional [target] transfer sequence with the control items dropped. + uncond_full_static (`dict`): + Unconditional [control..., target] transfer sequence for text CFG. + num_noisy_vision_tokens (`int`): + Number of noisy target vision tokens denoised each step. + timesteps (`Tensor`): + Scheduler timesteps for this chunk. + num_warmup_steps (`int`): + Number of scheduler warmup steps for this chunk. + vision_tokens_full (`list`): + Token list for the [control..., target] forward passes. + vision_tokens_target (`list`): + Token list for the target-only (no-control) forward pass. + vision_timesteps (`Tensor`): + Timesteps for the noisy target tokens. + velocity (`Tensor`): + Predicted (masked) transfer velocity. + previous_output (`Tensor`): + Decoded pixels of this chunk, used to seed the next chunk. + output_chunks (`list`): + Decoded pixel chunks accumulated so far (with this chunk appended). + """ + + model_name = "cosmos3-omni" + block_classes = [ + Cosmos3TransferChunkVaeEncoderStep, + Cosmos3TransferPrepareLatentsStep, + Cosmos3TransferPackSequenceStep, + Cosmos3TransferSetTimestepsStep, + Cosmos3TransferDenoiseStep, + Cosmos3TransferDecodeChunkStep, + ] + block_names = [ + "encode_transfer_chunk", + "prepare_transfer_latents", + "pack_transfer_sequence", + "set_timesteps", + "denoise", + "decode_chunk", + ] + + @property + def description(self) -> str: + return ( + "Autoregressive transfer chunk loop. Overrides __call__ to iterate chunks (the inner timestep loop is a " + "non-leaf LoopSequentialPipelineBlocks, so this outer loop cannot itself be a LoopSequentialPipelineBlocks). " + "Per-chunk cross-carry (previous_output, output_chunks) lives on PipelineState." + ) + + @property + def inputs(self) -> list[InputParam]: + return super().inputs + [ + InputParam(name="num_chunks", type_hint=int, required=True, description="Number of autoregressive chunks.") + ] + + @torch.no_grad() + def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState: + num_chunks = state.get("num_chunks") + state.set("output_chunks", []) + state.set("previous_output", None) + for chunk_id in range(num_chunks): + state.set("chunk_id", chunk_id) + for _, block in self.sub_blocks.items(): + components, state = block(components, state) + return components, state + + +# auto_docstring +class Cosmos3TransferCoreDenoiseStep(SequentialPipelineBlocks): + """ + Transfer denoise stage: prepare shared text segments once, then run the autoregressive chunk loop. + + Components: + transformer (`Cosmos3OmniTransformer`) + vae (`AutoencoderKLWan`) + video_processor (`VideoProcessor`) + scheduler (`UniPCMultistepScheduler`) + + Inputs: + cond_input_ids (`None`): + Token IDs for the conditional prompt. + uncond_input_ids (`None`): + Token IDs for the unconditional prompt. + chunk_id (`int`, *optional*, defaults to 0): + Index of the current chunk. + previous_output (`None`, *optional*): + Decoded pixels of the previous chunk, used to seed later chunks. + control_frames (`dict`): + Preprocessed, time-padded control maps in canonical hint order. + chunk_frames (`int`): + Pixel frames per chunk. + total_frames (`int`): + Total number of output frames. + stride (`int`): + Frame stride between chunks. + height (`int`): + Height of the generated video in pixels. + width (`int`): + Width of the generated video in pixels. + video (`None`, *optional*): + Optional input video that seeds the first chunk's conditioning. + num_first_chunk_conditional_frames (`int`, *optional*, defaults to 0): + Number of frames the first chunk reuses from the input video. + num_conditional_frames (`int`, *optional*, defaults to 1): + Number of frames each later chunk reuses from the previous chunk's tail. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + fps (`float`, *optional*, defaults to 24.0): + Frame rate of the generated video. + num_inference_steps (`int`): + The number of denoising steps. + **denoiser_input_fields (`None`, *optional*): + conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + guidance_scale (`float`, *optional*, defaults to 6.0): + Scale for text classifier-free guidance. + control_guidance (`float`, *optional*, defaults to 1.0): + Scale for the control (structural) guidance axis. + guidance_interval (`tuple`, *optional*): + Timestep interval [lo, hi] over which text guidance is active (None = always). + control_guidance_interval (`tuple`, *optional*): + Timestep interval [lo, hi] over which control guidance is active (None = always). + output_chunks (`list`): + Decoded pixel chunks accumulated so far. + num_chunks (`int`): + Number of autoregressive chunks. + + Outputs: + cond_text_segment (`dict`): + Conditional text segment for the denoiser. + uncond_text_segment (`dict`): + Unconditional text segment for the denoiser. + control_latents (`list`): + Clean control latents for this chunk, one per hint in canonical order. + x0_tokens_vision (`Tensor`): + Clean target vision latents encoded from the seeded target frames. + current_conditional_frames (`int`): + Number of pixel frames actually used to seed this chunk's target. + latents (`Tensor`): + Noisy target latents for this chunk. + velocity_mask (`Tensor`): + Mask that zeroes the velocity on conditioned (clean) latent frames. + condition_latents (`Tensor`): + Clean target latents on the conditioned frames (the autoregressive seed). + target_condition_indexes (`list`): + Latent-frame indexes fixed by the chunk's conditioning. + cond_full_static (`dict`): + Conditional [control..., target] transfer sequence carrying every control item. + cond_no_control_static (`dict`): + Conditional [target] transfer sequence with the control items dropped. + uncond_full_static (`dict`): + Unconditional [control..., target] transfer sequence for text CFG. + num_noisy_vision_tokens (`int`): + Number of noisy target vision tokens denoised each step. + timesteps (`Tensor`): + Scheduler timesteps for this chunk. + num_warmup_steps (`int`): + Number of scheduler warmup steps for this chunk. + vision_tokens_full (`list`): + Token list for the [control..., target] forward passes. + vision_tokens_target (`list`): + Token list for the target-only (no-control) forward pass. + vision_timesteps (`Tensor`): + Timesteps for the noisy target tokens. + velocity (`Tensor`): + Predicted (masked) transfer velocity. + previous_output (`Tensor`): + Decoded pixels of this chunk, used to seed the next chunk. + output_chunks (`list`): + Decoded pixel chunks accumulated so far (with this chunk appended). + """ + + model_name = "cosmos3-omni" + block_classes = [ + Cosmos3PrepareTextSegmentsStep, + Cosmos3TransferChunkDenoiseStep, + ] + block_names = ["prepare_text_segments", "chunk_denoise"] + + @property + def description(self) -> str: + return "Transfer denoise stage: prepare shared text segments once, then run the autoregressive chunk loop." + + # auto_docstring class Cosmos3AutoCoreDenoiseStep(ConditionalPipelineBlocks): """ Selects the Cosmos3 core denoising workflow. + - transfer runs the autoregressive control-video (ControlNet-style) chunk loop when control_videos are provided. - vision_sound_action runs when action and enable_sound are provided. - vision_action runs when action is provided. - vision_sound runs when enable_sound is true. - vision runs otherwise. Components: - transformer (`Cosmos3OmniTransformer`) scheduler (`UniPCMultistepScheduler`) + transformer (`Cosmos3OmniTransformer`) + vae (`AutoencoderKLWan`) + video_processor (`VideoProcessor`) + scheduler (`UniPCMultistepScheduler`) Inputs: cond_input_ids (`None`): Token IDs for the conditional prompt. uncond_input_ids (`None`): Token IDs for the unconditional prompt. - x0_tokens_vision (`Tensor`, *optional*): - Vision latents encoded from the conditioning image or video. - vision_condition_frames (`list`, *optional*): - Latent-frame indexes fixed by visual conditioning. - num_frames (`int`): - Number of frames to generate. + chunk_id (`int`, *optional*, defaults to 0): + Index of the current chunk. + previous_output (`None`, *optional*): + Decoded pixels of the previous chunk, used to seed later chunks. + control_frames (`dict`, *optional*): + Preprocessed, time-padded control maps in canonical hint order. + chunk_frames (`int`, *optional*): + Pixel frames per chunk. + total_frames (`int`, *optional*): + Total number of output frames. + stride (`int`, *optional*): + Frame stride between chunks. height (`int`): Height of the generated video in pixels. width (`int`): Width of the generated video in pixels. - fps (`float`, *optional*, defaults to 24.0): - Frame rate of the generated video. - latents (`Tensor`): - Pre-generated noisy vision latents. + video (`None`, *optional*): + Optional input video that seeds the first chunk's conditioning. + num_first_chunk_conditional_frames (`int`, *optional*, defaults to 0): + Number of frames the first chunk reuses from the input video. + num_conditional_frames (`int`, *optional*, defaults to 1): + Number of frames each later chunk reuses from the previous chunk's tail. generator (`Generator`, *optional*): Torch generator for deterministic generation. + fps (`float`, *optional*, defaults to 24.0): + Frame rate of the generated video. num_inference_steps (`int`): The number of denoising steps. + **denoiser_input_fields (`None`, *optional*): + conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + guidance_scale (`float`, *optional*, defaults to 6.0): + Scale for text classifier-free guidance. + control_guidance (`float`, *optional*, defaults to 1.0): + Scale for the control (structural) guidance axis. + guidance_interval (`tuple`, *optional*): + Timestep interval [lo, hi] over which text guidance is active (None = always). + control_guidance_interval (`tuple`, *optional*): + Timestep interval [lo, hi] over which control guidance is active (None = always). + output_chunks (`list`, *optional*): + Decoded pixel chunks accumulated so far. + num_chunks (`int`, *optional*): + Number of autoregressive chunks. + x0_tokens_vision (`Tensor`, *optional*): + Vision latents encoded from the conditioning image or video. + vision_condition_frames (`list`, *optional*): + Latent-frame indexes fixed by visual conditioning. + num_frames (`int`, *optional*): + Number of frames to generate. + latents (`Tensor`): + Pre-generated noisy vision latents. sound_latents (`Tensor`, *optional*): Pre-generated noisy sound latents. action (`CosmosActionCondition`, *optional*): @@ -609,16 +1036,52 @@ class Cosmos3AutoCoreDenoiseStep(ConditionalPipelineBlocks): Action-frame indexes fixed by action conditioning. action_latents (`Tensor`, *optional*): Pre-generated noisy action latents. - **denoiser_input_fields (`None`, *optional*): - conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. - guidance_scale (`float`, *optional*, defaults to 6.0): - Scale for classifier-free guidance. enable_sound (`bool`, *optional*, defaults to False): Whether to generate a synchronized sound track. Outputs: + cond_text_segment (`dict`): + Conditional text segment for the denoiser. + uncond_text_segment (`dict`): + Unconditional text segment for the denoiser. + control_latents (`list`): + Clean control latents for this chunk, one per hint in canonical order. + x0_tokens_vision (`Tensor`): + Clean target vision latents encoded from the seeded target frames. + current_conditional_frames (`int`): + Number of pixel frames actually used to seed this chunk's target. latents (`Tensor`): - Denoised latents. + Noisy target latents for this chunk. + velocity_mask (`Tensor`): + Mask that zeroes the velocity on conditioned (clean) latent frames. + condition_latents (`Tensor`): + Clean target latents on the conditioned frames (the autoregressive seed). + target_condition_indexes (`list`): + Latent-frame indexes fixed by the chunk's conditioning. + cond_full_static (`dict`): + Conditional [control..., target] transfer sequence carrying every control item. + cond_no_control_static (`dict`): + Conditional [target] transfer sequence with the control items dropped. + uncond_full_static (`dict`): + Unconditional [control..., target] transfer sequence for text CFG. + num_noisy_vision_tokens (`int`): + Number of noisy target vision tokens denoised each step. + timesteps (`Tensor`): + Scheduler timesteps for this chunk. + num_warmup_steps (`int`): + Number of scheduler warmup steps for this chunk. + vision_tokens_full (`list`): + Token list for the [control..., target] forward passes. + vision_tokens_target (`list`): + Token list for the target-only (no-control) forward pass. + vision_timesteps (`Tensor`): + Timesteps for the noisy target tokens. + velocity (`Tensor`): + Predicted (masked) transfer velocity. + previous_output (`Tensor`): + Decoded pixels of this chunk, used to seed the next chunk. + output_chunks (`list`): + Decoded pixel chunks accumulated so far (with this chunk appended). sound_latents (`Tensor`): Denoised sound latents. action_latents (`Tensor`): @@ -627,13 +1090,14 @@ class Cosmos3AutoCoreDenoiseStep(ConditionalPipelineBlocks): model_name = "cosmos3-omni" block_classes = [ + Cosmos3TransferCoreDenoiseStep, Cosmos3VisionSoundActionCoreDenoiseStep, Cosmos3VisionActionCoreDenoiseStep, Cosmos3VisionSoundCoreDenoiseStep, Cosmos3VisionCoreDenoiseStep, ] - block_names = ["vision_sound_action", "vision_action", "vision_sound", "vision"] - block_trigger_inputs = ["action", "enable_sound"] + block_names = ["transfer", "vision_sound_action", "vision_action", "vision_sound", "vision"] + block_trigger_inputs = ["action", "enable_sound", "control_videos"] default_block_name = "vision" @property @@ -652,6 +1116,8 @@ def inputs(self): def select_block(self, **kwargs) -> str | None: action = kwargs.get("action") enable_sound = kwargs.get("enable_sound") + if kwargs.get("control_videos") is not None: + return "transfer" if action is not None and enable_sound: return "vision_sound_action" if action is not None: @@ -664,6 +1130,7 @@ def select_block(self, **kwargs) -> str | None: def description(self): return ( "Selects the Cosmos3 core denoising workflow.\n" + + " - transfer runs the autoregressive control-video (ControlNet-style) chunk loop when control_videos are provided.\n" + " - vision_sound_action runs when action and enable_sound are provided.\n" + " - vision_action runs when action is provided.\n" + " - vision_sound runs when enable_sound is true.\n" @@ -689,27 +1156,36 @@ class Cosmos3OmniBlocks(SequentialPipelineBlocks): - `action_inverse_dynamics`: requires `prompt`, `action` Components: - text_tokenizer (`AutoTokenizer`) video_processor (`VideoProcessor`) vae (`AutoencoderKLWan`) transformer - (`Cosmos3OmniTransformer`) scheduler (`UniPCMultistepScheduler`) sound_tokenizer - (`Cosmos3AVAEAudioTokenizer`) + video_processor (`VideoProcessor`) + text_tokenizer (`AutoTokenizer`) + vae (`AutoencoderKLWan`) + transformer (`Cosmos3OmniTransformer`) + scheduler (`UniPCMultistepScheduler`) + sound_tokenizer (`Cosmos3AVAEAudioTokenizer`) Inputs: + control_videos (`dict`, *optional*): + Mapping of hint name (edge/blur/depth/seg/wsm) to the control video for that modality. + height (`int`, *optional*): + Height of the generated video in pixels. + width (`int`, *optional*): + Width of the generated video in pixels. + num_frames (`int`, *optional*): + Optional cap on the number of output frames (defaults to the control video length). + num_video_frames_per_chunk (`int`, *optional*): + Number of pixel frames generated per autoregressive chunk. + num_conditional_frames (`int`, *optional*, defaults to 1): + Number of frames each chunk reuses from the previous chunk's tail. prompt (`str`): The text prompt that guides Cosmos3 generation. negative_prompt (`str`, *optional*): The negative text prompt used for classifier-free guidance. + use_system_prompt (`bool`, *optional*, defaults to True): + Whether to prepend the Cosmos3 transfer system prompt. action (`CosmosActionCondition`, *optional*): Action-conditioning metadata and its reference visual input. - num_frames (`int`, *optional*): - Number of frames to generate. - height (`int`, *optional*): - Height of the generated video or image in pixels. - width (`int`, *optional*): - Width of the generated video or image in pixels. fps (`float`, *optional*, defaults to 24.0): Frame rate of the generated video. - use_system_prompt (`bool`, *optional*, defaults to True): - Whether to prepend the Cosmos3 system prompt. add_resolution_template (`bool`, *optional*, defaults to True): Whether to add resolution metadata to the prompt. add_duration_template (`bool`, *optional*, defaults to True): @@ -722,26 +1198,40 @@ class Cosmos3OmniBlocks(SequentialPipelineBlocks): Which end of a longer conditioning video to use: `first` or `last`. image (`None`, *optional*): Reference image for image-to-video conditioning. + chunk_id (`int`, *optional*, defaults to 0): + Index of the current chunk. + previous_output (`None`, *optional*): + Decoded pixels of the previous chunk, used to seed later chunks. + num_first_chunk_conditional_frames (`int`, *optional*, defaults to 0): + Number of frames the first chunk reuses from the input video. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + num_inference_steps (`int`): + The number of denoising steps. + **denoiser_input_fields (`None`, *optional*): + conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + guidance_scale (`float`, *optional*, defaults to 6.0): + Scale for text classifier-free guidance. + control_guidance (`float`, *optional*, defaults to 1.0): + Scale for the control (structural) guidance axis. + guidance_interval (`tuple`, *optional*): + Timestep interval [lo, hi] over which text guidance is active (None = always). + control_guidance_interval (`tuple`, *optional*): + Timestep interval [lo, hi] over which control guidance is active (None = always). + output_chunks (`list`, *optional*): + Decoded pixel chunks accumulated so far. x0_tokens_vision (`Tensor`, *optional*): Vision latents encoded from the conditioning image or video. vision_condition_frames (`list`, *optional*): Latent-frame indexes fixed by visual conditioning. latents (`Tensor`): Pre-generated noisy vision latents. - generator (`Generator`, *optional*): - Torch generator for deterministic generation. - num_inference_steps (`int`): - The number of denoising steps. sound_latents (`Tensor`, *optional*): Pre-generated noisy sound latents. action_condition_frame_indexes (`list`, *optional*): Action-frame indexes fixed by action conditioning. action_latents (`Tensor`, *optional*): Pre-generated noisy action latents. - **denoiser_input_fields (`None`, *optional*): - conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. - guidance_scale (`float`, *optional*, defaults to 6.0): - Scale for classifier-free guidance. enable_sound (`bool`, *optional*, defaults to False): Whether to generate a synchronized sound track. output_type (`str`, *optional*, defaults to pil): @@ -763,7 +1253,7 @@ class Cosmos3OmniBlocks(SequentialPipelineBlocks): Cosmos3AutoTextEncoderStep, Cosmos3AutoVaeEncoderStep, Cosmos3AutoCoreDenoiseStep, - Cosmos3DecodeStep, + Cosmos3AutoDecodeStep, Cosmos3ActionOutputStep, ] block_names = ["text_encoder", "vae_encoder", "denoise", "decode", "after_decode"] @@ -784,6 +1274,16 @@ class Cosmos3OmniBlocks(SequentialPipelineBlocks): def description(self): return "Modular pipeline blocks for Cosmos3 generation modes." + def get_workflow(self, workflow_name: str): + if workflow_name == "transfer": + raise NotImplementedError( + 'The standalone "transfer" workflow is temporarily unavailable because its nested autoregressive ' + "chunk and denoising loops cannot be preserved by the current workflow extraction logic. Transfer " + "remains available through the full Cosmos3OmniBlocks pipeline. The standalone workflow will be " + "enabled after migration to the upcoming composable nested-loop abstraction." + ) + return super().get_workflow(workflow_name) + @property def outputs(self): return [