Add tiny Cosmos3 modular pipeline fixture

README.md

@@ -0,0 +1,178 @@
+---
+library_name: diffusers
+tags:
+- modular-diffusers
+- diffusers
+- cosmos3-omni
+- text-to-image
+---
+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
+
+<details>
+<summary><strong>text2image</strong></summary>
+
+- `prompt` (`str`): The text prompt that guides Cosmos3 generation.
+- `num_frames` (`int`, *optional*): Number of frames to generate.
+
+</details>
+
+<details>
+<summary><strong>text2video</strong></summary>
+
+- `prompt` (`str`): The text prompt that guides Cosmos3 generation.
+
+</details>
+
+<details>
+<summary><strong>image2video</strong></summary>
+
+- `prompt` (`str`): The text prompt that guides Cosmos3 generation.
+- `image` (`None`, *optional*): Reference image for image-to-video conditioning.
+
+</details>
+
+<details>
+<summary><strong>video2video</strong></summary>
+
+- `prompt` (`str`): The text prompt that guides Cosmos3 generation.
+- `video` (`None`, *optional*): Reference video for video-to-video conditioning.
+
+</details>
+
+<details>
+<summary><strong>text2video_with_sound</strong></summary>
+
+- `prompt` (`str`): The text prompt that guides Cosmos3 generation.
+
+</details>
+
+<details>
+<summary><strong>image2video_with_sound</strong></summary>
+
+- `prompt` (`str`): The text prompt that guides Cosmos3 generation.
+- `image` (`None`, *optional*): Reference image for image-to-video conditioning.
+
+</details>
+
+<details>
+<summary><strong>video2video_with_sound</strong></summary>
+
+- `prompt` (`str`): The text prompt that guides Cosmos3 generation.
+- `video` (`None`, *optional*): Reference video for video-to-video conditioning.
+
+</details>
+
+<details>
+<summary><strong>action_policy</strong></summary>
+
+- `prompt` (`str`): The text prompt that guides Cosmos3 generation.
+- `action` (`CosmosActionCondition`): Action-conditioning metadata and its reference visual input.
+
+</details>
+
+<details>
+<summary><strong>action_forward_dynamics</strong></summary>
+
+- `prompt` (`str`): The text prompt that guides Cosmos3 generation.
+- `action` (`CosmosActionCondition`): Action-conditioning metadata and its reference visual input.
+
+</details>
+
+<details>
+<summary><strong>action_inverse_dynamics</strong></summary>
+
+- `prompt` (`str`): The text prompt that guides Cosmos3 generation.
+- `action` (`CosmosActionCondition`): Action-conditioning metadata and its reference visual input.
+
+</details>
+
+
+## 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.

modular_model_index.json

@@ -0,0 +1,78 @@
+{
+  "_blocks_class_name": "Cosmos3OmniBlocks",
+  "_class_name": "Cosmos3OmniModularPipeline",
+  "_diffusers_version": "0.40.0.dev0",
+  "default_use_system_prompt": true,
+  "enable_safety_checker": true,
+  "scheduler": [
+    "diffusers",
+    "UniPCMultistepScheduler",
+    {
+      "pretrained_model_name_or_path": "hf-internal-testing/tiny-cosmos3-modular-pipe",
+      "revision": null,
+      "subfolder": "scheduler",
+      "type_hint": [
+        "diffusers",
+        "UniPCMultistepScheduler"
+      ],
+      "variant": null
+    }
+  ],
+  "sound_tokenizer": [
+    "diffusers",
+    "Cosmos3AVAEAudioTokenizer",
+    {
+      "pretrained_model_name_or_path": "hf-internal-testing/tiny-cosmos3-modular-pipe",
+      "revision": null,
+      "subfolder": "sound_tokenizer",
+      "type_hint": [
+        "diffusers",
+        "Cosmos3AVAEAudioTokenizer"
+      ],
+      "variant": null
+    }
+  ],
+  "text_tokenizer": [
+    "transformers",
+    "Qwen2Tokenizer",
+    {
+      "pretrained_model_name_or_path": "hf-internal-testing/tiny-cosmos3-modular-pipe",
+      "revision": null,
+      "subfolder": "text_tokenizer",
+      "type_hint": [
+        "transformers",
+        "Qwen2Tokenizer"
+      ],
+      "variant": null
+    }
+  ],
+  "transformer": [
+    "diffusers",
+    "Cosmos3OmniTransformer",
+    {
+      "pretrained_model_name_or_path": "hf-internal-testing/tiny-cosmos3-modular-pipe",
+      "revision": null,
+      "subfolder": "transformer",
+      "type_hint": [
+        "diffusers",
+        "Cosmos3OmniTransformer"
+      ],
+      "variant": null
+    }
+  ],
+  "use_native_flow_schedule": false,
+  "vae": [
+    "diffusers",
+    "AutoencoderKLWan",
+    {
+      "pretrained_model_name_or_path": "hf-internal-testing/tiny-cosmos3-modular-pipe",
+      "revision": null,
+      "subfolder": "vae",
+      "type_hint": [
+        "diffusers",
+        "AutoencoderKLWan"
+      ],
+      "variant": null
+    }
+  ]
+}

