Decouple Cosmos3 transfer blocks from task-pipeline internals

Author: yzhautouskay

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,

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(

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",

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

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*):