import torch from ..modular_pipeline import ( AutoPipelineBlocks, ConditionalPipelineBlocks, PipelineState, SequentialPipelineBlocks, ) from ..modular_pipeline_utils import InputParam, OutputParam from .after_decode import Cosmos3ActionOutputStep from .before_denoise import ( Cosmos3ActionDenoiseInputStep, Cosmos3ActionPackSequenceStep, Cosmos3ActionPrepareLatentsStep, Cosmos3PrepareTextSegmentsStep, Cosmos3SetTimestepsStep, Cosmos3SoundDenoiseInputStep, Cosmos3SoundPackSequenceStep, Cosmos3SoundPrepareLatentsStep, Cosmos3TransferPackSequenceStep, Cosmos3TransferPrepareLatentsStep, Cosmos3TransferSetTimestepsStep, Cosmos3VisionDenoiseInputStep, Cosmos3VisionPackSequenceStep, Cosmos3VisionPrepareLatentsStep, ) from .before_encoder import Cosmos3TransferSetupStep from .decoders import ( Cosmos3SoundDecodeStep, Cosmos3TransferDecodeChunkStep, Cosmos3TransferStitchStep, Cosmos3VideoDecodeStep, ) from .denoise import ( Cosmos3TransferDenoiseStep, Cosmos3VisionActionDenoiseStep, Cosmos3VisionDenoiseStep, Cosmos3VisionSoundActionDenoiseStep, Cosmos3VisionSoundDenoiseStep, ) from .encoders import ( Cosmos3ActionTextStep, Cosmos3ActionVisionVaeEncoderStep, Cosmos3ImageVaeEncoderStep, Cosmos3TextEncoderStep, Cosmos3TransferChunkVaeEncoderStep, Cosmos3TransferTextStep, Cosmos3VideoVaeEncoderStep, ) from .modular_pipeline import Cosmos3OmniModularPipeline # auto_docstring class Cosmos3TransferTextBlocks(SequentialPipelineBlocks): """ Transfer text branch: resolves the control-video chunk geometry, then tokenizes the (pre-upsampled) prompt in transfer mode using the per-chunk frame count. Components: video_processor (`VideoProcessor`) text_tokenizer (`AutoTokenizer`) Inputs: control_videos (`dict`): Mapping of hint name (edge/blur/depth/seg/wsm) to the control video for that modality. height (`int`, *optional*): Height of the generated video in pixels. width (`int`, *optional*): Width of the generated video in pixels. num_frames (`int`, *optional*): Optional cap on the number of output frames (defaults to the control video length). num_video_frames_per_chunk (`int`, *optional*): Number of pixel frames generated per autoregressive chunk. num_conditional_frames (`int`, *optional*, defaults to 1): Number of frames each chunk reuses from the previous chunk's tail. prompt (`str`): The text prompt that guides Cosmos3 generation. negative_prompt (`str`, *optional*): The negative text prompt used for classifier-free guidance. use_system_prompt (`bool`, *optional*, defaults to True): Whether to prepend the Cosmos3 transfer system prompt. Outputs: height (`int`): Resolved output height in pixels. width (`int`): Resolved output width in pixels. control_frames (`dict`): Preprocessed, time-padded control maps in canonical hint order. total_frames (`int`): Total number of output frames to generate. chunk_frames (`int`): Number of pixel frames per autoregressive chunk. num_chunks (`int`): Number of autoregressive chunks. stride (`int`): Frame stride between consecutive chunks. cond_input_ids (`Tensor`): Token IDs for the conditional prompt. uncond_input_ids (`Tensor`): Token IDs for the unconditional prompt. """ model_name = "cosmos3-omni" block_classes = [Cosmos3TransferSetupStep, Cosmos3TransferTextStep] block_names = ["setup", "transfer_text"] @property def description(self): return ( "Transfer text branch: resolves the control-video chunk geometry, then tokenizes the (pre-upsampled) " "prompt in transfer mode using the per-chunk frame count." ) # auto_docstring class Cosmos3AutoTextEncoderStep(AutoPipelineBlocks): """ Auto text encoder block for Cosmos3. - Cosmos3TransferTextBlocks runs when control_videos are provided. - Cosmos3ActionTextStep runs when action is provided. - Cosmos3TextEncoderStep runs otherwise. Components: video_processor (`VideoProcessor`) text_tokenizer (`AutoTokenizer`) Configs: default_use_system_prompt (default: True) enable_safety_checker (default: True) Inputs: control_videos (`dict`, *optional*): Mapping of hint name (edge/blur/depth/seg/wsm) to the control video for that modality. height (`int`, *optional*): Height of the generated video in pixels. width (`int`, *optional*): Width of the generated video in pixels. num_frames (`int`, *optional*): Optional cap on the number of output frames (defaults to the control video length). num_video_frames_per_chunk (`int`, *optional*): Number of pixel frames generated per autoregressive chunk. num_conditional_frames (`int`, *optional*, defaults to 1): Number of frames each chunk reuses from the previous chunk's tail. prompt (`str`): The text prompt that guides Cosmos3 generation. negative_prompt (`str`, *optional*): The negative text prompt used for classifier-free guidance. use_system_prompt (`bool`, *optional*, defaults to True or None, depending on the workflow): 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. add_resolution_template (`bool`, *optional*, defaults to True): Whether to add resolution metadata to the prompt. add_duration_template (`bool`, *optional*, defaults to True): Whether to add duration metadata to the prompt. Outputs: 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. action_mode (`str`): Requested action-generation mode. num_frames (`int`): Number of frames to generate. """ model_name = "cosmos3-omni" 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." ) # auto_docstring class Cosmos3AutoVaeEncoderStep(ConditionalPipelineBlocks): """ Auto VAE conditioning block for Cosmos3. - Cosmos3ActionVisionVaeEncoderStep runs when action is provided. - Cosmos3VideoVaeEncoderStep runs for the non-action video path. - Cosmos3ImageVaeEncoderStep runs for the non-action image path. - when no action, image, or video conditioning is provided, this block is skipped. Components: vae (`AutoencoderKLWan`) video_processor (`VideoProcessor`) Inputs: 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)): Latent-frame indexes to preserve from the conditioning video. condition_video_keep (`str`, *optional*, defaults to first): Which end of a longer conditioning video to use: `first` or `last`. num_frames (`int`, *optional*): Number of frames to generate. height (`int`, *optional*): Height of the generated video in pixels. width (`int`, *optional*): Width of the generated video in pixels. image (`None`, *optional*): Reference image for image-to-video conditioning. Outputs: x0_tokens_vision (`Tensor`): Vision latents encoded from the conditioning image or video. vision_condition_frames (`list`): Latent-frame indexes fixed by visual conditioning. action_condition_frame_indexes (`list`): Action-frame indexes fixed by action conditioning. """ model_name = "cosmos3-omni" block_classes = [Cosmos3ActionVisionVaeEncoderStep, Cosmos3VideoVaeEncoderStep, Cosmos3ImageVaeEncoderStep] block_names = ["action_conditioning", "video_conditioning", "image_conditioning"] 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( "Pass action conditioning via `action.image` / `action.video`, not top-level image/video." ) return "action_conditioning" if image is not None and video is not None: raise ValueError("Pass either image or video, not both.") if video is not None: return "video_conditioning" if image is not None: return "image_conditioning" return None @property def description(self): return ( "Auto VAE conditioning block for Cosmos3.\n" + " - Cosmos3ActionVisionVaeEncoderStep runs when action is provided.\n" + " - Cosmos3VideoVaeEncoderStep runs for the non-action video path.\n" + " - Cosmos3ImageVaeEncoderStep runs for the non-action image path.\n" + " - when no action, image, or video conditioning is provided, this block is skipped." ) # auto_docstring class Cosmos3AutoSoundDecodeStep(AutoPipelineBlocks): """ Auto sound decoder block for Cosmos3. - Cosmos3SoundDecodeStep runs when sound_latents are present. - if sound_latents are not provided, this block is skipped. Components: sound_tokenizer (`Cosmos3AVAEAudioTokenizer`) Inputs: sound_latents (`Tensor`, *optional*): Denoised sound latents to decode. Outputs: sound (`Tensor`): Generated waveform. sampling_rate (`int`): Sample rate of the generated waveform in Hz. """ model_name = "cosmos3-omni" block_classes = [Cosmos3SoundDecodeStep] block_names = ["decode"] block_trigger_inputs = ["sound_latents"] @property def description(self): return ( "Auto sound decoder block for Cosmos3.\n" + " - Cosmos3SoundDecodeStep runs when sound_latents are present.\n" + " - if sound_latents are not provided, this block is skipped." ) # auto_docstring class Cosmos3DecodeStep(SequentialPipelineBlocks): """ Decodes denoised latents into modality outputs. Components: vae (`AutoencoderKLWan`) video_processor (`VideoProcessor`) sound_tokenizer (`Cosmos3AVAEAudioTokenizer`) Inputs: latents (`Tensor`): Denoised vision latents to decode. output_type (`str`, *optional*, defaults to pil): Output format: 'pil', 'np', 'pt'. sound_latents (`Tensor`, *optional*): Denoised sound latents to decode. Outputs: videos (`list`): The generated videos. sound (`Tensor`): Generated waveform. sampling_rate (`int`): Sample rate of the generated waveform in Hz. """ model_name = "cosmos3-omni" block_classes = [Cosmos3VideoDecodeStep, Cosmos3AutoSoundDecodeStep] block_names = ["video", "sound"] @property def description(self) -> str: return "Decodes denoised latents into modality outputs." class Cosmos3AutoDecodeStep(ConditionalPipelineBlocks): model_name = "cosmos3-omni" block_classes = [Cosmos3TransferStitchStep, Cosmos3DecodeStep] block_names = ["transfer", "standard"] block_trigger_inputs = ["control_videos"] default_block_name = "standard" def select_block(self, **kwargs) -> str | None: if kwargs.get("control_videos") is not None: return "transfer" return "standard" @property def description(self) -> str: return ( "Selects the Cosmos3 decode workflow.\n" + " - Cosmos3TransferStitchStep stitches the decoded transfer chunks when control_videos are provided.\n" + " - Cosmos3DecodeStep decodes the denoised latents otherwise." ) # auto_docstring class Cosmos3VisionCoreDenoiseStep(SequentialPipelineBlocks): """ Runs the text-and-vision Cosmos3 denoising workflow. Components: transformer (`Cosmos3OmniTransformer`) scheduler (`UniPCMultistepScheduler`) Configs: use_native_flow_schedule (default: False) 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. 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`, *optional*): Pre-generated noisy vision latents. generator (`Generator`, *optional*): Torch generator for deterministic generation. num_inference_steps (`int`): The number of denoising steps. **denoiser_input_fields (`None`, *optional*): conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. guidance_scale (`float`, *optional*, defaults to 6.0): Scale for classifier-free guidance. Outputs: latents (`Tensor`): Denoised latents. """ model_name = "cosmos3-omni" block_classes = [ Cosmos3PrepareTextSegmentsStep, Cosmos3VisionPrepareLatentsStep, Cosmos3VisionPackSequenceStep, Cosmos3VisionDenoiseInputStep, Cosmos3SetTimestepsStep, Cosmos3VisionDenoiseStep, ] block_names = [ "prepare_text_segments", "prepare_vision_latents", "pack_vision_sequence", "prepare_vision_denoiser_inputs", "set_timesteps", "denoise", ] @property def description(self): return "Runs the text-and-vision Cosmos3 denoising workflow." @property def outputs(self): return [OutputParam.template("latents")] # auto_docstring class Cosmos3VisionSoundCoreDenoiseStep(SequentialPipelineBlocks): """ Runs the text, vision, and sound Cosmos3 denoising workflow. Components: transformer (`Cosmos3OmniTransformer`) scheduler (`UniPCMultistepScheduler`) Configs: use_native_flow_schedule (default: False) 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. 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`, *optional*): 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. **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. Outputs: latents (`Tensor`): Denoised latents. sound_latents (`Tensor`): Denoised sound latents. """ model_name = "cosmos3-omni" block_classes = [ Cosmos3PrepareTextSegmentsStep, Cosmos3VisionPrepareLatentsStep, Cosmos3VisionPackSequenceStep, Cosmos3VisionDenoiseInputStep, Cosmos3SetTimestepsStep, Cosmos3SoundPrepareLatentsStep, Cosmos3SoundPackSequenceStep, Cosmos3SoundDenoiseInputStep, Cosmos3VisionSoundDenoiseStep, ] block_names = [ "prepare_text_segments", "prepare_vision_latents", "pack_vision_sequence", "prepare_vision_denoiser_inputs", "set_timesteps", "prepare_sound_latents", "pack_sound_sequence", "prepare_sound_denoiser_inputs", "denoise", ] @property def description(self): return "Runs the text, vision, and sound Cosmos3 denoising workflow." @property def outputs(self): return [ OutputParam.template("latents"), OutputParam("sound_latents", type_hint=torch.Tensor, description="Denoised sound latents."), ] # auto_docstring class Cosmos3VisionActionCoreDenoiseStep(SequentialPipelineBlocks): """ Runs the text, vision, and action Cosmos3 denoising workflow. Components: transformer (`Cosmos3OmniTransformer`) scheduler (`UniPCMultistepScheduler`) Configs: use_native_flow_schedule (default: False) 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. 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`, *optional*): Pre-generated noisy vision latents. generator (`Generator`, *optional*): Torch generator for deterministic generation. num_inference_steps (`int`): The number of denoising steps. action (`CosmosActionCondition`): Action-conditioning metadata. 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. Outputs: latents (`Tensor`): Denoised latents. action_latents (`Tensor`): Denoised action latents. """ model_name = "cosmos3-omni" block_classes = [ Cosmos3PrepareTextSegmentsStep, Cosmos3VisionPrepareLatentsStep, Cosmos3VisionPackSequenceStep, Cosmos3VisionDenoiseInputStep, Cosmos3SetTimestepsStep, Cosmos3ActionPrepareLatentsStep, Cosmos3ActionPackSequenceStep, Cosmos3ActionDenoiseInputStep, Cosmos3VisionActionDenoiseStep, ] block_names = [ "prepare_text_segments", "prepare_vision_latents", "pack_vision_sequence", "prepare_vision_denoiser_inputs", "set_timesteps", "prepare_action_latents", "pack_action_sequence", "prepare_action_denoiser_inputs", "denoise", ] @property def description(self): return "Runs the text, vision, and action Cosmos3 denoising workflow." @property def outputs(self): return [ OutputParam.template("latents"), OutputParam("action_latents", type_hint=torch.Tensor, description="Denoised action latents."), ] # auto_docstring class Cosmos3VisionSoundActionCoreDenoiseStep(SequentialPipelineBlocks): """ Runs the text, vision, sound, and action Cosmos3 denoising workflow. Components: transformer (`Cosmos3OmniTransformer`) scheduler (`UniPCMultistepScheduler`) Configs: use_native_flow_schedule (default: False) 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. 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`, *optional*): 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 (`CosmosActionCondition`): Action-conditioning metadata. 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. Outputs: latents (`Tensor`): Denoised latents. sound_latents (`Tensor`): Denoised sound latents. action_latents (`Tensor`): Denoised action latents. """ model_name = "cosmos3-omni" block_classes = [ Cosmos3PrepareTextSegmentsStep, Cosmos3VisionPrepareLatentsStep, Cosmos3VisionPackSequenceStep, Cosmos3VisionDenoiseInputStep, Cosmos3SetTimestepsStep, Cosmos3SoundPrepareLatentsStep, Cosmos3SoundPackSequenceStep, Cosmos3SoundDenoiseInputStep, Cosmos3ActionPrepareLatentsStep, Cosmos3ActionPackSequenceStep, Cosmos3ActionDenoiseInputStep, Cosmos3VisionSoundActionDenoiseStep, ] block_names = [ "prepare_text_segments", "prepare_vision_latents", "pack_vision_sequence", "prepare_vision_denoiser_inputs", "set_timesteps", "prepare_sound_latents", "pack_sound_sequence", "prepare_sound_denoiser_inputs", "prepare_action_latents", "pack_action_sequence", "prepare_action_denoiser_inputs", "denoise", ] @property def description(self): return "Runs the text, vision, sound, and action Cosmos3 denoising workflow." @property def outputs(self): return [ OutputParam.template("latents"), OutputParam("sound_latents", type_hint=torch.Tensor, description="Denoised sound latents."), OutputParam("action_latents", type_hint=torch.Tensor, description="Denoised action latents."), ] # auto_docstring class Cosmos3TransferChunkDenoiseStep(SequentialPipelineBlocks): """ Autoregressive transfer chunk loop. Overrides __call__ to iterate chunks (the inner timestep loop is a non-leaf LoopSequentialPipelineBlocks, so this outer loop cannot itself be a LoopSequentialPipelineBlocks). Per-chunk cross-carry (previous_output, output_chunks) lives on PipelineState. Components: vae (`AutoencoderKLWan`) video_processor (`VideoProcessor`) transformer (`Cosmos3OmniTransformer`) scheduler (`UniPCMultistepScheduler`) Inputs: chunk_id (`int`, *optional*, defaults to 0): Index of the current chunk. previous_output (`None`, *optional*): Decoded pixels of the previous chunk, used to seed later chunks. control_frames (`dict`): Preprocessed, time-padded control maps in canonical hint order. chunk_frames (`int`): Pixel frames per chunk. total_frames (`int`): Total number of output frames. stride (`int`): Frame stride between chunks. height (`int`): Height of the generated video in pixels. width (`int`): Width of the generated video in pixels. video (`None`, *optional*): Optional input video that seeds the first chunk's conditioning. num_first_chunk_conditional_frames (`int`, *optional*, defaults to 0): Number of frames the first chunk reuses from the input video. num_conditional_frames (`int`, *optional*, defaults to 1): Number of frames each later chunk reuses from the previous chunk's tail. generator (`Generator`, *optional*): Torch generator for deterministic generation. cond_text_segment (`dict`): Conditional text segment. uncond_text_segment (`dict`): Unconditional text segment. fps (`float`, *optional*, defaults to 24.0): Frame rate of the generated video. num_inference_steps (`int`): The number of denoising steps. **denoiser_input_fields (`None`, *optional*): conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. guidance_scale (`float`, *optional*, defaults to 6.0): Scale for text classifier-free guidance. control_guidance (`float`, *optional*, defaults to 1.0): Scale for the control (structural) guidance axis. guidance_interval (`tuple`, *optional*): Timestep interval [lo, hi] over which text guidance is active (None = always). control_guidance_interval (`tuple`, *optional*): Timestep interval [lo, hi] over which control guidance is active (None = always). output_chunks (`list`): Decoded pixel chunks accumulated so far. num_chunks (`int`): Number of autoregressive chunks. Outputs: control_latents (`list`): Clean control latents for this chunk, one per hint in canonical order. x0_tokens_vision (`Tensor`): Clean target vision latents encoded from the seeded target frames. current_conditional_frames (`int`): Number of pixel frames actually used to seed this chunk's target. latents (`Tensor`): Noisy target latents for this chunk. velocity_mask (`Tensor`): Mask that zeroes the velocity on conditioned (clean) latent frames. condition_latents (`Tensor`): Clean target latents on the conditioned frames (the autoregressive seed). target_condition_indexes (`list`): Latent-frame indexes fixed by the chunk's conditioning. cond_full_static (`dict`): Conditional [control..., target] transfer sequence carrying every control item. cond_no_control_static (`dict`): Conditional [target] transfer sequence with the control items dropped. uncond_full_static (`dict`): Unconditional [control..., target] transfer sequence for text CFG. num_noisy_vision_tokens (`int`): Number of noisy target vision tokens denoised each step. timesteps (`Tensor`): Scheduler timesteps for this chunk. num_warmup_steps (`int`): Number of scheduler warmup steps for this chunk. vision_tokens_full (`list`): Token list for the [control..., target] forward passes. vision_tokens_target (`list`): Token list for the target-only (no-control) forward pass. vision_timesteps (`Tensor`): Timesteps for the noisy target tokens. velocity (`Tensor`): Predicted (masked) transfer velocity. previous_output (`Tensor`): Decoded pixels of this chunk, used to seed the next chunk. output_chunks (`list`): Decoded pixel chunks accumulated so far (with this chunk appended). """ model_name = "cosmos3-omni" block_classes = [ Cosmos3TransferChunkVaeEncoderStep, Cosmos3TransferPrepareLatentsStep, Cosmos3TransferPackSequenceStep, Cosmos3TransferSetTimestepsStep, Cosmos3TransferDenoiseStep, Cosmos3TransferDecodeChunkStep, ] block_names = [ "encode_transfer_chunk", "prepare_transfer_latents", "pack_transfer_sequence", "set_timesteps", "denoise", "decode_chunk", ] @property def description(self) -> str: return ( "Autoregressive transfer chunk loop. Overrides __call__ to iterate chunks (the inner timestep loop is a " "non-leaf LoopSequentialPipelineBlocks, so this outer loop cannot itself be a LoopSequentialPipelineBlocks). " "Per-chunk cross-carry (previous_output, output_chunks) lives on PipelineState." ) @property def inputs(self) -> list[InputParam]: return super().inputs + [ InputParam(name="num_chunks", type_hint=int, required=True, description="Number of autoregressive chunks.") ] @torch.no_grad() def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState: num_chunks = state.get("num_chunks") state.set("output_chunks", []) state.set("previous_output", None) for chunk_id in range(num_chunks): state.set("chunk_id", chunk_id) for _, block in self.sub_blocks.items(): components, state = block(components, state) return components, state # auto_docstring class Cosmos3TransferCoreDenoiseStep(SequentialPipelineBlocks): """ Transfer denoise stage: prepare shared text segments once, then run the autoregressive chunk loop. Components: transformer (`Cosmos3OmniTransformer`) vae (`AutoencoderKLWan`) video_processor (`VideoProcessor`) scheduler (`UniPCMultistepScheduler`) Inputs: cond_input_ids (`None`): Token IDs for the conditional prompt. uncond_input_ids (`None`): Token IDs for the unconditional prompt. chunk_id (`int`, *optional*, defaults to 0): Index of the current chunk. previous_output (`None`, *optional*): Decoded pixels of the previous chunk, used to seed later chunks. control_frames (`dict`): Preprocessed, time-padded control maps in canonical hint order. chunk_frames (`int`): Pixel frames per chunk. total_frames (`int`): Total number of output frames. stride (`int`): Frame stride between chunks. height (`int`): Height of the generated video in pixels. width (`int`): Width of the generated video in pixels. video (`None`, *optional*): Optional input video that seeds the first chunk's conditioning. num_first_chunk_conditional_frames (`int`, *optional*, defaults to 0): Number of frames the first chunk reuses from the input video. num_conditional_frames (`int`, *optional*, defaults to 1): Number of frames each later chunk reuses from the previous chunk's tail. generator (`Generator`, *optional*): Torch generator for deterministic generation. fps (`float`, *optional*, defaults to 24.0): Frame rate of the generated video. num_inference_steps (`int`): The number of denoising steps. **denoiser_input_fields (`None`, *optional*): conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. guidance_scale (`float`, *optional*, defaults to 6.0): Scale for text classifier-free guidance. control_guidance (`float`, *optional*, defaults to 1.0): Scale for the control (structural) guidance axis. guidance_interval (`tuple`, *optional*): Timestep interval [lo, hi] over which text guidance is active (None = always). control_guidance_interval (`tuple`, *optional*): Timestep interval [lo, hi] over which control guidance is active (None = always). output_chunks (`list`): Decoded pixel chunks accumulated so far. num_chunks (`int`): Number of autoregressive chunks. Outputs: cond_text_segment (`dict`): Conditional text segment for the denoiser. uncond_text_segment (`dict`): Unconditional text segment for the denoiser. control_latents (`list`): Clean control latents for this chunk, one per hint in canonical order. x0_tokens_vision (`Tensor`): Clean target vision latents encoded from the seeded target frames. current_conditional_frames (`int`): Number of pixel frames actually used to seed this chunk's target. latents (`Tensor`): Noisy target latents for this chunk. velocity_mask (`Tensor`): Mask that zeroes the velocity on conditioned (clean) latent frames. condition_latents (`Tensor`): Clean target latents on the conditioned frames (the autoregressive seed). target_condition_indexes (`list`): Latent-frame indexes fixed by the chunk's conditioning. cond_full_static (`dict`): Conditional [control..., target] transfer sequence carrying every control item. cond_no_control_static (`dict`): Conditional [target] transfer sequence with the control items dropped. uncond_full_static (`dict`): Unconditional [control..., target] transfer sequence for text CFG. num_noisy_vision_tokens (`int`): Number of noisy target vision tokens denoised each step. timesteps (`Tensor`): Scheduler timesteps for this chunk. num_warmup_steps (`int`): Number of scheduler warmup steps for this chunk. vision_tokens_full (`list`): Token list for the [control..., target] forward passes. vision_tokens_target (`list`): Token list for the target-only (no-control) forward pass. vision_timesteps (`Tensor`): Timesteps for the noisy target tokens. velocity (`Tensor`): Predicted (masked) transfer velocity. previous_output (`Tensor`): Decoded pixels of this chunk, used to seed the next chunk. output_chunks (`list`): Decoded pixel chunks accumulated so far (with this chunk appended). """ model_name = "cosmos3-omni" block_classes = [ Cosmos3PrepareTextSegmentsStep, Cosmos3TransferChunkDenoiseStep, ] block_names = ["prepare_text_segments", "chunk_denoise"] @property def description(self) -> str: return "Transfer denoise stage: prepare shared text segments once, then run the autoregressive chunk loop." # auto_docstring class Cosmos3AutoCoreDenoiseStep(ConditionalPipelineBlocks): """ Selects the Cosmos3 core denoising workflow. - transfer runs the autoregressive control-video (ControlNet-style) chunk loop when control_videos are provided. - vision_sound_action runs when action and enable_sound are provided. - vision_action runs when action is provided. - vision_sound runs when enable_sound is true. - vision runs otherwise. Components: transformer (`Cosmos3OmniTransformer`) vae (`AutoencoderKLWan`) video_processor (`VideoProcessor`) scheduler (`UniPCMultistepScheduler`) Configs: use_native_flow_schedule (default: False) 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`, *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. video (`None`, *optional*): Optional input video that seeds the first chunk's conditioning. num_first_chunk_conditional_frames (`int`, *optional*, defaults to 0): Number of frames the first chunk reuses from the input video. num_conditional_frames (`int`, *optional*, defaults to 1): Number of frames each later chunk reuses from the previous chunk's tail. generator (`Generator`, *optional*): Torch generator for deterministic generation. fps (`float`, *optional*, defaults to 24.0): Frame rate of the generated video. num_inference_steps (`int`): The number of denoising steps. **denoiser_input_fields (`None`, *optional*): conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. guidance_scale (`float`, *optional*, defaults to 6.0): Scale for text classifier-free guidance. control_guidance (`float`, *optional*, defaults to 1.0): Scale for the control (structural) guidance axis. guidance_interval (`tuple`, *optional*): Timestep interval [lo, hi] over which text guidance is active (None = always). control_guidance_interval (`tuple`, *optional*): Timestep interval [lo, hi] over which control guidance is active (None = always). output_chunks (`list`, *optional*): Decoded pixel chunks accumulated so far. num_chunks (`int`, *optional*): Number of autoregressive chunks. x0_tokens_vision (`Tensor`, *optional*): Vision latents encoded from the conditioning image or video. vision_condition_frames (`list`, *optional*): Latent-frame indexes fixed by visual conditioning. num_frames (`int`, *optional*): Number of frames to generate. latents (`Tensor`): Pre-generated noisy vision latents. sound_latents (`Tensor`, *optional*): Pre-generated noisy sound latents. action (`CosmosActionCondition`, *optional*): Action-conditioning metadata. action_condition_frame_indexes (`list`, *optional*): Action-frame indexes fixed by action conditioning. action_latents (`Tensor`, *optional*): Pre-generated noisy action latents. 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`): 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`): Denoised action latents. """ model_name = "cosmos3-omni" block_classes = [ Cosmos3TransferCoreDenoiseStep, Cosmos3VisionSoundActionCoreDenoiseStep, Cosmos3VisionActionCoreDenoiseStep, Cosmos3VisionSoundCoreDenoiseStep, Cosmos3VisionCoreDenoiseStep, ] block_names = ["transfer", "vision_sound_action", "vision_action", "vision_sound", "vision"] block_trigger_inputs = ["action", "enable_sound", "control_videos"] default_block_name = "vision" @property def inputs(self): inputs = super().inputs inputs.append( InputParam( name="enable_sound", type_hint=bool, default=False, description="Whether to generate a synchronized sound track.", ) ) return inputs 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: return "vision_action" if enable_sound: return "vision_sound" return "vision" @property 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" + " - vision runs otherwise." ) # auto_docstring class Cosmos3OmniBlocks(SequentialPipelineBlocks): """ Modular pipeline blocks for Cosmos3 generation modes. Supported workflows: - `text2image`: requires `prompt`, `num_frames` - `text2video`: requires `prompt` - `image2video`: requires `prompt`, `image` - `video2video`: requires `prompt`, `video` - `text2video_with_sound`: requires `prompt`, `enable_sound` - `image2video_with_sound`: requires `prompt`, `image`, `enable_sound` - `video2video_with_sound`: requires `prompt`, `video`, `enable_sound` - `action_policy`: requires `prompt`, `action` - `action_forward_dynamics`: requires `prompt`, `action` - `action_inverse_dynamics`: requires `prompt`, `action` Components: video_processor (`VideoProcessor`) text_tokenizer (`AutoTokenizer`) vae (`AutoencoderKLWan`) transformer (`Cosmos3OmniTransformer`) scheduler (`UniPCMultistepScheduler`) sound_tokenizer (`Cosmos3AVAEAudioTokenizer`) Configs: default_use_system_prompt (default: True) enable_safety_checker (default: True) use_native_flow_schedule (default: False) Inputs: control_videos (`dict`, *optional*): Mapping of hint name (edge/blur/depth/seg/wsm) to the control video for that modality. height (`int`, *optional*): Height of the generated video in pixels. width (`int`, *optional*): Width of the generated video in pixels. num_frames (`int`, *optional*): Optional cap on the number of output frames (defaults to the control video length). num_video_frames_per_chunk (`int`, *optional*): Number of pixel frames generated per autoregressive chunk. num_conditional_frames (`int`, *optional*, defaults to 1): Number of frames each chunk reuses from the previous chunk's tail. prompt (`str`): The text prompt that guides Cosmos3 generation. negative_prompt (`str`, *optional*): The negative text prompt used for classifier-free guidance. use_system_prompt (`bool`, *optional*, defaults to True or None, depending on the workflow): 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. 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. video (`None`, *optional*): Reference video for video-to-video conditioning. condition_frame_indexes_vision (`tuple | list`, *optional*, defaults to (0, 1)): Latent-frame indexes to preserve from the conditioning video. condition_video_keep (`str`, *optional*, defaults to first): Which end of a longer conditioning video to use: `first` or `last`. image (`None`, *optional*): Reference image for image-to-video conditioning. chunk_id (`int`, *optional*, defaults to 0): Index of the current chunk. previous_output (`None`, *optional*): Decoded pixels of the previous chunk, used to seed later chunks. num_first_chunk_conditional_frames (`int`, *optional*, defaults to 0): Number of frames the first chunk reuses from the input video. generator (`Generator`, *optional*): Torch generator for deterministic generation. num_inference_steps (`int`): The number of denoising steps. **denoiser_input_fields (`None`, *optional*): conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. guidance_scale (`float`, *optional*, defaults to 6.0): Scale for text classifier-free guidance. control_guidance (`float`, *optional*, defaults to 1.0): Scale for the control (structural) guidance axis. guidance_interval (`tuple`, *optional*): Timestep interval [lo, hi] over which text guidance is active (None = always). control_guidance_interval (`tuple`, *optional*): Timestep interval [lo, hi] over which control guidance is active (None = always). output_chunks (`list`, *optional*): Decoded pixel chunks accumulated so far. x0_tokens_vision (`Tensor`, *optional*): Vision latents encoded from the conditioning image or video. vision_condition_frames (`list`, *optional*): Latent-frame indexes fixed by visual conditioning. latents (`Tensor`): Pre-generated noisy vision latents. 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. enable_sound (`bool`, *optional*, defaults to False): Whether to generate a synchronized sound track. output_type (`str`, *optional*, defaults to pil): Output format: 'pil', 'np', 'pt'. Outputs: videos (`list`): The generated videos. sound (`Tensor`): Generated waveform. sampling_rate (`int`): Sample rate of the generated waveform in Hz. action (`list`): Generated action vectors. """ model_name = "cosmos3-omni" block_classes = [ Cosmos3AutoTextEncoderStep, Cosmos3AutoVaeEncoderStep, Cosmos3AutoCoreDenoiseStep, Cosmos3AutoDecodeStep, Cosmos3ActionOutputStep, ] block_names = ["text_encoder", "vae_encoder", "denoise", "decode", "after_decode"] _workflow_map = { "text2image": {"prompt": True, "num_frames": 1}, "text2video": {"prompt": True}, "image2video": {"prompt": True, "image": True}, "video2video": {"prompt": True, "video": True}, "text2video_with_sound": {"prompt": True, "enable_sound": True}, "image2video_with_sound": {"prompt": True, "image": True, "enable_sound": True}, "video2video_with_sound": {"prompt": True, "video": True, "enable_sound": True}, "action_policy": {"prompt": True, "action": True}, "action_forward_dynamics": {"prompt": True, "action": True}, "action_inverse_dynamics": {"prompt": True, "action": 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 [ OutputParam.template("videos"), OutputParam("sound", type_hint=torch.Tensor, description="Generated waveform."), OutputParam("sampling_rate", type_hint=int, description="Sample rate of the generated waveform in Hz."), OutputParam("action", type_hint=list[torch.Tensor], description="Generated action vectors."), ]