"""
# Cosmos 2 Predict

Download checkpoint
```bash
hf download nvidia/Cosmos-Predict2-2B-Text2Image
```

convert checkpoint
```bash
transformer_ckpt_path=~/.cache/huggingface/hub/models--nvidia--Cosmos-Predict2-2B-Text2Image/snapshots/acdb5fde992a73ef0355f287977d002cbfd127e0/model.pt

python scripts/convert_cosmos_to_diffusers.py \
    --transformer_ckpt_path $transformer_ckpt_path \
    --transformer_type Cosmos-2.0-Diffusion-2B-Text2Image \
    --text_encoder_path google-t5/t5-11b \
    --tokenizer_path google-t5/t5-11b \
    --vae_type wan2.1 \
    --output_path converted/cosmos-p2-t2i-2b \
    --save_pipeline
```

# Cosmos 2.5 Predict

Download checkpoint
```bash
hf download nvidia/Cosmos-Predict2.5-2B
```

Convert checkpoint
```bash
# pre-trained
transformer_ckpt_path=~/.cache/huggingface/hub/models--nvidia--Cosmos-Predict2.5-2B/snapshots/865baf084d4c9e850eac59a021277d5a9b9e8b63/base/pre-trained/d20b7120-df3e-4911-919d-db6e08bad31c_ema_bf16.pt

python scripts/convert_cosmos_to_diffusers.py \
    --transformer_type Cosmos-2.5-Predict-Base-2B \
    --transformer_ckpt_path $transformer_ckpt_path \
    --vae_type wan2.1 \
    --output_path converted/2b/d20b7120-df3e-4911-919d-db6e08bad31c \
    --save_pipeline

# post-trained
transformer_ckpt_path=~/.cache/huggingface/hub/models--nvidia--Cosmos-Predict2.5-2B/snapshots/865baf084d4c9e850eac59a021277d5a9b9e8b63/base/post-trained/81edfebe-bd6a-4039-8c1d-737df1a790bf_ema_bf16.pt

python scripts/convert_cosmos_to_diffusers.py \
    --transformer_type Cosmos-2.5-Predict-Base-2B \
    --transformer_ckpt_path $transformer_ckpt_path \
    --vae_type wan2.1 \
    --output_path converted/2b/81edfebe-bd6a-4039-8c1d-737df1a790bf \
    --save_pipeline
```

## 14B

```bash
hf download nvidia/Cosmos-Predict2.5-14B
```

```bash
# pre-trained
transformer_ckpt_path=~/.cache/huggingface/hub/models--nvidia--Cosmos-Predict2.5-14B/snapshots/71ebf3e8af30ecfe440bf0481115975fcc052b46/base/pre-trained/54937b8c-29de-4f04-862c-e67b04ec41e8_ema_bf16.pt

python scripts/convert_cosmos_to_diffusers.py \
    --transformer_type Cosmos-2.5-Predict-Base-14B \
    --transformer_ckpt_path $transformer_ckpt_path \
    --vae_type wan2.1 \
    --output_path converted/14b/54937b8c-29de-4f04-862c-e67b04ec41e8/ \
    --save_pipeline

# post-trained
transformer_ckpt_path=~/.cache/huggingface/hub/models--nvidia--Cosmos-Predict2.5-14B/snapshots/71ebf3e8af30ecfe440bf0481115975fcc052b46/base/post-trained/e21d2a49-4747-44c8-ba44-9f6f9243715f_ema_bf16.pt

python scripts/convert_cosmos_to_diffusers.py \
    --transformer_type Cosmos-2.5-Predict-Base-14B \
    --transformer_ckpt_path $transformer_ckpt_path \
    --vae_type wan2.1 \
    --output_path converted/14b/e21d2a49-4747-44c8-ba44-9f6f9243715f/ \
    --save_pipeline
```

# Cosmos 2.5 Transfer

Download checkpoint
```bash
hf download nvidia/Cosmos-Transfer2.5-2B
```

Convert checkpoint
```bash
# depth
transformer_ckpt_path=~/.cache/huggingface/hub/models--nvidia--Cosmos-Transfer2.5-2B/snapshots/eb5325b77d358944da58a690157dd2b8071bbf85/general/depth/626e6618-bfcd-4d9a-a077-1409e2ce353f_ema_bf16.pt

python scripts/convert_cosmos_to_diffusers.py \
    --transformer_type Cosmos-2.5-Transfer-General-2B \
    --transformer_ckpt_path $transformer_ckpt_path \
    --vae_type wan2.1 \
    --output_path converted/transfer/2b/general/depth/pipeline \
    --save_pipeline

python scripts/convert_cosmos_to_diffusers.py \
    --transformer_type Cosmos-2.5-Transfer-General-2B \
    --transformer_ckpt_path $transformer_ckpt_path \
    --vae_type wan2.1 \
    --output_path converted/transfer/2b/general/depth/models

# edge
transformer_ckpt_path=~/.cache/huggingface/hub/models--nvidia--Cosmos-Transfer2.5-2B/snapshots/eb5325b77d358944da58a690157dd2b8071bbf85/general/edge/61f5694b-0ad5-4ecd-8ad7-c8545627d125_ema_bf16.pt

python scripts/convert_cosmos_to_diffusers.py \
    --transformer_type Cosmos-2.5-Transfer-General-2B \
    --transformer_ckpt_path $transformer_ckpt_path \
    --vae_type wan2.1 \
    --output_path converted/transfer/2b/general/edge/pipeline \
    --save_pipeline

python scripts/convert_cosmos_to_diffusers.py \
    --transformer_type Cosmos-2.5-Transfer-General-2B \
    --transformer_ckpt_path $transformer_ckpt_path \
    --vae_type wan2.1 \
    --output_path converted/transfer/2b/general/edge/models

# blur
transformer_ckpt_path=~/.cache/huggingface/hub/models--nvidia--Cosmos-Transfer2.5-2B/snapshots/eb5325b77d358944da58a690157dd2b8071bbf85/general/blur/ba2f44f2-c726-4fe7-949f-597069d9b91c_ema_bf16.pt

python scripts/convert_cosmos_to_diffusers.py \
    --transformer_type Cosmos-2.5-Transfer-General-2B \
    --transformer_ckpt_path $transformer_ckpt_path \
    --vae_type wan2.1 \
    --output_path converted/transfer/2b/general/blur/pipeline \
    --save_pipeline

python scripts/convert_cosmos_to_diffusers.py \
    --transformer_type Cosmos-2.5-Transfer-General-2B \
    --transformer_ckpt_path $transformer_ckpt_path \
    --vae_type wan2.1 \
    --output_path converted/transfer/2b/general/blur/models

# seg
transformer_ckpt_path=~/.cache/huggingface/hub/models--nvidia--Cosmos-Transfer2.5-2B/snapshots/eb5325b77d358944da58a690157dd2b8071bbf85/general/seg/5136ef49-6d8d-42e8-8abf-7dac722a304a_ema_bf16.pt

python scripts/convert_cosmos_to_diffusers.py \
    --transformer_type Cosmos-2.5-Transfer-General-2B \
    --transformer_ckpt_path $transformer_ckpt_path \
    --vae_type wan2.1 \
    --output_path converted/transfer/2b/general/seg/pipeline \
    --save_pipeline

python scripts/convert_cosmos_to_diffusers.py \
    --transformer_type Cosmos-2.5-Transfer-General-2B \
    --transformer_ckpt_path $transformer_ckpt_path \
    --vae_type wan2.1 \
    --output_path converted/transfer/2b/general/seg/models
```
"""