scheduler/scheduler_config.json

@@ -0,0 +1,33 @@
+{
+  "_class_name": "UniPCMultistepScheduler",
+  "_diffusers_version": "0.40.0.dev0",
+  "beta_end": 0.02,
+  "beta_schedule": "linear",
+  "beta_start": 0.0001,
+  "disable_corrector": [],
+  "dynamic_thresholding_ratio": 0.995,
+  "final_sigmas_type": "zero",
+  "flow_shift": 1.0,
+  "lower_order_final": true,
+  "num_train_timesteps": 1000,
+  "predict_x0": true,
+  "prediction_type": "epsilon",
+  "rescale_betas_zero_snr": false,
+  "sample_max_value": 1.0,
+  "shift_terminal": null,
+  "sigma_max": null,
+  "sigma_min": null,
+  "solver_order": 2,
+  "solver_p": null,
+  "solver_type": "bh2",
+  "steps_offset": 0,
+  "thresholding": false,
+  "time_shift_type": "exponential",
+  "timestep_spacing": "linspace",
+  "trained_betas": null,
+  "use_beta_sigmas": false,
+  "use_dynamic_shifting": false,
+  "use_exponential_sigmas": false,
+  "use_flow_sigmas": false,
+  "use_karras_sigmas": false
+}

sound_tokenizer/config.json

@@ -0,0 +1,56 @@
+{
+  "_class_name": "Cosmos3AVAEAudioTokenizer",
+  "_diffusers_version": "0.40.0.dev0",
+  "_name_or_path": "/tmp/cosmos3-tiny-modular-prototype-20260715",
+  "activation": "snakebeta",
+  "anti_aliasing": false,
+  "bottleneck": null,
+  "bottleneck_type": "vae",
+  "causal": false,
+  "dec_anti_aliasing": false,
+  "dec_c_mults": [
+    1,
+    2
+  ],
+  "dec_dim": 4,
+  "dec_final_tanh": false,
+  "dec_out_channels": 2,
+  "dec_strides": [
+    2,
+    2
+  ],
+  "dec_type": "oobleck",
+  "dec_use_nearest_upsample": false,
+  "dec_use_snake": true,
+  "dec_use_tanh_at_final": false,
+  "enc_c_mults": [
+    1
+  ],
+  "enc_dim": 4,
+  "enc_hop_length": 2,
+  "enc_identity_init": false,
+  "enc_intermediate_dim": 768,
+  "enc_latent_dim": 8,
+  "enc_n_fft": 8,
+  "enc_num_blocks": 1,
+  "enc_num_layers": 12,
+  "enc_strides": [
+    2
+  ],
+  "enc_type": "spec_convnext",
+  "enc_use_snake": true,
+  "encoder_enabled": true,
+  "hop_size": 4,
+  "input_channels": 1,
+  "latent_mean": null,
+  "latent_std": null,
+  "model_type": "autoencoder_v2",
+  "normalize_volume": false,
+  "padding_mode": "zeros",
+  "sampling_rate": 16,
+  "snake_logscale": true,
+  "stereo": true,
+  "use_cuda_kernel": false,
+  "use_wav_as_input": true,
+  "vocoder_input_dim": 4
+}

sound_tokenizer/diffusion_pytorch_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:09e100dd66c9a2a16bd479178016e133ac0d3f8fcc5ced7a9bed1d0a0c390ede
+size 16536

text_tokenizer/added_tokens.json

@@ -0,0 +1,16 @@
+{
+  "<|box_end|>": 151649,
+  "<|box_start|>": 151648,
+  "<|endoftext|>": 151643,
+  "<|im_end|>": 151645,
+  "<|im_start|>": 151644,
+  "<|image_pad|>": 151655,
+  "<|object_ref_end|>": 151647,
+  "<|object_ref_start|>": 151646,
+  "<|quad_end|>": 151651,
+  "<|quad_start|>": 151650,
+  "<|video_pad|>": 151656,
+  "<|vision_end|>": 151653,
+  "<|vision_pad|>": 151654,
+  "<|vision_start|>": 151652
+}

text_tokenizer/chat_template.jinja

@@ -0,0 +1,7 @@
+{% set image_count = namespace(value=0) %}{% set video_count = namespace(value=0) %}{% for message in messages %}{% if loop.first and message['role'] != 'system' %}<|im_start|>system
+You are a helpful assistant.<|im_end|>
+{% endif %}<|im_start|>{{ message['role'] }}
+{% if message['content'] is string %}{{ message['content'] }}<|im_end|>
+{% else %}{% for content in message['content'] %}{% if content['type'] == 'image' or 'image' in content or 'image_url' in content %}{% set image_count.value = image_count.value + 1 %}{% if add_vision_id %}Picture {{ image_count.value }}: {% endif %}<|vision_start|><|image_pad|><|vision_end|>{% elif content['type'] == 'video' or 'video' in content %}{% set video_count.value = video_count.value + 1 %}{% if add_vision_id %}Video {{ video_count.value }}: {% endif %}<|vision_start|><|video_pad|><|vision_end|>{% elif 'text' in content %}{{ content['text'] }}{% endif %}{% endfor %}<|im_end|>
+{% endif %}{% endfor %}{% if add_generation_prompt %}<|im_start|>assistant
+{% endif %}

