From 4fea4dab4035e150da7e4623d6d398a11854d807 Mon Sep 17 00:00:00 2001 From: Yuliya Zhautouskaya Date: Wed, 8 Jul 2026 07:45:44 -0700 Subject: [PATCH 1/7] Add transfer support to the Cosmos3 modular pipeline --- docs/source/en/api/pipelines/cosmos3.md | 59 ++- .../cosmos/before_denoise.py | 375 ++++++++++++++++++ .../modular_pipelines/cosmos/decoders.py | 113 ++++++ .../modular_pipelines/cosmos/denoise.py | 199 +++++++++- .../modular_pipelines/cosmos/encoders.py | 97 +++++ .../cosmos/modular_blocks_cosmos3.py | 133 ++++++- 6 files changed, 964 insertions(+), 12 deletions(-) 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..b0485df359de 100644 --- a/src/diffusers/modular_pipelines/cosmos/before_denoise.py +++ b/src/diffusers/modular_pipelines/cosmos/before_denoise.py @@ -1,11 +1,15 @@ import copy +import math import torch +from ...configuration_utils import FrozenDict +from ...models.autoencoders.autoencoder_kl_wan import AutoencoderKLWan from ...models.transformers.transformer_cosmos3 import Cosmos3OmniTransformer from ...pipelines.cosmos.pipeline_cosmos3_omni import _EMBODIMENT_TO_DOMAIN_ID, CosmosActionCondition from ...schedulers import UniPCMultistepScheduler from ...utils.torch_utils import randn_tensor +from ...video_processor import VideoProcessor from ..modular_pipeline import ModularPipelineBlocks, PipelineState from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam from .modular_pipeline import Cosmos3OmniModularPipeline @@ -960,3 +964,374 @@ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) ) self.set_block_state(state, block_state) return components, state + + +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("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="control_videos", required=True), + InputParam(name="height", default=None), + InputParam(name="width", default=None), + InputParam(name="num_frames", default=None), + InputParam(name="num_video_frames_per_chunk", default=None), + InputParam(name="num_conditional_frames", type_hint=int, default=1), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam("height"), + OutputParam("width"), + OutputParam("control_frames"), + OutputParam("total_frames"), + OutputParam("chunk_frames"), + OutputParam("num_chunks"), + OutputParam("stride"), + ] + + @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 = int(components.vae.config.scale_factor_temporal) + sf = int(components.vae.config.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 + + +class Cosmos3TransferPrepareLatentsStep(ModularPipelineBlocks): + model_name = "cosmos3-omni" + + @property + def description(self) -> str: + return ( + "Per-chunk transfer latent prep: slice + pad the chunk's control maps, seed the target's conditioning " + "frames (first chunk from the input video, later chunks from the previous chunk's tail), encode the " + "controls as clean latents and build the noisy target latents, velocity mask and condition latents." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("transformer", Cosmos3OmniTransformer), + 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 [ + # Loop carry (set by the chunk-loop wrapper): + InputParam(name="chunk_id", type_hint=int, default=0), + InputParam(name="previous_output", default=None), + # Setup artifacts the loop can't cheaply re-derive (preprocessed controls + chunk geometry): + InputParam(name="control_frames", required=True), + InputParam(name="chunk_frames", required=True), + InputParam(name="total_frames", required=True), + InputParam(name="stride", required=True), + # User inputs (mirrors how the other prepare-latents steps declare height/width/generator/etc.): + InputParam(name="height", required=True), + InputParam(name="width", required=True), + InputParam(name="video", default=None), + InputParam(name="num_first_chunk_conditional_frames", type_hint=int, default=0), + InputParam(name="num_conditional_frames", type_hint=int, default=1), + InputParam(name="generator", default=None), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam("latents"), + OutputParam("control_latents"), + OutputParam("velocity_mask"), + OutputParam("condition_latents"), + OutputParam("target_condition_indexes"), + OutputParam("current_conditional_frames"), + ] + + @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 + tcf = int(components.vae.config.scale_factor_temporal) + + # 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 + ) + + # Encode controls as clean latents and build the noisy target latents + conditioning mask. + block_state.control_latents = [components._encode_video(ctrl).contiguous().float() for ctrl in chunk_controls] + target_x0 = components._encode_video(target).contiguous().float() + 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)) + block_state.current_conditional_frames = current_conditional_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", required=True), + InputParam(name="uncond_text_segment", required=True), + InputParam(name="control_latents", required=True), + InputParam(name="latents", required=True), + InputParam(name="target_condition_indexes", required=True), + InputParam(name="fps", type_hint=float, default=24.0), + InputParam(name="share_vision_temporal_positions", type_hint=bool, default=True), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam("cond_full_static"), + OutputParam("cond_no_control_static"), + OutputParam("uncond_full_static"), + OutputParam("num_noisy_vision_tokens"), + ] + + @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] + vision_segment = components._prepare_vision_segment( + input_vision_tokens=vision_items, + has_image_condition=False, + mrope_offset=text_segment["vision_start_temporal_offset"], + vision_fps=block_state.fps, + curr=text_segment["und_len"], + device=device, + condition_frame_indexes=condition_indexes, + clean_item_flags=clean_flags, + share_vision_temporal_positions=block_state.share_vision_temporal_positions, + ) + 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"), + OutputParam("num_warmup_steps"), + ] + + @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/decoders.py b/src/diffusers/modular_pipelines/cosmos/decoders.py index c557e98ab9cc..a52b5e9d45ef 100644 --- a/src/diffusers/modular_pipelines/cosmos/decoders.py +++ b/src/diffusers/modular_pipelines/cosmos/decoders.py @@ -114,3 +114,116 @@ 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", required=True), + InputParam(name="chunk_id", type_hint=int, default=0), + InputParam(name="current_conditional_frames", required=True), + InputParam(name="output_chunks", required=True), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam("previous_output"), + OutputParam("output_chunks"), + ] + + @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", required=True), + InputParam(name="total_frames", required=True), + InputParam.template("output_type", default="pil"), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam("videos"), + OutputParam("sound"), + OutputParam("sampling_rate"), + ] + + @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..15b6d65ac000 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,195 @@ 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", required=True), + InputParam(name="latents", required=True), + InputParam(name="num_noisy_vision_tokens", required=True), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam("vision_tokens_full"), + OutputParam("vision_tokens_target"), + OutputParam("vision_timesteps"), + ] + + @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): + 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 [ + InputParam(name="cond_full_static", required=True), + InputParam(name="cond_no_control_static", required=True), + InputParam(name="uncond_full_static", required=True), + InputParam(name="vision_tokens_full", required=True), + InputParam(name="vision_tokens_target", required=True), + InputParam(name="vision_timesteps", required=True), + InputParam(name="velocity_mask", required=True), + InputParam(name="guidance_scale", default=6.0), + InputParam(name="control_guidance", type_hint=float, default=1.0), + InputParam(name="guidance_interval", default=None), + InputParam(name="control_guidance_interval", default=None), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [OutputParam("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 + + cond_full = self._forward( + components, block_state.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, + block_state.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, + block_state.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", required=True), + InputParam(name="velocity", required=True), + InputParam(name="velocity_mask", required=True), + InputParam(name="condition_latents", required=True), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [OutputParam("latents")] + + @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 + + +class Cosmos3TransferDenoiseStep(Cosmos3DenoiseLoopWrapper): + 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..67bbd70b4d0b 100644 --- a/src/diffusers/modular_pipelines/cosmos/encoders.py +++ b/src/diffusers/modular_pipelines/cosmos/encoders.py @@ -144,6 +144,103 @@ 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 in transfer mode using the per-chunk frame count. Transfer prompts are " + "pre-upsampled JSON captions passed through verbatim (the metadata templates are skipped)." + ) + + @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), + InputParam(name="negative_prompt", default=None), + InputParam(name="chunk_frames", type_hint=int, required=True), + InputParam(name="height", default=None), + InputParam(name="width", default=None), + InputParam(name="fps", type_hint=float, default=24.0), + InputParam(name="use_system_prompt", type_hint=bool, default=True), + InputParam(name="add_resolution_template", type_hint=bool, default=True), + InputParam(name="add_duration_template", type_hint=bool, default=True), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam("cond_input_ids"), + OutputParam("uncond_input_ids"), + ] + + @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") + + block_state.cond_input_ids, block_state.uncond_input_ids = components.tokenize_prompt( + block_state.prompt, + block_state.negative_prompt, + num_frames=block_state.chunk_frames, + height=block_state.height, + width=block_state.width, + fps=block_state.fps, + use_system_prompt=block_state.use_system_prompt, + add_resolution_template=block_state.add_resolution_template, + add_duration_template=block_state.add_duration_template, + action_mode=None, + action_view_point=None, + transfer_mode=True, + ) + + self.set_block_state(state, block_state) + return components, state + + class Cosmos3ActionTextStep(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..f9be51581213 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, + Cosmos3TransferSetupStep, Cosmos3VisionDenoiseInputStep, Cosmos3VisionPackSequenceStep, Cosmos3VisionPrepareLatentsStep, ) -from .decoders import Cosmos3SoundDecodeStep, Cosmos3VideoDecodeStep +from .decoders import ( + Cosmos3SoundDecodeStep, + Cosmos3TransferDecodeChunkStep, + Cosmos3TransferStitchStep, + Cosmos3VideoDecodeStep, +) from .denoise import ( + Cosmos3TransferDenoiseStep, Cosmos3VisionActionDenoiseStep, Cosmos3VisionDenoiseStep, Cosmos3VisionSoundActionDenoiseStep, @@ -28,8 +43,23 @@ Cosmos3ActionVisionVaeEncoderStep, Cosmos3ImageVaeEncoderStep, Cosmos3TextEncoderStep, + Cosmos3TransferTextStep, Cosmos3VideoVaeEncoderStep, ) +from .modular_pipeline import Cosmos3OmniModularPipeline + + +class Cosmos3TransferTextBlocks(SequentialPipelineBlocks): + 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 @@ -80,14 +110,15 @@ class Cosmos3AutoTextEncoderStep(AutoPipelineBlocks): """ 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." ) @@ -135,13 +166,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( @@ -236,6 +271,27 @@ 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): """ @@ -566,6 +622,60 @@ def outputs(self): ] +class Cosmos3TransferChunkDenoiseStep(SequentialPipelineBlocks): + model_name = "cosmos3-omni" + block_classes = [ + Cosmos3TransferPrepareLatentsStep, + Cosmos3TransferPackSequenceStep, + Cosmos3TransferSetTimestepsStep, + Cosmos3TransferDenoiseStep, + Cosmos3TransferDecodeChunkStep, + ] + block_names = [ + "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", required=True)] + + @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 + + +class Cosmos3TransferCoreDenoiseStep(SequentialPipelineBlocks): + 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): """ @@ -627,13 +737,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 +763,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 +777,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" @@ -763,7 +877,7 @@ class Cosmos3OmniBlocks(SequentialPipelineBlocks): Cosmos3AutoTextEncoderStep, Cosmos3AutoVaeEncoderStep, Cosmos3AutoCoreDenoiseStep, - Cosmos3DecodeStep, + Cosmos3AutoDecodeStep, Cosmos3ActionOutputStep, ] block_names = ["text_encoder", "vae_encoder", "denoise", "decode", "after_decode"] @@ -778,6 +892,7 @@ class Cosmos3OmniBlocks(SequentialPipelineBlocks): "action_policy": {"prompt": True, "action": True}, "action_forward_dynamics": {"prompt": True, "action": True}, "action_inverse_dynamics": {"prompt": True, "action": True}, + "transfer": {"prompt": True, "control_videos": True}, } @property From ef3335e53cdcd14f45bf024ede12f0aca3f353f4 Mon Sep 17 00:00:00 2001 From: Yuliya Zhautouskaya Date: Thu, 9 Jul 2026 04:12:34 -0700 Subject: [PATCH 2/7] Convention matches fixes --- .../cosmos/before_denoise.py | 61 +++++++++++++------ 1 file changed, 44 insertions(+), 17 deletions(-) diff --git a/src/diffusers/modular_pipelines/cosmos/before_denoise.py b/src/diffusers/modular_pipelines/cosmos/before_denoise.py index b0485df359de..c80dd84dfee5 100644 --- a/src/diffusers/modular_pipelines/cosmos/before_denoise.py +++ b/src/diffusers/modular_pipelines/cosmos/before_denoise.py @@ -979,7 +979,6 @@ def description(self) -> str: @property def expected_components(self) -> list[ComponentSpec]: return [ - ComponentSpec("vae", AutoencoderKLWan), ComponentSpec( "video_processor", VideoProcessor, @@ -991,24 +990,52 @@ def expected_components(self) -> list[ComponentSpec]: @property def inputs(self) -> list[InputParam]: return [ - InputParam(name="control_videos", required=True), - InputParam(name="height", default=None), - InputParam(name="width", default=None), - InputParam(name="num_frames", default=None), - InputParam(name="num_video_frames_per_chunk", default=None), - InputParam(name="num_conditional_frames", type_hint=int, default=1), + 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"), - OutputParam("width"), - OutputParam("control_frames"), - OutputParam("total_frames"), - OutputParam("chunk_frames"), - OutputParam("num_chunks"), - OutputParam("stride"), + 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() @@ -1033,8 +1060,8 @@ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) if any(v is None for v in control_videos.values()): raise ValueError("`control_videos` entries must be loaded videos, not None.") - tcf = int(components.vae.config.scale_factor_temporal) - sf = int(components.vae.config.scale_factor_spatial) + 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})." @@ -1161,7 +1188,7 @@ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) chunk_frames = block_state.chunk_frames height = block_state.height width = block_state.width - tcf = int(components.vae.config.scale_factor_temporal) + tcf = components.vae_scale_factor_temporal # 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. From 2181295218f5288cff6c50e6541adc48206aa9af Mon Sep 17 00:00:00 2001 From: Yuliya Zhautouskaya Date: Thu, 9 Jul 2026 04:46:11 -0700 Subject: [PATCH 3/7] Align Cosmos3 transfer blocks with modular conventions --- .../cosmos/before_denoise.py | 188 +++++++------- .../modular_pipelines/cosmos/decoders.py | 49 +++- .../modular_pipelines/cosmos/denoise.py | 132 ++++++++-- .../modular_pipelines/cosmos/encoders.py | 209 ++++++++++++++- .../cosmos/modular_blocks_cosmos3.py | 244 ++++++++++++++---- 5 files changed, 618 insertions(+), 204 deletions(-) diff --git a/src/diffusers/modular_pipelines/cosmos/before_denoise.py b/src/diffusers/modular_pipelines/cosmos/before_denoise.py index c80dd84dfee5..534365b4ef89 100644 --- a/src/diffusers/modular_pipelines/cosmos/before_denoise.py +++ b/src/diffusers/modular_pipelines/cosmos/before_denoise.py @@ -4,7 +4,6 @@ import torch from ...configuration_utils import FrozenDict -from ...models.autoencoders.autoencoder_kl_wan import AutoencoderKLWan from ...models.transformers.transformer_cosmos3 import Cosmos3OmniTransformer from ...pipelines.cosmos.pipeline_cosmos3_omni import _EMBODIMENT_TO_DOMAIN_ID, CosmosActionCondition from ...schedulers import UniPCMultistepScheduler @@ -1129,53 +1128,52 @@ class Cosmos3TransferPrepareLatentsStep(ModularPipelineBlocks): @property def description(self) -> str: return ( - "Per-chunk transfer latent prep: slice + pad the chunk's control maps, seed the target's conditioning " - "frames (first chunk from the input video, later chunks from the previous chunk's tail), encode the " - "controls as clean latents and build the noisy target latents, velocity mask and condition latents." + "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), - ComponentSpec("vae", AutoencoderKLWan), - ComponentSpec( - "video_processor", - VideoProcessor, - config=FrozenDict({"vae_scale_factor": 16, "resample": "bilinear"}), - default_creation_method="from_config", - ), - ] + return [ComponentSpec("transformer", Cosmos3OmniTransformer)] @property def inputs(self) -> list[InputParam]: return [ - # Loop carry (set by the chunk-loop wrapper): - InputParam(name="chunk_id", type_hint=int, default=0), - InputParam(name="previous_output", default=None), - # Setup artifacts the loop can't cheaply re-derive (preprocessed controls + chunk geometry): - InputParam(name="control_frames", required=True), - InputParam(name="chunk_frames", required=True), - InputParam(name="total_frames", required=True), - InputParam(name="stride", required=True), - # User inputs (mirrors how the other prepare-latents steps declare height/width/generator/etc.): - InputParam(name="height", required=True), - InputParam(name="width", required=True), - InputParam(name="video", default=None), - InputParam(name="num_first_chunk_conditional_frames", type_hint=int, default=0), - InputParam(name="num_conditional_frames", type_hint=int, default=1), - InputParam(name="generator", default=None), + 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"), - OutputParam("control_latents"), - OutputParam("velocity_mask"), - OutputParam("condition_latents"), - OutputParam("target_condition_indexes"), - OutputParam("current_conditional_frames"), + 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() @@ -1183,59 +1181,12 @@ def __call__(self, components: Cosmos3OmniModularPipeline, state: 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 tcf = components.vae_scale_factor_temporal - # 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 - ) + target_x0 = block_state.x0_tokens_vision.to(device=device) + current_conditional_frames = block_state.current_conditional_frames - # Encode controls as clean latents and build the noisy target latents + conditioning mask. - block_state.control_latents = [components._encode_video(ctrl).contiguous().float() for ctrl in chunk_controls] - target_x0 = components._encode_video(target).contiguous().float() + # 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 @@ -1247,7 +1198,6 @@ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) 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)) - block_state.current_conditional_frames = current_conditional_frames self.set_block_state(state, block_state) return components, state @@ -1266,22 +1216,60 @@ def description(self) -> str: @property def inputs(self) -> list[InputParam]: return [ - InputParam(name="cond_text_segment", required=True), - InputParam(name="uncond_text_segment", required=True), - InputParam(name="control_latents", required=True), - InputParam(name="latents", required=True), - InputParam(name="target_condition_indexes", required=True), - InputParam(name="fps", type_hint=float, default=24.0), - InputParam(name="share_vision_temporal_positions", type_hint=bool, default=True), + 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."), + InputParam( + name="share_vision_temporal_positions", + type_hint=bool, + default=True, + description="Whether control and target items share vision temporal positions.", + ), ] @property def intermediate_outputs(self) -> list[OutputParam]: return [ - OutputParam("cond_full_static"), - OutputParam("cond_no_control_static"), - OutputParam("uncond_full_static"), - OutputParam("num_noisy_vision_tokens"), + 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() @@ -1347,8 +1335,10 @@ def inputs(self) -> list[InputParam]: @property def intermediate_outputs(self) -> list[OutputParam]: return [ - OutputParam("timesteps"), - OutputParam("num_warmup_steps"), + 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() diff --git a/src/diffusers/modular_pipelines/cosmos/decoders.py b/src/diffusers/modular_pipelines/cosmos/decoders.py index a52b5e9d45ef..7f3e4fa6393b 100644 --- a/src/diffusers/modular_pipelines/cosmos/decoders.py +++ b/src/diffusers/modular_pipelines/cosmos/decoders.py @@ -134,17 +134,37 @@ def expected_components(self) -> list[ComponentSpec]: @property def inputs(self) -> list[InputParam]: return [ - InputParam(name="latents", required=True), - InputParam(name="chunk_id", type_hint=int, default=0), - InputParam(name="current_conditional_frames", required=True), - InputParam(name="output_chunks", required=True), + 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"), - OutputParam("output_chunks"), + 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() @@ -190,17 +210,24 @@ def expected_components(self) -> list[ComponentSpec]: @property def inputs(self) -> list[InputParam]: return [ - InputParam(name="output_chunks", required=True), - InputParam(name="total_frames", required=True), + 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"), - OutputParam("sound"), - OutputParam("sampling_rate"), + 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() diff --git a/src/diffusers/modular_pipelines/cosmos/denoise.py b/src/diffusers/modular_pipelines/cosmos/denoise.py index 15b6d65ac000..06f21894f236 100644 --- a/src/diffusers/modular_pipelines/cosmos/denoise.py +++ b/src/diffusers/modular_pipelines/cosmos/denoise.py @@ -496,17 +496,39 @@ def expected_components(self) -> list[ComponentSpec]: @property def inputs(self) -> list[InputParam]: return [ - InputParam(name="control_latents", required=True), - InputParam(name="latents", required=True), - InputParam(name="num_noisy_vision_tokens", required=True), + 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"), - OutputParam("vision_tokens_target"), - OutputParam("vision_timesteps"), + 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() @@ -522,6 +544,8 @@ def __call__(self, components: Cosmos3OmniModularPipeline, block_state: BlockSta 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 @@ -538,22 +562,59 @@ def expected_components(self) -> list[ComponentSpec]: @property def inputs(self) -> list[InputParam]: return [ - InputParam(name="cond_full_static", required=True), - InputParam(name="cond_no_control_static", required=True), - InputParam(name="uncond_full_static", required=True), - InputParam(name="vision_tokens_full", required=True), - InputParam(name="vision_tokens_target", required=True), - InputParam(name="vision_timesteps", required=True), - InputParam(name="velocity_mask", required=True), - InputParam(name="guidance_scale", default=6.0), - InputParam(name="control_guidance", type_hint=float, default=1.0), - InputParam(name="guidance_interval", default=None), - InputParam(name="control_guidance_interval", default=None), + # 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")] + return [OutputParam("velocity", type_hint=torch.Tensor, description="Predicted (masked) transfer velocity.")] @staticmethod def _forward(components, static, vision_tokens, vision_timesteps): @@ -588,15 +649,20 @@ def __call__(self, components: Cosmos3OmniModularPipeline, block_state: BlockSta 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, block_state.cond_full_static, block_state.vision_tokens_full, block_state.vision_timesteps + 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, - block_state.cond_no_control_static, + cond_no_control_static, block_state.vision_tokens_target, block_state.vision_timesteps, ) @@ -605,7 +671,7 @@ def __call__(self, components: Cosmos3OmniModularPipeline, block_state: BlockSta if needs_text_cfg: uncond_full = self._forward( components, - block_state.uncond_full_static, + uncond_full_static, block_state.vision_tokens_full, block_state.vision_timesteps, ) @@ -638,15 +704,29 @@ def expected_components(self) -> list[ComponentSpec]: @property def inputs(self) -> list[InputParam]: return [ - InputParam(name="latents", required=True), - InputParam(name="velocity", required=True), - InputParam(name="velocity_mask", required=True), - InputParam(name="condition_latents", required=True), + 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")] + 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): diff --git a/src/diffusers/modular_pipelines/cosmos/encoders.py b/src/diffusers/modular_pipelines/cosmos/encoders.py index 67bbd70b4d0b..9b3f0e5c8dcf 100644 --- a/src/diffusers/modular_pipelines/cosmos/encoders.py +++ b/src/diffusers/modular_pipelines/cosmos/encoders.py @@ -177,22 +177,58 @@ def expected_components(self) -> list[ComponentSpec]: @property def inputs(self) -> list[InputParam]: return [ - InputParam(name="prompt", type_hint=str, required=True), - InputParam(name="negative_prompt", default=None), - InputParam(name="chunk_frames", type_hint=int, required=True), - InputParam(name="height", default=None), - InputParam(name="width", default=None), - InputParam(name="fps", type_hint=float, default=24.0), - InputParam(name="use_system_prompt", type_hint=bool, default=True), - InputParam(name="add_resolution_template", type_hint=bool, default=True), - InputParam(name="add_duration_template", type_hint=bool, default=True), + 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="chunk_frames", type_hint=int, required=True, description="Number of pixel frames in this chunk." + ), + 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="fps", type_hint=float, default=24.0, description="Frame rate of the generated video."), + InputParam( + name="use_system_prompt", + type_hint=bool, + default=True, + description="Whether to prepend the Cosmos3 system prompt.", + ), + InputParam( + name="add_resolution_template", + type_hint=bool, + default=True, + description="Whether to add resolution metadata to the prompt.", + ), + InputParam( + name="add_duration_template", + type_hint=bool, + default=True, + description="Whether to add duration metadata to the prompt.", + ), ] @property def intermediate_outputs(self) -> list[OutputParam]: return [ - OutputParam("cond_input_ids"), - OutputParam("uncond_input_ids"), + 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() @@ -628,6 +664,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 f9be51581213..726eb656f71b 100644 --- a/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3.py +++ b/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3.py @@ -43,6 +43,7 @@ Cosmos3ActionVisionVaeEncoderStep, Cosmos3ImageVaeEncoderStep, Cosmos3TextEncoderStep, + Cosmos3TransferChunkVaeEncoderStep, Cosmos3TransferTextStep, Cosmos3VideoVaeEncoderStep, ) @@ -66,25 +67,31 @@ def description(self): 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. - 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): @@ -93,20 +100,32 @@ class Cosmos3AutoTextEncoderStep(AutoPipelineBlocks): Whether to add resolution metadata to the prompt. add_duration_template (`bool`, *optional*, defaults to True): Whether to add duration metadata to the prompt. + action (`CosmosActionCondition`, *optional*): + Action-conditioning metadata and its reference visual input. 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" @@ -134,7 +153,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*): @@ -243,7 +263,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`): @@ -298,7 +320,8 @@ 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`): @@ -366,7 +389,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`): @@ -447,7 +471,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`): @@ -532,7 +557,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`): @@ -625,6 +651,7 @@ def outputs(self): class Cosmos3TransferChunkDenoiseStep(SequentialPipelineBlocks): model_name = "cosmos3-omni" block_classes = [ + Cosmos3TransferChunkVaeEncoderStep, Cosmos3TransferPrepareLatentsStep, Cosmos3TransferPackSequenceStep, Cosmos3TransferSetTimestepsStep, @@ -632,6 +659,7 @@ class Cosmos3TransferChunkDenoiseStep(SequentialPipelineBlocks): Cosmos3TransferDecodeChunkStep, ] block_names = [ + "encode_transfer_chunk", "prepare_transfer_latents", "pack_transfer_sequence", "set_timesteps", @@ -649,7 +677,9 @@ def description(self) -> str: @property def inputs(self) -> list[InputParam]: - return super().inputs + [InputParam(name="num_chunks", required=True)] + 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: @@ -680,37 +710,75 @@ def description(self) -> str: 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. + share_vision_temporal_positions (`bool`, *optional*, defaults to True): + Whether control and target items share vision temporal positions. 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*): @@ -719,16 +787,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`): @@ -801,25 +905,33 @@ class Cosmos3OmniBlocks(SequentialPipelineBlocks): - `action_policy`: requires `prompt`, `action` - `action_forward_dynamics`: requires `prompt`, `action` - `action_inverse_dynamics`: requires `prompt`, `action` + - `transfer`: requires `prompt`, `control_videos` 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. - 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): @@ -828,6 +940,8 @@ class Cosmos3OmniBlocks(SequentialPipelineBlocks): Whether to add resolution metadata to the prompt. add_duration_template (`bool`, *optional*, defaults to True): Whether to add duration metadata to the prompt. + action (`CosmosActionCondition`, *optional*): + Action-conditioning metadata and its reference visual input. video (`None`, *optional*): Reference video for video-to-video conditioning. condition_frame_indexes_vision (`tuple | list`, *optional*, defaults to (0, 1)): @@ -836,26 +950,42 @@ 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. + share_vision_temporal_positions (`bool`, *optional*, defaults to True): + Whether control and target items share vision temporal positions. + 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): From 0cfc55c59b9d273eb77d642b8154c5ca64b6add7 Mon Sep 17 00:00:00 2001 From: Yuliya Zhautouskaya Date: Fri, 10 Jul 2026 05:12:38 -0700 Subject: [PATCH 4/7] Move Cosmos3TransferSetupStep to before_encoder.py --- .../cosmos/before_denoise.py | 160 ----------------- .../cosmos/before_encoder.py | 166 ++++++++++++++++++ .../cosmos/modular_blocks_cosmos3.py | 2 +- 3 files changed, 167 insertions(+), 161 deletions(-) create mode 100644 src/diffusers/modular_pipelines/cosmos/before_encoder.py diff --git a/src/diffusers/modular_pipelines/cosmos/before_denoise.py b/src/diffusers/modular_pipelines/cosmos/before_denoise.py index 534365b4ef89..6049a66406a6 100644 --- a/src/diffusers/modular_pipelines/cosmos/before_denoise.py +++ b/src/diffusers/modular_pipelines/cosmos/before_denoise.py @@ -1,14 +1,11 @@ import copy -import math import torch -from ...configuration_utils import FrozenDict from ...models.transformers.transformer_cosmos3 import Cosmos3OmniTransformer from ...pipelines.cosmos.pipeline_cosmos3_omni import _EMBODIMENT_TO_DOMAIN_ID, CosmosActionCondition from ...schedulers import UniPCMultistepScheduler from ...utils.torch_utils import randn_tensor -from ...video_processor import VideoProcessor from ..modular_pipeline import ModularPipelineBlocks, PipelineState from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam from .modular_pipeline import Cosmos3OmniModularPipeline @@ -965,163 +962,6 @@ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) return components, state -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 - - class Cosmos3TransferPrepareLatentsStep(ModularPipelineBlocks): model_name = "cosmos3-omni" 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/modular_blocks_cosmos3.py b/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3.py index 726eb656f71b..d0aa596c7646 100644 --- a/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3.py +++ b/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3.py @@ -20,11 +20,11 @@ Cosmos3TransferPackSequenceStep, Cosmos3TransferPrepareLatentsStep, Cosmos3TransferSetTimestepsStep, - Cosmos3TransferSetupStep, Cosmos3VisionDenoiseInputStep, Cosmos3VisionPackSequenceStep, Cosmos3VisionPrepareLatentsStep, ) +from .before_encoder import Cosmos3TransferSetupStep from .decoders import ( Cosmos3SoundDecodeStep, Cosmos3TransferDecodeChunkStep, From 5191d32c926f60f1fd5491b213cd8b24f1b0e553 Mon Sep 17 00:00:00 2001 From: Yuliya Zhautouskaya Date: Fri, 10 Jul 2026 05:14:01 -0700 Subject: [PATCH 5/7] Decouple Cosmos3 transfer blocks from task-pipeline internals --- .../cosmos/before_denoise.py | 70 ++++++++++++----- .../modular_pipelines/cosmos/decoders.py | 5 +- .../modular_pipelines/cosmos/denoise.py | 10 ++- .../modular_pipelines/cosmos/encoders.py | 78 +++++++++---------- .../cosmos/modular_blocks_cosmos3.py | 4 - 5 files changed, 99 insertions(+), 68 deletions(-) diff --git a/src/diffusers/modular_pipelines/cosmos/before_denoise.py b/src/diffusers/modular_pipelines/cosmos/before_denoise.py index 6049a66406a6..aa9fd94d32b3 100644 --- a/src/diffusers/modular_pipelines/cosmos/before_denoise.py +++ b/src/diffusers/modular_pipelines/cosmos/before_denoise.py @@ -1056,7 +1056,9 @@ def description(self) -> str: @property def inputs(self) -> list[InputParam]: return [ - InputParam(name="cond_text_segment", type_hint=dict, required=True, description="Conditional text segment."), + 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." ), @@ -1067,7 +1069,10 @@ def inputs(self) -> list[InputParam]: 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." + name="latents", + type_hint=torch.Tensor, + required=True, + description="Noisy target latents for this chunk.", ), InputParam( name="target_condition_indexes", @@ -1076,12 +1081,6 @@ def inputs(self) -> list[InputParam]: 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."), - InputParam( - name="share_vision_temporal_positions", - type_hint=bool, - default=True, - description="Whether control and target items share vision temporal positions.", - ), ] @property @@ -1127,17 +1126,50 @@ def _vision_pack(text_segment: dict, include_controls: bool) -> dict: vision_items = [block_state.latents] condition_indexes = [block_state.target_condition_indexes] clean_flags = [False] - vision_segment = components._prepare_vision_segment( - input_vision_tokens=vision_items, - has_image_condition=False, - mrope_offset=text_segment["vision_start_temporal_offset"], - vision_fps=block_state.fps, - curr=text_segment["und_len"], - device=device, - condition_frame_indexes=condition_indexes, - clean_item_flags=clean_flags, - share_vision_temporal_positions=block_state.share_vision_temporal_positions, - ) + + # 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, diff --git a/src/diffusers/modular_pipelines/cosmos/decoders.py b/src/diffusers/modular_pipelines/cosmos/decoders.py index 7f3e4fa6393b..a76e48501d85 100644 --- a/src/diffusers/modular_pipelines/cosmos/decoders.py +++ b/src/diffusers/modular_pipelines/cosmos/decoders.py @@ -135,7 +135,10 @@ def expected_components(self) -> list[ComponentSpec]: def inputs(self) -> list[InputParam]: return [ InputParam( - name="latents", type_hint=torch.Tensor, required=True, description="Denoised target latents for this chunk." + 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( diff --git a/src/diffusers/modular_pipelines/cosmos/denoise.py b/src/diffusers/modular_pipelines/cosmos/denoise.py index 06f21894f236..f603082131bc 100644 --- a/src/diffusers/modular_pipelines/cosmos/denoise.py +++ b/src/diffusers/modular_pipelines/cosmos/denoise.py @@ -590,7 +590,10 @@ def inputs(self) -> list[InputParam]: 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." + name="guidance_scale", + type_hint=float, + default=6.0, + description="Scale for text classifier-free guidance.", ), InputParam( name="control_guidance", @@ -708,7 +711,10 @@ def inputs(self) -> list[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." + name="velocity", + type_hint=torch.Tensor, + required=True, + description="Predicted (masked) transfer velocity.", ), InputParam( name="velocity_mask", diff --git a/src/diffusers/modular_pipelines/cosmos/encoders.py b/src/diffusers/modular_pipelines/cosmos/encoders.py index 9b3f0e5c8dcf..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" @@ -150,8 +158,9 @@ class Cosmos3TransferTextStep(ModularPipelineBlocks): @property def description(self) -> str: return ( - "Tokenizes the transfer prompt in transfer mode using the per-chunk frame count. Transfer prompts are " - "pre-upsampled JSON captions passed through verbatim (the metadata templates are skipped)." + "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 @@ -189,36 +198,11 @@ def inputs(self) -> list[InputParam]: default=None, description="The negative text prompt used for classifier-free guidance.", ), - InputParam( - name="chunk_frames", type_hint=int, required=True, description="Number of pixel frames in this chunk." - ), - 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="fps", type_hint=float, default=24.0, description="Frame rate of the generated video."), InputParam( name="use_system_prompt", type_hint=bool, default=True, - description="Whether to prepend the Cosmos3 system prompt.", - ), - InputParam( - name="add_resolution_template", - type_hint=bool, - default=True, - description="Whether to add resolution metadata to the prompt.", - ), - InputParam( - name="add_duration_template", - type_hint=bool, - default=True, - description="Whether to add duration metadata to the prompt.", + description="Whether to prepend the Cosmos3 transfer system prompt.", ), ] @@ -258,20 +242,30 @@ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) finally: components.safety_checker.to("cpu") - block_state.cond_input_ids, block_state.uncond_input_ids = components.tokenize_prompt( - block_state.prompt, - block_state.negative_prompt, - num_frames=block_state.chunk_frames, - height=block_state.height, - width=block_state.width, - fps=block_state.fps, - use_system_prompt=block_state.use_system_prompt, - add_resolution_template=block_state.add_resolution_template, - add_duration_template=block_state.add_duration_template, - action_mode=None, - action_view_point=None, - transfer_mode=True, - ) + # 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 diff --git a/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3.py b/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3.py index d0aa596c7646..88c2f8aee307 100644 --- a/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3.py +++ b/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3.py @@ -753,8 +753,6 @@ class Cosmos3AutoCoreDenoiseStep(ConditionalPipelineBlocks): Torch generator for deterministic generation. fps (`float`, *optional*, defaults to 24.0): Frame rate of the generated video. - share_vision_temporal_positions (`bool`, *optional*, defaults to True): - Whether control and target items share vision temporal positions. num_inference_steps (`int`): The number of denoising steps. **denoiser_input_fields (`None`, *optional*): @@ -958,8 +956,6 @@ class Cosmos3OmniBlocks(SequentialPipelineBlocks): Number of frames the first chunk reuses from the input video. generator (`Generator`, *optional*): Torch generator for deterministic generation. - share_vision_temporal_positions (`bool`, *optional*, defaults to True): - Whether control and target items share vision temporal positions. num_inference_steps (`int`): The number of denoising steps. **denoiser_input_fields (`None`, *optional*): From b4408aed3c15fb1dca40cc7bdc005cc480cf0b88 Mon Sep 17 00:00:00 2001 From: Yuliya Zhautouskaya Date: Sat, 11 Jul 2026 08:00:58 -0700 Subject: [PATCH 6/7] Temporarily disable standalone transfer workflow --- .../cosmos/modular_blocks_cosmos3.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3.py b/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3.py index 88c2f8aee307..2f950dd78166 100644 --- a/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3.py +++ b/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3.py @@ -903,8 +903,6 @@ class Cosmos3OmniBlocks(SequentialPipelineBlocks): - `action_policy`: requires `prompt`, `action` - `action_forward_dynamics`: requires `prompt`, `action` - `action_inverse_dynamics`: requires `prompt`, `action` - - `transfer`: requires `prompt`, `control_videos` - Components: video_processor (`VideoProcessor`) text_tokenizer (`AutoTokenizer`) @@ -1018,13 +1016,22 @@ class Cosmos3OmniBlocks(SequentialPipelineBlocks): "action_policy": {"prompt": True, "action": True}, "action_forward_dynamics": {"prompt": True, "action": True}, "action_inverse_dynamics": {"prompt": True, "action": True}, - "transfer": {"prompt": True, "control_videos": True}, } @property 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 [ From 98deee4e8a39d825772580f8e0032e980c00fd8f Mon Sep 17 00:00:00 2001 From: Yuliya Zhautouskaya Date: Mon, 13 Jul 2026 01:23:56 -0700 Subject: [PATCH 7/7] Auto docstring run fix --- .../modular_pipelines/cosmos/denoise.py | 43 +++ .../cosmos/modular_blocks_cosmos3.py | 268 +++++++++++++++++- 2 files changed, 303 insertions(+), 8 deletions(-) diff --git a/src/diffusers/modular_pipelines/cosmos/denoise.py b/src/diffusers/modular_pipelines/cosmos/denoise.py index f603082131bc..7d2cdfe16fff 100644 --- a/src/diffusers/modular_pipelines/cosmos/denoise.py +++ b/src/diffusers/modular_pipelines/cosmos/denoise.py @@ -747,7 +747,50 @@ def __call__(self, components: Cosmos3OmniModularPipeline, block_state: BlockSta 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, diff --git a/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3.py b/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3.py index 2f950dd78166..f04226c79150 100644 --- a/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3.py +++ b/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3.py @@ -50,7 +50,56 @@ 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"] @@ -92,16 +141,16 @@ class Cosmos3AutoTextEncoderStep(AutoPipelineBlocks): 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. 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. - action (`CosmosActionCondition`, *optional*): - Action-conditioning metadata and its reference visual input. Outputs: height (`int`): @@ -648,7 +697,106 @@ 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, @@ -693,7 +841,110 @@ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) 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, @@ -903,6 +1154,7 @@ class Cosmos3OmniBlocks(SequentialPipelineBlocks): - `action_policy`: requires `prompt`, `action` - `action_forward_dynamics`: requires `prompt`, `action` - `action_inverse_dynamics`: requires `prompt`, `action` + Components: video_processor (`VideoProcessor`) text_tokenizer (`AutoTokenizer`) @@ -928,16 +1180,16 @@ class Cosmos3OmniBlocks(SequentialPipelineBlocks): 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. 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. - action (`CosmosActionCondition`, *optional*): - Action-conditioning metadata and its reference visual input. video (`None`, *optional*): Reference video for video-to-video conditioning. condition_frame_indexes_vision (`tuple | list`, *optional*, defaults to (0, 1)):