import argparse
import pathlib
import sys
from typing import Any, Dict, Optional

import torch
from accelerate import init_empty_weights
from huggingface_hub import snapshot_download
from transformers import AutoTokenizer, Qwen2_5_VLForConditionalGeneration, T5EncoderModel, T5TokenizerFast

from diffusers import (
    AutoencoderKLCosmos,
    AutoencoderKLWan,
    Cosmos2TextToImagePipeline,
    Cosmos2VideoToWorldPipeline,
    CosmosControlNetModel,
    CosmosTextToWorldPipeline,
    CosmosTransformer3DModel,
    CosmosVideoToWorldPipeline,
    EDMEulerScheduler,
    FlowMatchEulerDiscreteScheduler,
    UniPCMultistepScheduler,
)
from diffusers.pipelines.cosmos.pipeline_cosmos2_5_predict import Cosmos2_5_PredictBasePipeline
from diffusers.pipelines.cosmos.pipeline_cosmos2_5_transfer import Cosmos2_5_TransferPipeline


def remove_keys_(key: str, state_dict: Dict[str, Any]):
    state_dict.pop(key)


def update_state_dict_(state_dict: Dict[str, Any], old_key: str, new_key: str) -> dict[str, Any]:
    state_dict[new_key] = state_dict.pop(old_key)


def rename_transformer_blocks_(key: str, state_dict: Dict[str, Any]):
    block_index = int(key.split(".")[1].removeprefix("block"))
    new_key = key

    old_prefix = f"blocks.block{block_index}"
    new_prefix = f"transformer_blocks.{block_index}"
    new_key = new_prefix + new_key.removeprefix(old_prefix)

    state_dict[new_key] = state_dict.pop(key)


TRANSFORMER_KEYS_RENAME_DICT_COSMOS_1_0 = {
    "t_embedder.1": "time_embed.t_embedder",
    "affline_norm": "time_embed.norm",
    ".blocks.0.block.attn": ".attn1",
    ".blocks.1.block.attn": ".attn2",
    ".blocks.2.block": ".ff",
    ".blocks.0.adaLN_modulation.1": ".norm1.linear_1",
    ".blocks.0.adaLN_modulation.2": ".norm1.linear_2",
    ".blocks.1.adaLN_modulation.1": ".norm2.linear_1",
    ".blocks.1.adaLN_modulation.2": ".norm2.linear_2",
    ".blocks.2.adaLN_modulation.1": ".norm3.linear_1",
    ".blocks.2.adaLN_modulation.2": ".norm3.linear_2",
    "to_q.0": "to_q",
    "to_q.1": "norm_q",
    "to_k.0": "to_k",
    "to_k.1": "norm_k",
    "to_v.0": "to_v",
    "layer1": "net.0.proj",
    "layer2": "net.2",
    "proj.1": "proj",
    "x_embedder": "patch_embed",
    "extra_pos_embedder": "learnable_pos_embed",
    "final_layer.adaLN_modulation.1": "norm_out.linear_1",
    "final_layer.adaLN_modulation.2": "norm_out.linear_2",
    "final_layer.linear": "proj_out",
}

