Transfer support for Cosmos3 ModularPipeline

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

src/diffusers/modular_pipelines/cosmos/before_denoise.py

@@ -960,3 +960,267 @@ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState)
         )
         self.set_block_state(state, block_state)
         return components, state
+
+
+class Cosmos3TransferPrepareLatentsStep(ModularPipelineBlocks):
+    model_name = "cosmos3-omni"
+
+    @property
+    def description(self) -> str:
+        return (
+            "Per-chunk transfer latent prep: takes the clean target latents encoded by "
+            "Cosmos3TransferChunkVaeEncoderStep and builds the noisy target latents, velocity mask, condition latents "
+            "and conditioned-frame indexes for this chunk."
+        )
+
+    @property
+    def expected_components(self) -> list[ComponentSpec]:
+        return [ComponentSpec("transformer", Cosmos3OmniTransformer)]
+
+    @property
+    def inputs(self) -> list[InputParam]:
+        return [
+            InputParam(
+                name="x0_tokens_vision",
+                type_hint=torch.Tensor,
+                required=True,
+                description="Clean target vision latents encoded from the seeded target frames.",
+            ),
+            InputParam(
+                name="current_conditional_frames",
+                type_hint=int,
+                required=True,
+                description="Number of pixel frames used to seed this chunk's target.",
+            ),
+            InputParam.template("generator"),
+        ]
+
+    @property
+    def intermediate_outputs(self) -> list[OutputParam]:
+        return [
+            OutputParam("latents", type_hint=torch.Tensor, description="Noisy target latents for this chunk."),
+            OutputParam(
+                "velocity_mask",
+                type_hint=torch.Tensor,
+                description="Mask that zeroes the velocity on conditioned (clean) latent frames.",
+            ),
+            OutputParam(
+                "condition_latents",
+                type_hint=torch.Tensor,
+                description="Clean target latents on the conditioned frames (the autoregressive seed).",
+            ),
+            OutputParam(
+                "target_condition_indexes",
+                type_hint=list[int],
+                description="Latent-frame indexes fixed by the chunk's conditioning.",
+            ),
+        ]
+
+    @torch.no_grad()
+    def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
+        block_state = self.get_block_state(state)
+        device = components._execution_device
+        dtype = components.transformer.dtype
+        tcf = components.vae_scale_factor_temporal
+
+        target_x0 = block_state.x0_tokens_vision.to(device=device)
+        current_conditional_frames = block_state.current_conditional_frames
+
+        # Build the noisy target latents + conditioning mask from the clean target latents.
+        latent_t = target_x0.shape[2]
+        condition_mask = torch.zeros((latent_t, 1, 1), device=device, dtype=dtype)
+        latent_condition_frames = 0
+        if current_conditional_frames > 0:
+            latent_condition_frames = (current_conditional_frames - 1) // tcf + 1
+            condition_mask[:latent_condition_frames] = 1.0
+        noise = randn_tensor(tuple(target_x0.shape), generator=block_state.generator, device=device, dtype=dtype)
+        block_state.latents = condition_mask * target_x0 + (1.0 - condition_mask) * noise
+        block_state.velocity_mask = 1.0 - condition_mask
+        block_state.condition_latents = condition_mask * target_x0
+        block_state.target_condition_indexes = list(range(latent_condition_frames))
+
+        self.set_block_state(state, block_state)
+        return components, state
+
+
+class Cosmos3TransferPackSequenceStep(ModularPipelineBlocks):
+    model_name = "cosmos3-omni"
+
+    @property
+    def description(self) -> str:
+        return (
+            "Pre-packs the three transfer CFG sequence variants: cond_full / uncond_full carry every control item, "
+            "the no-control branch drops them (only [text, target]) so the control axis can be amplified."
+        )
+
+    @property
+    def inputs(self) -> list[InputParam]:
+        return [
+            InputParam(
+                name="cond_text_segment", type_hint=dict, required=True, description="Conditional text segment."
+            ),
+            InputParam(
+                name="uncond_text_segment", type_hint=dict, required=True, description="Unconditional text segment."
+            ),
+            InputParam(
+                name="control_latents",
+                type_hint=list[torch.Tensor],
+                required=True,
+                description="Clean control latents for this chunk, one per hint in canonical order.",
+            ),
+            InputParam(
+                name="latents",
+                type_hint=torch.Tensor,
+                required=True,
+                description="Noisy target latents for this chunk.",
+            ),
+            InputParam(
+                name="target_condition_indexes",
+                type_hint=list[int],
+                required=True,
+                description="Latent-frame indexes fixed by the chunk's conditioning.",
+            ),
+            InputParam(name="fps", type_hint=float, default=24.0, description="Frame rate of the generated video."),
+        ]
+
+    @property
+    def intermediate_outputs(self) -> list[OutputParam]:
+        return [
+            OutputParam(
+                "cond_full_static",
+                type_hint=dict,
+                kwargs_type="denoiser_input_fields",
+                description="Conditional [control..., target] transfer sequence carrying every control item.",
+            ),
+            OutputParam(
+                "cond_no_control_static",
+                type_hint=dict,
+                kwargs_type="denoiser_input_fields",
+                description="Conditional [target] transfer sequence with the control items dropped.",
+            ),
+            OutputParam(
+                "uncond_full_static",
+                type_hint=dict,
+                kwargs_type="denoiser_input_fields",
+                description="Unconditional [control..., target] transfer sequence for text CFG.",
+            ),
+            OutputParam(
+                "num_noisy_vision_tokens",
+                type_hint=int,
+                description="Number of noisy target vision tokens denoised each step.",
+            ),
+        ]
+
+    @torch.no_grad()
+    def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
+        block_state = self.get_block_state(state)
+        device = components._execution_device
+        num_hints = len(block_state.control_latents)
+
+        def _vision_pack(text_segment: dict, include_controls: bool) -> dict:
+            if include_controls:
+                vision_items = [*block_state.control_latents, block_state.latents]
+                condition_indexes = [None] * num_hints + [block_state.target_condition_indexes]
+                clean_flags = [True] * num_hints + [False]
+            else:
+                vision_items = [block_state.latents]
+                condition_indexes = [block_state.target_condition_indexes]
+                clean_flags = [False]
+
+            # Transfer packs [ctrl_1, ..., ctrl_N, target] into one vision segment
+            mrope_offset = text_segment["vision_start_temporal_offset"]
+            item_curr = text_segment["und_len"]
+            token_shapes = []
+            sequence_index_parts = []
+            mse_loss_index_parts = []
+            noisy_frame_indexes_per_item = []
+            mrope_id_parts = []
+            num_vision_tokens = 0
+            num_noisy_vision_tokens = 0
+            for item, item_condition, is_clean in zip(vision_items, condition_indexes, clean_flags):
+                latent_t = item.shape[2]
+                if is_clean:
+                    frame_condition = list(range(latent_t))
+                else:
+                    frame_condition = item_condition if item_condition is not None else []
+                item_segment = components._prepare_vision_segment(
+                    input_vision_tokens=item,
+                    has_image_condition=False,
+                    mrope_offset=mrope_offset,
+                    vision_fps=block_state.fps,
+                    curr=item_curr,
+                    device=device,
+                    condition_frame_indexes=frame_condition,
+                )
+                token_shapes.extend(item_segment["vision_token_shapes"])
+                sequence_index_parts.append(item_segment["vision_sequence_indexes"])
+                mse_loss_index_parts.append(item_segment["vision_mse_loss_indexes"])
+                noisy_frame_indexes_per_item.extend(item_segment["vision_noisy_frame_indexes"])
+                mrope_id_parts.append(item_segment["vision_mrope_ids"])
+                num_vision_tokens += item_segment["num_vision_tokens"]
+                num_noisy_vision_tokens += item_segment["num_noisy_vision_tokens"]
+                item_curr += item_segment["num_vision_tokens"]
+
+            vision_segment = {
+                "vision_token_shapes": token_shapes,
+                "vision_sequence_indexes": torch.cat(sequence_index_parts, dim=0),
+                "vision_mse_loss_indexes": torch.cat(mse_loss_index_parts, dim=0),
+                "vision_noisy_frame_indexes": noisy_frame_indexes_per_item,
+                "vision_mrope_ids": torch.cat(mrope_id_parts, dim=1),
+                "num_vision_tokens": num_vision_tokens,
+                "num_noisy_vision_tokens": num_noisy_vision_tokens,
+            }
+            return {
+                **text_segment,
+                **vision_segment,
+                "position_ids": torch.cat([text_segment["text_mrope_ids"], vision_segment["vision_mrope_ids"]], dim=1),
+                "sequence_length": text_segment["und_len"] + vision_segment["num_vision_tokens"],
+            }
+
+        block_state.cond_full_static = _vision_pack(block_state.cond_text_segment, include_controls=True)
+        block_state.cond_no_control_static = _vision_pack(block_state.cond_text_segment, include_controls=False)
+        block_state.uncond_full_static = _vision_pack(block_state.uncond_text_segment, include_controls=True)
+        block_state.num_noisy_vision_tokens = block_state.cond_full_static["num_noisy_vision_tokens"]
+
+        self.set_block_state(state, block_state)
+        return components, state
+
+
+class Cosmos3TransferSetTimestepsStep(ModularPipelineBlocks):
+    model_name = "cosmos3-omni"
+
+    @property
+    def description(self) -> str:
+        return (
+            "Resets the scheduler and computes timesteps for a single transfer chunk. UniPCMultistepScheduler keeps "
+            "per-step state on the instance, so it is reset per chunk (each autoregressive chunk is a full denoise)."
+        )
+
+    @property
+    def expected_components(self) -> list[ComponentSpec]:
+        return [ComponentSpec("scheduler", UniPCMultistepScheduler)]
+
+    @property
+    def inputs(self) -> list[InputParam]:
+        return [InputParam.template("num_inference_steps", required=True)]
+
+    @property
+    def intermediate_outputs(self) -> list[OutputParam]:
+        return [
+            OutputParam("timesteps", type_hint=torch.Tensor, description="Scheduler timesteps for this chunk."),
+            OutputParam(
+                "num_warmup_steps", type_hint=int, description="Number of scheduler warmup steps for this chunk."
+            ),
+        ]
+
+    @torch.no_grad()
+    def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
+        block_state = self.get_block_state(state)
+        device = components._execution_device
+        components.scheduler.set_timesteps(block_state.num_inference_steps, device=device)
+        block_state.timesteps = components.scheduler.timesteps
+        block_state.num_warmup_steps = (
+            len(block_state.timesteps) - block_state.num_inference_steps * components.scheduler.order
+        )
+        self.set_block_state(state, block_state)
+        return components, state