def convert_megatron_transformer_checkpoint_to_diffusers(
    ckpt_path: str,
    num_layers: int,
    num_heads: int,
    hidden_size: int,
):
    """
    Convert a Megatron Transformer checkpoint to Diffusers format.

    Args:
        ckpt_path (str): Path to the Megatron Transformer checkpoint.
        num_layers (int): Number of Transformer layers.
        num_heads (int): Number of attention heads.
        hidden_size (int): Hidden size of the Transformer.

    Returns:
        dict: The converted state dictionary compatible with Diffusers.
    """
    ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
    mega = ckpt["model"]

    new_state_dict = {}

    # Patch Embedding
    new_state_dict["patch_embed.proj.weight"] = mega["encoder_expand_linear.weight"].reshape(
        hidden_size, 128 if args.control else 64
    )
    new_state_dict["patch_embed.proj.bias"] = mega["encoder_expand_linear.bias"]
    new_state_dict["patch_embed.text_proj.weight"] = mega["text_projector.weight"]
    new_state_dict["patch_embed.text_proj.bias"] = mega["text_projector.bias"]

    # Time Condition Embedding
    new_state_dict["time_condition_embed.timestep_embedder.linear_1.weight"] = mega[
        "time_embedding.time_embed.0.weight"
    ]
    new_state_dict["time_condition_embed.timestep_embedder.linear_1.bias"] = mega["time_embedding.time_embed.0.bias"]
    new_state_dict["time_condition_embed.timestep_embedder.linear_2.weight"] = mega[
        "time_embedding.time_embed.2.weight"
    ]
    new_state_dict["time_condition_embed.timestep_embedder.linear_2.bias"] = mega["time_embedding.time_embed.2.bias"]

    new_state_dict["time_condition_embed.condition_embedder.linear_1.weight"] = mega[
        "label_embedding.label_embed.0.weight"
    ]
    new_state_dict["time_condition_embed.condition_embedder.linear_1.bias"] = mega[
        "label_embedding.label_embed.0.bias"
    ]
    new_state_dict["time_condition_embed.condition_embedder.linear_2.weight"] = mega[
        "label_embedding.label_embed.2.weight"
    ]
    new_state_dict["time_condition_embed.condition_embedder.linear_2.bias"] = mega[
        "label_embedding.label_embed.2.bias"
    ]

    # Convert each Transformer layer
    for i in tqdm(range(num_layers), desc="Converting layers (Megatron->Diffusers)"):
        block_prefix = f"transformer_blocks.{i}."

        # AdaLayerNorm
        new_state_dict[block_prefix + "norm1.linear.weight"] = mega[f"decoder.layers.{i}.adaln.weight"]
        new_state_dict[block_prefix + "norm1.linear.bias"] = mega[f"decoder.layers.{i}.adaln.bias"]
        qkv_weight = mega[f"decoder.layers.{i}.self_attention.linear_qkv.weight"]
        qkv_bias = mega[f"decoder.layers.{i}.self_attention.linear_qkv.bias"]

        # Reshape to match SAT logic
        qkv_weight = qkv_weight.view(num_heads, 3, hidden_size // num_heads, hidden_size)
        qkv_weight = qkv_weight.permute(1, 0, 2, 3).reshape(3 * hidden_size, hidden_size)

        qkv_bias = qkv_bias.view(num_heads, 3, hidden_size // num_heads)
        qkv_bias = qkv_bias.permute(1, 0, 2).reshape(3 * hidden_size)

        # Assign to Diffusers keys
        q, k, v = torch.chunk(qkv_weight, 3, dim=0)
        qb, kb, vb = torch.chunk(qkv_bias, 3, dim=0)

        new_state_dict[block_prefix + "attn1.to_q.weight"] = q
        new_state_dict[block_prefix + "attn1.to_q.bias"] = qb
        new_state_dict[block_prefix + "attn1.to_k.weight"] = k
        new_state_dict[block_prefix + "attn1.to_k.bias"] = kb
        new_state_dict[block_prefix + "attn1.to_v.weight"] = v
        new_state_dict[block_prefix + "attn1.to_v.bias"] = vb

        # Attention Output
        new_state_dict[block_prefix + "attn1.to_out.0.weight"] = mega[
            f"decoder.layers.{i}.self_attention.linear_proj.weight"
        ]
        new_state_dict[block_prefix + "attn1.to_out.0.bias"] = mega[
            f"decoder.layers.{i}.self_attention.linear_proj.bias"
        ]

        # MLP
        new_state_dict[block_prefix + "ff.net.0.proj.weight"] = mega[f"decoder.layers.{i}.mlp.linear_fc1.weight"]
        new_state_dict[block_prefix + "ff.net.0.proj.bias"] = mega[f"decoder.layers.{i}.mlp.linear_fc1.bias"]
        new_state_dict[block_prefix + "ff.net.2.weight"] = mega[f"decoder.layers.{i}.mlp.linear_fc2.weight"]
        new_state_dict[block_prefix + "ff.net.2.bias"] = mega[f"decoder.layers.{i}.mlp.linear_fc2.bias"]

    # Final Layers
    new_state_dict["norm_out.linear.weight"] = swap_scale_shift(mega["adaln_final.weight"], dim=0)
    new_state_dict["norm_out.linear.bias"] = swap_scale_shift(mega["adaln_final.bias"], dim=0)
    new_state_dict["proj_out.weight"] = mega["output_projector.weight"]
    new_state_dict["proj_out.bias"] = mega["output_projector.bias"]

    return new_state_dict

    Convert a Megatron Transformer checkpoint to Diffusers format.

    Args:
        ckpt_path (str): Path to the Megatron Transformer checkpoint.
        num_layers (int): Number of Transformer layers.
        num_heads (int): Number of attention heads.
        hidden_size (int): Hidden size of the Transformer.

    Returns:
        dict: The converted state dictionary compatible with Diffusers.
    
scripts/convert_cogview4_to_diffusers_megatron.py