TRANSFORMER_SPECIAL_KEYS_REMAP_COSMOS_1_0 = {
    "blocks.block": rename_transformer_blocks_,
    "logvar.0.freqs": remove_keys_,
    "logvar.0.phases": remove_keys_,
    "logvar.1.weight": remove_keys_,
    "pos_embedder.seq": remove_keys_,
}

TRANSFORMER_KEYS_RENAME_DICT_COSMOS_2_0 = {
    "t_embedder.1": "time_embed.t_embedder",
    "t_embedding_norm": "time_embed.norm",
    "blocks": "transformer_blocks",
    "adaln_modulation_self_attn.1": "norm1.linear_1",
    "adaln_modulation_self_attn.2": "norm1.linear_2",
    "adaln_modulation_cross_attn.1": "norm2.linear_1",
    "adaln_modulation_cross_attn.2": "norm2.linear_2",
    "adaln_modulation_mlp.1": "norm3.linear_1",
    "adaln_modulation_mlp.2": "norm3.linear_2",
    "self_attn": "attn1",
    "cross_attn": "attn2",
    "q_proj": "to_q",
    "k_proj": "to_k",
    "v_proj": "to_v",
    "output_proj": "to_out.0",
    "q_norm": "norm_q",
    "k_norm": "norm_k",
    "mlp.layer1": "ff.net.0.proj",
    "mlp.layer2": "ff.net.2",
    "x_embedder.proj.1": "patch_embed.proj",
    "final_layer.adaln_modulation.1": "norm_out.linear_1",
    "final_layer.adaln_modulation.2": "norm_out.linear_2",
    "final_layer.linear": "proj_out",
}

TRANSFORMER_SPECIAL_KEYS_REMAP_COSMOS_2_0 = {
    "accum_video_sample_counter": remove_keys_,
    "accum_image_sample_counter": remove_keys_,
    "accum_iteration": remove_keys_,
    "accum_train_in_hours": remove_keys_,
    "pos_embedder.seq": remove_keys_,
    "pos_embedder.dim_spatial_range": remove_keys_,
    "pos_embedder.dim_temporal_range": remove_keys_,
    "_extra_state": remove_keys_,
}


