import copy import numpy as np import torch from ...models.transformers.transformer_cosmos3 import Cosmos3OmniTransformer from ...pipelines.cosmos.pipeline_cosmos3_omni import _EMBODIMENT_TO_DOMAIN_ID, CosmosActionCondition from ...schedulers import FlowMatchEulerDiscreteScheduler, UniPCMultistepScheduler from ...utils.torch_utils import randn_tensor from ..modular_pipeline import ModularPipelineBlocks, PipelineState from ..modular_pipeline_utils import ComponentSpec, ConfigSpec, InputParam, OutputParam from .modular_pipeline import Cosmos3OmniModularPipeline class Cosmos3PrepareTextSegmentsStep(ModularPipelineBlocks): model_name = "cosmos3-omni" @property def description(self) -> str: return "Builds cond/uncond text segments before denoising." @property def expected_components(self) -> list[ComponentSpec]: return [ComponentSpec("transformer", Cosmos3OmniTransformer)] @property def inputs(self) -> list[InputParam]: return [ InputParam(name="cond_input_ids", required=True, description="Token IDs for the conditional prompt."), InputParam(name="uncond_input_ids", required=True, description="Token IDs for the unconditional prompt."), ] @property def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam( "cond_text_segment", type_hint=dict, kwargs_type="denoiser_input_fields", description="Conditional text segment for the denoiser.", ), OutputParam( "uncond_text_segment", type_hint=dict, kwargs_type="denoiser_input_fields", description="Unconditional text segment for the denoiser.", ), ] @torch.no_grad() def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) device = components._execution_device block_state.cond_text_segment = components._prepare_text_segment(block_state.cond_input_ids, device=device) block_state.uncond_text_segment = components._prepare_text_segment(block_state.uncond_input_ids, device=device) self.set_block_state(state, block_state) return components, state class Cosmos3VisionPrepareLatentsStep(ModularPipelineBlocks): model_name = "cosmos3-omni" @property def description(self) -> str: return "Prepares noisy vision latents and the vision conditioning mask." @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, default=None, description="Vision latents encoded from the conditioning image or video.", ), InputParam( name="vision_condition_frames", type_hint=list[int], default=None, description="Latent-frame indexes fixed by visual conditioning.", ), InputParam(name="num_frames", type_hint=int, required=True, description="Number of frames to generate."), 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="fps", type_hint=float, default=24.0, description="Frame rate of the generated video."), InputParam( name="latents", type_hint=torch.Tensor, default=None, description="Pre-generated noisy vision latents.", ), InputParam.template("generator"), ] @property def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam("latents", type_hint=torch.Tensor, description="Noisy vision latents for denoising."), OutputParam("fps_vision", type_hint=float, description="Frame rate used to pack vision latents."), OutputParam( "vision_condition_mask", type_hint=torch.Tensor, kwargs_type="denoiser_input_fields", description="Mask marking conditioned vision latent frames.", ), OutputParam( "vision_condition_indexes_for_pack", type_hint=list[int], description="Indexes of conditioned vision latent frames.", ), OutputParam( "vision_conditioning_latents", type_hint=torch.Tensor, description="Clean encoded vision latents used to re-anchor image conditioning each step.", ), ] @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 x0_tokens_vision = block_state.x0_tokens_vision if x0_tokens_vision is None: if block_state.num_frames < 1: raise ValueError(f"num_frames must be >= 1, got {block_state.num_frames}.") sf_spatial = components.vae_scale_factor_spatial if block_state.height % sf_spatial != 0 or block_state.width % sf_spatial != 0: raise ValueError( f"height and width must be multiples of {sf_spatial}, got ({block_state.height}, {block_state.width})." ) latent_shape = ( 1, components.num_channels_latents, (block_state.num_frames - 1) // components.vae_scale_factor_temporal + 1, block_state.height // sf_spatial, block_state.width // sf_spatial, ) x0_tokens_vision = torch.zeros(latent_shape, device=device, dtype=torch.float32) else: x0_tokens_vision = x0_tokens_vision.to(device=device, dtype=torch.float32) block_state.fps_vision = float(block_state.fps) condition_frames = block_state.vision_condition_frames or [] block_state.vision_condition_mask = torch.zeros((x0_tokens_vision.shape[2], 1, 1), device=device, dtype=dtype) for frame_idx in condition_frames: if 0 <= frame_idx < block_state.vision_condition_mask.shape[0]: block_state.vision_condition_mask[frame_idx, 0, 0] = 1.0 if block_state.latents is None: pure_noise = randn_tensor( tuple(x0_tokens_vision.shape), generator=block_state.generator, device=device, dtype=dtype ) block_state.latents = ( block_state.vision_condition_mask * x0_tokens_vision.to(device=device, dtype=dtype) + (1.0 - block_state.vision_condition_mask) * pure_noise ) else: block_state.latents = block_state.latents.to(device=device, dtype=dtype) vision_condition_indexes = torch.nonzero( block_state.vision_condition_mask[:, 0, 0] > 0, as_tuple=False ).flatten() block_state.vision_condition_indexes_for_pack = [int(idx.item()) for idx in vision_condition_indexes] block_state.vision_conditioning_latents = x0_tokens_vision self.set_block_state(state, block_state) return components, state class Cosmos3SoundPrepareLatentsStep(ModularPipelineBlocks): model_name = "cosmos3-omni" @property def description(self) -> str: return "Prepares noisy sound latents and the sound conditioning mask." @property def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec("transformer", Cosmos3OmniTransformer), ComponentSpec("scheduler", UniPCMultistepScheduler), ] @property def inputs(self) -> list[InputParam]: return [ InputParam(name="num_frames", type_hint=int, required=True, description="Number of frames to generate."), InputParam(name="fps", type_hint=float, default=24.0, description="Frame rate of the generated video."), InputParam( name="sound_latents", type_hint=torch.Tensor, default=None, description="Pre-generated noisy sound latents.", ), InputParam.template("generator"), ] @property def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam("sound_latents", type_hint=torch.Tensor, description="Noisy sound latents for denoising."), OutputParam("fps_sound", type_hint=float, description="Frame rate of the sound latent sequence."), OutputParam( "sound_condition_mask", type_hint=torch.Tensor, kwargs_type="denoiser_input_fields", description="Mask marking conditioned sound latent frames.", ), OutputParam("sound_scheduler", description="Scheduler used to update sound latents."), ] @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 not components.transformer.config.sound_gen: raise ValueError("Sound generation requires a transformer trained with sound_gen=True.") sound_dim = components.transformer.config.sound_dim block_state.fps_sound = float(components.transformer.config.sound_latent_fps) n_audio_samples = int(block_state.num_frames / block_state.fps * components.sound_sampling_rate) hop_size = components.sound_hop_size t_sound = (n_audio_samples + hop_size - 1) // hop_size x0_tokens_sound = torch.zeros(sound_dim, t_sound, device=device, dtype=dtype) block_state.sound_condition_mask = torch.zeros((x0_tokens_sound.shape[1], 1), device=device, dtype=dtype) if block_state.sound_latents is None: pure_noise = randn_tensor( tuple(x0_tokens_sound.shape), generator=block_state.generator, device=device, dtype=dtype ) block_state.sound_latents = ( block_state.sound_condition_mask.T * x0_tokens_sound + (1.0 - block_state.sound_condition_mask.T) * pure_noise ) else: block_state.sound_latents = block_state.sound_latents.to(device=device, dtype=dtype) block_state.sound_scheduler = copy.deepcopy(components.scheduler) self.set_block_state(state, block_state) return components, state class Cosmos3ActionPrepareLatentsStep(ModularPipelineBlocks): model_name = "cosmos3-omni" @property def description(self) -> str: return "Prepares noisy action latents and the action conditioning mask." @property def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec("transformer", Cosmos3OmniTransformer), ComponentSpec("scheduler", UniPCMultistepScheduler), ] @property def inputs(self) -> list[InputParam]: return [ InputParam( name="action", type_hint=CosmosActionCondition, required=True, description="Action-conditioning metadata.", ), InputParam( name="action_condition_frame_indexes", type_hint=list[int], default=None, description="Action-frame indexes fixed by action conditioning.", ), InputParam( name="action_latents", type_hint=torch.Tensor, default=None, description="Pre-generated noisy action latents.", ), InputParam.template("generator"), ] @property def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam("action_latents", type_hint=torch.Tensor, description="Noisy action latents for denoising."), OutputParam( "action_condition_mask", type_hint=torch.Tensor, kwargs_type="denoiser_input_fields", description="Mask marking conditioned action latent frames.", ), OutputParam( "action_domain_ids", type_hint=list[torch.Tensor], kwargs_type="denoiser_input_fields", description="Embodiment domain IDs for action conditioning.", ), OutputParam( "raw_action_dim_resolved", type_hint=int, kwargs_type="denoiser_input_fields", description="Unpadded action-vector dimension.", ), OutputParam("action_scheduler", description="Scheduler used to update action latents."), ] @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 action = block_state.action if not components.transformer.config.action_gen: raise ValueError("action requires a transformer trained with action_gen=True.") block_state.raw_action_dim_resolved = int(action.raw_action_dim) if action.raw_action_dim is not None else None if ( block_state.raw_action_dim_resolved is not None and block_state.raw_action_dim_resolved > components.transformer.config.action_dim ): raise ValueError( f"raw_action_dim={block_state.raw_action_dim_resolved} exceeds the model action_dim=" f"{components.transformer.config.action_dim}." ) action_chunk_size = action.chunk_size action_dim = components.transformer.action_dim if action.mode == "forward_dynamics": raw_actions = action.raw_actions if raw_actions is None: raise ValueError("action_mode='forward_dynamics' requires an action tensor.") raw_actions = raw_actions.to(device=device, dtype=dtype) if raw_actions.shape[-1] > action_dim: raise ValueError( f"Cosmos3 action dimension {raw_actions.shape[-1]} exceeds model action_dim={action_dim}." ) if raw_actions.shape[0] < action_chunk_size: raw_actions = torch.cat( [raw_actions, raw_actions[-1:].expand(action_chunk_size - raw_actions.shape[0], -1)], dim=0, ) raw_actions = raw_actions[:action_chunk_size] if raw_actions.shape[-1] < action_dim: action_padding = torch.zeros( raw_actions.shape[0], action_dim - raw_actions.shape[-1], dtype=raw_actions.dtype, device=raw_actions.device, ) raw_actions = torch.cat([raw_actions, action_padding], dim=-1) x0_tokens_action = raw_actions else: x0_tokens_action = torch.zeros(action_chunk_size, action_dim, device=device, dtype=dtype) if action.domain_name not in _EMBODIMENT_TO_DOMAIN_ID: raise ValueError( f"Unknown Cosmos3 action domain_name={action.domain_name!r}; expected one of {sorted(_EMBODIMENT_TO_DOMAIN_ID)}." ) block_state.action_domain_ids = [ torch.tensor([_EMBODIMENT_TO_DOMAIN_ID[action.domain_name]], dtype=torch.long, device=device) ] condition_frames = block_state.action_condition_frame_indexes or [] block_state.action_condition_mask = torch.zeros((x0_tokens_action.shape[0], 1), device=device, dtype=dtype) for frame_idx in condition_frames: if 0 <= frame_idx < block_state.action_condition_mask.shape[0]: block_state.action_condition_mask[frame_idx, 0] = 1.0 if block_state.action_latents is None: pure_noise = randn_tensor( tuple(x0_tokens_action.shape), generator=block_state.generator, device=device, dtype=dtype ) block_state.action_latents = ( block_state.action_condition_mask * x0_tokens_action + (1.0 - block_state.action_condition_mask) * pure_noise ) if block_state.raw_action_dim_resolved is not None: block_state.action_latents[:, block_state.raw_action_dim_resolved :] = 0 else: block_state.action_latents = block_state.action_latents.to(device=device, dtype=dtype) block_state.action_scheduler = copy.deepcopy(components.scheduler) self.set_block_state(state, block_state) return components, state class Cosmos3VisionPackSequenceStep(ModularPipelineBlocks): model_name = "cosmos3-omni" @property def description(self) -> str: return "Builds separate cond/uncond vision sequence segments." @property def expected_components(self) -> list[ComponentSpec]: return [ComponentSpec("transformer", Cosmos3OmniTransformer)] @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="latents", type_hint=torch.Tensor, required=True, description="Noisy vision latents to pack." ), InputParam( name="fps_vision", type_hint=float, required=True, description="Frame rate used to pack vision latents.", ), InputParam( name="vision_condition_indexes_for_pack", type_hint=list[int], required=True, description="Indexes of conditioned vision latent frames.", ), ] @property def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam( "cond_vision_segment", type_hint=dict, kwargs_type="denoiser_input_fields", description="Conditional vision segment for the denoiser.", ), OutputParam( "uncond_vision_segment", type_hint=dict, kwargs_type="denoiser_input_fields", description="Unconditional vision segment for the denoiser.", ), ] @torch.no_grad() def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) device = components._execution_device has_image_condition = bool(block_state.vision_condition_indexes_for_pack) block_state.cond_vision_segment = components._prepare_vision_segment( input_vision_tokens=block_state.latents, has_image_condition=has_image_condition, mrope_offset=block_state.cond_text_segment["vision_start_temporal_offset"], vision_fps=block_state.fps_vision, curr=block_state.cond_text_segment["und_len"], device=device, condition_frame_indexes=block_state.vision_condition_indexes_for_pack, ) block_state.uncond_vision_segment = components._prepare_vision_segment( input_vision_tokens=block_state.latents, has_image_condition=has_image_condition, mrope_offset=block_state.uncond_text_segment["vision_start_temporal_offset"], vision_fps=block_state.fps_vision, curr=block_state.uncond_text_segment["und_len"], device=device, condition_frame_indexes=block_state.vision_condition_indexes_for_pack, ) self.set_block_state(state, block_state) return components, state class Cosmos3SoundPackSequenceStep(ModularPipelineBlocks): model_name = "cosmos3-omni" @property def description(self) -> str: return "Builds separate cond/uncond sound sequence segments." @property def expected_components(self) -> list[ComponentSpec]: return [ComponentSpec("transformer", Cosmos3OmniTransformer)] @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="cond_sequence_length", type_hint=int, required=True, description="Conditional multimodal sequence length.", ), InputParam( name="uncond_sequence_length", type_hint=int, required=True, description="Unconditional multimodal sequence length.", ), InputParam( name="sound_latents", type_hint=torch.Tensor, required=True, description="Noisy sound latents to pack." ), InputParam( name="fps_sound", type_hint=float, required=True, description="Frame rate of the sound latent sequence.", ), ] @property def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam( "cond_sound_segment", type_hint=dict, kwargs_type="denoiser_input_fields", description="Conditional sound segment for the denoiser.", ), OutputParam( "uncond_sound_segment", type_hint=dict, kwargs_type="denoiser_input_fields", description="Unconditional sound segment for the denoiser.", ), ] @torch.no_grad() def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) device = components._execution_device block_state.cond_sound_segment = components._prepare_sound_segment( input_sound_tokens=block_state.sound_latents, mrope_offset=block_state.cond_text_segment["vision_start_temporal_offset"], sound_fps=block_state.fps_sound, curr=block_state.cond_sequence_length, device=device, ) block_state.uncond_sound_segment = components._prepare_sound_segment( input_sound_tokens=block_state.sound_latents, mrope_offset=block_state.uncond_text_segment["vision_start_temporal_offset"], sound_fps=block_state.fps_sound, curr=block_state.uncond_sequence_length, device=device, ) self.set_block_state(state, block_state) return components, state class Cosmos3ActionPackSequenceStep(ModularPipelineBlocks): model_name = "cosmos3-omni" @property def description(self) -> str: return "Builds separate cond/uncond action sequence segments." @property def expected_components(self) -> list[ComponentSpec]: return [ComponentSpec("transformer", Cosmos3OmniTransformer)] @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="cond_sequence_length", type_hint=int, required=True, description="Conditional multimodal sequence length.", ), InputParam( name="uncond_sequence_length", type_hint=int, required=True, description="Unconditional multimodal sequence length.", ), InputParam( name="action_latents", type_hint=torch.Tensor, required=True, description="Noisy action latents to pack.", ), InputParam( name="action_condition_frame_indexes", type_hint=list[int], default=None, description="Action-frame indexes fixed by action conditioning.", ), InputParam( name="fps_vision", type_hint=float, required=True, description="Frame rate used to pack vision latents.", ), ] @property def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam( "cond_action_segment", type_hint=dict, kwargs_type="denoiser_input_fields", description="Conditional action segment for the denoiser.", ), OutputParam( "uncond_action_segment", type_hint=dict, kwargs_type="denoiser_input_fields", description="Unconditional action segment for the denoiser.", ), ] @torch.no_grad() def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) device = components._execution_device block_state.cond_action_segment = components._prepare_action_segment( input_action_tokens=block_state.action_latents, condition_frame_indexes=block_state.action_condition_frame_indexes, mrope_offset=block_state.cond_text_segment["vision_start_temporal_offset"], action_fps=block_state.fps_vision, curr=block_state.cond_sequence_length, device=device, ) block_state.uncond_action_segment = components._prepare_action_segment( input_action_tokens=block_state.action_latents, condition_frame_indexes=block_state.action_condition_frame_indexes, mrope_offset=block_state.uncond_text_segment["vision_start_temporal_offset"], action_fps=block_state.fps_vision, curr=block_state.uncond_sequence_length, device=device, ) self.set_block_state(state, block_state) return components, state class Cosmos3VisionDenoiseInputStep(ModularPipelineBlocks): model_name = "cosmos3-omni" @property def description(self) -> str: return "Assembles text and vision sequence metadata for the denoising loop." @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="cond_vision_segment", type_hint=dict, required=True, description="Conditional vision segment." ), InputParam( name="uncond_vision_segment", type_hint=dict, required=True, description="Unconditional vision segment.", ), ] @property def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam( "cond_position_ids", type_hint=torch.Tensor, kwargs_type="denoiser_input_fields", description="Conditional multimodal RoPE position IDs.", ), OutputParam( "uncond_position_ids", type_hint=torch.Tensor, kwargs_type="denoiser_input_fields", description="Unconditional multimodal RoPE position IDs.", ), OutputParam( "cond_sequence_length", type_hint=int, kwargs_type="denoiser_input_fields", description="Conditional multimodal sequence length.", ), OutputParam( "uncond_sequence_length", type_hint=int, kwargs_type="denoiser_input_fields", description="Unconditional multimodal sequence length.", ), ] @torch.no_grad() def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) block_state.cond_position_ids = torch.cat( [ block_state.cond_text_segment["text_mrope_ids"], block_state.cond_vision_segment["vision_mrope_ids"], ], dim=1, ) block_state.uncond_position_ids = torch.cat( [ block_state.uncond_text_segment["text_mrope_ids"], block_state.uncond_vision_segment["vision_mrope_ids"], ], dim=1, ) block_state.cond_sequence_length = ( block_state.cond_text_segment["und_len"] + block_state.cond_vision_segment["num_vision_tokens"] ) block_state.uncond_sequence_length = ( block_state.uncond_text_segment["und_len"] + block_state.uncond_vision_segment["num_vision_tokens"] ) self.set_block_state(state, block_state) return components, state class Cosmos3SoundDenoiseInputStep(ModularPipelineBlocks): model_name = "cosmos3-omni" @property def description(self) -> str: return "Appends sound sequence metadata to the denoising-loop inputs." @property def inputs(self) -> list[InputParam]: return [ InputParam( name="cond_position_ids", type_hint=torch.Tensor, required=True, description="Conditional multimodal RoPE position IDs.", ), InputParam( name="uncond_position_ids", type_hint=torch.Tensor, required=True, description="Unconditional multimodal RoPE position IDs.", ), InputParam( name="cond_sequence_length", type_hint=int, required=True, description="Conditional multimodal sequence length.", ), InputParam( name="uncond_sequence_length", type_hint=int, required=True, description="Unconditional multimodal sequence length.", ), InputParam( name="cond_sound_segment", type_hint=dict, required=True, description="Conditional sound segment." ), InputParam( name="uncond_sound_segment", type_hint=dict, required=True, description="Unconditional sound segment.", ), ] @property def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam( "cond_position_ids", type_hint=torch.Tensor, kwargs_type="denoiser_input_fields", description="Conditional multimodal RoPE position IDs.", ), OutputParam( "uncond_position_ids", type_hint=torch.Tensor, kwargs_type="denoiser_input_fields", description="Unconditional multimodal RoPE position IDs.", ), OutputParam( "cond_sequence_length", type_hint=int, kwargs_type="denoiser_input_fields", description="Conditional multimodal sequence length.", ), OutputParam( "uncond_sequence_length", type_hint=int, kwargs_type="denoiser_input_fields", description="Unconditional multimodal sequence length.", ), ] @torch.no_grad() def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) block_state.cond_position_ids = torch.cat( [block_state.cond_position_ids, block_state.cond_sound_segment["sound_mrope_ids"]], dim=1 ) block_state.uncond_position_ids = torch.cat( [block_state.uncond_position_ids, block_state.uncond_sound_segment["sound_mrope_ids"]], dim=1 ) block_state.cond_sequence_length += block_state.cond_sound_segment["sound_len"] block_state.uncond_sequence_length += block_state.uncond_sound_segment["sound_len"] self.set_block_state(state, block_state) return components, state class Cosmos3ActionDenoiseInputStep(ModularPipelineBlocks): model_name = "cosmos3-omni" @property def description(self) -> str: return "Appends action sequence metadata to the denoising-loop inputs." @property def inputs(self) -> list[InputParam]: return [ InputParam( name="cond_position_ids", type_hint=torch.Tensor, required=True, description="Conditional multimodal RoPE position IDs.", ), InputParam( name="uncond_position_ids", type_hint=torch.Tensor, required=True, description="Unconditional multimodal RoPE position IDs.", ), InputParam( name="cond_sequence_length", type_hint=int, required=True, description="Conditional multimodal sequence length.", ), InputParam( name="uncond_sequence_length", type_hint=int, required=True, description="Unconditional multimodal sequence length.", ), InputParam( name="cond_action_segment", type_hint=dict, required=True, description="Conditional action segment." ), InputParam( name="uncond_action_segment", type_hint=dict, required=True, description="Unconditional action segment.", ), ] @property def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam( "cond_position_ids", type_hint=torch.Tensor, kwargs_type="denoiser_input_fields", description="Conditional multimodal RoPE position IDs.", ), OutputParam( "uncond_position_ids", type_hint=torch.Tensor, kwargs_type="denoiser_input_fields", description="Unconditional multimodal RoPE position IDs.", ), OutputParam( "cond_sequence_length", type_hint=int, kwargs_type="denoiser_input_fields", description="Conditional multimodal sequence length.", ), OutputParam( "uncond_sequence_length", type_hint=int, kwargs_type="denoiser_input_fields", description="Unconditional multimodal sequence length.", ), ] @torch.no_grad() def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) block_state.cond_position_ids = torch.cat( [block_state.cond_position_ids, block_state.cond_action_segment["action_mrope_ids"]], dim=1 ) block_state.uncond_position_ids = torch.cat( [block_state.uncond_position_ids, block_state.uncond_action_segment["action_mrope_ids"]], dim=1 ) block_state.cond_sequence_length += block_state.cond_action_segment["action_len"] block_state.uncond_sequence_length += block_state.uncond_action_segment["action_len"] self.set_block_state(state, block_state) return components, state class Cosmos3SetTimestepsStep(ModularPipelineBlocks): model_name = "cosmos3-omni" @property def description(self) -> str: return "Initializes scheduler timesteps." @property def expected_components(self) -> list[ComponentSpec]: return [ComponentSpec("scheduler", UniPCMultistepScheduler)] @property def expected_configs(self) -> list[ConfigSpec]: return [ConfigSpec(name="use_native_flow_schedule", default=False)] @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 denoising."), OutputParam("num_warmup_steps", type_hint=int, description="Number of scheduler warmup steps."), ] @torch.no_grad() def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) device = components._execution_device if components.config.use_native_flow_schedule: sigmas = np.linspace( 1.0 - 1.0 / components.scheduler.config.num_train_timesteps, 0.0, block_state.num_inference_steps + 1, )[:-1] components.scheduler.set_timesteps(block_state.num_inference_steps, device=device, sigmas=sigmas) else: 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 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 class Cosmos3DistilledSetTimestepsStep(ModularPipelineBlocks): model_name = "cosmos3-omni" @property def description(self) -> str: return "Initializes the fixed distilled sampling schedule from the pipeline's `distilled_sigmas` config." @property def expected_components(self) -> list[ComponentSpec]: return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)] @property def expected_configs(self) -> list[ConfigSpec]: return [ ConfigSpec(name="is_distilled", default=True), ConfigSpec(name="distilled_sigmas", default=None), ] @property def inputs(self) -> list[InputParam]: return [ InputParam.template("num_inference_steps", required=False, default=None), InputParam( name="guidance_scale", type_hint=float, default=None, description=( "Unused for distilled checkpoints; classifier-free guidance is baked into the weights and the " "scale is forced to 1.0. Passing a value other than 1.0 raises an error." ), ), ] @property def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam("timesteps", type_hint=torch.Tensor, description="Scheduler timesteps for denoising."), OutputParam("num_warmup_steps", type_hint=int, description="Number of scheduler warmup steps."), OutputParam( "num_inference_steps", type_hint=int, description="Resolved number of denoising steps (fixed by the distilled schedule).", ), OutputParam( name="guidance_scale", type_hint=float, description="Resolved classifier-free guidance scale (always 1.0 for distilled checkpoints).", ), ] @torch.no_grad() def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) device = components._execution_device sigmas = components.config.distilled_sigmas if not sigmas: raise ValueError( "Cosmos3DistilledSetTimestepsStep requires the pipeline config `distilled_sigmas` to be set " "(populated from the distilled checkpoint's `modular_model_index.json`). Load a distilled Cosmos3 " "checkpoint or use `Cosmos3OmniModularPipeline` for base checkpoints." ) sigmas = [float(s) for s in sigmas] distilled_steps = len(sigmas) if block_state.num_inference_steps is not None and block_state.num_inference_steps != distilled_steps: raise ValueError( "This is a distilled checkpoint; the step count is fixed by the pipeline's " f"`distilled_sigmas` config ({distilled_steps} steps). " f"`num_inference_steps` must be {distilled_steps} or left unset (got {block_state.num_inference_steps})." ) if block_state.guidance_scale is not None and block_state.guidance_scale != 1.0: raise ValueError( "This is a distilled checkpoint; classifier-free guidance is baked into the weights. " f"`guidance_scale` must be 1.0 or left unset (got {block_state.guidance_scale})." ) components.scheduler.set_timesteps(sigmas=sigmas, device=device) block_state.num_inference_steps = distilled_steps block_state.guidance_scale = 1.0 block_state.timesteps = components.scheduler.timesteps block_state.num_warmup_steps = len(block_state.timesteps) - distilled_steps * components.scheduler.order self.set_block_state(state, block_state) return components, state