Add files using upload-large-folder tool

README.md

@@ -0,0 +1,69 @@
+---
+license: other
+license_name: openmdw1.1-license
+license_link: https://openmdw.ai/license/1-1/
+library_name: cosmos
+base_model: nvidia/Cosmos3-Nano-Policy-DROID
+tags:
+  - robotics
+  - droid
+  - policy
+  - world-model
+  - cosmos
+---
+
+# Cosmos-droid
+
+Intermediate checkpoints from a full fine-tune of
+[nvidia/Cosmos3-Nano-Policy-DROID](https://huggingface.co/nvidia/Cosmos3-Nano-Policy-DROID)
+(revision `805c0d6d46196ecdc789213a6de72303c383895b`) on five DROID manipulation tasks.
+The model jointly generates video and actions (joint video/action "WAM" policy).
+
+**Work in progress:** the run targets 5000 optimizer updates. These are intermediate
+checkpoints and have not been evaluated on a robot yet.
+
+## Checkpoints
+
+Each folder is a self-contained checkpoint in the same layout as the base model
+(`transformer/`, `vae/`, `vision_encoder/`, `text_tokenizer/`, `scheduler/`, `model_index.json`, ...).
+They were produced with `cosmos_framework.scripts.export_model` using the **EMA weights**,
+followed by `cosmos_framework.scripts.convert_model_to_diffusers`.
+
+| Folder | Optimizer step | Samples seen | Training loss (median over steps ±50) |
+|---|---:|---:|---:|
+| `iter_000500/` | 500 | 256k | ~1.14 |
+| `iter_001000/` | 1000 | 512k | ~1.01 |
+| `iter_001500/` | 1500 | 768k | ~0.97 |
+
+Load one by downloading its folder and passing it as `--checkpoint-path`, for example:
+
+```shell
+hf download Cosmos-Data/Cosmos-droid --include "iter_001500/*" --local-dir Cosmos-droid
+python -m cosmos_framework.scripts.inference --checkpoint-path Cosmos-droid/iter_001500 ...
+```
+
+## Tasks and data
+
+Five tabletop tasks: stack blue bowl on green bowl; serve bread on plate; put grapes into
+bowl; insert marker into mug; pick cherries and put on plate.
+
+The training data was 96 successful, relabeled DROID-format episodes (19,627 training windows)
+at 640x360, 15 FPS. Each window has 33 frames and 32 action targets, plus three cameras: wrist
+on top, two exterior views below. Actions are 7 absolute joint commands plus gripper.
+
+## Training
+
+- Full fine-tune (not LoRA), using the official generation/action optimizer selectors.
+- Global batch 512 (8x H100 80GB, 16 samples/GPU x grad-accum 4), BF16, full activation
+  checkpointing.
+- FusedAdam, base LR 2e-4 with scheduler multiplier 0.4, action-head LR multiplier 5,
+  weight decay 0.05, betas (0.9, 0.99), no warmup. EMA is enabled (power EMA, rate 0.1)
+  and re-initialized from the base model's regular (non-EMA) weights.
+- Recipe: `examples/toml/sft_config/droid_five_task_8h100_gbs512.toml` and
+  `docs/droid_five_task_policy.md` in the Cosmos-Framework fork used for training.
+
+## License
+
+Derived from nvidia/Cosmos3-Nano-Policy-DROID and released under the same
+[OpenMDW-1.1](https://openmdw.ai/license/1-1/) license. See the base model card for the
+bias, explainability, privacy and safety notes, which apply to this model as well.

iter_000500/chat_template.json

@@ -0,0 +1,3 @@
+{
+    "chat_template": "{%- if tools %}\n    {{- '<|im_start|>system\\n' }}\n    {%- if messages[0].role == 'system' %}\n        {%- if messages[0].content is string %}\n            {{- messages[0].content }}\n        {%- else %}\n            {%- for content in messages[0].content %}\n                {%- if 'text' in content %}\n                    {{- content.text }}\n                {%- endif %}\n            {%- endfor %}\n        {%- endif %}\n        {{- '\\n\\n' }}\n    {%- endif %}\n    {{- \"# Tools\\n\\nYou may call one or more functions to assist with the user query.\\n\\nYou are provided with function signatures within <tools></tools> XML tags:\\n<tools>\" }}\n    {%- for tool in tools %}\n        {{- \"\\n\" }}\n        {{- tool | tojson }}\n    {%- endfor %}\n    {{- \"\\n</tools>\\n\\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\\n<tool_call>\\n{\\\"name\\\": <function-name>, \\\"arguments\\\": <args-json-object>}\\n</tool_call><|im_end|>\\n\" }}\n{%- else %}\n    {%- if messages[0].role == 'system' %}\n        {{- '<|im_start|>system\\n' }}\n        {%- if messages[0].content is string %}\n            {{- messages[0].content }}\n        {%- else %}\n            {%- for content in messages[0].content %}\n                {%- if 'text' in content %}\n                    {{- content.text }}\n                {%- endif %}\n            {%- endfor %}\n        {%- endif %}\n        {{- '<|im_end|>\\n' }}\n    {%- endif %}\n{%- endif %}\n{%- set image_count = namespace(value=0) %}\n{%- set video_count = namespace(value=0) %}\n{%- for message in messages %}\n    {%- if message.role == \"user\" %}\n        {{- '<|im_start|>' + message.role + '\\n' }}\n        {%- if message.content is string %}\n            {{- message.content }}\n        {%- else %}\n            {%- for content in message.content %}\n                {%- if content.type == 'image' or 'image' in content or 'image_url' in content %}\n                    {%- set image_count.value = image_count.value + 1 %}\n                    {%- if add_vision_id %}Picture {{ image_count.value }}: {% endif -%}\n                    <|vision_start|><|image_pad|><|vision_end|>\n                {%- elif content.type == 'video' or 'video' in content %}\n                    {%- set video_count.value = video_count.value + 1 %}\n                    {%- if add_vision_id %}Video {{ video_count.value }}: {% endif -%}\n                    <|vision_start|><|video_pad|><|vision_end|>\n                {%- elif 'text' in content %}\n                    {{- content.text }}\n                {%- endif %}\n            {%- endfor %}\n        {%- endif %}\n        {{- '<|im_end|>\\n' }}\n    {%- elif message.role == \"assistant\" %}\n        {{- '<|im_start|>' + message.role + '\\n' }}\n        {%- if message.content is string %}\n            {{- message.content }}\n        {%- else %}\n            {%- for content_item in message.content %}\n                {%- if 'text' in content_item %}\n                    {{- content_item.text }}\n                {%- endif %}\n            {%- endfor %}\n        {%- endif %}\n        {%- if message.tool_calls %}\n            {%- for tool_call in message.tool_calls %}\n                {%- if (loop.first and message.content) or (not loop.first) %}\n                    {{- '\\n' }}\n                {%- endif %}\n                {%- if tool_call.function %}\n                    {%- set tool_call = tool_call.function %}\n                {%- endif %}\n                {{- '<tool_call>\\n{\"name\": \"' }}\n                {{- tool_call.name }}\n                {{- '\", \"arguments\": ' }}\n                {%- if tool_call.arguments is string %}\n                    {{- tool_call.arguments }}\n                {%- else %}\n                    {{- tool_call.arguments | tojson }}\n                {%- endif %}\n                {{- '}\\n</tool_call>' }}\n            {%- endfor %}\n        {%- endif %}\n        {{- '<|im_end|>\\n' }}\n    {%- elif message.role == \"tool\" %}\n        {%- if loop.first or (messages[loop.index0 - 1].role != \"tool\") %}\n            {{- '<|im_start|>user' }}\n        {%- endif %}\n        {{- '\\n<tool_response>\\n' }}\n        {%- if message.content is string %}\n            {{- message.content }}\n        {%- else %}\n            {%- for content in message.content %}\n                {%- if content.type == 'image' or 'image' in content or 'image_url' in content %}\n                    {%- set image_count.value = image_count.value + 1 %}\n                    {%- if add_vision_id %}Picture {{ image_count.value }}: {% endif -%}\n                    <|vision_start|><|image_pad|><|vision_end|>\n                {%- elif content.type == 'video' or 'video' in content %}\n                    {%- set video_count.value = video_count.value + 1 %}\n                    {%- if add_vision_id %}Video {{ video_count.value }}: {% endif -%}\n                    <|vision_start|><|video_pad|><|vision_end|>\n                {%- elif 'text' in content %}\n                    {{- content.text }}\n                {%- endif %}\n            {%- endfor %}\n        {%- endif %}\n        {{- '\\n</tool_response>' }}\n        {%- if loop.last or (messages[loop.index0 + 1].role != \"tool\") %}\n            {{- '<|im_end|>\\n' }}\n        {%- endif %}\n    {%- endif %}\n{%- endfor %}\n{%- if add_generation_prompt %}\n    {{- '<|im_start|>assistant\\n' }}\n{%- endif %}\n"
+}

iter_000500/checkpoint.json

@@ -0,0 +1 @@
+{}

iter_000500/config.json

@@ -0,0 +1,240 @@
+{
+  "allow_patterns_overrides": [
+    "*/*.safetensors"
+  ],
+  "architectures": [
+    "Cosmos3ForConditionalGeneration"
+  ],
+  "image_token_id": 151655,
+  "model": {
+    "_recursive_": false,
+    "_target": "omni_mot_model",
+    "config": {
+      "_type": "omni_mot_model_config",
+      "action_gen": true,
+      "activation_checkpointing": {
+        "determinism_check": "default",
+        "mode": "full",
+        "preserve_rng_state": true,
+        "save_ops_regex": [
+          "fmha"
+        ]
+      },
+      "causal_training_strategy": "none",
+      "compile": {
+        "ar_post_saturation_mode": "default",
+        "compile_dynamic": true,
+        "compiled_region": "language",
+        "coordinate_descent_tuning": false,
+        "enabled": false,
+        "max_autotune_pointwise": false,
+        "use_cuda_graphs": false
+      },
+      "diffusion_expert_config": {
+        "base_fps": 24,
+        "enable_fps_modulation": true,
+        "load_weights_from_pretrained": false,
+        "max_vae_latent_side_after_patchify": 20,
+        "patch_spatial": 2,
+        "sound_base_temporal_compression_factor": null,
+        "timestep_range": 1.0,
+        "unified_3d_mrope_reset_spatial_ids": true,
+        "unified_3d_mrope_temporal_modality_margin": 15000,
+        "vision_temporal_position_mode": "latent_index"
+      },
+      "ema": {
+        "enabled": false,
+        "iteration_shift": 0,
+        "rate": 0.1
+      },
+      "enable_input_bias": true,
+      "exclude_reasoner_weights_from_checkpoint": false,
+      "fixed_step_sampler_config": null,
+      "input_caption_key": "ai_caption",
+      "input_image_key": "images",
+      "input_video_key": "video",
+      "joint_attn_implementation": "two_way",
+      "latent_downsample_factor": 16,
+      "lbl": {
+        "coeff_gen": null,
+        "coeff_und": null,
+        "method": "local"
+      },
+      "load_vision_tokenizer": true,
+      "log_enc_time_every_n": 100,
+      "lora_alpha": 32,
+      "lora_enabled": false,
+      "lora_extra_trainable_modules": "",
+      "lora_rank": 16,
+      "lora_target_modules": "q_proj_moe_gen,k_proj_moe_gen,v_proj_moe_gen,o_proj_moe_gen",
+      "max_action_dim": 64,
+      "max_num_tokens_after_packing": -1,
+      "natten_parameter_list": null,
+      "net": null,
+      "num_embodiment_domains": 32,
+      "parallelism": {
+        "attention_io_layout": "sequence_sharded",
+        "cfg_parallel_shard_degree": 1,
+        "context_parallel_shard_degree": 1,
+        "data_parallel_replicate_degree": 1,
+        "data_parallel_shard_degree": 8,
+        "enable_inference_mode": false,
+        "fsdp_master_dtype": "float32"
+      },
+      "precision": "bfloat16",
+      "rectified_flow_inference_config": {
+        "num_train_timesteps": 1000,
+        "scheduler_type": "unipc",
+        "shift": 1,
+        "use_dynamic_shifting": false
+      },
+      "rectified_flow_training_config": {
+        "action_loss_weight": 10.0,
+        "align_shared_action_schedule": false,
+        "image_loss_scale": 1.0,
+        "independent_action_schedule": false,
+        "independent_sound_schedule": false,
+        "loss_scale": 10.0,
+        "normalize_loss_by_active": false,
+        "sample_level_loss_averaging": false,
+        "shift": {
+          "256": 3,
+          "480": 5,
+          "720": 10
+        },
+        "shift_action": null,
+        "shift_sound": null,
+        "sound_loss_scale": null,
+        "train_time_action_distribution": "logitnormal",
+        "train_time_image_distribution": "logitnormal",
+        "train_time_sound_distribution": "logitnormal",
+        "train_time_video_distribution": "waver",
+        "train_time_weight": "uniform",
+        "use_discrete_rf": false,
+        "use_dynamic_shift": false
+      },
+      "resolution": "720",
+      "sound_dim": null,
+      "sound_gen": false,
+      "sound_latent_fps": 25,
+      "sound_tokenizer": null,
+      "state_ch": 48,
+      "state_t": 300,
+      "tokenizer": {
+        "_target": "wan2pt2_vae_interface",
+        "bucket_name": "",
+        "causal": true,
+        "chunk_duration": 93,
+        "encode_bucket_multiple": null,
+        "encode_chunk_frames": {
+          "256": 68,
+          "480": 24,
+          "720": 8,
+          "768": 8
+        },
+        "encode_exact_durations": [
+          33
+        ],
+        "keep_decoder_cache": false,
+        "object_store_credential_path_pretrained": "",
+        "spatial_compression_factor": 16,
+        "temporal_compression_factor": 4,
+        "temporal_window": null,
+        "use_streaming_encode": false,
+        "vae_path": "pretrained/tokenizers/video/wan2pt2/Wan2.2_VAE.pth"
+      },
+      "video_temporal_causal": false,
+      "vision_gen": true,
+      "vlm_config": {
+        "layer_module": "Qwen2MoTDecoderLayer",
+        "model_instance": {
+          "_target": "qwen3_vl_text_for_causal_lm",
+          "config": {
+            "_target": "create_vlm_config",
+            "base_config": {
+              "_target": "qwen3_vl_mot_config_from_json_file",
+              "json_file": "cosmos3://vfm/models/vlm/qwen3_vl/configs/Qwen3-VL-8B-Instruct.json"
+            },
+            "freeze_und": false,
+            "layer_module": "MoTDecoderLayer",
+            "qk_norm_for_text": true,
+            "tie_word_embeddings": true
+          }
+        },
+        "model_name": "Qwen/Qwen3-VL-8B-Instruct",
+        "pretrained_weights": {
+          "backbone_path": "s3://bucket0/cosmos3/pretrained/huggingface/Qwen/Qwen3-VL-8B-Instruct/",
+          "checkpoint_format": null,
+          "credentials_path": "",
+          "enable_gcs_patch_in_boto3": true,
+          "enabled": false
+        },
+        "qk_norm": false,
+        "safetensors_path": "",
+        "tie_word_embeddings": false,
+        "tokenizer": {
+          "_target": "create_qwen2_tokenizer_with_download",
+          "config_variant": "hf",
+          "pretrained_model_name": "nvidia/Cosmos3-Nano-Policy-DROID"
+        },
+        "use_system_prompt": false
+      }
+    }
+  },
+  "model_type": "cosmos3_omni",
+  "text_config": {
+    "attention_bias": false,
+    "attention_dropout": 0.0,
+    "bos_token_id": 151643,
+    "dtype": "bfloat16",
+    "eos_token_id": 151645,
+    "head_dim": 128,
+    "hidden_act": "silu",
+    "hidden_size": 4096,
+    "initializer_range": 0.02,
+    "intermediate_size": 12288,
+    "max_position_embeddings": 262144,
+    "model_type": "qwen3_vl_text",
+    "num_attention_heads": 32,
+    "num_hidden_layers": 36,
+    "num_key_value_heads": 8,
+    "rms_norm_eps": 1e-06,
+    "rope_scaling": {
+      "mrope_interleaved": true,
+      "mrope_section": [
+        24,
+        20,
+        20
+      ],
+      "rope_type": "default"
+    },
+    "rope_theta": 5000000,
+    "use_cache": true,
+    "vocab_size": 151936
+  },
+  "tie_word_embeddings": false,
+  "transformers_version": "4.57.0.dev0",
+  "video_token_id": 151656,
+  "vision_config": {
+    "deepstack_visual_indexes": [
+      8,
+      16,
+      24
+    ],
+    "depth": 27,
+    "hidden_act": "gelu_pytorch_tanh",
+    "hidden_size": 1152,
+    "in_channels": 3,
+    "initializer_range": 0.02,
+    "intermediate_size": 4304,
+    "model_type": "qwen3_vl",
+    "num_heads": 16,
+    "num_position_embeddings": 2304,
+    "out_hidden_size": 4096,
+    "patch_size": 16,
+    "spatial_merge_size": 2,
+    "temporal_patch_size": 2
+  },
+  "vision_end_token_id": 151653,
+  "vision_start_token_id": 151652
+}

iter_000500/generation_config.json

@@ -0,0 +1,14 @@
+{
+    "bos_token_id": 151643,
+    "pad_token_id": 151643,
+    "do_sample": true,
+    "eos_token_id": [
+        151645,
+        151643
+    ],
+    "top_k": 20,
+    "top_p": 0.8,
+    "repetition_penalty": 1.0,
+    "temperature": 0.7,
+    "transformers_version": "4.56.0"
+}

iter_000500/model_index.json

@@ -0,0 +1,32 @@
+{
+  "_class_name": "Cosmos3OmniPipeline",
+  "_diffusers_version": "0.39.0",
+  "safety_checker": [
+    null,
+    null
+  ],
+  "scheduler": [
+    "diffusers",
+    "UniPCMultistepScheduler"
+  ],
+  "sound_tokenizer": [
+    null,
+    null
+  ],
+  "text_tokenizer": [
+    "transformers",
+    "Qwen2TokenizerFast"
+  ],
+  "transformer": [
+    "diffusers",
+    "Cosmos3OmniTransformer"
+  ],
+  "vae": [
+    "diffusers",
+    "AutoencoderKLWan"
+  ],
+  "vision_encoder": [
+    "transformers",
+    "Qwen3VLVisionModel"
+  ]
+}

iter_000500/modular_model_index.json

@@ -0,0 +1,58 @@
+{
+  "_blocks_class_name": "Cosmos3OmniBlocks",
+  "_class_name": "Cosmos3OmniModularPipeline",
+  "_diffusers_version": "0.39.0",
+  "is_distilled": false,
+  "scheduler": [
+    "diffusers",
+    "UniPCMultistepScheduler",
+    {
+      "pretrained_model_name_or_path": "Cosmos-Data/Cosmos-droid",
+      "subfolder": "iter_000500/scheduler",
+      "type_hint": [
+        "diffusers",
+        "UniPCMultistepScheduler"
+      ],
+      "variant": null
+    }
+  ],
+  "text_tokenizer": [
+    "transformers",
+    "Qwen2TokenizerFast",
+    {
+      "pretrained_model_name_or_path": "Cosmos-Data/Cosmos-droid",
+      "subfolder": "iter_000500/text_tokenizer",
+      "type_hint": [
+        "transformers",
+        "Qwen2TokenizerFast"
+      ],
+      "variant": null
+    }
+  ],
+  "transformer": [
+    "diffusers",
+    "Cosmos3OmniTransformer",
+    {
+      "pretrained_model_name_or_path": "Cosmos-Data/Cosmos-droid",
+      "subfolder": "iter_000500/transformer",
+      "type_hint": [
+        "diffusers",
+        "Cosmos3OmniTransformer"
+      ],
+      "variant": null
+    }
+  ],
+  "vae": [
+    "diffusers",
+    "AutoencoderKLWan",
+    {
+      "pretrained_model_name_or_path": "Cosmos-Data/Cosmos-droid",
+      "subfolder": "iter_000500/vae",
+      "type_hint": [
+        "diffusers",
+        "AutoencoderKLWan"
+      ],
+      "variant": null
+    }
+  ]
+}

iter_000500/preprocessor_config.json

@@ -0,0 +1,21 @@
+{
+    "size": {
+        "longest_edge": 16777216,
+        "shortest_edge": 65536
+    },
+    "patch_size": 16,
+    "temporal_patch_size": 2,
+    "merge_size": 2,
+    "image_mean": [
+        0.5,
+        0.5,
+        0.5
+    ],
+    "image_std": [
+        0.5,
+        0.5,
+        0.5
+    ],
+    "processor_class": "Qwen3VLProcessor",
+    "image_processor_type": "Qwen2VLImageProcessorFast"
+}

iter_000500/tokenizer_config.json

@@ -0,0 +1,239 @@
+{
+  "add_bos_token": false,
+  "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
+    },
+    "151657": {
+      "content": "<tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151658": {
+      "content": "</tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151659": {
+      "content": "<|fim_prefix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151660": {
+      "content": "<|fim_middle|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151661": {
+      "content": "<|fim_suffix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151662": {
+      "content": "<|fim_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151663": {
+      "content": "<|repo_name|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151664": {
+      "content": "<|file_sep|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151665": {
+      "content": "<tool_response>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151666": {
+      "content": "</tool_response>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151667": {
+      "content": "<think>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151668": {
+      "content": "</think>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    }
+  },
+  "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,
+  "chat_template": "{%- if tools %}\n    {{- '<|im_start|>system\\n' }}\n    {%- if messages[0].role == 'system' %}\n        {%- if messages[0].content is string %}\n            {{- messages[0].content }}\n        {%- else %}\n            {%- for content in messages[0].content %}\n                {%- if 'text' in content %}\n                    {{- content.text }}\n                {%- endif %}\n            {%- endfor %}\n        {%- endif %}\n        {{- '\\n\\n' }}\n    {%- endif %}\n    {{- \"# Tools\\n\\nYou may call one or more functions to assist with the user query.\\n\\nYou are provided with function signatures within <tools></tools> XML tags:\\n<tools>\" }}\n    {%- for tool in tools %}\n        {{- \"\\n\" }}\n        {{- tool | tojson }}\n    {%- endfor %}\n    {{- \"\\n</tools>\\n\\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\\n<tool_call>\\n{\\\"name\\\": <function-name>, \\\"arguments\\\": <args-json-object>}\\n</tool_call><|im_end|>\\n\" }}\n{%- else %}\n    {%- if messages[0].role == 'system' %}\n        {{- '<|im_start|>system\\n' }}\n        {%- if messages[0].content is string %}\n            {{- messages[0].content }}\n        {%- else %}\n            {%- for content in messages[0].content %}\n                {%- if 'text' in content %}\n                    {{- content.text }}\n                {%- endif %}\n            {%- endfor %}\n        {%- endif %}\n        {{- '<|im_end|>\\n' }}\n    {%- endif %}\n{%- endif %}\n{%- set image_count = namespace(value=0) %}\n{%- set video_count = namespace(value=0) %}\n{%- for message in messages %}\n    {%- if message.role == \"user\" %}\n        {{- '<|im_start|>' + message.role + '\\n' }}\n        {%- if message.content is string %}\n            {{- message.content }}\n        {%- else %}\n            {%- for content in message.content %}\n                {%- if content.type == 'image' or 'image' in content or 'image_url' in content %}\n                    {%- set image_count.value = image_count.value + 1 %}\n                    {%- if add_vision_id %}Picture {{ image_count.value }}: {% endif -%}\n                    <|vision_start|><|image_pad|><|vision_end|>\n                {%- elif content.type == 'video' or 'video' in content %}\n                    {%- set video_count.value = video_count.value + 1 %}\n                    {%- if add_vision_id %}Video {{ video_count.value }}: {% endif -%}\n                    <|vision_start|><|video_pad|><|vision_end|>\n                {%- elif 'text' in content %}\n                    {{- content.text }}\n                {%- endif %}\n            {%- endfor %}\n        {%- endif %}\n        {{- '<|im_end|>\\n' }}\n    {%- elif message.role == \"assistant\" %}\n        {{- '<|im_start|>' + message.role + '\\n' }}\n        {%- if message.content is string %}\n            {{- message.content }}\n        {%- else %}\n            {%- for content_item in message.content %}\n                {%- if 'text' in content_item %}\n                    {{- content_item.text }}\n                {%- endif %}\n            {%- endfor %}\n        {%- endif %}\n        {%- if message.tool_calls %}\n            {%- for tool_call in message.tool_calls %}\n                {%- if (loop.first and message.content) or (not loop.first) %}\n                    {{- '\\n' }}\n                {%- endif %}\n                {%- if tool_call.function %}\n                    {%- set tool_call = tool_call.function %}\n                {%- endif %}\n                {{- '<tool_call>\\n{\"name\": \"' }}\n                {{- tool_call.name }}\n                {{- '\", \"arguments\": ' }}\n                {%- if tool_call.arguments is string %}\n                    {{- tool_call.arguments }}\n                {%- else %}\n                    {{- tool_call.arguments | tojson }}\n                {%- endif %}\n                {{- '}\\n</tool_call>' }}\n            {%- endfor %}\n        {%- endif %}\n        {{- '<|im_end|>\\n' }}\n    {%- elif message.role == \"tool\" %}\n        {%- if loop.first or (messages[loop.index0 - 1].role != \"tool\") %}\n            {{- '<|im_start|>user' }}\n        {%- endif %}\n        {{- '\\n<tool_response>\\n' }}\n        {%- if message.content is string %}\n            {{- message.content }}\n        {%- else %}\n            {%- for content in message.content %}\n                {%- if content.type == 'image' or 'image' in content or 'image_url' in content %}\n                    {%- set image_count.value = image_count.value + 1 %}\n                    {%- if add_vision_id %}Picture {{ image_count.value }}: {% endif -%}\n                    <|vision_start|><|image_pad|><|vision_end|>\n                {%- elif content.type == 'video' or 'video' in content %}\n                    {%- set video_count.value = video_count.value + 1 %}\n                    {%- if add_vision_id %}Video {{ video_count.value }}: {% endif -%}\n                    <|vision_start|><|video_pad|><|vision_end|>\n                {%- elif 'text' in content %}\n                    {{- content.text }}\n                {%- endif %}\n            {%- endfor %}\n        {%- endif %}\n        {{- '\\n</tool_response>' }}\n        {%- if loop.last or (messages[loop.index0 + 1].role != \"tool\") %}\n            {{- '<|im_end|>\\n' }}\n        {%- endif %}\n    {%- endif %}\n{%- endfor %}\n{%- if add_generation_prompt %}\n    {{- '<|im_start|>assistant\\n' }}\n{%- endif %}\n",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "errors": "replace",
+  "model_max_length": 262144,
+  "pad_token": "<|endoftext|>",
+  "split_special_tokens": false,
+  "tokenizer_class": "Qwen2Tokenizer",
+  "unk_token": null
+}

iter_000500/video_preprocessor_config.json

@@ -0,0 +1,21 @@
+{
+    "size": {
+        "longest_edge": 25165824,
+        "shortest_edge": 4096
+    },
+    "patch_size": 16,
+    "temporal_patch_size": 2,
+    "merge_size": 2,
+    "image_mean": [
+        0.5,
+        0.5,
+        0.5
+    ],
+    "image_std": [
+        0.5,
+        0.5,
+        0.5
+    ],
+    "processor_class": "Qwen3VLProcessor",
+    "video_processor_type": "Qwen3VLVideoProcessor"
+}

iter_001000/chat_template.json

@@ -0,0 +1,3 @@
+{
+    "chat_template": "{%- if tools %}\n    {{- '<|im_start|>system\\n' }}\n    {%- if messages[0].role == 'system' %}\n        {%- if messages[0].content is string %}\n            {{- messages[0].content }}\n        {%- else %}\n            {%- for content in messages[0].content %}\n                {%- if 'text' in content %}\n                    {{- content.text }}\n                {%- endif %}\n            {%- endfor %}\n        {%- endif %}\n        {{- '\\n\\n' }}\n    {%- endif %}\n    {{- \"# Tools\\n\\nYou may call one or more functions to assist with the user query.\\n\\nYou are provided with function signatures within <tools></tools> XML tags:\\n<tools>\" }}\n    {%- for tool in tools %}\n        {{- \"\\n\" }}\n        {{- tool | tojson }}\n    {%- endfor %}\n    {{- \"\\n</tools>\\n\\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\\n<tool_call>\\n{\\\"name\\\": <function-name>, \\\"arguments\\\": <args-json-object>}\\n</tool_call><|im_end|>\\n\" }}\n{%- else %}\n    {%- if messages[0].role == 'system' %}\n        {{- '<|im_start|>system\\n' }}\n        {%- if messages[0].content is string %}\n            {{- messages[0].content }}\n        {%- else %}\n            {%- for content in messages[0].content %}\n                {%- if 'text' in content %}\n                    {{- content.text }}\n                {%- endif %}\n            {%- endfor %}\n        {%- endif %}\n        {{- '<|im_end|>\\n' }}\n    {%- endif %}\n{%- endif %}\n{%- set image_count = namespace(value=0) %}\n{%- set video_count = namespace(value=0) %}\n{%- for message in messages %}\n    {%- if message.role == \"user\" %}\n        {{- '<|im_start|>' + message.role + '\\n' }}\n        {%- if message.content is string %}\n            {{- message.content }}\n        {%- else %}\n            {%- for content in message.content %}\n                {%- if content.type == 'image' or 'image' in content or 'image_url' in content %}\n                    {%- set image_count.value = image_count.value + 1 %}\n                    {%- if add_vision_id %}Picture {{ image_count.value }}: {% endif -%}\n                    <|vision_start|><|image_pad|><|vision_end|>\n                {%- elif content.type == 'video' or 'video' in content %}\n                    {%- set video_count.value = video_count.value + 1 %}\n                    {%- if add_vision_id %}Video {{ video_count.value }}: {% endif -%}\n                    <|vision_start|><|video_pad|><|vision_end|>\n                {%- elif 'text' in content %}\n                    {{- content.text }}\n                {%- endif %}\n            {%- endfor %}\n        {%- endif %}\n        {{- '<|im_end|>\\n' }}\n    {%- elif message.role == \"assistant\" %}\n        {{- '<|im_start|>' + message.role + '\\n' }}\n        {%- if message.content is string %}\n            {{- message.content }}\n        {%- else %}\n            {%- for content_item in message.content %}\n                {%- if 'text' in content_item %}\n                    {{- content_item.text }}\n                {%- endif %}\n            {%- endfor %}\n        {%- endif %}\n        {%- if message.tool_calls %}\n            {%- for tool_call in message.tool_calls %}\n                {%- if (loop.first and message.content) or (not loop.first) %}\n                    {{- '\\n' }}\n                {%- endif %}\n                {%- if tool_call.function %}\n                    {%- set tool_call = tool_call.function %}\n                {%- endif %}\n                {{- '<tool_call>\\n{\"name\": \"' }}\n                {{- tool_call.name }}\n                {{- '\", \"arguments\": ' }}\n                {%- if tool_call.arguments is string %}\n                    {{- tool_call.arguments }}\n                {%- else %}\n                    {{- tool_call.arguments | tojson }}\n                {%- endif %}\n                {{- '}\\n</tool_call>' }}\n            {%- endfor %}\n        {%- endif %}\n        {{- '<|im_end|>\\n' }}\n    {%- elif message.role == \"tool\" %}\n        {%- if loop.first or (messages[loop.index0 - 1].role != \"tool\") %}\n            {{- '<|im_start|>user' }}\n        {%- endif %}\n        {{- '\\n<tool_response>\\n' }}\n        {%- if message.content is string %}\n            {{- message.content }}\n        {%- else %}\n            {%- for content in message.content %}\n                {%- if content.type == 'image' or 'image' in content or 'image_url' in content %}\n                    {%- set image_count.value = image_count.value + 1 %}\n                    {%- if add_vision_id %}Picture {{ image_count.value }}: {% endif -%}\n                    <|vision_start|><|image_pad|><|vision_end|>\n                {%- elif content.type == 'video' or 'video' in content %}\n                    {%- set video_count.value = video_count.value + 1 %}\n                    {%- if add_vision_id %}Video {{ video_count.value }}: {% endif -%}\n                    <|vision_start|><|video_pad|><|vision_end|>\n                {%- elif 'text' in content %}\n                    {{- content.text }}\n                {%- endif %}\n            {%- endfor %}\n        {%- endif %}\n        {{- '\\n</tool_response>' }}\n        {%- if loop.last or (messages[loop.index0 + 1].role != \"tool\") %}\n            {{- '<|im_end|>\\n' }}\n        {%- endif %}\n    {%- endif %}\n{%- endfor %}\n{%- if add_generation_prompt %}\n    {{- '<|im_start|>assistant\\n' }}\n{%- endif %}\n"
+}

iter_001000/checkpoint.json

@@ -0,0 +1 @@
+{}

iter_001000/config.json

@@ -0,0 +1,240 @@
+{
+  "allow_patterns_overrides": [
+    "*/*.safetensors"
+  ],
+  "architectures": [
+    "Cosmos3ForConditionalGeneration"
+  ],
+  "image_token_id": 151655,
+  "model": {
+    "_recursive_": false,
+    "_target": "omni_mot_model",
+    "config": {
+      "_type": "omni_mot_model_config",
+      "action_gen": true,
+      "activation_checkpointing": {
+        "determinism_check": "default",
+        "mode": "full",
+        "preserve_rng_state": true,
+        "save_ops_regex": [
+          "fmha"
+        ]
+      },
+      "causal_training_strategy": "none",
+      "compile": {
+        "ar_post_saturation_mode": "default",
+        "compile_dynamic": true,
+        "compiled_region": "language",
+        "coordinate_descent_tuning": false,
+        "enabled": false,
+        "max_autotune_pointwise": false,
+        "use_cuda_graphs": false
+      },
+      "diffusion_expert_config": {
+        "base_fps": 24,
+        "enable_fps_modulation": true,
+        "load_weights_from_pretrained": false,
+        "max_vae_latent_side_after_patchify": 20,
+        "patch_spatial": 2,
+        "sound_base_temporal_compression_factor": null,
+        "timestep_range": 1.0,
+        "unified_3d_mrope_reset_spatial_ids": true,
+        "unified_3d_mrope_temporal_modality_margin": 15000,
+        "vision_temporal_position_mode": "latent_index"
+      },
+      "ema": {
+        "enabled": false,
+        "iteration_shift": 0,
+        "rate": 0.1
+      },
+      "enable_input_bias": true,
+      "exclude_reasoner_weights_from_checkpoint": false,
+      "fixed_step_sampler_config": null,
+      "input_caption_key": "ai_caption",
+      "input_image_key": "images",
+      "input_video_key": "video",
+      "joint_attn_implementation": "two_way",
+      "latent_downsample_factor": 16,
+      "lbl": {
+        "coeff_gen": null,
+        "coeff_und": null,
+        "method": "local"
+      },
+      "load_vision_tokenizer": true,
+      "log_enc_time_every_n": 100,
+      "lora_alpha": 32,
+      "lora_enabled": false,
+      "lora_extra_trainable_modules": "",
+      "lora_rank": 16,
+      "lora_target_modules": "q_proj_moe_gen,k_proj_moe_gen,v_proj_moe_gen,o_proj_moe_gen",
+      "max_action_dim": 64,
+      "max_num_tokens_after_packing": -1,
+      "natten_parameter_list": null,
+      "net": null,
+      "num_embodiment_domains": 32,
+      "parallelism": {
+        "attention_io_layout": "sequence_sharded",
+        "cfg_parallel_shard_degree": 1,
+        "context_parallel_shard_degree": 1,
+        "data_parallel_replicate_degree": 1,
+        "data_parallel_shard_degree": 8,
+        "enable_inference_mode": false,
+        "fsdp_master_dtype": "float32"
+      },
+      "precision": "bfloat16",
+      "rectified_flow_inference_config": {
+        "num_train_timesteps": 1000,
+        "scheduler_type": "unipc",
+        "shift": 1,
+        "use_dynamic_shifting": false
+      },
+      "rectified_flow_training_config": {
+        "action_loss_weight": 10.0,
+        "align_shared_action_schedule": false,
+        "image_loss_scale": 1.0,
+        "independent_action_schedule": false,
+        "independent_sound_schedule": false,
+        "loss_scale": 10.0,
+        "normalize_loss_by_active": false,
+        "sample_level_loss_averaging": false,
+        "shift": {
+          "256": 3,
+          "480": 5,
+          "720": 10
+        },
+        "shift_action": null,
+        "shift_sound": null,
+        "sound_loss_scale": null,
+        "train_time_action_distribution": "logitnormal",
+        "train_time_image_distribution": "logitnormal",
+        "train_time_sound_distribution": "logitnormal",
+        "train_time_video_distribution": "waver",
+        "train_time_weight": "uniform",
+        "use_discrete_rf": false,
+        "use_dynamic_shift": false
+      },
+      "resolution": "720",
+      "sound_dim": null,
+      "sound_gen": false,
+      "sound_latent_fps": 25,
+      "sound_tokenizer": null,
+      "state_ch": 48,
+      "state_t": 300,
+      "tokenizer": {
+        "_target": "wan2pt2_vae_interface",
+        "bucket_name": "",
+        "causal": true,
+        "chunk_duration": 93,
+        "encode_bucket_multiple": null,
+        "encode_chunk_frames": {
+          "256": 68,
+          "480": 24,
+          "720": 8,
+          "768": 8
+        },
+        "encode_exact_durations": [
+          33
+        ],
+        "keep_decoder_cache": false,
+        "object_store_credential_path_pretrained": "",
+        "spatial_compression_factor": 16,
+        "temporal_compression_factor": 4,
+        "temporal_window": null,
+        "use_streaming_encode": false,
+        "vae_path": "pretrained/tokenizers/video/wan2pt2/Wan2.2_VAE.pth"
+      },
+      "video_temporal_causal": false,
+      "vision_gen": true,
+      "vlm_config": {
+        "layer_module": "Qwen2MoTDecoderLayer",
+        "model_instance": {
+          "_target": "qwen3_vl_text_for_causal_lm",
+          "config": {
+            "_target": "create_vlm_config",
+            "base_config": {
+              "_target": "qwen3_vl_mot_config_from_json_file",
+              "json_file": "cosmos3://vfm/models/vlm/qwen3_vl/configs/Qwen3-VL-8B-Instruct.json"
+            },
+            "freeze_und": false,
+            "layer_module": "MoTDecoderLayer",
+            "qk_norm_for_text": true,
+            "tie_word_embeddings": true
+          }
+        },
+        "model_name": "Qwen/Qwen3-VL-8B-Instruct",
+        "pretrained_weights": {
+          "backbone_path": "s3://bucket0/cosmos3/pretrained/huggingface/Qwen/Qwen3-VL-8B-Instruct/",
+          "checkpoint_format": null,
+          "credentials_path": "",
+          "enable_gcs_patch_in_boto3": true,
+          "enabled": false
+        },
+        "qk_norm": false,
+        "safetensors_path": "",
+        "tie_word_embeddings": false,
+        "tokenizer": {
+          "_target": "create_qwen2_tokenizer_with_download",
+          "config_variant": "hf",
+          "pretrained_model_name": "nvidia/Cosmos3-Nano-Policy-DROID"
+        },
+        "use_system_prompt": false
+      }
+    }
+  },
+  "model_type": "cosmos3_omni",
+  "text_config": {
+    "attention_bias": false,
+    "attention_dropout": 0.0,
+    "bos_token_id": 151643,
+    "dtype": "bfloat16",
+    "eos_token_id": 151645,
+    "head_dim": 128,
+    "hidden_act": "silu",
+    "hidden_size": 4096,
+    "initializer_range": 0.02,
+    "intermediate_size": 12288,
+    "max_position_embeddings": 262144,
+    "model_type": "qwen3_vl_text",
+    "num_attention_heads": 32,
+    "num_hidden_layers": 36,
+    "num_key_value_heads": 8,
+    "rms_norm_eps": 1e-06,
+    "rope_scaling": {
+      "mrope_interleaved": true,
+      "mrope_section": [
+        24,
+        20,
+        20
+      ],
+      "rope_type": "default"
+    },
+    "rope_theta": 5000000,
+    "use_cache": true,
+    "vocab_size": 151936
+  },
+  "tie_word_embeddings": false,
+  "transformers_version": "4.57.0.dev0",
+  "video_token_id": 151656,
+  "vision_config": {
+    "deepstack_visual_indexes": [
+      8,
+      16,
+      24
+    ],
+    "depth": 27,
+    "hidden_act": "gelu_pytorch_tanh",
+    "hidden_size": 1152,
+    "in_channels": 3,
+    "initializer_range": 0.02,
+    "intermediate_size": 4304,
+    "model_type": "qwen3_vl",
+    "num_heads": 16,
+    "num_position_embeddings": 2304,
+    "out_hidden_size": 4096,
+    "patch_size": 16,
+    "spatial_merge_size": 2,
+    "temporal_patch_size": 2
+  },
+  "vision_end_token_id": 151653,
+  "vision_start_token_id": 151652
+}

iter_001000/generation_config.json

@@ -0,0 +1,14 @@
+{
+    "bos_token_id": 151643,
+    "pad_token_id": 151643,
+    "do_sample": true,
+    "eos_token_id": [
+        151645,
+        151643
+    ],
+    "top_k": 20,
+    "top_p": 0.8,
+    "repetition_penalty": 1.0,
+    "temperature": 0.7,
+    "transformers_version": "4.56.0"
+}

iter_001000/model_index.json

@@ -0,0 +1,32 @@
+{
+  "_class_name": "Cosmos3OmniPipeline",
+  "_diffusers_version": "0.39.0",
+  "safety_checker": [
+    null,
+    null
+  ],
+  "scheduler": [
+    "diffusers",
+    "UniPCMultistepScheduler"
+  ],
+  "sound_tokenizer": [
+    null,
+    null
+  ],
+  "text_tokenizer": [
+    "transformers",
+    "Qwen2TokenizerFast"
+  ],
+  "transformer": [
+    "diffusers",
+    "Cosmos3OmniTransformer"
+  ],
+  "vae": [
+    "diffusers",
+    "AutoencoderKLWan"
+  ],
+  "vision_encoder": [
+    "transformers",
+    "Qwen3VLVisionModel"
+  ]
+}

iter_001000/modular_model_index.json

@@ -0,0 +1,58 @@
+{
+  "_blocks_class_name": "Cosmos3OmniBlocks",
+  "_class_name": "Cosmos3OmniModularPipeline",
+  "_diffusers_version": "0.39.0",
+  "is_distilled": false,
+  "scheduler": [
+    "diffusers",
+    "UniPCMultistepScheduler",
+    {
+      "pretrained_model_name_or_path": "Cosmos-Data/Cosmos-droid",
+      "subfolder": "iter_001000/scheduler",
+      "type_hint": [
+        "diffusers",
+        "UniPCMultistepScheduler"
+      ],
+      "variant": null
+    }
+  ],
+  "text_tokenizer": [
+    "transformers",
+    "Qwen2TokenizerFast",
+    {
+      "pretrained_model_name_or_path": "Cosmos-Data/Cosmos-droid",
+      "subfolder": "iter_001000/text_tokenizer",
+      "type_hint": [
+        "transformers",
+        "Qwen2TokenizerFast"
+      ],
+      "variant": null
+    }
+  ],
+  "transformer": [
+    "diffusers",
+    "Cosmos3OmniTransformer",
+    {
+      "pretrained_model_name_or_path": "Cosmos-Data/Cosmos-droid",
+      "subfolder": "iter_001000/transformer",
+      "type_hint": [
+        "diffusers",
+        "Cosmos3OmniTransformer"
+      ],
+      "variant": null
+    }
+  ],
+  "vae": [
+    "diffusers",
+    "AutoencoderKLWan",
+    {
+      "pretrained_model_name_or_path": "Cosmos-Data/Cosmos-droid",
+      "subfolder": "iter_001000/vae",
+      "type_hint": [
+        "diffusers",
+        "AutoencoderKLWan"
+      ],
+      "variant": null
+    }
+  ]
+}

iter_001000/preprocessor_config.json

@@ -0,0 +1,21 @@
+{
+    "size": {
+        "longest_edge": 16777216,
+        "shortest_edge": 65536
+    },
+    "patch_size": 16,
+    "temporal_patch_size": 2,
+    "merge_size": 2,
+    "image_mean": [
+        0.5,
+        0.5,
+        0.5
+    ],
+    "image_std": [
+        0.5,
+        0.5,
+        0.5
+    ],
+    "processor_class": "Qwen3VLProcessor",
+    "image_processor_type": "Qwen2VLImageProcessorFast"
+}

iter_001000/tokenizer_config.json

@@ -0,0 +1,239 @@
+{
+  "add_bos_token": false,
+  "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
+    },
+    "151657": {
+      "content": "<tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151658": {
+      "content": "</tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151659": {
+      "content": "<|fim_prefix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151660": {
+      "content": "<|fim_middle|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151661": {
+      "content": "<|fim_suffix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151662": {
+      "content": "<|fim_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151663": {
+      "content": "<|repo_name|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151664": {
+      "content": "<|file_sep|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151665": {
+      "content": "<tool_response>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151666": {
+      "content": "</tool_response>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151667": {
+      "content": "<think>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151668": {
+      "content": "</think>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    }
+  },
+  "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,
+  "chat_template": "{%- if tools %}\n    {{- '<|im_start|>system\\n' }}\n    {%- if messages[0].role == 'system' %}\n        {%- if messages[0].content is string %}\n            {{- messages[0].content }}\n        {%- else %}\n            {%- for content in messages[0].content %}\n                {%- if 'text' in content %}\n                    {{- content.text }}\n                {%- endif %}\n            {%- endfor %}\n        {%- endif %}\n        {{- '\\n\\n' }}\n    {%- endif %}\n    {{- \"# Tools\\n\\nYou may call one or more functions to assist with the user query.\\n\\nYou are provided with function signatures within <tools></tools> XML tags:\\n<tools>\" }}\n    {%- for tool in tools %}\n        {{- \"\\n\" }}\n        {{- tool | tojson }}\n    {%- endfor %}\n    {{- \"\\n</tools>\\n\\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\\n<tool_call>\\n{\\\"name\\\": <function-name>, \\\"arguments\\\": <args-json-object>}\\n</tool_call><|im_end|>\\n\" }}\n{%- else %}\n    {%- if messages[0].role == 'system' %}\n        {{- '<|im_start|>system\\n' }}\n        {%- if messages[0].content is string %}\n            {{- messages[0].content }}\n        {%- else %}\n            {%- for content in messages[0].content %}\n                {%- if 'text' in content %}\n                    {{- content.text }}\n                {%- endif %}\n            {%- endfor %}\n        {%- endif %}\n        {{- '<|im_end|>\\n' }}\n    {%- endif %}\n{%- endif %}\n{%- set image_count = namespace(value=0) %}\n{%- set video_count = namespace(value=0) %}\n{%- for message in messages %}\n    {%- if message.role == \"user\" %}\n        {{- '<|im_start|>' + message.role + '\\n' }}\n        {%- if message.content is string %}\n            {{- message.content }}\n        {%- else %}\n            {%- for content in message.content %}\n                {%- if content.type == 'image' or 'image' in content or 'image_url' in content %}\n                    {%- set image_count.value = image_count.value + 1 %}\n                    {%- if add_vision_id %}Picture {{ image_count.value }}: {% endif -%}\n                    <|vision_start|><|image_pad|><|vision_end|>\n                {%- elif content.type == 'video' or 'video' in content %}\n                    {%- set video_count.value = video_count.value + 1 %}\n                    {%- if add_vision_id %}Video {{ video_count.value }}: {% endif -%}\n                    <|vision_start|><|video_pad|><|vision_end|>\n                {%- elif 'text' in content %}\n                    {{- content.text }}\n                {%- endif %}\n            {%- endfor %}\n        {%- endif %}\n        {{- '<|im_end|>\\n' }}\n    {%- elif message.role == \"assistant\" %}\n        {{- '<|im_start|>' + message.role + '\\n' }}\n        {%- if message.content is string %}\n            {{- message.content }}\n        {%- else %}\n            {%- for content_item in message.content %}\n                {%- if 'text' in content_item %}\n                    {{- content_item.text }}\n                {%- endif %}\n            {%- endfor %}\n        {%- endif %}\n        {%- if message.tool_calls %}\n            {%- for tool_call in message.tool_calls %}\n                {%- if (loop.first and message.content) or (not loop.first) %}\n                    {{- '\\n' }}\n                {%- endif %}\n                {%- if tool_call.function %}\n                    {%- set tool_call = tool_call.function %}\n                {%- endif %}\n                {{- '<tool_call>\\n{\"name\": \"' }}\n                {{- tool_call.name }}\n                {{- '\", \"arguments\": ' }}\n                {%- if tool_call.arguments is string %}\n                    {{- tool_call.arguments }}\n                {%- else %}\n                    {{- tool_call.arguments | tojson }}\n                {%- endif %}\n                {{- '}\\n</tool_call>' }}\n            {%- endfor %}\n        {%- endif %}\n        {{- '<|im_end|>\\n' }}\n    {%- elif message.role == \"tool\" %}\n        {%- if loop.first or (messages[loop.index0 - 1].role != \"tool\") %}\n            {{- '<|im_start|>user' }}\n        {%- endif %}\n        {{- '\\n<tool_response>\\n' }}\n        {%- if message.content is string %}\n            {{- message.content }}\n        {%- else %}\n            {%- for content in message.content %}\n                {%- if content.type == 'image' or 'image' in content or 'image_url' in content %}\n                    {%- set image_count.value = image_count.value + 1 %}\n                    {%- if add_vision_id %}Picture {{ image_count.value }}: {% endif -%}\n                    <|vision_start|><|image_pad|><|vision_end|>\n                {%- elif content.type == 'video' or 'video' in content %}\n                    {%- set video_count.value = video_count.value + 1 %}\n                    {%- if add_vision_id %}Video {{ video_count.value }}: {% endif -%}\n                    <|vision_start|><|video_pad|><|vision_end|>\n                {%- elif 'text' in content %}\n                    {{- content.text }}\n                {%- endif %}\n            {%- endfor %}\n        {%- endif %}\n        {{- '\\n</tool_response>' }}\n        {%- if loop.last or (messages[loop.index0 + 1].role != \"tool\") %}\n            {{- '<|im_end|>\\n' }}\n        {%- endif %}\n    {%- endif %}\n{%- endfor %}\n{%- if add_generation_prompt %}\n    {{- '<|im_start|>assistant\\n' }}\n{%- endif %}\n",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "errors": "replace",
+  "model_max_length": 262144,
+  "pad_token": "<|endoftext|>",
+  "split_special_tokens": false,
+  "tokenizer_class": "Qwen2Tokenizer",
+  "unk_token": null
+}

iter_001000/video_preprocessor_config.json

@@ -0,0 +1,21 @@
+{
+    "size": {
+        "longest_edge": 25165824,
+        "shortest_edge": 4096
+    },
+    "patch_size": 16,
+    "temporal_patch_size": 2,
+    "merge_size": 2,
+    "image_mean": [
+        0.5,
+        0.5,
+        0.5
+    ],
+    "image_std": [
+        0.5,
+        0.5,
+        0.5
+    ],
+    "processor_class": "Qwen3VLProcessor",
+    "video_processor_type": "Qwen3VLVideoProcessor"
+}

iter_001500/chat_template.json

@@ -0,0 +1,3 @@
+{
+    "chat_template": "{%- if tools %}\n    {{- '<|im_start|>system\\n' }}\n    {%- if messages[0].role == 'system' %}\n        {%- if messages[0].content is string %}\n            {{- messages[0].content }}\n        {%- else %}\n            {%- for content in messages[0].content %}\n                {%- if 'text' in content %}\n                    {{- content.text }}\n                {%- endif %}\n            {%- endfor %}\n        {%- endif %}\n        {{- '\\n\\n' }}\n    {%- endif %}\n    {{- \"# Tools\\n\\nYou may call one or more functions to assist with the user query.\\n\\nYou are provided with function signatures within <tools></tools> XML tags:\\n<tools>\" }}\n    {%- for tool in tools %}\n        {{- \"\\n\" }}\n        {{- tool | tojson }}\n    {%- endfor %}\n    {{- \"\\n</tools>\\n\\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\\n<tool_call>\\n{\\\"name\\\": <function-name>, \\\"arguments\\\": <args-json-object>}\\n</tool_call><|im_end|>\\n\" }}\n{%- else %}\n    {%- if messages[0].role == 'system' %}\n        {{- '<|im_start|>system\\n' }}\n        {%- if messages[0].content is string %}\n            {{- messages[0].content }}\n        {%- else %}\n            {%- for content in messages[0].content %}\n                {%- if 'text' in content %}\n                    {{- content.text }}\n                {%- endif %}\n            {%- endfor %}\n        {%- endif %}\n        {{- '<|im_end|>\\n' }}\n    {%- endif %}\n{%- endif %}\n{%- set image_count = namespace(value=0) %}\n{%- set video_count = namespace(value=0) %}\n{%- for message in messages %}\n    {%- if message.role == \"user\" %}\n        {{- '<|im_start|>' + message.role + '\\n' }}\n        {%- if message.content is string %}\n            {{- message.content }}\n        {%- else %}\n            {%- for content in message.content %}\n                {%- if content.type == 'image' or 'image' in content or 'image_url' in content %}\n                    {%- set image_count.value = image_count.value + 1 %}\n                    {%- if add_vision_id %}Picture {{ image_count.value }}: {% endif -%}\n                    <|vision_start|><|image_pad|><|vision_end|>\n                {%- elif content.type == 'video' or 'video' in content %}\n                    {%- set video_count.value = video_count.value + 1 %}\n                    {%- if add_vision_id %}Video {{ video_count.value }}: {% endif -%}\n                    <|vision_start|><|video_pad|><|vision_end|>\n                {%- elif 'text' in content %}\n                    {{- content.text }}\n                {%- endif %}\n            {%- endfor %}\n        {%- endif %}\n        {{- '<|im_end|>\\n' }}\n    {%- elif message.role == \"assistant\" %}\n        {{- '<|im_start|>' + message.role + '\\n' }}\n        {%- if message.content is string %}\n            {{- message.content }}\n        {%- else %}\n            {%- for content_item in message.content %}\n                {%- if 'text' in content_item %}\n                    {{- content_item.text }}\n                {%- endif %}\n            {%- endfor %}\n        {%- endif %}\n        {%- if message.tool_calls %}\n            {%- for tool_call in message.tool_calls %}\n                {%- if (loop.first and message.content) or (not loop.first) %}\n                    {{- '\\n' }}\n                {%- endif %}\n                {%- if tool_call.function %}\n                    {%- set tool_call = tool_call.function %}\n                {%- endif %}\n                {{- '<tool_call>\\n{\"name\": \"' }}\n                {{- tool_call.name }}\n                {{- '\", \"arguments\": ' }}\n                {%- if tool_call.arguments is string %}\n                    {{- tool_call.arguments }}\n                {%- else %}\n                    {{- tool_call.arguments | tojson }}\n                {%- endif %}\n                {{- '}\\n</tool_call>' }}\n            {%- endfor %}\n        {%- endif %}\n        {{- '<|im_end|>\\n' }}\n    {%- elif message.role == \"tool\" %}\n        {%- if loop.first or (messages[loop.index0 - 1].role != \"tool\") %}\n            {{- '<|im_start|>user' }}\n        {%- endif %}\n        {{- '\\n<tool_response>\\n' }}\n        {%- if message.content is string %}\n            {{- message.content }}\n        {%- else %}\n            {%- for content in message.content %}\n                {%- if content.type == 'image' or 'image' in content or 'image_url' in content %}\n                    {%- set image_count.value = image_count.value + 1 %}\n                    {%- if add_vision_id %}Picture {{ image_count.value }}: {% endif -%}\n                    <|vision_start|><|image_pad|><|vision_end|>\n                {%- elif content.type == 'video' or 'video' in content %}\n                    {%- set video_count.value = video_count.value + 1 %}\n                    {%- if add_vision_id %}Video {{ video_count.value }}: {% endif -%}\n                    <|vision_start|><|video_pad|><|vision_end|>\n                {%- elif 'text' in content %}\n                    {{- content.text }}\n                {%- endif %}\n            {%- endfor %}\n        {%- endif %}\n        {{- '\\n</tool_response>' }}\n        {%- if loop.last or (messages[loop.index0 + 1].role != \"tool\") %}\n            {{- '<|im_end|>\\n' }}\n        {%- endif %}\n    {%- endif %}\n{%- endfor %}\n{%- if add_generation_prompt %}\n    {{- '<|im_start|>assistant\\n' }}\n{%- endif %}\n"
+}

iter_001500/checkpoint.json

@@ -0,0 +1 @@
+{}

iter_001500/config.json

@@ -0,0 +1,240 @@
+{
+  "allow_patterns_overrides": [
+    "*/*.safetensors"
+  ],
+  "architectures": [
+    "Cosmos3ForConditionalGeneration"
+  ],
+  "image_token_id": 151655,
+  "model": {
+    "_recursive_": false,
+    "_target": "omni_mot_model",
+    "config": {
+      "_type": "omni_mot_model_config",
+      "action_gen": true,
+      "activation_checkpointing": {
+        "determinism_check": "default",
+        "mode": "full",
+        "preserve_rng_state": true,
+        "save_ops_regex": [
+          "fmha"
+        ]
+      },
+      "causal_training_strategy": "none",
+      "compile": {
+        "ar_post_saturation_mode": "default",
+        "compile_dynamic": true,
+        "compiled_region": "language",
+        "coordinate_descent_tuning": false,
+        "enabled": false,
+        "max_autotune_pointwise": false,
+        "use_cuda_graphs": false
+      },
+      "diffusion_expert_config": {
+        "base_fps": 24,
+        "enable_fps_modulation": true,
+        "load_weights_from_pretrained": false,
+        "max_vae_latent_side_after_patchify": 20,
+        "patch_spatial": 2,
+        "sound_base_temporal_compression_factor": null,
+        "timestep_range": 1.0,
+        "unified_3d_mrope_reset_spatial_ids": true,
+        "unified_3d_mrope_temporal_modality_margin": 15000,
+        "vision_temporal_position_mode": "latent_index"
+      },
+      "ema": {
+        "enabled": false,
+        "iteration_shift": 0,
+        "rate": 0.1
+      },
+      "enable_input_bias": true,
+      "exclude_reasoner_weights_from_checkpoint": false,
+      "fixed_step_sampler_config": null,
+      "input_caption_key": "ai_caption",
+      "input_image_key": "images",
+      "input_video_key": "video",
+      "joint_attn_implementation": "two_way",
+      "latent_downsample_factor": 16,
+      "lbl": {
+        "coeff_gen": null,
+        "coeff_und": null,
+        "method": "local"
+      },
+      "load_vision_tokenizer": true,
+      "log_enc_time_every_n": 100,
+      "lora_alpha": 32,
+      "lora_enabled": false,
+      "lora_extra_trainable_modules": "",
+      "lora_rank": 16,
+      "lora_target_modules": "q_proj_moe_gen,k_proj_moe_gen,v_proj_moe_gen,o_proj_moe_gen",
+      "max_action_dim": 64,
+      "max_num_tokens_after_packing": -1,
+      "natten_parameter_list": null,
+      "net": null,
+      "num_embodiment_domains": 32,
+      "parallelism": {
+        "attention_io_layout": "sequence_sharded",
+        "cfg_parallel_shard_degree": 1,
+        "context_parallel_shard_degree": 1,
+        "data_parallel_replicate_degree": 1,
+        "data_parallel_shard_degree": 8,
+        "enable_inference_mode": false,
+        "fsdp_master_dtype": "float32"
+      },
+      "precision": "bfloat16",
+      "rectified_flow_inference_config": {
+        "num_train_timesteps": 1000,
+        "scheduler_type": "unipc",
+        "shift": 1,
+        "use_dynamic_shifting": false
+      },
+      "rectified_flow_training_config": {
+        "action_loss_weight": 10.0,
+        "align_shared_action_schedule": false,
+        "image_loss_scale": 1.0,
+        "independent_action_schedule": false,
+        "independent_sound_schedule": false,
+        "loss_scale": 10.0,
+        "normalize_loss_by_active": false,
+        "sample_level_loss_averaging": false,
+        "shift": {
+          "256": 3,
+          "480": 5,
+          "720": 10
+        },
+        "shift_action": null,
+        "shift_sound": null,
+        "sound_loss_scale": null,
+        "train_time_action_distribution": "logitnormal",
+        "train_time_image_distribution": "logitnormal",
+        "train_time_sound_distribution": "logitnormal",
+        "train_time_video_distribution": "waver",
+        "train_time_weight": "uniform",
+        "use_discrete_rf": false,
+        "use_dynamic_shift": false
+      },
+      "resolution": "720",
+      "sound_dim": null,
+      "sound_gen": false,
+      "sound_latent_fps": 25,
+      "sound_tokenizer": null,
+      "state_ch": 48,
+      "state_t": 300,
+      "tokenizer": {
+        "_target": "wan2pt2_vae_interface",
+        "bucket_name": "",
+        "causal": true,
+        "chunk_duration": 93,
+        "encode_bucket_multiple": null,
+        "encode_chunk_frames": {
+          "256": 68,
+          "480": 24,
+          "720": 8,
+          "768": 8
+        },
+        "encode_exact_durations": [
+          33
+        ],
+        "keep_decoder_cache": false,
+        "object_store_credential_path_pretrained": "",
+        "spatial_compression_factor": 16,
+        "temporal_compression_factor": 4,
+        "temporal_window": null,
+        "use_streaming_encode": false,
+        "vae_path": "pretrained/tokenizers/video/wan2pt2/Wan2.2_VAE.pth"
+      },
+      "video_temporal_causal": false,
+      "vision_gen": true,
+      "vlm_config": {
+        "layer_module": "Qwen2MoTDecoderLayer",
+        "model_instance": {
+          "_target": "qwen3_vl_text_for_causal_lm",
+          "config": {
+            "_target": "create_vlm_config",
+            "base_config": {
+              "_target": "qwen3_vl_mot_config_from_json_file",
+              "json_file": "cosmos3://vfm/models/vlm/qwen3_vl/configs/Qwen3-VL-8B-Instruct.json"
+            },
+            "freeze_und": false,
+            "layer_module": "MoTDecoderLayer",
+            "qk_norm_for_text": true,
+            "tie_word_embeddings": true
+          }
+        },
+        "model_name": "Qwen/Qwen3-VL-8B-Instruct",
+        "pretrained_weights": {
+          "backbone_path": "s3://bucket0/cosmos3/pretrained/huggingface/Qwen/Qwen3-VL-8B-Instruct/",
+          "checkpoint_format": null,
+          "credentials_path": "",
+          "enable_gcs_patch_in_boto3": true,
+          "enabled": false
+        },
+        "qk_norm": false,
+        "safetensors_path": "",
+        "tie_word_embeddings": false,
+        "tokenizer": {
+          "_target": "create_qwen2_tokenizer_with_download",
+          "config_variant": "hf",
+          "pretrained_model_name": "nvidia/Cosmos3-Nano-Policy-DROID"
+        },
+        "use_system_prompt": false
+      }
+    }
+  },
+  "model_type": "cosmos3_omni",
+  "text_config": {
+    "attention_bias": false,
+    "attention_dropout": 0.0,
+    "bos_token_id": 151643,
+    "dtype": "bfloat16",
+    "eos_token_id": 151645,
+    "head_dim": 128,
+    "hidden_act": "silu",
+    "hidden_size": 4096,
+    "initializer_range": 0.02,
+    "intermediate_size": 12288,
+    "max_position_embeddings": 262144,
+    "model_type": "qwen3_vl_text",
+    "num_attention_heads": 32,
+    "num_hidden_layers": 36,
+    "num_key_value_heads": 8,
+    "rms_norm_eps": 1e-06,
+    "rope_scaling": {
+      "mrope_interleaved": true,
+      "mrope_section": [
+        24,
+        20,
+        20
+      ],
+      "rope_type": "default"
+    },
+    "rope_theta": 5000000,
+    "use_cache": true,
+    "vocab_size": 151936
+  },
+  "tie_word_embeddings": false,
+  "transformers_version": "4.57.0.dev0",
+  "video_token_id": 151656,
+  "vision_config": {
+    "deepstack_visual_indexes": [
+      8,
+      16,
+      24
+    ],
+    "depth": 27,
+    "hidden_act": "gelu_pytorch_tanh",
+    "hidden_size": 1152,
+    "in_channels": 3,
+    "initializer_range": 0.02,
+    "intermediate_size": 4304,
+    "model_type": "qwen3_vl",
+    "num_heads": 16,
+    "num_position_embeddings": 2304,
+    "out_hidden_size": 4096,
+    "patch_size": 16,
+    "spatial_merge_size": 2,
+    "temporal_patch_size": 2
+  },
+  "vision_end_token_id": 151653,
+  "vision_start_token_id": 151652
+}

iter_001500/generation_config.json

@@ -0,0 +1,14 @@
+{
+    "bos_token_id": 151643,
+    "pad_token_id": 151643,
+    "do_sample": true,
+    "eos_token_id": [
+        151645,
+        151643
+    ],
+    "top_k": 20,
+    "top_p": 0.8,
+    "repetition_penalty": 1.0,
+    "temperature": 0.7,
+    "transformers_version": "4.56.0"
+}

iter_001500/model_index.json

@@ -0,0 +1,32 @@
+{
+  "_class_name": "Cosmos3OmniPipeline",
+  "_diffusers_version": "0.39.0",
+  "safety_checker": [
+    null,
+    null
+  ],
+  "scheduler": [
+    "diffusers",
+    "UniPCMultistepScheduler"
+  ],
+  "sound_tokenizer": [
+    null,
+    null
+  ],
+  "text_tokenizer": [
+    "transformers",
+    "Qwen2TokenizerFast"
+  ],
+  "transformer": [
+    "diffusers",
+    "Cosmos3OmniTransformer"
+  ],
+  "vae": [
+    "diffusers",
+    "AutoencoderKLWan"
+  ],
+  "vision_encoder": [
+    "transformers",
+    "Qwen3VLVisionModel"
+  ]
+}

iter_001500/modular_model_index.json

@@ -0,0 +1,58 @@
+{
+  "_blocks_class_name": "Cosmos3OmniBlocks",
+  "_class_name": "Cosmos3OmniModularPipeline",
+  "_diffusers_version": "0.39.0",
+  "is_distilled": false,
+  "scheduler": [
+    "diffusers",
+    "UniPCMultistepScheduler",
+    {
+      "pretrained_model_name_or_path": "Cosmos-Data/Cosmos-droid",
+      "subfolder": "iter_001500/scheduler",
+      "type_hint": [
+        "diffusers",
+        "UniPCMultistepScheduler"
+      ],
+      "variant": null
+    }
+  ],
+  "text_tokenizer": [
+    "transformers",
+    "Qwen2TokenizerFast",
+    {
+      "pretrained_model_name_or_path": "Cosmos-Data/Cosmos-droid",
+      "subfolder": "iter_001500/text_tokenizer",
+      "type_hint": [
+        "transformers",
+        "Qwen2TokenizerFast"
+      ],
+      "variant": null
+    }
+  ],
+  "transformer": [
+    "diffusers",
+    "Cosmos3OmniTransformer",
+    {
+      "pretrained_model_name_or_path": "Cosmos-Data/Cosmos-droid",
+      "subfolder": "iter_001500/transformer",
+      "type_hint": [
+        "diffusers",
+        "Cosmos3OmniTransformer"
+      ],
+      "variant": null
+    }
+  ],
+  "vae": [
+    "diffusers",
+    "AutoencoderKLWan",
+    {
+      "pretrained_model_name_or_path": "Cosmos-Data/Cosmos-droid",
+      "subfolder": "iter_001500/vae",
+      "type_hint": [
+        "diffusers",
+        "AutoencoderKLWan"
+      ],
+      "variant": null
+    }
+  ]
+}

iter_001500/preprocessor_config.json

@@ -0,0 +1,21 @@
+{
+    "size": {
+        "longest_edge": 16777216,
+        "shortest_edge": 65536
+    },
+    "patch_size": 16,
+    "temporal_patch_size": 2,
+    "merge_size": 2,
+    "image_mean": [
+        0.5,
+        0.5,
+        0.5
+    ],
+    "image_std": [
+        0.5,
+        0.5,
+        0.5
+    ],
+    "processor_class": "Qwen3VLProcessor",
+    "image_processor_type": "Qwen2VLImageProcessorFast"
+}

iter_001500/scheduler/scheduler_config.json

@@ -0,0 +1,33 @@
+{
+  "_class_name": "UniPCMultistepScheduler",
+  "_diffusers_version": "0.39.0",
+  "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": "flow_prediction",
+  "rescale_betas_zero_snr": false,
+  "sample_max_value": 1.0,
+  "shift_terminal": null,
+  "sigma_max": 200.0,
+  "sigma_min": 0.147,
+  "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": true,
+  "use_karras_sigmas": true
+}

iter_001500/text_tokenizer/added_tokens.json

@@ -0,0 +1,28 @@
+{
+  "</think>": 151668,
+  "</tool_call>": 151658,
+  "</tool_response>": 151666,
+  "<think>": 151667,
+  "<tool_call>": 151657,
+  "<tool_response>": 151665,
+  "<|box_end|>": 151649,
+  "<|box_start|>": 151648,
+  "<|endoftext|>": 151643,
+  "<|file_sep|>": 151664,
+  "<|fim_middle|>": 151660,
+  "<|fim_pad|>": 151662,
+  "<|fim_prefix|>": 151659,
+  "<|fim_suffix|>": 151661,
+  "<|im_end|>": 151645,
+  "<|im_start|>": 151644,
+  "<|image_pad|>": 151655,
+  "<|object_ref_end|>": 151647,
+  "<|object_ref_start|>": 151646,
+  "<|quad_end|>": 151651,
+  "<|quad_start|>": 151650,
+  "<|repo_name|>": 151663,
+  "<|video_pad|>": 151656,
+  "<|vision_end|>": 151653,
+  "<|vision_pad|>": 151654,
+  "<|vision_start|>": 151652
+}

iter_001500/text_tokenizer/chat_template.jinja

@@ -0,0 +1,120 @@
+{%- if tools %}
+    {{- '<|im_start|>system\n' }}
+    {%- if messages[0].role == 'system' %}
+        {%- if messages[0].content is string %}
+            {{- messages[0].content }}
+        {%- else %}
+            {%- for content in messages[0].content %}
+                {%- if 'text' in content %}
+                    {{- content.text }}
+                {%- endif %}
+            {%- endfor %}
+        {%- endif %}
+        {{- '\n\n' }}
+    {%- endif %}
+    {{- "# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>" }}
+    {%- for tool in tools %}
+        {{- "\n" }}
+        {{- tool | tojson }}
+    {%- endfor %}
+    {{- "\n</tools>\n\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\n<tool_call>\n{\"name\": <function-name>, \"arguments\": <args-json-object>}\n</tool_call><|im_end|>\n" }}
+{%- else %}
+    {%- if messages[0].role == 'system' %}
+        {{- '<|im_start|>system\n' }}
+        {%- if messages[0].content is string %}
+            {{- messages[0].content }}
+        {%- else %}
+            {%- for content in messages[0].content %}
+                {%- if 'text' in content %}
+                    {{- content.text }}
+                {%- endif %}
+            {%- endfor %}
+        {%- endif %}
+        {{- '<|im_end|>\n' }}
+    {%- endif %}
+{%- endif %}
+{%- set image_count = namespace(value=0) %}
+{%- set video_count = namespace(value=0) %}
+{%- for message in messages %}
+    {%- if message.role == "user" %}
+        {{- '<|im_start|>' + message.role + '\n' }}
+        {%- if message.content is string %}
+            {{- message.content }}
+        {%- 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 %}
+        {%- endif %}
+        {{- '<|im_end|>\n' }}
+    {%- elif message.role == "assistant" %}
+        {{- '<|im_start|>' + message.role + '\n' }}
+        {%- if message.content is string %}
+            {{- message.content }}
+        {%- else %}
+            {%- for content_item in message.content %}
+                {%- if 'text' in content_item %}
+                    {{- content_item.text }}
+                {%- endif %}
+            {%- endfor %}
+        {%- endif %}
+        {%- if message.tool_calls %}
+            {%- for tool_call in message.tool_calls %}
+                {%- if (loop.first and message.content) or (not loop.first) %}
+                    {{- '\n' }}
+                {%- endif %}
+                {%- if tool_call.function %}
+                    {%- set tool_call = tool_call.function %}
+                {%- endif %}
+                {{- '<tool_call>\n{"name": "' }}
+                {{- tool_call.name }}
+                {{- '", "arguments": ' }}
+                {%- if tool_call.arguments is string %}
+                    {{- tool_call.arguments }}
+                {%- else %}
+                    {{- tool_call.arguments | tojson }}
+                {%- endif %}
+                {{- '}\n</tool_call>' }}
+            {%- endfor %}
+        {%- endif %}
+        {{- '<|im_end|>\n' }}
+    {%- elif message.role == "tool" %}
+        {%- if loop.first or (messages[loop.index0 - 1].role != "tool") %}
+            {{- '<|im_start|>user' }}
+        {%- endif %}
+        {{- '\n<tool_response>\n' }}
+        {%- if message.content is string %}
+            {{- message.content }}
+        {%- 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 %}
+        {%- endif %}
+        {{- '\n</tool_response>' }}
+        {%- if loop.last or (messages[loop.index0 + 1].role != "tool") %}
+            {{- '<|im_end|>\n' }}
+        {%- endif %}
+    {%- endif %}
+{%- endfor %}
+{%- if add_generation_prompt %}
+    {{- '<|im_start|>assistant\n' }}
+{%- endif %}

iter_001500/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
+  }
+}

iter_001500/text_tokenizer/tokenizer_config.json

@@ -0,0 +1,239 @@
+{
+  "add_bos_token": false,
+  "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
+    },
+    "151657": {
+      "content": "<tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151658": {
+      "content": "</tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151659": {
+      "content": "<|fim_prefix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151660": {
+      "content": "<|fim_middle|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151661": {
+      "content": "<|fim_suffix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151662": {
+      "content": "<|fim_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151663": {
+      "content": "<|repo_name|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151664": {
+      "content": "<|file_sep|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151665": {
+      "content": "<tool_response>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151666": {
+      "content": "</tool_response>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151667": {
+      "content": "<think>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151668": {
+      "content": "</think>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    }
+  },
+  "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": 262144,
+  "pad_token": "<|endoftext|>",
+  "split_special_tokens": false,
+  "tokenizer_class": "Qwen2Tokenizer",
+  "unk_token": null
+}

iter_001500/tokenizer_config.json

@@ -0,0 +1,239 @@
+{
+  "add_bos_token": false,
+  "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
+    },
+    "151657": {
+      "content": "<tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151658": {
+      "content": "</tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151659": {
+      "content": "<|fim_prefix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151660": {
+      "content": "<|fim_middle|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151661": {
+      "content": "<|fim_suffix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151662": {
+      "content": "<|fim_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151663": {
+      "content": "<|repo_name|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151664": {
+      "content": "<|file_sep|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151665": {
+      "content": "<tool_response>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151666": {
+      "content": "</tool_response>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151667": {
+      "content": "<think>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151668": {
+      "content": "</think>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    }
+  },
+  "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,
+  "chat_template": "{%- if tools %}\n    {{- '<|im_start|>system\\n' }}\n    {%- if messages[0].role == 'system' %}\n        {%- if messages[0].content is string %}\n            {{- messages[0].content }}\n        {%- else %}\n            {%- for content in messages[0].content %}\n                {%- if 'text' in content %}\n                    {{- content.text }}\n                {%- endif %}\n            {%- endfor %}\n        {%- endif %}\n        {{- '\\n\\n' }}\n    {%- endif %}\n    {{- \"# Tools\\n\\nYou may call one or more functions to assist with the user query.\\n\\nYou are provided with function signatures within <tools></tools> XML tags:\\n<tools>\" }}\n    {%- for tool in tools %}\n        {{- \"\\n\" }}\n        {{- tool | tojson }}\n    {%- endfor %}\n    {{- \"\\n</tools>\\n\\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\\n<tool_call>\\n{\\\"name\\\": <function-name>, \\\"arguments\\\": <args-json-object>}\\n</tool_call><|im_end|>\\n\" }}\n{%- else %}\n    {%- if messages[0].role == 'system' %}\n        {{- '<|im_start|>system\\n' }}\n        {%- if messages[0].content is string %}\n            {{- messages[0].content }}\n        {%- else %}\n            {%- for content in messages[0].content %}\n                {%- if 'text' in content %}\n                    {{- content.text }}\n                {%- endif %}\n            {%- endfor %}\n        {%- endif %}\n        {{- '<|im_end|>\\n' }}\n    {%- endif %}\n{%- endif %}\n{%- set image_count = namespace(value=0) %}\n{%- set video_count = namespace(value=0) %}\n{%- for message in messages %}\n    {%- if message.role == \"user\" %}\n        {{- '<|im_start|>' + message.role + '\\n' }}\n        {%- if message.content is string %}\n            {{- message.content }}\n        {%- else %}\n            {%- for content in message.content %}\n                {%- if content.type == 'image' or 'image' in content or 'image_url' in content %}\n                    {%- set image_count.value = image_count.value + 1 %}\n                    {%- if add_vision_id %}Picture {{ image_count.value }}: {% endif -%}\n                    <|vision_start|><|image_pad|><|vision_end|>\n                {%- elif content.type == 'video' or 'video' in content %}\n                    {%- set video_count.value = video_count.value + 1 %}\n                    {%- if add_vision_id %}Video {{ video_count.value }}: {% endif -%}\n                    <|vision_start|><|video_pad|><|vision_end|>\n                {%- elif 'text' in content %}\n                    {{- content.text }}\n                {%- endif %}\n            {%- endfor %}\n        {%- endif %}\n        {{- '<|im_end|>\\n' }}\n    {%- elif message.role == \"assistant\" %}\n        {{- '<|im_start|>' + message.role + '\\n' }}\n        {%- if message.content is string %}\n            {{- message.content }}\n        {%- else %}\n            {%- for content_item in message.content %}\n                {%- if 'text' in content_item %}\n                    {{- content_item.text }}\n                {%- endif %}\n            {%- endfor %}\n        {%- endif %}\n        {%- if message.tool_calls %}\n            {%- for tool_call in message.tool_calls %}\n                {%- if (loop.first and message.content) or (not loop.first) %}\n                    {{- '\\n' }}\n                {%- endif %}\n                {%- if tool_call.function %}\n                    {%- set tool_call = tool_call.function %}\n                {%- endif %}\n                {{- '<tool_call>\\n{\"name\": \"' }}\n                {{- tool_call.name }}\n                {{- '\", \"arguments\": ' }}\n                {%- if tool_call.arguments is string %}\n                    {{- tool_call.arguments }}\n                {%- else %}\n                    {{- tool_call.arguments | tojson }}\n                {%- endif %}\n                {{- '}\\n</tool_call>' }}\n            {%- endfor %}\n        {%- endif %}\n        {{- '<|im_end|>\\n' }}\n    {%- elif message.role == \"tool\" %}\n        {%- if loop.first or (messages[loop.index0 - 1].role != \"tool\") %}\n            {{- '<|im_start|>user' }}\n        {%- endif %}\n        {{- '\\n<tool_response>\\n' }}\n        {%- if message.content is string %}\n            {{- message.content }}\n        {%- else %}\n            {%- for content in message.content %}\n                {%- if content.type == 'image' or 'image' in content or 'image_url' in content %}\n                    {%- set image_count.value = image_count.value + 1 %}\n                    {%- if add_vision_id %}Picture {{ image_count.value }}: {% endif -%}\n                    <|vision_start|><|image_pad|><|vision_end|>\n                {%- elif content.type == 'video' or 'video' in content %}\n                    {%- set video_count.value = video_count.value + 1 %}\n                    {%- if add_vision_id %}Video {{ video_count.value }}: {% endif -%}\n                    <|vision_start|><|video_pad|><|vision_end|>\n                {%- elif 'text' in content %}\n                    {{- content.text }}\n                {%- endif %}\n            {%- endfor %}\n        {%- endif %}\n        {{- '\\n</tool_response>' }}\n        {%- if loop.last or (messages[loop.index0 + 1].role != \"tool\") %}\n            {{- '<|im_end|>\\n' }}\n        {%- endif %}\n    {%- endif %}\n{%- endfor %}\n{%- if add_generation_prompt %}\n    {{- '<|im_start|>assistant\\n' }}\n{%- endif %}\n",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "errors": "replace",
+  "model_max_length": 262144,
+  "pad_token": "<|endoftext|>",
+  "split_special_tokens": false,
+  "tokenizer_class": "Qwen2Tokenizer",
+  "unk_token": null
+}

iter_001500/transformer/config.json

@@ -0,0 +1,43 @@
+{
+  "_class_name": "Cosmos3OmniTransformer",
+  "_diffusers_version": "0.39.0",
+  "action_dim": 64,
+  "action_gen": true,
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "base_fps": 24,
+  "enable_fps_modulation": true,
+  "head_dim": 128,
+  "hidden_size": 4096,
+  "intermediate_size": 12288,
+  "latent_channel": 48,
+  "latent_patch_size": 2,
+  "num_attention_heads": 32,
+  "num_embodiment_domains": 32,
+  "num_hidden_layers": 36,
+  "num_key_value_heads": 8,
+  "patch_latent_dim": 192,
+  "rms_norm_eps": 1e-06,
+  "rope_axes_dim": [
+    24,
+    20,
+    20
+  ],
+  "rope_scaling": {
+    "mrope_interleaved": true,
+    "mrope_section": [
+      24,
+      20,
+      20
+    ],
+    "rope_type": "default"
+  },
+  "rope_theta": 5000000,
+  "sound_dim": null,
+  "sound_gen": false,
+  "sound_latent_fps": 25,
+  "timestep_scale": 0.001,
+  "unified_3d_mrope_reset_spatial_ids": true,
+  "unified_3d_mrope_temporal_modality_margin": 15000,
+  "vocab_size": 151936
+}

iter_001500/transformer/diffusion_pytorch_model.safetensors.index.json

@@ -0,0 +1,816 @@
+{
+  "metadata": {
+    "total_size": 30346273152
+  },
+  "weight_map": {
+    "action_modality_embed": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "action_proj_in.bias.weight": "diffusion_pytorch_model-00007-of-00007.safetensors",
+    "action_proj_in.fc.weight": "diffusion_pytorch_model-00007-of-00007.safetensors",
+    "action_proj_out.bias.weight": "diffusion_pytorch_model-00007-of-00007.safetensors",
+    "action_proj_out.fc.weight": "diffusion_pytorch_model-00007-of-00007.safetensors",
+    "embed_tokens.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.0.input_layernorm.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.0.input_layernorm_moe_gen.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.0.mlp.down_proj.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.0.mlp.gate_proj.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.0.mlp.up_proj.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.0.mlp_moe_gen.down_proj.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.0.mlp_moe_gen.gate_proj.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.0.mlp_moe_gen.up_proj.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.0.post_attention_layernorm.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.0.post_attention_layernorm_moe_gen.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.0.self_attn.add_k_proj.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.0.self_attn.add_q_proj.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.0.self_attn.add_v_proj.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.0.self_attn.norm_added_k.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.0.self_attn.norm_added_q.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.0.self_attn.norm_k.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.0.self_attn.norm_q.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.0.self_attn.to_add_out.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.0.self_attn.to_k.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.0.self_attn.to_out.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.0.self_attn.to_q.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.0.self_attn.to_v.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.1.input_layernorm.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.1.input_layernorm_moe_gen.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.1.mlp.down_proj.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.1.mlp.gate_proj.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.1.mlp.up_proj.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.1.mlp_moe_gen.down_proj.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.1.mlp_moe_gen.gate_proj.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.1.mlp_moe_gen.up_proj.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.1.post_attention_layernorm.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.1.post_attention_layernorm_moe_gen.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.1.self_attn.add_k_proj.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.1.self_attn.add_q_proj.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.1.self_attn.add_v_proj.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.1.self_attn.norm_added_k.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.1.self_attn.norm_added_q.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.1.self_attn.norm_k.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.1.self_attn.norm_q.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.1.self_attn.to_add_out.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.1.self_attn.to_k.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.1.self_attn.to_out.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.1.self_attn.to_q.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.1.self_attn.to_v.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.10.input_layernorm.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.10.input_layernorm_moe_gen.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.10.mlp.down_proj.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.10.mlp.gate_proj.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.10.mlp.up_proj.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.10.mlp_moe_gen.down_proj.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.10.mlp_moe_gen.gate_proj.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.10.mlp_moe_gen.up_proj.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.10.post_attention_layernorm.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.10.post_attention_layernorm_moe_gen.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.10.self_attn.add_k_proj.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.10.self_attn.add_q_proj.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.10.self_attn.add_v_proj.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.10.self_attn.norm_added_k.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.10.self_attn.norm_added_q.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.10.self_attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.10.self_attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.10.self_attn.to_add_out.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.10.self_attn.to_k.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.10.self_attn.to_out.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.10.self_attn.to_q.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.10.self_attn.to_v.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.11.input_layernorm.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.11.input_layernorm_moe_gen.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.11.mlp.down_proj.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.11.mlp.gate_proj.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.11.mlp.up_proj.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.11.mlp_moe_gen.down_proj.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.11.mlp_moe_gen.gate_proj.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.11.mlp_moe_gen.up_proj.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.11.post_attention_layernorm.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.11.post_attention_layernorm_moe_gen.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.11.self_attn.add_k_proj.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.11.self_attn.add_q_proj.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.11.self_attn.add_v_proj.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.11.self_attn.norm_added_k.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.11.self_attn.norm_added_q.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.11.self_attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.11.self_attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.11.self_attn.to_add_out.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.11.self_attn.to_k.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.11.self_attn.to_out.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.11.self_attn.to_q.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.11.self_attn.to_v.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.12.input_layernorm.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.12.input_layernorm_moe_gen.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.12.mlp.down_proj.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.12.mlp.gate_proj.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.12.mlp.up_proj.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.12.mlp_moe_gen.down_proj.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.12.mlp_moe_gen.gate_proj.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.12.mlp_moe_gen.up_proj.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.12.post_attention_layernorm.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.12.post_attention_layernorm_moe_gen.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.12.self_attn.add_k_proj.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.12.self_attn.add_q_proj.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.12.self_attn.add_v_proj.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.12.self_attn.norm_added_k.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.12.self_attn.norm_added_q.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.12.self_attn.norm_k.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.12.self_attn.norm_q.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.12.self_attn.to_add_out.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.12.self_attn.to_k.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.12.self_attn.to_out.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.12.self_attn.to_q.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.12.self_attn.to_v.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.13.input_layernorm.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.13.input_layernorm_moe_gen.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.13.mlp.down_proj.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.13.mlp.gate_proj.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.13.mlp.up_proj.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.13.mlp_moe_gen.down_proj.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.13.mlp_moe_gen.gate_proj.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.13.mlp_moe_gen.up_proj.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.13.post_attention_layernorm.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.13.post_attention_layernorm_moe_gen.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.13.self_attn.add_k_proj.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.13.self_attn.add_q_proj.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.13.self_attn.add_v_proj.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.13.self_attn.norm_added_k.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.13.self_attn.norm_added_q.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.13.self_attn.norm_k.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.13.self_attn.norm_q.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.13.self_attn.to_add_out.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.13.self_attn.to_k.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.13.self_attn.to_out.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.13.self_attn.to_q.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.13.self_attn.to_v.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.14.input_layernorm.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.14.input_layernorm_moe_gen.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.14.mlp.down_proj.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.14.mlp.gate_proj.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.14.mlp.up_proj.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.14.mlp_moe_gen.down_proj.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.14.mlp_moe_gen.gate_proj.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.14.mlp_moe_gen.up_proj.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.14.post_attention_layernorm.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.14.post_attention_layernorm_moe_gen.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.14.self_attn.add_k_proj.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.14.self_attn.add_q_proj.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.14.self_attn.add_v_proj.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.14.self_attn.norm_added_k.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.14.self_attn.norm_added_q.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.14.self_attn.norm_k.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.14.self_attn.norm_q.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.14.self_attn.to_add_out.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.14.self_attn.to_k.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.14.self_attn.to_out.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.14.self_attn.to_q.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.14.self_attn.to_v.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.15.input_layernorm.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.15.input_layernorm_moe_gen.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.15.mlp.down_proj.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.15.mlp.gate_proj.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.15.mlp.up_proj.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.15.mlp_moe_gen.down_proj.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.15.mlp_moe_gen.gate_proj.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.15.mlp_moe_gen.up_proj.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.15.post_attention_layernorm.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.15.post_attention_layernorm_moe_gen.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.15.self_attn.add_k_proj.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.15.self_attn.add_q_proj.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.15.self_attn.add_v_proj.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.15.self_attn.norm_added_k.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.15.self_attn.norm_added_q.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.15.self_attn.norm_k.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.15.self_attn.norm_q.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.15.self_attn.to_add_out.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.15.self_attn.to_k.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.15.self_attn.to_out.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.15.self_attn.to_q.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.15.self_attn.to_v.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.16.input_layernorm.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.16.input_layernorm_moe_gen.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.16.mlp.down_proj.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.16.mlp.gate_proj.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.16.mlp.up_proj.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.16.mlp_moe_gen.down_proj.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.16.mlp_moe_gen.gate_proj.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.16.mlp_moe_gen.up_proj.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.16.post_attention_layernorm.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.16.post_attention_layernorm_moe_gen.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.16.self_attn.add_k_proj.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.16.self_attn.add_q_proj.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.16.self_attn.add_v_proj.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.16.self_attn.norm_added_k.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.16.self_attn.norm_added_q.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.16.self_attn.norm_k.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.16.self_attn.norm_q.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.16.self_attn.to_add_out.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.16.self_attn.to_k.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.16.self_attn.to_out.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.16.self_attn.to_q.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.16.self_attn.to_v.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.17.input_layernorm.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.17.input_layernorm_moe_gen.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.17.mlp.down_proj.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.17.mlp.gate_proj.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.17.mlp.up_proj.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.17.mlp_moe_gen.down_proj.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.17.mlp_moe_gen.gate_proj.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.17.mlp_moe_gen.up_proj.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.17.post_attention_layernorm.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.17.post_attention_layernorm_moe_gen.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.17.self_attn.add_k_proj.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.17.self_attn.add_q_proj.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.17.self_attn.add_v_proj.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.17.self_attn.norm_added_k.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.17.self_attn.norm_added_q.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.17.self_attn.norm_k.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.17.self_attn.norm_q.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.17.self_attn.to_add_out.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.17.self_attn.to_k.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.17.self_attn.to_out.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.17.self_attn.to_q.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.17.self_attn.to_v.weight": "diffusion_pytorch_model-00003-of-00007.safetensors",
+    "layers.18.input_layernorm.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.18.input_layernorm_moe_gen.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.18.mlp.down_proj.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.18.mlp.gate_proj.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.18.mlp.up_proj.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.18.mlp_moe_gen.down_proj.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.18.mlp_moe_gen.gate_proj.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.18.mlp_moe_gen.up_proj.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.18.post_attention_layernorm.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.18.post_attention_layernorm_moe_gen.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.18.self_attn.add_k_proj.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.18.self_attn.add_q_proj.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.18.self_attn.add_v_proj.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.18.self_attn.norm_added_k.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.18.self_attn.norm_added_q.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.18.self_attn.norm_k.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.18.self_attn.norm_q.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.18.self_attn.to_add_out.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.18.self_attn.to_k.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.18.self_attn.to_out.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.18.self_attn.to_q.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.18.self_attn.to_v.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.19.input_layernorm.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.19.input_layernorm_moe_gen.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.19.mlp.down_proj.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.19.mlp.gate_proj.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.19.mlp.up_proj.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.19.mlp_moe_gen.down_proj.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.19.mlp_moe_gen.gate_proj.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.19.mlp_moe_gen.up_proj.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.19.post_attention_layernorm.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.19.post_attention_layernorm_moe_gen.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.19.self_attn.add_k_proj.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.19.self_attn.add_q_proj.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.19.self_attn.add_v_proj.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.19.self_attn.norm_added_k.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.19.self_attn.norm_added_q.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.19.self_attn.norm_k.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.19.self_attn.norm_q.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.19.self_attn.to_add_out.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.19.self_attn.to_k.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.19.self_attn.to_out.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.19.self_attn.to_q.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.19.self_attn.to_v.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.2.input_layernorm.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.2.input_layernorm_moe_gen.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.2.mlp.down_proj.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.2.mlp.gate_proj.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.2.mlp.up_proj.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.2.mlp_moe_gen.down_proj.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.2.mlp_moe_gen.gate_proj.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.2.mlp_moe_gen.up_proj.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.2.post_attention_layernorm.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.2.post_attention_layernorm_moe_gen.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.2.self_attn.add_k_proj.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.2.self_attn.add_q_proj.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.2.self_attn.add_v_proj.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.2.self_attn.norm_added_k.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.2.self_attn.norm_added_q.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.2.self_attn.norm_k.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.2.self_attn.norm_q.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.2.self_attn.to_add_out.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.2.self_attn.to_k.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.2.self_attn.to_out.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.2.self_attn.to_q.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.2.self_attn.to_v.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.20.input_layernorm.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.20.input_layernorm_moe_gen.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.20.mlp.down_proj.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.20.mlp.gate_proj.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.20.mlp.up_proj.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.20.mlp_moe_gen.down_proj.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.20.mlp_moe_gen.gate_proj.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.20.mlp_moe_gen.up_proj.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.20.post_attention_layernorm.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.20.post_attention_layernorm_moe_gen.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.20.self_attn.add_k_proj.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.20.self_attn.add_q_proj.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.20.self_attn.add_v_proj.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.20.self_attn.norm_added_k.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.20.self_attn.norm_added_q.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.20.self_attn.norm_k.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.20.self_attn.norm_q.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.20.self_attn.to_add_out.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.20.self_attn.to_k.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.20.self_attn.to_out.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.20.self_attn.to_q.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.20.self_attn.to_v.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.21.input_layernorm.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.21.input_layernorm_moe_gen.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.21.mlp.down_proj.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.21.mlp.gate_proj.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.21.mlp.up_proj.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.21.mlp_moe_gen.down_proj.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.21.mlp_moe_gen.gate_proj.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.21.mlp_moe_gen.up_proj.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.21.post_attention_layernorm.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.21.post_attention_layernorm_moe_gen.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.21.self_attn.add_k_proj.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.21.self_attn.add_q_proj.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.21.self_attn.add_v_proj.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.21.self_attn.norm_added_k.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.21.self_attn.norm_added_q.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.21.self_attn.norm_k.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.21.self_attn.norm_q.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.21.self_attn.to_add_out.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.21.self_attn.to_k.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.21.self_attn.to_out.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.21.self_attn.to_q.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.21.self_attn.to_v.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.22.input_layernorm.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.22.input_layernorm_moe_gen.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.22.mlp.down_proj.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.22.mlp.gate_proj.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.22.mlp.up_proj.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.22.mlp_moe_gen.down_proj.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.22.mlp_moe_gen.gate_proj.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.22.mlp_moe_gen.up_proj.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.22.post_attention_layernorm.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.22.post_attention_layernorm_moe_gen.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.22.self_attn.add_k_proj.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.22.self_attn.add_q_proj.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.22.self_attn.add_v_proj.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.22.self_attn.norm_added_k.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.22.self_attn.norm_added_q.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.22.self_attn.norm_k.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.22.self_attn.norm_q.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.22.self_attn.to_add_out.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.22.self_attn.to_k.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.22.self_attn.to_out.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.22.self_attn.to_q.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.22.self_attn.to_v.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.23.input_layernorm.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.23.input_layernorm_moe_gen.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.23.mlp.down_proj.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.23.mlp.gate_proj.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.23.mlp.up_proj.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.23.mlp_moe_gen.down_proj.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.23.mlp_moe_gen.gate_proj.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.23.mlp_moe_gen.up_proj.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.23.post_attention_layernorm.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.23.post_attention_layernorm_moe_gen.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.23.self_attn.add_k_proj.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.23.self_attn.add_q_proj.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.23.self_attn.add_v_proj.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.23.self_attn.norm_added_k.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.23.self_attn.norm_added_q.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.23.self_attn.norm_k.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.23.self_attn.norm_q.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.23.self_attn.to_add_out.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.23.self_attn.to_k.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.23.self_attn.to_out.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.23.self_attn.to_q.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.23.self_attn.to_v.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.24.input_layernorm.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.24.input_layernorm_moe_gen.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.24.mlp.down_proj.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.24.mlp.gate_proj.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.24.mlp.up_proj.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.24.mlp_moe_gen.down_proj.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.24.mlp_moe_gen.gate_proj.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.24.mlp_moe_gen.up_proj.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.24.post_attention_layernorm.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.24.post_attention_layernorm_moe_gen.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.24.self_attn.add_k_proj.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.24.self_attn.add_q_proj.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.24.self_attn.add_v_proj.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.24.self_attn.norm_added_k.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.24.self_attn.norm_added_q.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.24.self_attn.norm_k.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.24.self_attn.norm_q.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.24.self_attn.to_add_out.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.24.self_attn.to_k.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.24.self_attn.to_out.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.24.self_attn.to_q.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.24.self_attn.to_v.weight": "diffusion_pytorch_model-00004-of-00007.safetensors",
+    "layers.25.input_layernorm.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.25.input_layernorm_moe_gen.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.25.mlp.down_proj.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.25.mlp.gate_proj.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.25.mlp.up_proj.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.25.mlp_moe_gen.down_proj.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.25.mlp_moe_gen.gate_proj.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.25.mlp_moe_gen.up_proj.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.25.post_attention_layernorm.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.25.post_attention_layernorm_moe_gen.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.25.self_attn.add_k_proj.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.25.self_attn.add_q_proj.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.25.self_attn.add_v_proj.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.25.self_attn.norm_added_k.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.25.self_attn.norm_added_q.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.25.self_attn.norm_k.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.25.self_attn.norm_q.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.25.self_attn.to_add_out.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.25.self_attn.to_k.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.25.self_attn.to_out.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.25.self_attn.to_q.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.25.self_attn.to_v.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.26.input_layernorm.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.26.input_layernorm_moe_gen.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.26.mlp.down_proj.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.26.mlp.gate_proj.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.26.mlp.up_proj.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.26.mlp_moe_gen.down_proj.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.26.mlp_moe_gen.gate_proj.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.26.mlp_moe_gen.up_proj.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.26.post_attention_layernorm.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.26.post_attention_layernorm_moe_gen.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.26.self_attn.add_k_proj.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.26.self_attn.add_q_proj.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.26.self_attn.add_v_proj.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.26.self_attn.norm_added_k.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.26.self_attn.norm_added_q.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.26.self_attn.norm_k.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.26.self_attn.norm_q.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.26.self_attn.to_add_out.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.26.self_attn.to_k.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.26.self_attn.to_out.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.26.self_attn.to_q.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.26.self_attn.to_v.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.27.input_layernorm.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.27.input_layernorm_moe_gen.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.27.mlp.down_proj.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.27.mlp.gate_proj.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.27.mlp.up_proj.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.27.mlp_moe_gen.down_proj.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.27.mlp_moe_gen.gate_proj.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.27.mlp_moe_gen.up_proj.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.27.post_attention_layernorm.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.27.post_attention_layernorm_moe_gen.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.27.self_attn.add_k_proj.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.27.self_attn.add_q_proj.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.27.self_attn.add_v_proj.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.27.self_attn.norm_added_k.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.27.self_attn.norm_added_q.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.27.self_attn.norm_k.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.27.self_attn.norm_q.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.27.self_attn.to_add_out.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.27.self_attn.to_k.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.27.self_attn.to_out.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.27.self_attn.to_q.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.27.self_attn.to_v.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.28.input_layernorm.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.28.input_layernorm_moe_gen.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.28.mlp.down_proj.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.28.mlp.gate_proj.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.28.mlp.up_proj.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.28.mlp_moe_gen.down_proj.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.28.mlp_moe_gen.gate_proj.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.28.mlp_moe_gen.up_proj.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.28.post_attention_layernorm.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.28.post_attention_layernorm_moe_gen.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.28.self_attn.add_k_proj.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.28.self_attn.add_q_proj.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.28.self_attn.add_v_proj.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.28.self_attn.norm_added_k.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.28.self_attn.norm_added_q.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.28.self_attn.norm_k.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.28.self_attn.norm_q.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.28.self_attn.to_add_out.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.28.self_attn.to_k.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.28.self_attn.to_out.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.28.self_attn.to_q.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.28.self_attn.to_v.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.29.input_layernorm.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.29.input_layernorm_moe_gen.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.29.mlp.down_proj.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.29.mlp.gate_proj.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.29.mlp.up_proj.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.29.mlp_moe_gen.down_proj.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.29.mlp_moe_gen.gate_proj.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.29.mlp_moe_gen.up_proj.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.29.post_attention_layernorm.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.29.post_attention_layernorm_moe_gen.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.29.self_attn.add_k_proj.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.29.self_attn.add_q_proj.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.29.self_attn.add_v_proj.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.29.self_attn.norm_added_k.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.29.self_attn.norm_added_q.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.29.self_attn.norm_k.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.29.self_attn.norm_q.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.29.self_attn.to_add_out.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.29.self_attn.to_k.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.29.self_attn.to_out.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.29.self_attn.to_q.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.29.self_attn.to_v.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.3.input_layernorm.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.3.input_layernorm_moe_gen.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.3.mlp.down_proj.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.3.mlp.gate_proj.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.3.mlp.up_proj.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.3.mlp_moe_gen.down_proj.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.3.mlp_moe_gen.gate_proj.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.3.mlp_moe_gen.up_proj.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.3.post_attention_layernorm.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.3.post_attention_layernorm_moe_gen.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.3.self_attn.add_k_proj.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.3.self_attn.add_q_proj.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.3.self_attn.add_v_proj.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.3.self_attn.norm_added_k.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.3.self_attn.norm_added_q.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.3.self_attn.norm_k.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.3.self_attn.norm_q.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.3.self_attn.to_add_out.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.3.self_attn.to_k.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.3.self_attn.to_out.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.3.self_attn.to_q.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.3.self_attn.to_v.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.30.input_layernorm.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.30.input_layernorm_moe_gen.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.30.mlp.down_proj.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.30.mlp.gate_proj.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.30.mlp.up_proj.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.30.mlp_moe_gen.down_proj.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.30.mlp_moe_gen.gate_proj.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.30.mlp_moe_gen.up_proj.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.30.post_attention_layernorm.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.30.post_attention_layernorm_moe_gen.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.30.self_attn.add_k_proj.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.30.self_attn.add_q_proj.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.30.self_attn.add_v_proj.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.30.self_attn.norm_added_k.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.30.self_attn.norm_added_q.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.30.self_attn.norm_k.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.30.self_attn.norm_q.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.30.self_attn.to_add_out.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.30.self_attn.to_k.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.30.self_attn.to_out.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.30.self_attn.to_q.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.30.self_attn.to_v.weight": "diffusion_pytorch_model-00005-of-00007.safetensors",
+    "layers.31.input_layernorm.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.31.input_layernorm_moe_gen.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.31.mlp.down_proj.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.31.mlp.gate_proj.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.31.mlp.up_proj.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.31.mlp_moe_gen.down_proj.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.31.mlp_moe_gen.gate_proj.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.31.mlp_moe_gen.up_proj.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.31.post_attention_layernorm.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.31.post_attention_layernorm_moe_gen.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.31.self_attn.add_k_proj.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.31.self_attn.add_q_proj.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.31.self_attn.add_v_proj.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.31.self_attn.norm_added_k.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.31.self_attn.norm_added_q.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.31.self_attn.norm_k.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.31.self_attn.norm_q.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.31.self_attn.to_add_out.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.31.self_attn.to_k.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.31.self_attn.to_out.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.31.self_attn.to_q.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.31.self_attn.to_v.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.32.input_layernorm.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.32.input_layernorm_moe_gen.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.32.mlp.down_proj.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.32.mlp.gate_proj.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.32.mlp.up_proj.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.32.mlp_moe_gen.down_proj.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.32.mlp_moe_gen.gate_proj.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.32.mlp_moe_gen.up_proj.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.32.post_attention_layernorm.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.32.post_attention_layernorm_moe_gen.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.32.self_attn.add_k_proj.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.32.self_attn.add_q_proj.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.32.self_attn.add_v_proj.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.32.self_attn.norm_added_k.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.32.self_attn.norm_added_q.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.32.self_attn.norm_k.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.32.self_attn.norm_q.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.32.self_attn.to_add_out.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.32.self_attn.to_k.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.32.self_attn.to_out.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.32.self_attn.to_q.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.32.self_attn.to_v.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.33.input_layernorm.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.33.input_layernorm_moe_gen.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.33.mlp.down_proj.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.33.mlp.gate_proj.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.33.mlp.up_proj.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.33.mlp_moe_gen.down_proj.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.33.mlp_moe_gen.gate_proj.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.33.mlp_moe_gen.up_proj.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.33.post_attention_layernorm.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.33.post_attention_layernorm_moe_gen.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.33.self_attn.add_k_proj.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.33.self_attn.add_q_proj.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.33.self_attn.add_v_proj.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.33.self_attn.norm_added_k.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.33.self_attn.norm_added_q.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.33.self_attn.norm_k.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.33.self_attn.norm_q.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.33.self_attn.to_add_out.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.33.self_attn.to_k.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.33.self_attn.to_out.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.33.self_attn.to_q.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.33.self_attn.to_v.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.34.input_layernorm.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.34.input_layernorm_moe_gen.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.34.mlp.down_proj.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.34.mlp.gate_proj.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.34.mlp.up_proj.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.34.mlp_moe_gen.down_proj.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.34.mlp_moe_gen.gate_proj.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.34.mlp_moe_gen.up_proj.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.34.post_attention_layernorm.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.34.post_attention_layernorm_moe_gen.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.34.self_attn.add_k_proj.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.34.self_attn.add_q_proj.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.34.self_attn.add_v_proj.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.34.self_attn.norm_added_k.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.34.self_attn.norm_added_q.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.34.self_attn.norm_k.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.34.self_attn.norm_q.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.34.self_attn.to_add_out.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.34.self_attn.to_k.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.34.self_attn.to_out.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.34.self_attn.to_q.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.34.self_attn.to_v.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.35.input_layernorm.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.35.input_layernorm_moe_gen.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.35.mlp.down_proj.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.35.mlp.gate_proj.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.35.mlp.up_proj.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.35.mlp_moe_gen.down_proj.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.35.mlp_moe_gen.gate_proj.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.35.mlp_moe_gen.up_proj.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.35.post_attention_layernorm.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.35.post_attention_layernorm_moe_gen.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.35.self_attn.add_k_proj.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.35.self_attn.add_q_proj.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.35.self_attn.add_v_proj.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.35.self_attn.norm_added_k.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.35.self_attn.norm_added_q.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.35.self_attn.norm_k.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.35.self_attn.norm_q.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.35.self_attn.to_add_out.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.35.self_attn.to_k.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.35.self_attn.to_out.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.35.self_attn.to_q.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.35.self_attn.to_v.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "layers.4.input_layernorm.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.4.input_layernorm_moe_gen.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.4.mlp.down_proj.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.4.mlp.gate_proj.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.4.mlp.up_proj.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.4.mlp_moe_gen.down_proj.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.4.mlp_moe_gen.gate_proj.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.4.mlp_moe_gen.up_proj.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.4.post_attention_layernorm.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.4.post_attention_layernorm_moe_gen.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.4.self_attn.add_k_proj.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.4.self_attn.add_q_proj.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.4.self_attn.add_v_proj.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.4.self_attn.norm_added_k.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.4.self_attn.norm_added_q.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.4.self_attn.norm_k.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.4.self_attn.norm_q.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.4.self_attn.to_add_out.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.4.self_attn.to_k.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.4.self_attn.to_out.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.4.self_attn.to_q.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.4.self_attn.to_v.weight": "diffusion_pytorch_model-00001-of-00007.safetensors",
+    "layers.5.input_layernorm.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.5.input_layernorm_moe_gen.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.5.mlp.down_proj.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.5.mlp.gate_proj.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.5.mlp.up_proj.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.5.mlp_moe_gen.down_proj.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.5.mlp_moe_gen.gate_proj.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.5.mlp_moe_gen.up_proj.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.5.post_attention_layernorm.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.5.post_attention_layernorm_moe_gen.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.5.self_attn.add_k_proj.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.5.self_attn.add_q_proj.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.5.self_attn.add_v_proj.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.5.self_attn.norm_added_k.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.5.self_attn.norm_added_q.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.5.self_attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.5.self_attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.5.self_attn.to_add_out.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.5.self_attn.to_k.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.5.self_attn.to_out.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.5.self_attn.to_q.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.5.self_attn.to_v.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.6.input_layernorm.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.6.input_layernorm_moe_gen.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.6.mlp.down_proj.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.6.mlp.gate_proj.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.6.mlp.up_proj.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.6.mlp_moe_gen.down_proj.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.6.mlp_moe_gen.gate_proj.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.6.mlp_moe_gen.up_proj.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.6.post_attention_layernorm.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.6.post_attention_layernorm_moe_gen.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.6.self_attn.add_k_proj.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.6.self_attn.add_q_proj.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.6.self_attn.add_v_proj.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.6.self_attn.norm_added_k.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.6.self_attn.norm_added_q.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.6.self_attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.6.self_attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.6.self_attn.to_add_out.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.6.self_attn.to_k.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.6.self_attn.to_out.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.6.self_attn.to_q.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.6.self_attn.to_v.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.7.input_layernorm.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.7.input_layernorm_moe_gen.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.7.mlp.down_proj.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.7.mlp.gate_proj.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.7.mlp.up_proj.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.7.mlp_moe_gen.down_proj.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.7.mlp_moe_gen.gate_proj.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.7.mlp_moe_gen.up_proj.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.7.post_attention_layernorm.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.7.post_attention_layernorm_moe_gen.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.7.self_attn.add_k_proj.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.7.self_attn.add_q_proj.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.7.self_attn.add_v_proj.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.7.self_attn.norm_added_k.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.7.self_attn.norm_added_q.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.7.self_attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.7.self_attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.7.self_attn.to_add_out.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.7.self_attn.to_k.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.7.self_attn.to_out.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.7.self_attn.to_q.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.7.self_attn.to_v.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.8.input_layernorm.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.8.input_layernorm_moe_gen.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.8.mlp.down_proj.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.8.mlp.gate_proj.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.8.mlp.up_proj.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.8.mlp_moe_gen.down_proj.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.8.mlp_moe_gen.gate_proj.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.8.mlp_moe_gen.up_proj.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.8.post_attention_layernorm.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.8.post_attention_layernorm_moe_gen.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.8.self_attn.add_k_proj.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.8.self_attn.add_q_proj.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.8.self_attn.add_v_proj.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.8.self_attn.norm_added_k.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.8.self_attn.norm_added_q.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.8.self_attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.8.self_attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.8.self_attn.to_add_out.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.8.self_attn.to_k.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.8.self_attn.to_out.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.8.self_attn.to_q.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.8.self_attn.to_v.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.9.input_layernorm.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.9.input_layernorm_moe_gen.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.9.mlp.down_proj.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.9.mlp.gate_proj.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.9.mlp.up_proj.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.9.mlp_moe_gen.down_proj.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.9.mlp_moe_gen.gate_proj.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.9.mlp_moe_gen.up_proj.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.9.post_attention_layernorm.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.9.post_attention_layernorm_moe_gen.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.9.self_attn.add_k_proj.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.9.self_attn.add_q_proj.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.9.self_attn.add_v_proj.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.9.self_attn.norm_added_k.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.9.self_attn.norm_added_q.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.9.self_attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.9.self_attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.9.self_attn.to_add_out.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.9.self_attn.to_k.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.9.self_attn.to_out.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.9.self_attn.to_q.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "layers.9.self_attn.to_v.weight": "diffusion_pytorch_model-00002-of-00007.safetensors",
+    "lm_head.weight": "diffusion_pytorch_model-00007-of-00007.safetensors",
+    "norm.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "norm_moe_gen.weight": "diffusion_pytorch_model-00006-of-00007.safetensors",
+    "proj_in.bias": "diffusion_pytorch_model-00007-of-00007.safetensors",
+    "proj_in.weight": "diffusion_pytorch_model-00007-of-00007.safetensors",
+    "proj_out.bias": "diffusion_pytorch_model-00007-of-00007.safetensors",
+    "proj_out.weight": "diffusion_pytorch_model-00007-of-00007.safetensors",
+    "time_embedder.linear_1.bias": "diffusion_pytorch_model-00007-of-00007.safetensors",
+    "time_embedder.linear_1.weight": "diffusion_pytorch_model-00007-of-00007.safetensors",
+    "time_embedder.linear_2.bias": "diffusion_pytorch_model-00007-of-00007.safetensors",
+    "time_embedder.linear_2.weight": "diffusion_pytorch_model-00007-of-00007.safetensors"
+  }
+}

iter_001500/vae/config.json

@@ -0,0 +1,129 @@
+{
+  "_class_name": "AutoencoderKLWan",
+  "_diffusers_version": "0.39.0",
+  "_name_or_path": "Wan-AI/Wan2.2-TI2V-5B-Diffusers",
+  "attn_scales": [],
+  "base_dim": 160,
+  "clip_output": false,
+  "decoder_base_dim": 256,
+  "dim_mult": [
+    1,
+    2,
+    4,
+    4
+  ],
+  "dropout": 0.0,
+  "in_channels": 12,
+  "is_residual": true,
+  "latents_mean": [
+    -0.2289,
+    -0.0052,
+    -0.1323,
+    -0.2339,
+    -0.2799,
+    0.0174,
+    0.1838,
+    0.1557,
+    -0.1382,
+    0.0542,
+    0.2813,
+    0.0891,
+    0.157,
+    -0.0098,
+    0.0375,
+    -0.1825,
+    -0.2246,
+    -0.1207,
+    -0.0698,
+    0.5109,
+    0.2665,
+    -0.2108,
+    -0.2158,
+    0.2502,
+    -0.2055,
+    -0.0322,
+    0.1109,
+    0.1567,
+    -0.0729,
+    0.0899,
+    -0.2799,
+    -0.123,
+    -0.0313,
+    -0.1649,
+    0.0117,
+    0.0723,
+    -0.2839,
+    -0.2083,
+    -0.052,
+    0.3748,
+    0.0152,
+    0.1957,
+    0.1433,
+    -0.2944,
+    0.3573,
+    -0.0548,
+    -0.1681,
+    -0.0667
+  ],
+  "latents_std": [
+    0.4765,
+    1.0364,
+    0.4514,
+    1.1677,
+    0.5313,
+    0.499,
+    0.4818,
+    0.5013,
+    0.8158,
+    1.0344,
+    0.5894,
+    1.0901,
+    0.6885,
+    0.6165,
+    0.8454,
+    0.4978,
+    0.5759,
+    0.3523,
+    0.7135,
+    0.6804,
+    0.5833,
+    1.4146,
+    0.8986,
+    0.5659,
+    0.7069,
+    0.5338,
+    0.4889,
+    0.4917,
+    0.4069,
+    0.4999,
+    0.6866,
+    0.4093,
+    0.5709,
+    0.6065,
+    0.6415,
+    0.4944,
+    0.5726,
+    1.2042,
+    0.5458,
+    1.6887,
+    0.3971,
+    1.06,
+    0.3943,
+    0.5537,
+    0.5444,
+    0.4089,
+    0.7468,
+    0.7744
+  ],
+  "num_res_blocks": 2,
+  "out_channels": 12,
+  "patch_size": 2,
+  "scale_factor_spatial": 16,
+  "scale_factor_temporal": 4,
+  "temperal_downsample": [
+    false,
+    true,
+    true
+  ],
+  "z_dim": 48
+}

iter_001500/video_preprocessor_config.json

@@ -0,0 +1,21 @@
+{
+    "size": {
+        "longest_edge": 25165824,
+        "shortest_edge": 4096
+    },
+    "patch_size": 16,
+    "temporal_patch_size": 2,
+    "merge_size": 2,
+    "image_mean": [
+        0.5,
+        0.5,
+        0.5
+    ],
+    "image_std": [
+        0.5,
+        0.5,
+        0.5
+    ],
+    "processor_class": "Qwen3VLProcessor",
+    "video_processor_type": "Qwen3VLVideoProcessor"
+}

iter_001500/vision_encoder/config.json

@@ -0,0 +1,25 @@
+{
+  "architectures": [
+    "Qwen3VLVisionModel"
+  ],
+  "deepstack_visual_indexes": [
+    8,
+    16,
+    24
+  ],
+  "depth": 27,
+  "dtype": "bfloat16",
+  "hidden_act": "gelu_pytorch_tanh",
+  "hidden_size": 1152,
+  "in_channels": 3,
+  "initializer_range": 0.02,
+  "intermediate_size": 4304,
+  "model_type": "qwen3_vl",
+  "num_heads": 16,
+  "num_position_embeddings": 2304,
+  "out_hidden_size": 4096,
+  "patch_size": 16,
+  "spatial_merge_size": 2,
+  "temporal_patch_size": 2,
+  "transformers_version": "4.57.6"
+}