TRANSFORMER_CONFIGS = {
    "Cosmos-1.0-Diffusion-7B-Text2World": {
        "in_channels": 16,
        "out_channels": 16,
        "num_attention_heads": 32,
        "attention_head_dim": 128,
        "num_layers": 28,
        "mlp_ratio": 4.0,
        "text_embed_dim": 1024,
        "adaln_lora_dim": 256,
        "max_size": (128, 240, 240),
        "patch_size": (1, 2, 2),
        "rope_scale": (2.0, 1.0, 1.0),
        "concat_padding_mask": True,
        "extra_pos_embed_type": "learnable",
    },
    "Cosmos-1.0-Diffusion-7B-Video2World": {
        "in_channels": 16 + 1,
        "out_channels": 16,
        "num_attention_heads": 32,
        "attention_head_dim": 128,
        "num_layers": 28,
        "mlp_ratio": 4.0,
        "text_embed_dim": 1024,
        "adaln_lora_dim": 256,
        "max_size": (128, 240, 240),
        "patch_size": (1, 2, 2),
        "rope_scale": (2.0, 1.0, 1.0),
        "concat_padding_mask": True,
        "extra_pos_embed_type": "learnable",
    },
    "Cosmos-1.0-Diffusion-14B-Text2World": {
        "in_channels": 16,
        "out_channels": 16,
        "num_attention_heads": 40,
        "attention_head_dim": 128,
        "num_layers": 36,
        "mlp_ratio": 4.0,
        "text_embed_dim": 1024,
        "adaln_lora_dim": 256,
        "max_size": (128, 240, 240),
        "patch_size": (1, 2, 2),
        "rope_scale": (2.0, 2.0, 2.0),
        "concat_padding_mask": True,
        "extra_pos_embed_type": "learnable",
    },
    "Cosmos-1.0-Diffusion-14B-Video2World": {
        "in_channels": 16 + 1,
        "out_channels": 16,
        "num_attention_heads": 40,
        "attention_head_dim": 128,
        "num_layers": 36,
        "mlp_ratio": 4.0,
        "text_embed_dim": 1024,
        "adaln_lora_dim": 256,
        "max_size": (128, 240, 240),
        "patch_size": (1, 2, 2),
        "rope_scale": (2.0, 2.0, 2.0),
        "concat_padding_mask": True,
        "extra_pos_embed_type": "learnable",
    },
    "Cosmos-2.0-Diffusion-2B-Text2Image": {
        "in_channels": 16,
        "out_channels": 16,
        "num_attention_heads": 16,
        "attention_head_dim": 128,
        "num_layers": 28,
        "mlp_ratio": 4.0,
        "text_embed_dim": 1024,
        "adaln_lora_dim": 256,
        "max_size": (128, 240, 240),
        "patch_size": (1, 2, 2),
        "rope_scale": (1.0, 4.0, 4.0),
        "concat_padding_mask": True,
        "extra_pos_embed_type": None,
    },
    "Cosmos-2.0-Diffusion-14B-Text2Image": {
        "in_channels": 16,
        "out_channels": 16,
        "num_attention_heads": 40,
        "attention_head_dim": 128,
        "num_layers": 36,
        "mlp_ratio": 4.0,
        "text_embed_dim": 1024,
        "adaln_lora_dim": 256,
        "max_size": (128, 240, 240),
        "patch_size": (1, 2, 2),
        "rope_scale": (1.0, 4.0, 4.0),
        "concat_padding_mask": True,
        "extra_pos_embed_type": None,
    },
    "Cosmos-2.0-Diffusion-2B-Video2World": {
        "in_channels": 16 + 1,
        "out_channels": 16,
        "num_attention_heads": 16,
        "attention_head_dim": 128,
        "num_layers": 28,
        "mlp_ratio": 4.0,
        "text_embed_dim": 1024,
        "adaln_lora_dim": 256,
        "max_size": (128, 240, 240),
        "patch_size": (1, 2, 2),
        "rope_scale": (1.0, 3.0, 3.0),
        "concat_padding_mask": True,
        "extra_pos_embed_type": None,
    },
    "Cosmos-2.0-Diffusion-14B-Video2World": {
        "in_channels": 16 + 1,
        "out_channels": 16,
        "num_attention_heads": 40,
        "attention_head_dim": 128,
        "num_layers": 36,
        "mlp_ratio": 4.0,
        "text_embed_dim": 1024,
        "adaln_lora_dim": 256,
        "max_size": (128, 240, 240),
        "patch_size": (1, 2, 2),
        "rope_scale": (20 / 24, 2.0, 2.0),
        "concat_padding_mask": True,
        "extra_pos_embed_type": None,
    },
    "Cosmos-2.5-Predict-Base-2B": {
        "in_channels": 16 + 1,
        "out_channels": 16,
        "num_attention_heads": 16,
        "attention_head_dim": 128,
        "num_layers": 28,
        "mlp_ratio": 4.0,
        "text_embed_dim": 1024,
        "adaln_lora_dim": 256,
        "max_size": (128, 240, 240),
        "patch_size": (1, 2, 2),
        "rope_scale": (1.0, 3.0, 3.0),
        "concat_padding_mask": True,
        # NOTE: source config has pos_emb_learnable: 'True' - but params are missing
        "extra_pos_embed_type": None,
        "use_crossattn_projection": True,
        "crossattn_proj_in_channels": 100352,
        "encoder_hidden_states_channels": 1024,
    },
    "Cosmos-2.5-Predict-Base-14B": {
        "in_channels": 16 + 1,
        "out_channels": 16,
        "num_attention_heads": 40,
        "attention_head_dim": 128,
        "num_layers": 36,
        "mlp_ratio": 4.0,
        "text_embed_dim": 1024,
        "adaln_lora_dim": 256,
        "max_size": (128, 240, 240),
        "patch_size": (1, 2, 2),
        "rope_scale": (1.0, 3.0, 3.0),
        "concat_padding_mask": True,
        # NOTE: source config has pos_emb_learnable: 'True' - but params are missing
        "extra_pos_embed_type": None,
        "use_crossattn_projection": True,
        "crossattn_proj_in_channels": 100352,
        "encoder_hidden_states_channels": 1024,
    },
    "Cosmos-2.5-Transfer-General-2B": {
        "in_channels": 16 + 1,
        "out_channels": 16,
        "num_attention_heads": 16,
        "attention_head_dim": 128,
        "num_layers": 28,
        "mlp_ratio": 4.0,
        "text_embed_dim": 1024,
        "adaln_lora_dim": 256,
        "max_size": (128, 240, 240),
        "patch_size": (1, 2, 2),
        "rope_scale": (1.0, 3.0, 3.0),
        "concat_padding_mask": True,
        "extra_pos_embed_type": None,
        "use_crossattn_projection": True,
        "crossattn_proj_in_channels": 100352,
        "encoder_hidden_states_channels": 1024,
        "controlnet_block_every_n": 7,
        "img_context_dim_in": 1152,
        "img_context_dim_out": 2048,
        "img_context_num_tokens": 256,
    },
}