text_tokenizer/special_tokens_map.json

@@ -0,0 +1,31 @@
+{
+  "additional_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|object_ref_start|>",
+    "<|object_ref_end|>",
+    "<|box_start|>",
+    "<|box_end|>",
+    "<|quad_start|>",
+    "<|quad_end|>",
+    "<|vision_start|>",
+    "<|vision_end|>",
+    "<|vision_pad|>",
+    "<|image_pad|>",
+    "<|video_pad|>"
+  ],
+  "eos_token": {
+    "content": "<|im_end|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

text_tokenizer/tokenizer_config.json

@@ -0,0 +1,145 @@
+{
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "151643": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151644": {
+      "content": "<|im_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151645": {
+      "content": "<|im_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151646": {
+      "content": "<|object_ref_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151647": {
+      "content": "<|object_ref_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151648": {
+      "content": "<|box_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151649": {
+      "content": "<|box_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151650": {
+      "content": "<|quad_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151651": {
+      "content": "<|quad_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151652": {
+      "content": "<|vision_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151653": {
+      "content": "<|vision_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151654": {
+      "content": "<|vision_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151655": {
+      "content": "<|image_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151656": {
+      "content": "<|video_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "additional_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|object_ref_start|>",
+    "<|object_ref_end|>",
+    "<|box_start|>",
+    "<|box_end|>",
+    "<|quad_start|>",
+    "<|quad_end|>",
+    "<|vision_start|>",
+    "<|vision_end|>",
+    "<|vision_pad|>",
+    "<|image_pad|>",
+    "<|video_pad|>"
+  ],
+  "bos_token": null,
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "errors": "replace",
+  "extra_special_tokens": {},
+  "model_max_length": 32768,
+  "pad_token": "<|endoftext|>",
+  "padding_side": "left",
+  "processor_class": "Qwen2VLProcessor",
+  "split_special_tokens": false,
+  "tokenizer_class": "Qwen2Tokenizer",
+  "unk_token": null,
+  "variant": null
+}

transformer/config.json

@@ -0,0 +1,46 @@
+{
+  "_class_name": "Cosmos3OmniTransformer",
+  "_diffusers_version": "0.40.0.dev0",
+  "_name_or_path": "/tmp/cosmos3-tiny-modular-prototype-20260715",
+  "action_dim": 10,
+  "action_gen": true,
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "backbone_type": "cosmos3_qwen3vl",
+  "base_fps": 24,
+  "enable_fps_modulation": true,
+  "head_dim": 8,
+  "hidden_act": "silu",
+  "hidden_size": 32,
+  "intermediate_size": 64,
+  "latent_channel": 4,
+  "latent_patch_size": 2,
+  "num_attention_heads": 4,
+  "num_embodiment_domains": 32,
+  "num_hidden_layers": 2,
+  "num_key_value_heads": 2,
+  "patch_latent_dim": 16,
+  "qk_norm_for_text": true,
+  "rms_norm_eps": 1e-06,
+  "rope_axes_dim": [
+    2,
+    1,
+    1
+  ],
+  "rope_scaling": {
+    "mrope_section": [
+      2,
+      1,
+      1
+    ]
+  },
+  "rope_theta": 5000000.0,
+  "sound_dim": 4,
+  "sound_gen": true,
+  "sound_latent_fps": 5.0,
+  "temporal_compression_factor": 4,
+  "timestep_scale": 0.001,
+  "unified_3d_mrope_reset_spatial_ids": true,
+  "unified_3d_mrope_temporal_modality_margin": 15000,
+  "vocab_size": 151657
+}

transformer/diffusion_pytorch_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:db737dd3c20b7b31d0e71972de97177cb4839b16c199b213c8f380ba144b5fe5
+size 39109792

vae/config.json

@@ -0,0 +1,36 @@
+{
+  "_class_name": "AutoencoderKLWan",
+  "_diffusers_version": "0.40.0.dev0",
+  "_name_or_path": "/tmp/cosmos3-tiny-modular-prototype-20260715",
+  "attn_scales": [],
+  "base_dim": 8,
+  "decoder_base_dim": 8,
+  "dim_mult": [
+    1,
+    1
+  ],
+  "dropout": 0.0,
+  "in_channels": 3,
+  "is_residual": false,
+  "latents_mean": [
+    0.0,
+    0.0,
+    0.0,
+    0.0
+  ],
+  "latents_std": [
+    1.0,
+    1.0,
+    1.0,
+    1.0
+  ],
+  "num_res_blocks": 1,
+  "out_channels": 3,
+  "patch_size": null,
+  "scale_factor_spatial": 8,
+  "scale_factor_temporal": 4,
+  "temperal_downsample": [
+    false
+  ],
+  "z_dim": 4
+}

vae/diffusion_pytorch_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9baafae0091064db9eeeb21d2bdd7645cb4c97bc590d5eabba4780bf50237560
+size 167348