hf-internal-testing/tiny-cosmos3-modular-pipe

This is a modular diffusion pipeline built with 🧨 Diffusers' modular pipeline framework.

Pipeline Type: Cosmos3OmniBlocks

Description: Modular pipeline blocks for Cosmos3 generation modes.

This pipeline uses a 5-block architecture that can be customized and extended.

Example Usage

[TODO]

Pipeline Architecture

This modular pipeline is composed of the following blocks:

1. text_encoder (Cosmos3AutoTextEncoderStep)
   • Auto text encoder block for Cosmos3.
2. vae_encoder (Cosmos3AutoVaeEncoderStep)
   • Auto VAE conditioning block for Cosmos3.
3. denoise (Cosmos3AutoCoreDenoiseStep)
   • Selects the Cosmos3 core denoising workflow.
4. decode (Cosmos3AutoDecodeStep)
   • Selects the Cosmos3 decode workflow.
5. after_decode (Cosmos3ActionOutputStep)
   • Post-processes action latents into action outputs.

Model Components

1. video_processor (VideoProcessor)
2. text_tokenizer (AutoTokenizer)
3. vae (AutoencoderKLWan)
4. transformer (Cosmos3OmniTransformer)
5. scheduler (UniPCMultistepScheduler)
6. sound_tokenizer (Cosmos3AVAEAudioTokenizer)

Configuration Parameters

default_use_system_prompt (default: True) enable_safety_checker (default: True) use_native_flow_schedule (default: False)

Workflow Input Specification

text2image

• prompt (str): The text prompt that guides Cosmos3 generation.
• num_frames (int, optional): Number of frames to generate.

text2video

• prompt (str): The text prompt that guides Cosmos3 generation.

image2video

• prompt (str): The text prompt that guides Cosmos3 generation.
• image (None, optional): Reference image for image-to-video conditioning.

video2video

• prompt (str): The text prompt that guides Cosmos3 generation.
• video (None, optional): Reference video for video-to-video conditioning.

text2video_with_sound

• prompt (str): The text prompt that guides Cosmos3 generation.

image2video_with_sound

• prompt (str): The text prompt that guides Cosmos3 generation.
• image (None, optional): Reference image for image-to-video conditioning.

video2video_with_sound

• prompt (str): The text prompt that guides Cosmos3 generation.
• video (None, optional): Reference video for video-to-video conditioning.

action_policy

• prompt (str): The text prompt that guides Cosmos3 generation.
• action (CosmosActionCondition): Action-conditioning metadata and its reference visual input.

action_forward_dynamics

• prompt (str): The text prompt that guides Cosmos3 generation.
• action (CosmosActionCondition): Action-conditioning metadata and its reference visual input.

action_inverse_dynamics

• prompt (str): The text prompt that guides Cosmos3 generation.
• action (CosmosActionCondition): Action-conditioning metadata and its reference visual input.

Input/Output Specification

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