CONTROLNET_CONFIGS = {
    "Cosmos-2.5-Transfer-General-2B": {
        "n_controlnet_blocks": 4,
        "model_channels": 2048,
        "in_channels": 130,
        "latent_channels": 18,  # (16 latent + 1 condition_mask) + 1 padding_mask = 18
        "num_attention_heads": 16,
        "attention_head_dim": 128,
        "mlp_ratio": 4.0,
        "text_embed_dim": 1024,
        "adaln_lora_dim": 256,
        "patch_size": (1, 2, 2),
        "max_size": (128, 240, 240),
        "rope_scale": (1.0, 3.0, 3.0),
        "extra_pos_embed_type": None,
        "img_context_dim_in": 1152,
        "img_context_dim_out": 2048,
        "use_crossattn_projection": True,
        "crossattn_proj_in_channels": 100352,
        "encoder_hidden_states_channels": 1024,
    },
}

CONTROLNET_KEYS_RENAME_DICT = {
    **TRANSFORMER_KEYS_RENAME_DICT_COSMOS_2_0,
    "blocks": "blocks",
    "control_embedder.proj.1": "patch_embed.proj",
}


CONTROLNET_SPECIAL_KEYS_REMAP = {**TRANSFORMER_SPECIAL_KEYS_REMAP_COSMOS_2_0}

VAE_KEYS_RENAME_DICT = {
    "down.0": "down_blocks.0",
    "down.1": "down_blocks.1",
    "down.2": "down_blocks.2",
    "up.0": "up_blocks.2",
    "up.1": "up_blocks.1",
    "up.2": "up_blocks.0",
    ".block.": ".resnets.",
    "downsample": "downsamplers.0",
    "upsample": "upsamplers.0",
    "mid.block_1": "mid_block.resnets.0",
    "mid.attn_1.0": "mid_block.attentions.0",
    "mid.attn_1.1": "mid_block.temp_attentions.0",
    "mid.block_2": "mid_block.resnets.1",
    ".q.conv3d": ".to_q",
    ".k.conv3d": ".to_k",
    ".v.conv3d": ".to_v",
    ".proj_out.conv3d": ".to_out.0",
    ".0.conv3d": ".conv_s",
    ".1.conv3d": ".conv_t",
    "conv1.conv3d": "conv1",
    "conv2.conv3d": "conv2",
    "conv3.conv3d": "conv3",
    "nin_shortcut.conv3d": "conv_shortcut",
    "quant_conv.conv3d": "quant_conv",
    "post_quant_conv.conv3d": "post_quant_conv",
}

VAE_SPECIAL_KEYS_REMAP = {
    "wavelets": remove_keys_,
    "_arange": remove_keys_,
    "patch_size_buffer": remove_keys_,
}

VAE_CONFIGS = {
    "CV8x8x8-0.1": {
        "name": "nvidia/Cosmos-0.1-Tokenizer-CV8x8x8",
        "diffusers_config": {
            "in_channels": 3,
            "out_channels": 3,
            "latent_channels": 16,
            "encoder_block_out_channels": (128, 256, 512, 512),
            "decode_block_out_channels": (256, 512, 512, 512),
            "attention_resolutions": (32,),
            "resolution": 1024,
            "num_layers": 2,
            "patch_size": 4,
            "patch_type": "haar",
            "scaling_factor": 1.0,
            "spatial_compression_ratio": 8,
            "temporal_compression_ratio": 8,
            "latents_mean": None,
            "latents_std": None,
        },
    },
    "CV8x8x8-1.0": {
        "name": "nvidia/Cosmos-1.0-Tokenizer-CV8x8x8",
        "diffusers_config": {
            "in_channels": 3,
            "out_channels": 3,
            "latent_channels": 16,
            "encoder_block_out_channels": (128, 256, 512, 512),
            "decode_block_out_channels": (256, 512, 512, 512),
            "attention_resolutions": (32,),
            "resolution": 1024,
            "num_layers": 2,
            "patch_size": 4,
            "patch_type": "haar",
            "scaling_factor": 1.0,
            "spatial_compression_ratio": 8,
            "temporal_compression_ratio": 8,
            "latents_mean": None,
            "latents_std": None,
        },
    },
}


def get_state_dict(saved_dict: Dict[str, Any]) -> dict[str, Any]:
    state_dict = saved_dict
    if "model" in saved_dict.keys():
        state_dict = state_dict["model"]
    if "module" in saved_dict.keys():
        state_dict = state_dict["module"]
    if "state_dict" in saved_dict.keys():
        state_dict = state_dict["state_dict"]
    return state_dict


