Add tiny Cosmos3 distilled modular pipeline fixture (#1)

- Add tiny Cosmos3 distilled modular pipeline fixture (2059b596353c32d2da7560ad8bcf18ced291ff89)


Co-authored-by: Yuliya Zhautouskaya <yzhautouskay@users.noreply.huggingface.co>

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+text_tokenizer/tokenizer.json filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,108 @@
+---
+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**: Cosmos3DistilledBlocks
+
+**Description**: Modular pipeline blocks for distilled (few-step) Cosmos3 generation modes.
+
+This pipeline uses a 4-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** (`Cosmos3DistilledTextEncoderStep`)
+   - Prepares distilled prompt token IDs. Classifier-free guidance is baked into the weights, so `negative_prompt` is not exposed and the unconditional branch is derived from an empty prompt.
+2. **vae_encoder** (`Cosmos3DistilledAutoVaeEncoderStep`)
+   - Auto VAE conditioning block for distilled Cosmos3.
+3. **denoise** (`Cosmos3DistilledVisionCoreDenoiseStep`)
+   - Runs the text-and-vision distilled Cosmos3 denoising workflow.
+4. **decode** (`Cosmos3VideoDecodeStep`)
+   - Decodes denoised vision latents into video outputs. 
+
+## Model Components
+
+1. text_tokenizer (`AutoTokenizer`)
+2. vae (`AutoencoderKLWan`)
+3. video_processor (`VideoProcessor`)
+4. transformer (`Cosmos3OmniTransformer`)
+5. scheduler (`FlowMatchEulerDiscreteScheduler`) 
+
+## Configuration Parameters
+
+is_distilled (default: True)
+    distilled_sigmas (default: None)
+
+## 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>
+
+
+## Input/Output Specification
+
+**Inputs:**
+
+- `prompt` (`str`): The text prompt that guides Cosmos3 generation.
+- `num_frames` (`int`, *optional*): Number of frames to generate.
+- `height` (`int`, *optional*): Height of the generated video or image in pixels.
+- `width` (`int`, *optional*): Width of the generated video or image in pixels.
+- `fps` (`float`, *optional*, defaults to `24.0`): Frame rate of the generated video.
+- `use_system_prompt` (`bool`, *optional*, defaults to `True`): Whether to prepend the Cosmos3 system prompt.
+- `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.
+- `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`, *optional*): Pre-generated noisy vision latents.
+- `generator` (`Generator`, *optional*): Torch generator for deterministic generation.
+- `num_inference_steps` (`int`, *optional*): The number of denoising steps.
+- `guidance_scale` (`float`, *optional*): 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.
+- `**denoiser_input_fields` (`None`, *optional*): conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
+- `output_type` (`str`, *optional*, defaults to `pil`): Output format: 'pil', 'np', 'pt'.
+
+**Outputs:**
+
+- `videos` (`list`): The generated videos.

modular_model_index.json

@@ -0,0 +1,71 @@
+{
+  "_blocks_class_name": "Cosmos3DistilledBlocks",
+  "_class_name": "Cosmos3DistilledModularPipeline",
+  "_diffusers_version": "0.40.0.dev0",
+  "default_use_system_prompt": true,
+  "distilled_sigmas": [
+    1.0,
+    0.9375,
+    0.8333333333333334,
+    0.625
+  ],
+  "enable_safety_checker": true,
+  "is_distilled": true,
+  "scheduler": [
+    "diffusers",
+    "FlowMatchEulerDiscreteScheduler",
+    {
+      "pretrained_model_name_or_path": "hf-internal-testing/tiny-cosmos3-distilled-modular-pipe",
+      "revision": null,
+      "subfolder": "scheduler",
+      "type_hint": [
+        "diffusers",
+        "FlowMatchEulerDiscreteScheduler"
+      ],
+      "variant": null
+    }
+  ],
+  "text_tokenizer": [
+    "transformers",
+    "Qwen2Tokenizer",
+    {
+      "pretrained_model_name_or_path": "hf-internal-testing/tiny-cosmos3-distilled-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-distilled-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-distilled-modular-pipe",
+      "revision": null,
+      "subfolder": "vae",
+      "type_hint": [
+        "diffusers",
+        "AutoencoderKLWan"
+      ],
+      "variant": null
+    }
+  ]
+}

scheduler/scheduler_config.json

@@ -0,0 +1,27 @@
+{
+  "_class_name": "FlowMatchEulerDiscreteScheduler",
+  "_diffusers_version": "0.40.0.dev0",
+  "base_image_seq_len": 256,
+  "base_shift": 0.5,
+  "fixed_step_sampler_config": {
+    "sample_type": "sde",
+    "t_list": [
+      1.0,
+      0.9375,
+      0.8333333333333334,
+      0.625
+    ]
+  },
+  "invert_sigmas": false,
+  "max_image_seq_len": 4096,
+  "max_shift": 1.15,
+  "num_train_timesteps": 1000,
+  "shift": 1.0,
+  "shift_terminal": null,
+  "stochastic_sampling": true,
+  "time_shift_type": "exponential",
+  "use_beta_sigmas": false,
+  "use_dynamic_shifting": false,
+  "use_exponential_sigmas": false,
+  "use_karras_sigmas": false
+}

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/tokenizer.json

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

text_tokenizer/tokenizer_config.json

@@ -0,0 +1,18 @@
+{
+  "add_prefix_space": false,
+  "backend": "tokenizers",
+  "bos_token": null,
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "errors": "replace",
+  "is_local": true,
+  "local_files_only": false,
+  "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": "/root/.cache/huggingface/hub/models--hf-internal-testing--tiny-cosmos3-modular-pipe/snapshots/9b0268fcede397cca83377c4a41c8f3efd2702fd/transformer",
+  "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": 6,
+  "hidden_act": "silu",
+  "hidden_size": 6,
+  "intermediate_size": 12,
+  "latent_channel": 4,
+  "latent_patch_size": 2,
+  "num_attention_heads": 1,
+  "num_embodiment_domains": 32,
+  "num_hidden_layers": 1,
+  "num_key_value_heads": 1,
+  "patch_latent_dim": 16,
+  "qk_norm_for_text": true,
+  "rms_norm_eps": 1e-06,
+  "rope_axes_dim": [
+    1,
+    1,
+    1
+  ],
+  "rope_scaling": {
+    "mrope_section": [
+      1,
+      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:382866740d429a75e2ca32857f115ab8d5ce3bc430ff0e4d0191c9afb6a7f1ce
+size 7311624

vae/config.json

@@ -0,0 +1,36 @@
+{
+  "_class_name": "AutoencoderKLWan",
+  "_diffusers_version": "0.40.0.dev0",
+  "_name_or_path": "/root/.cache/huggingface/hub/models--hf-internal-testing--tiny-cosmos3-modular-pipe/snapshots/9b0268fcede397cca83377c4a41c8f3efd2702fd/vae",
+  "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