def convert_transformer(
    transformer_type: str,
    state_dict: Optional[Dict[str, Any]] = None,
    weights_only: bool = True,
):
    PREFIX_KEY = "net."

    if "Cosmos-1.0" in transformer_type:
        TRANSFORMER_KEYS_RENAME_DICT = TRANSFORMER_KEYS_RENAME_DICT_COSMOS_1_0
        TRANSFORMER_SPECIAL_KEYS_REMAP = TRANSFORMER_SPECIAL_KEYS_REMAP_COSMOS_1_0
    elif "Cosmos-2.0" in transformer_type:
        TRANSFORMER_KEYS_RENAME_DICT = TRANSFORMER_KEYS_RENAME_DICT_COSMOS_2_0
        TRANSFORMER_SPECIAL_KEYS_REMAP = TRANSFORMER_SPECIAL_KEYS_REMAP_COSMOS_2_0
    elif "Cosmos-2.5" in transformer_type:
        TRANSFORMER_KEYS_RENAME_DICT = TRANSFORMER_KEYS_RENAME_DICT_COSMOS_2_0
        TRANSFORMER_SPECIAL_KEYS_REMAP = TRANSFORMER_SPECIAL_KEYS_REMAP_COSMOS_2_0
    else:
        assert False

    with init_empty_weights():
        config = TRANSFORMER_CONFIGS[transformer_type]
        transformer = CosmosTransformer3DModel(**config)

    old2new = {}
    new2old = {}
    for key in list(state_dict.keys()):
        new_key = key[:]
        if new_key.startswith(PREFIX_KEY):
            new_key = new_key.removeprefix(PREFIX_KEY)
        for replace_key, rename_key in TRANSFORMER_KEYS_RENAME_DICT.items():
            new_key = new_key.replace(replace_key, rename_key)
        print(key, "->", new_key, flush=True)
        assert new_key not in new2old, f"new key {new_key} already mapped"
        assert key not in old2new, f"old key {key} already mapped"
        old2new[key] = new_key
        new2old[new_key] = key
        update_state_dict_(state_dict, key, new_key)

    for key in list(state_dict.keys()):
        for special_key, handler_fn_inplace in TRANSFORMER_SPECIAL_KEYS_REMAP.items():
            if special_key not in key:
                continue
            handler_fn_inplace(key, state_dict)

    expected_keys = set(transformer.state_dict().keys())
    mapped_keys = set(state_dict.keys())
    missing_keys = expected_keys - mapped_keys
    unexpected_keys = mapped_keys - expected_keys
    if missing_keys:
        print(f"ERROR: missing keys ({len(missing_keys)} from state_dict:", flush=True, file=sys.stderr)
        for k in missing_keys:
            print(k)
        sys.exit(1)
    if unexpected_keys:
        print(f"ERROR: unexpected keys ({len(unexpected_keys)}) from state_dict:", flush=True, file=sys.stderr)
        for k in unexpected_keys:
            print(k)
        sys.exit(2)

    transformer.load_state_dict(state_dict, strict=True, assign=True)
    return transformer


def convert_controlnet(
    transformer_type: str,
    control_state_dict: Dict[str, Any],
    base_state_dict: Dict[str, Any],
    weights_only: bool = True,
):
    """
    Convert controlnet weights.

    Args:
        transformer_type: The type of transformer/controlnet
        control_state_dict: State dict containing controlnet-specific weights
        base_state_dict: State dict containing base transformer weights (for shared modules)
        weights_only: Whether to use weights_only loading
    """
    if transformer_type not in CONTROLNET_CONFIGS:
        raise AssertionError(f"{transformer_type} does not define a ControlNet config")

    PREFIX_KEY = "net."

    # Process control-specific keys
    for key in list(control_state_dict.keys()):
        new_key = key[:]
        if new_key.startswith(PREFIX_KEY):
            new_key = new_key.removeprefix(PREFIX_KEY)
        for replace_key, rename_key in CONTROLNET_KEYS_RENAME_DICT.items():
            new_key = new_key.replace(replace_key, rename_key)
        update_state_dict_(control_state_dict, key, new_key)

    for key in list(control_state_dict.keys()):
        for special_key, handler_fn_inplace in CONTROLNET_SPECIAL_KEYS_REMAP.items():
            if special_key not in key:
                continue
            handler_fn_inplace(key, control_state_dict)

    # Copy shared weights from base transformer to controlnet
    # These are the duplicated modules: patch_embed_base, time_embed, learnable_pos_embed, img_context_proj, crossattn_proj
    shared_module_mappings = {
        # transformer key prefix -> controlnet key prefix
        "patch_embed.": "patch_embed_base.",
        "time_embed.": "time_embed.",
        "learnable_pos_embed.": "learnable_pos_embed.",
        "img_context_proj.": "img_context_proj.",
        "crossattn_proj.": "crossattn_proj.",
    }

    for key in list(base_state_dict.keys()):
        for transformer_prefix, controlnet_prefix in shared_module_mappings.items():
            if key.startswith(transformer_prefix):
                controlnet_key = controlnet_prefix + key[len(transformer_prefix) :]
                control_state_dict[controlnet_key] = base_state_dict[key].clone()
                print(f"Copied shared weight: {key} -> {controlnet_key}", flush=True)
                break

    cfg = CONTROLNET_CONFIGS[transformer_type]
    controlnet = CosmosControlNetModel(**cfg)

    expected_keys = set(controlnet.state_dict().keys())
    mapped_keys = set(control_state_dict.keys())
    missing_keys = expected_keys - mapped_keys
    unexpected_keys = mapped_keys - expected_keys
    if missing_keys:
        print(f"WARNING: missing controlnet keys ({len(missing_keys)}):", file=sys.stderr, flush=True)
        for k in sorted(missing_keys):
            print(k, file=sys.stderr)
        sys.exit(3)
    if unexpected_keys:
        print(f"WARNING: unexpected controlnet keys ({len(unexpected_keys)}):", file=sys.stderr, flush=True)
        for k in sorted(unexpected_keys):
            print(k, file=sys.stderr)
        sys.exit(4)

    controlnet.load_state_dict(control_state_dict, strict=True, assign=True)
    return controlnet


def convert_vae(vae_type: str):
    model_name = VAE_CONFIGS[vae_type]["name"]
    snapshot_directory = snapshot_download(model_name, repo_type="model")
    directory = pathlib.Path(snapshot_directory)

    autoencoder_file = directory / "autoencoder.jit"
    mean_std_file = directory / "mean_std.pt"

    original_state_dict = torch.jit.load(autoencoder_file.as_posix()).state_dict()
    if mean_std_file.exists():
        mean_std = torch.load(mean_std_file, map_location="cpu", weights_only=True)
    else:
        mean_std = (None, None)

    config = VAE_CONFIGS[vae_type]["diffusers_config"]
    config.update(
        {
            "latents_mean": mean_std[0].detach().cpu().numpy().tolist(),
            "latents_std": mean_std[1].detach().cpu().numpy().tolist(),
        }
    )
    vae = AutoencoderKLCosmos(**config)

    for key in list(original_state_dict.keys()):
        new_key = key[:]
        for replace_key, rename_key in VAE_KEYS_RENAME_DICT.items():
            new_key = new_key.replace(replace_key, rename_key)
        update_state_dict_(original_state_dict, key, new_key)

    for key in list(original_state_dict.keys()):
        for special_key, handler_fn_inplace in VAE_SPECIAL_KEYS_REMAP.items():
            if special_key not in key:
                continue
            handler_fn_inplace(key, original_state_dict)

    vae.load_state_dict(original_state_dict, strict=True, assign=True)
    return vae


def save_pipeline_cosmos_1_0(args, transformer, vae):
    text_encoder = T5EncoderModel.from_pretrained(args.text_encoder_path, torch_dtype=torch.bfloat16)
    tokenizer = T5TokenizerFast.from_pretrained(args.tokenizer_path)
    # The original code initializes EDM config with sigma_min=0.0002, but does not make use of it anywhere directly.
    # So, the sigma_min values that is used is the default value of 0.002.
    scheduler = EDMEulerScheduler(
        sigma_min=0.002,
        sigma_max=80,
        sigma_data=0.5,
        sigma_schedule="karras",
        num_train_timesteps=1000,
        prediction_type="epsilon",
        rho=7.0,
        final_sigmas_type="sigma_min",
    )

    pipe_cls = CosmosTextToWorldPipeline if "Text2World" in args.transformer_type else CosmosVideoToWorldPipeline
    pipe = pipe_cls(
        text_encoder=text_encoder,
        tokenizer=tokenizer,
        transformer=transformer,
        vae=vae,
        scheduler=scheduler,
        safety_checker=lambda *args, **kwargs: None,
    )
    pipe.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB")


def save_pipeline_cosmos_2_0(args, transformer, vae):
    text_encoder = T5EncoderModel.from_pretrained(args.text_encoder_path, torch_dtype=torch.bfloat16)
    tokenizer = T5TokenizerFast.from_pretrained(args.tokenizer_path)

    scheduler = FlowMatchEulerDiscreteScheduler(use_karras_sigmas=True)

    pipe_cls = Cosmos2TextToImagePipeline if "Text2Image" in args.transformer_type else Cosmos2VideoToWorldPipeline
    pipe = pipe_cls(
        text_encoder=text_encoder,
        tokenizer=tokenizer,
        transformer=transformer,
        vae=vae,
        scheduler=scheduler,
        safety_checker=lambda *args, **kwargs: None,
    )
    pipe.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB")


def save_pipeline_cosmos2_5_predict(args, transformer, vae):
    text_encoder_path = args.text_encoder_path or "nvidia/Cosmos-Reason1-7B"
    tokenizer_path = args.tokenizer_path or "Qwen/Qwen2.5-VL-7B-Instruct"

    text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained(
        text_encoder_path, torch_dtype="auto", device_map="cpu"
    )
    tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)

    scheduler = UniPCMultistepScheduler(
        use_karras_sigmas=True,
        use_flow_sigmas=True,
        prediction_type="flow_prediction",
        sigma_max=200.0,
        sigma_min=0.01,
    )

    pipe = Cosmos2_5_PredictBasePipeline(
        text_encoder=text_encoder,
        tokenizer=tokenizer,
        transformer=transformer,
        vae=vae,
        scheduler=scheduler,
        safety_checker=lambda *args, **kwargs: None,
    )
    pipe.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB")


def save_pipeline_cosmos2_5_transfer(args, transformer, controlnet, vae):
    text_encoder_path = args.text_encoder_path or "nvidia/Cosmos-Reason1-7B"
    tokenizer_path = args.tokenizer_path or "Qwen/Qwen2.5-VL-7B-Instruct"

    text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained(
        text_encoder_path, torch_dtype="auto", device_map="cpu"
    )
    tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)

    scheduler = UniPCMultistepScheduler(
        use_karras_sigmas=True,
        use_flow_sigmas=True,
        prediction_type="flow_prediction",
        sigma_max=200.0,
        sigma_min=0.01,
    )

    pipe = Cosmos2_5_TransferPipeline(
        text_encoder=text_encoder,
        tokenizer=tokenizer,
        transformer=transformer,
        controlnet=controlnet,
        vae=vae,
        scheduler=scheduler,
        safety_checker=lambda *args, **kwargs: None,
    )
    pipe.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB")


def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--transformer_type", type=str, default=None, choices=list(TRANSFORMER_CONFIGS.keys()))
    parser.add_argument(
        "--transformer_ckpt_path", type=str, default=None, help="Path to original transformer checkpoint"
    )
    parser.add_argument(
        "--vae_type", type=str, default="wan2.1", choices=["wan2.1", *list(VAE_CONFIGS.keys())], help="Type of VAE"
    )
    parser.add_argument("--text_encoder_path", type=str, default=None)
    parser.add_argument("--tokenizer_path", type=str, default=None)
    parser.add_argument("--save_pipeline", action="store_true")
    parser.add_argument("--output_path", type=str, required=True, help="Path where converted model should be saved")
    parser.add_argument("--dtype", default="bf16", help="Torch dtype to save the transformer in.")
    return parser.parse_args()


DTYPE_MAPPING = {
    "fp32": torch.float32,
    "fp16": torch.float16,
    "bf16": torch.bfloat16,
}


if __name__ == "__main__":
    args = get_args()

    transformer = None
    controlnet = None
    dtype = DTYPE_MAPPING[args.dtype]

    if args.save_pipeline:
        assert args.transformer_ckpt_path is not None
        assert args.vae_type is not None

    raw_state_dict = None
    if args.transformer_ckpt_path is not None:
        weights_only = "Cosmos-1.0" in args.transformer_type
        raw_state_dict = get_state_dict(
            torch.load(args.transformer_ckpt_path, map_location="cpu", weights_only=weights_only)
        )

    if raw_state_dict is not None:
        if "Transfer" in args.transformer_type:
            base_state_dict = {}
            control_state_dict = {}
            for k, v in raw_state_dict.items():
                plain_key = k.removeprefix("net.") if k.startswith("net.") else k
                if "control" in plain_key.lower():
                    control_state_dict[k] = v
                else:
                    base_state_dict[k] = v
            assert len(base_state_dict.keys() & control_state_dict.keys()) == 0

            # Convert transformer first to get the processed base state dict
            transformer = convert_transformer(
                args.transformer_type, state_dict=base_state_dict, weights_only=weights_only
            )
            transformer = transformer.to(dtype=dtype)

            # Get converted transformer state dict to copy shared weights to controlnet
            converted_base_state_dict = transformer.state_dict()

            # Convert controlnet with both control-specific and shared weights from transformer
            controlnet = convert_controlnet(
                args.transformer_type, control_state_dict, converted_base_state_dict, weights_only=weights_only
            )
            controlnet = controlnet.to(dtype=dtype)

            if not args.save_pipeline:
                transformer.save_pretrained(
                    pathlib.Path(args.output_path) / "transformer", safe_serialization=True, max_shard_size="5GB"
                )
                controlnet.save_pretrained(
                    pathlib.Path(args.output_path) / "controlnet", safe_serialization=True, max_shard_size="5GB"
                )
        else:
            transformer = convert_transformer(
                args.transformer_type, state_dict=raw_state_dict, weights_only=weights_only
            )
            transformer = transformer.to(dtype=dtype)
            if not args.save_pipeline:
                transformer.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB")

    if args.vae_type is not None:
        if "Cosmos-1.0" in args.transformer_type:
            vae = convert_vae(args.vae_type)
        elif "Cosmos-2.0" in args.transformer_type or "Cosmos-2.5" in args.transformer_type:
            vae = AutoencoderKLWan.from_pretrained(
                "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", subfolder="vae", torch_dtype=torch.float32
            )
        else:
            raise AssertionError(f"{args.transformer_type} not supported")

        if not args.save_pipeline:
            vae.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB")
    else:
        vae = None

    if args.save_pipeline:
        if "Cosmos-1.0" in args.transformer_type:
            assert args.text_encoder_path is not None
            assert args.tokenizer_path is not None
            save_pipeline_cosmos_1_0(args, transformer, vae)
        elif "Cosmos-2.0" in args.transformer_type:
            assert args.text_encoder_path is not None
            assert args.tokenizer_path is not None
            save_pipeline_cosmos_2_0(args, transformer, vae)
        elif "Cosmos-2.5" in args.transformer_type:
            if "Predict" in args.transformer_type:
                save_pipeline_cosmos2_5_predict(args, transformer, vae)
            elif "Transfer" in args.transformer_type:
                save_pipeline_cosmos2_5_transfer(args, transformer, None, vae)
            else:
                raise AssertionError(f"{args.transformer_type} not supported")
        else:
            raise AssertionError(f"{args.transformer_type} not supported")
