from typing import List

from mflux.models.common.weights.mapping.weight_mapping import WeightMapping, WeightTarget
from mflux.models.common.weights.mapping.weight_transforms import WeightTransforms


class QwenWeightMapping(WeightMapping):
    @staticmethod
    def get_transformer_mapping() -> List[WeightTarget]:
        return [
            WeightTarget(
                to_pattern="img_in.weight",
                from_pattern=["img_in.weight"],
            ),
            WeightTarget(
                to_pattern="img_in.bias",
                from_pattern=["img_in.bias"],
            ),
            WeightTarget(
                to_pattern="txt_norm.weight",
                from_pattern=["txt_norm.weight"],
            ),
            WeightTarget(
                to_pattern="txt_in.weight",
                from_pattern=["txt_in.weight"],
            ),
            WeightTarget(
                to_pattern="txt_in.bias",
                from_pattern=["txt_in.bias"],
            ),
            WeightTarget(
                to_pattern="time_text_embed.timestep_embedder.linear_1.weight",
                from_pattern=["time_text_embed.timestep_embedder.linear_1.weight"],
            ),
            WeightTarget(
                to_pattern="time_text_embed.timestep_embedder.linear_1.bias",
                from_pattern=["time_text_embed.timestep_embedder.linear_1.bias"],
            ),
            WeightTarget(
                to_pattern="time_text_embed.timestep_embedder.linear_2.weight",
                from_pattern=["time_text_embed.timestep_embedder.linear_2.weight"],
            ),
            WeightTarget(
                to_pattern="time_text_embed.timestep_embedder.linear_2.bias",
                from_pattern=["time_text_embed.timestep_embedder.linear_2.bias"],
            ),
            WeightTarget(
                to_pattern="norm_out.linear.weight",
                from_pattern=["norm_out.linear.weight"],
            ),
            WeightTarget(
                to_pattern="norm_out.linear.bias",
                from_pattern=["norm_out.linear.bias"],
            ),
            WeightTarget(
                to_pattern="proj_out.weight",
                from_pattern=["proj_out.weight"],
            ),
            WeightTarget(
                to_pattern="proj_out.bias",
                from_pattern=["proj_out.bias"],
            ),
            WeightTarget(
                to_pattern="transformer_blocks.{block}.attn.to_q.weight",
                from_pattern=["transformer_blocks.{block}.attn.to_q.weight"],
            ),
            WeightTarget(
                to_pattern="transformer_blocks.{block}.attn.to_q.bias",
                from_pattern=["transformer_blocks.{block}.attn.to_q.bias"],
            ),
            WeightTarget(
                to_pattern="transformer_blocks.{block}.attn.to_k.weight",
                from_pattern=["transformer_blocks.{block}.attn.to_k.weight"],
            ),
            WeightTarget(
                to_pattern="transformer_blocks.{block}.attn.to_k.bias",
                from_pattern=["transformer_blocks.{block}.attn.to_k.bias"],
            ),
            WeightTarget(
                to_pattern="transformer_blocks.{block}.attn.to_v.weight",
                from_pattern=["transformer_blocks.{block}.attn.to_v.weight"],
            ),
            WeightTarget(
                to_pattern="transformer_blocks.{block}.attn.to_v.bias",
                from_pattern=["transformer_blocks.{block}.attn.to_v.bias"],
            ),
            WeightTarget(
                to_pattern="transformer_blocks.{block}.attn.add_q_proj.weight",
                from_pattern=["transformer_blocks.{block}.attn.add_q_proj.weight"],
            ),
            WeightTarget(
                to_pattern="transformer_blocks.{block}.attn.add_q_proj.bias",
                from_pattern=["transformer_blocks.{block}.attn.add_q_proj.bias"],
            ),
            WeightTarget(
                to_pattern="transformer_blocks.{block}.attn.add_k_proj.weight",
                from_pattern=["transformer_blocks.{block}.attn.add_k_proj.weight"],
            ),
            WeightTarget(
                to_pattern="transformer_blocks.{block}.attn.add_k_proj.bias",
                from_pattern=["transformer_blocks.{block}.attn.add_k_proj.bias"],
            ),
            WeightTarget(
                to_pattern="transformer_blocks.{block}.attn.add_v_proj.weight",
                from_pattern=["transformer_blocks.{block}.attn.add_v_proj.weight"],
            ),
            WeightTarget(
                to_pattern="transformer_blocks.{block}.attn.add_v_proj.bias",
                from_pattern=["transformer_blocks.{block}.attn.add_v_proj.bias"],
            ),
            WeightTarget(
                to_pattern="transformer_blocks.{block}.attn.norm_q.weight",
                from_pattern=["transformer_blocks.{block}.attn.norm_q.weight"],
            ),
            WeightTarget(
                to_pattern="transformer_blocks.{block}.attn.norm_k.weight",
                from_pattern=["transformer_blocks.{block}.attn.norm_k.weight"],
            ),
            WeightTarget(
                to_pattern="transformer_blocks.{block}.attn.norm_added_q.weight",
                from_pattern=["transformer_blocks.{block}.attn.norm_added_q.weight"],
            ),
            WeightTarget(
                to_pattern="transformer_blocks.{block}.attn.norm_added_k.weight",
                from_pattern=["transformer_blocks.{block}.attn.norm_added_k.weight"],
            ),
            WeightTarget(
                to_pattern="transformer_blocks.{block}.attn.attn_to_out.0.weight",
                from_pattern=["transformer_blocks.{block}.attn.to_out.0.weight"],
            ),
            WeightTarget(
                to_pattern="transformer_blocks.{block}.attn.attn_to_out.0.bias",
                from_pattern=["transformer_blocks.{block}.attn.to_out.0.bias"],
            ),
            WeightTarget(
                to_pattern="transformer_blocks.{block}.attn.to_add_out.weight",
                from_pattern=["transformer_blocks.{block}.attn.to_add_out.weight"],
            ),
            WeightTarget(
                to_pattern="transformer_blocks.{block}.attn.to_add_out.bias",
                from_pattern=["transformer_blocks.{block}.attn.to_add_out.bias"],
            ),
            WeightTarget(
                to_pattern="transformer_blocks.{block}.img_mod_linear.weight",
                from_pattern=["transformer_blocks.{block}.img_mod.1.weight"],
            ),
            WeightTarget(
                to_pattern="transformer_blocks.{block}.img_mod_linear.bias",
                from_pattern=["transformer_blocks.{block}.img_mod.1.bias"],
            ),
            WeightTarget(
                to_pattern="transformer_blocks.{block}.txt_mod_linear.weight",
                from_pattern=["transformer_blocks.{block}.txt_mod.1.weight"],
            ),
            WeightTarget(
                to_pattern="transformer_blocks.{block}.txt_mod_linear.bias",
                from_pattern=["transformer_blocks.{block}.txt_mod.1.bias"],
            ),
            WeightTarget(
                to_pattern="transformer_blocks.{block}.img_ff.mlp_in.weight",
                from_pattern=["transformer_blocks.{block}.img_mlp.net.0.proj.weight"],
            ),
            WeightTarget(
                to_pattern="transformer_blocks.{block}.img_ff.mlp_in.bias",
                from_pattern=["transformer_blocks.{block}.img_mlp.net.0.proj.bias"],
            ),
            WeightTarget(
                to_pattern="transformer_blocks.{block}.img_ff.mlp_out.weight",
                from_pattern=["transformer_blocks.{block}.img_mlp.net.2.weight"],
            ),
            WeightTarget(
                to_pattern="transformer_blocks.{block}.img_ff.mlp_out.bias",
                from_pattern=["transformer_blocks.{block}.img_mlp.net.2.bias"],
            ),
            WeightTarget(
                to_pattern="transformer_blocks.{block}.txt_ff.mlp_in.weight",
                from_pattern=["transformer_blocks.{block}.txt_mlp.net.0.proj.weight"],
            ),
            WeightTarget(
                to_pattern="transformer_blocks.{block}.txt_ff.mlp_in.bias",
                from_pattern=["transformer_blocks.{block}.txt_mlp.net.0.proj.bias"],
            ),
            WeightTarget(
                to_pattern="transformer_blocks.{block}.txt_ff.mlp_out.weight",
                from_pattern=["transformer_blocks.{block}.txt_mlp.net.2.weight"],
            ),
            WeightTarget(
                to_pattern="transformer_blocks.{block}.txt_ff.mlp_out.bias",
                from_pattern=["transformer_blocks.{block}.txt_mlp.net.2.bias"],
            ),
        ]

    @staticmethod
    def get_vae_mapping() -> List[WeightTarget]:
        return [
            WeightTarget(
                to_pattern="decoder.conv_in.conv3d.weight",
                from_pattern=["decoder.conv_in.weight"],
                transform=WeightTransforms.transpose_conv3d_weight,
            ),
            WeightTarget(
                to_pattern="decoder.conv_in.conv3d.bias",
                from_pattern=["decoder.conv_in.bias"],
            ),
            WeightTarget(
                to_pattern="decoder.conv_out.conv3d.weight",
                from_pattern=["decoder.conv_out.weight"],
                transform=WeightTransforms.transpose_conv3d_weight,
            ),
            WeightTarget(
                to_pattern="decoder.conv_out.conv3d.bias",
                from_pattern=["decoder.conv_out.bias"],
            ),
            WeightTarget(
                to_pattern="decoder.norm_out.weight",
                from_pattern=["decoder.norm_out.gamma"],
                transform=WeightTransforms.reshape_gamma_to_1d,
            ),
            WeightTarget(
                to_pattern="post_quant_conv.conv3d.weight",
                from_pattern=["post_quant_conv.weight"],
                transform=WeightTransforms.transpose_conv3d_weight,
            ),
            WeightTarget(
                to_pattern="post_quant_conv.conv3d.bias",
                from_pattern=["post_quant_conv.bias"],
            ),
            WeightTarget(
                to_pattern="decoder.mid_block.resnets.{i}.conv1.conv3d.weight",
                from_pattern=["decoder.mid_block.resnets.{i}.conv1.weight"],
                transform=WeightTransforms.transpose_conv3d_weight,
            ),
            WeightTarget(
                to_pattern="decoder.mid_block.resnets.{i}.conv1.conv3d.bias",
                from_pattern=["decoder.mid_block.resnets.{i}.conv1.bias"],
            ),
            WeightTarget(
                to_pattern="decoder.mid_block.resnets.{i}.conv2.conv3d.weight",
                from_pattern=["decoder.mid_block.resnets.{i}.conv2.weight"],
                transform=WeightTransforms.transpose_conv3d_weight,
            ),
            WeightTarget(
                to_pattern="decoder.mid_block.resnets.{i}.conv2.conv3d.bias",
                from_pattern=["decoder.mid_block.resnets.{i}.conv2.bias"],
            ),
            WeightTarget(
                to_pattern="decoder.mid_block.resnets.{i}.norm1.weight",
                from_pattern=["decoder.mid_block.resnets.{i}.norm1.gamma"],
                transform=WeightTransforms.reshape_gamma_to_1d,
            ),
            WeightTarget(
                to_pattern="decoder.mid_block.resnets.{i}.norm2.weight",
                from_pattern=["decoder.mid_block.resnets.{i}.norm2.gamma"],
                transform=WeightTransforms.reshape_gamma_to_1d,
            ),
            WeightTarget(
                to_pattern="decoder.mid_block.attentions.0.norm.weight",
                from_pattern=["decoder.mid_block.attentions.0.norm.gamma"],
                transform=WeightTransforms.reshape_gamma_to_1d,
            ),
            WeightTarget(
                to_pattern="decoder.mid_block.attentions.0.to_qkv.weight",
                from_pattern=["decoder.mid_block.attentions.0.to_qkv.weight"],
                transform=WeightTransforms.transpose_conv2d_weight,
            ),
            WeightTarget(
                to_pattern="decoder.mid_block.attentions.0.to_qkv.bias",
                from_pattern=["decoder.mid_block.attentions.0.to_qkv.bias"],
            ),
            WeightTarget(
                to_pattern="decoder.mid_block.attentions.0.proj.weight",
                from_pattern=["decoder.mid_block.attentions.0.proj.weight"],
                transform=WeightTransforms.transpose_conv2d_weight,
            ),
            WeightTarget(
                to_pattern="decoder.mid_block.attentions.0.proj.bias",
                from_pattern=["decoder.mid_block.attentions.0.proj.bias"],
            ),
            WeightTarget(
                to_pattern="decoder.up_block{block}.resnets.{res}.conv1.conv3d.weight",
                from_pattern=["decoder.up_blocks.{block}.resnets.{res}.conv1.weight"],
                transform=WeightTransforms.transpose_conv3d_weight,
            ),
            WeightTarget(
                to_pattern="decoder.up_block{block}.resnets.{res}.conv1.conv3d.bias",
                from_pattern=["decoder.up_blocks.{block}.resnets.{res}.conv1.bias"],
            ),
            WeightTarget(
                to_pattern="decoder.up_block{block}.resnets.{res}.conv2.conv3d.weight",
                from_pattern=["decoder.up_blocks.{block}.resnets.{res}.conv2.weight"],
                transform=WeightTransforms.transpose_conv3d_weight,
            ),
            WeightTarget(
                to_pattern="decoder.up_block{block}.resnets.{res}.conv2.conv3d.bias",
                from_pattern=["decoder.up_blocks.{block}.resnets.{res}.conv2.bias"],
            ),
            WeightTarget(
                to_pattern="decoder.up_block{block}.resnets.{res}.norm1.weight",
                from_pattern=["decoder.up_blocks.{block}.resnets.{res}.norm1.gamma"],
                transform=WeightTransforms.reshape_gamma_to_1d,
            ),
            WeightTarget(
                to_pattern="decoder.up_block{block}.resnets.{res}.norm2.weight",
                from_pattern=["decoder.up_blocks.{block}.resnets.{res}.norm2.gamma"],
                transform=WeightTransforms.reshape_gamma_to_1d,
            ),
            WeightTarget(
                to_pattern="decoder.up_block{block}.resnets.{res}.skip_conv.conv3d.weight",
                from_pattern=["decoder.up_blocks.{block}.resnets.{res}.conv_shortcut.weight"],
                transform=WeightTransforms.transpose_conv3d_weight,
                required=False,
            ),
            WeightTarget(
                to_pattern="decoder.up_block{block}.resnets.{res}.skip_conv.conv3d.bias",
                from_pattern=["decoder.up_blocks.{block}.resnets.{res}.conv_shortcut.bias"],
                required=False,
            ),
            WeightTarget(
                to_pattern="decoder.up_block{block}.upsamplers.0.resample_conv.weight",
                from_pattern=["decoder.up_blocks.{block}.upsamplers.0.resample.1.weight"],
                transform=WeightTransforms.transpose_conv2d_weight,
            ),
            WeightTarget(
                to_pattern="decoder.up_block{block}.upsamplers.0.resample_conv.bias",
                from_pattern=["decoder.up_blocks.{block}.upsamplers.0.resample.1.bias"],
            ),
            WeightTarget(
                to_pattern="decoder.up_block{block}.upsamplers.0.time_conv.conv3d.weight",
                from_pattern=["decoder.up_blocks.{block}.upsamplers.0.time_conv.weight"],
                transform=WeightTransforms.transpose_conv3d_weight,
            ),
            WeightTarget(
                to_pattern="decoder.up_block{block}.upsamplers.0.time_conv.conv3d.bias",
                from_pattern=["decoder.up_blocks.{block}.upsamplers.0.time_conv.bias"],
            ),
            WeightTarget(
                to_pattern="encoder.conv_in.conv3d.weight",
                from_pattern=["encoder.conv_in.weight"],
                transform=WeightTransforms.transpose_conv3d_weight,
            ),
            WeightTarget(
                to_pattern="encoder.conv_in.conv3d.bias",
                from_pattern=["encoder.conv_in.bias"],
            ),
            WeightTarget(
                to_pattern="encoder.conv_out.conv3d.weight",
                from_pattern=["encoder.conv_out.weight"],
                transform=WeightTransforms.transpose_conv3d_weight,
            ),
            WeightTarget(
                to_pattern="encoder.conv_out.conv3d.bias",
                from_pattern=["encoder.conv_out.bias"],
            ),
            WeightTarget(
                to_pattern="encoder.norm_out.weight",
                from_pattern=["encoder.norm_out.gamma"],
                transform=WeightTransforms.reshape_gamma_to_1d,
            ),
            WeightTarget(
                to_pattern="encoder.mid_block.attentions.0.norm.weight",
                from_pattern=["encoder.mid_block.attentions.0.norm.gamma"],
                transform=WeightTransforms.reshape_gamma_to_1d,
            ),
            WeightTarget(
                to_pattern="encoder.mid_block.attentions.0.to_qkv.weight",
                from_pattern=["encoder.mid_block.attentions.0.to_qkv.weight"],
                transform=WeightTransforms.transpose_conv2d_weight,
            ),
            WeightTarget(
                to_pattern="encoder.mid_block.attentions.0.to_qkv.bias",
                from_pattern=["encoder.mid_block.attentions.0.to_qkv.bias"],
            ),
            WeightTarget(
                to_pattern="encoder.mid_block.attentions.0.proj.weight",
                from_pattern=["encoder.mid_block.attentions.0.proj.weight"],
                transform=WeightTransforms.transpose_conv2d_weight,
            ),
            WeightTarget(
                to_pattern="encoder.mid_block.attentions.0.proj.bias",
                from_pattern=["encoder.mid_block.attentions.0.proj.bias"],
            ),
            WeightTarget(
                to_pattern="encoder.mid_block.resnets.{i}.conv1.conv3d.weight",
                from_pattern=["encoder.mid_block.resnets.{i}.conv1.weight"],
                transform=WeightTransforms.transpose_conv3d_weight,
            ),
            WeightTarget(
                to_pattern="encoder.mid_block.resnets.{i}.conv1.conv3d.bias",
                from_pattern=["encoder.mid_block.resnets.{i}.conv1.bias"],
            ),
            WeightTarget(
                to_pattern="encoder.mid_block.resnets.{i}.conv2.conv3d.weight",
                from_pattern=["encoder.mid_block.resnets.{i}.conv2.weight"],
                transform=WeightTransforms.transpose_conv3d_weight,
            ),
            WeightTarget(
                to_pattern="encoder.mid_block.resnets.{i}.conv2.conv3d.bias",
                from_pattern=["encoder.mid_block.resnets.{i}.conv2.bias"],
            ),
            WeightTarget(
                to_pattern="encoder.mid_block.resnets.{i}.norm1.weight",
                from_pattern=["encoder.mid_block.resnets.{i}.norm1.gamma"],
                transform=WeightTransforms.reshape_gamma_to_1d,
            ),
            WeightTarget(
                to_pattern="encoder.mid_block.resnets.{i}.norm2.weight",
                from_pattern=["encoder.mid_block.resnets.{i}.norm2.gamma"],
                transform=WeightTransforms.reshape_gamma_to_1d,
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.0.resnets.0.conv1.conv3d.weight",
                from_pattern=["encoder.down_blocks.0.conv1.weight"],
                transform=WeightTransforms.transpose_conv3d_weight,
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.0.resnets.0.conv1.conv3d.bias",
                from_pattern=["encoder.down_blocks.0.conv1.bias"],
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.0.resnets.0.conv2.conv3d.weight",
                from_pattern=["encoder.down_blocks.0.conv2.weight"],
                transform=WeightTransforms.transpose_conv3d_weight,
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.0.resnets.0.conv2.conv3d.bias",
                from_pattern=["encoder.down_blocks.0.conv2.bias"],
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.0.resnets.0.norm1.weight",
                from_pattern=["encoder.down_blocks.0.norm1.gamma"],
                transform=WeightTransforms.reshape_gamma_to_1d,
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.0.resnets.0.norm2.weight",
                from_pattern=["encoder.down_blocks.0.norm2.gamma"],
                transform=WeightTransforms.reshape_gamma_to_1d,
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.0.resnets.1.conv1.conv3d.weight",
                from_pattern=["encoder.down_blocks.1.conv1.weight"],
                transform=WeightTransforms.transpose_conv3d_weight,
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.0.resnets.1.conv1.conv3d.bias",
                from_pattern=["encoder.down_blocks.1.conv1.bias"],
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.0.resnets.1.conv2.conv3d.weight",
                from_pattern=["encoder.down_blocks.1.conv2.weight"],
                transform=WeightTransforms.transpose_conv3d_weight,
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.0.resnets.1.conv2.conv3d.bias",
                from_pattern=["encoder.down_blocks.1.conv2.bias"],
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.0.resnets.1.norm1.weight",
                from_pattern=["encoder.down_blocks.1.norm1.gamma"],
                transform=WeightTransforms.reshape_gamma_to_1d,
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.0.resnets.1.norm2.weight",
                from_pattern=["encoder.down_blocks.1.norm2.gamma"],
                transform=WeightTransforms.reshape_gamma_to_1d,
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.0.downsamplers.0.resample_conv.weight",
                from_pattern=["encoder.down_blocks.2.resample.1.weight"],
                transform=WeightTransforms.transpose_conv2d_weight,
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.0.downsamplers.0.resample_conv.bias",
                from_pattern=["encoder.down_blocks.2.resample.1.bias"],
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.1.resnets.0.conv1.conv3d.weight",
                from_pattern=["encoder.down_blocks.3.conv1.weight"],
                transform=WeightTransforms.transpose_conv3d_weight,
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.1.resnets.0.conv1.conv3d.bias",
                from_pattern=["encoder.down_blocks.3.conv1.bias"],
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.1.resnets.0.conv2.conv3d.weight",
                from_pattern=["encoder.down_blocks.3.conv2.weight"],
                transform=WeightTransforms.transpose_conv3d_weight,
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.1.resnets.0.conv2.conv3d.bias",
                from_pattern=["encoder.down_blocks.3.conv2.bias"],
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.1.resnets.0.norm1.weight",
                from_pattern=["encoder.down_blocks.3.norm1.gamma"],
                transform=WeightTransforms.reshape_gamma_to_1d,
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.1.resnets.0.norm2.weight",
                from_pattern=["encoder.down_blocks.3.norm2.gamma"],
                transform=WeightTransforms.reshape_gamma_to_1d,
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.1.resnets.0.skip_conv.conv3d.weight",
                from_pattern=["encoder.down_blocks.3.conv_shortcut.weight"],
                transform=WeightTransforms.transpose_conv3d_weight,
                required=False,
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.1.resnets.0.skip_conv.conv3d.bias",
                from_pattern=["encoder.down_blocks.3.conv_shortcut.bias"],
                required=False,
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.1.resnets.1.conv1.conv3d.weight",
                from_pattern=["encoder.down_blocks.4.conv1.weight"],
                transform=WeightTransforms.transpose_conv3d_weight,
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.1.resnets.1.conv1.conv3d.bias",
                from_pattern=["encoder.down_blocks.4.conv1.bias"],
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.1.resnets.1.conv2.conv3d.weight",
                from_pattern=["encoder.down_blocks.4.conv2.weight"],
                transform=WeightTransforms.transpose_conv3d_weight,
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.1.resnets.1.conv2.conv3d.bias",
                from_pattern=["encoder.down_blocks.4.conv2.bias"],
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.1.resnets.1.norm1.weight",
                from_pattern=["encoder.down_blocks.4.norm1.gamma"],
                transform=WeightTransforms.reshape_gamma_to_1d,
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.1.resnets.1.norm2.weight",
                from_pattern=["encoder.down_blocks.4.norm2.gamma"],
                transform=WeightTransforms.reshape_gamma_to_1d,
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.1.downsamplers.0.resample_conv.weight",
                from_pattern=["encoder.down_blocks.5.resample.1.weight"],
                transform=WeightTransforms.transpose_conv2d_weight,
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.1.downsamplers.0.resample_conv.bias",
                from_pattern=["encoder.down_blocks.5.resample.1.bias"],
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.2.resnets.0.conv1.conv3d.weight",
                from_pattern=["encoder.down_blocks.6.conv1.weight"],
                transform=WeightTransforms.transpose_conv3d_weight,
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.2.resnets.0.conv1.conv3d.bias",
                from_pattern=["encoder.down_blocks.6.conv1.bias"],
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.2.resnets.0.conv2.conv3d.weight",
                from_pattern=["encoder.down_blocks.6.conv2.weight"],
                transform=WeightTransforms.transpose_conv3d_weight,
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.2.resnets.0.conv2.conv3d.bias",
                from_pattern=["encoder.down_blocks.6.conv2.bias"],
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.2.resnets.0.norm1.weight",
                from_pattern=["encoder.down_blocks.6.norm1.gamma"],
                transform=WeightTransforms.reshape_gamma_to_1d,
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.2.resnets.0.norm2.weight",
                from_pattern=["encoder.down_blocks.6.norm2.gamma"],
                transform=WeightTransforms.reshape_gamma_to_1d,
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.2.resnets.0.skip_conv.conv3d.weight",
                from_pattern=["encoder.down_blocks.6.conv_shortcut.weight"],
                transform=WeightTransforms.transpose_conv3d_weight,
                required=False,
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.2.resnets.0.skip_conv.conv3d.bias",
                from_pattern=["encoder.down_blocks.6.conv_shortcut.bias"],
                required=False,
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.2.resnets.1.conv1.conv3d.weight",
                from_pattern=["encoder.down_blocks.7.conv1.weight"],
                transform=WeightTransforms.transpose_conv3d_weight,
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.2.resnets.1.conv1.conv3d.bias",
                from_pattern=["encoder.down_blocks.7.conv1.bias"],
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.2.resnets.1.conv2.conv3d.weight",
                from_pattern=["encoder.down_blocks.7.conv2.weight"],
                transform=WeightTransforms.transpose_conv3d_weight,
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.2.resnets.1.conv2.conv3d.bias",
                from_pattern=["encoder.down_blocks.7.conv2.bias"],
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.2.resnets.1.norm1.weight",
                from_pattern=["encoder.down_blocks.7.norm1.gamma"],
                transform=WeightTransforms.reshape_gamma_to_1d,
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.2.resnets.1.norm2.weight",
                from_pattern=["encoder.down_blocks.7.norm2.gamma"],
                transform=WeightTransforms.reshape_gamma_to_1d,
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.2.downsamplers.0.resample_conv.weight",
                from_pattern=["encoder.down_blocks.8.resample.1.weight"],
                transform=WeightTransforms.transpose_conv2d_weight,
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.2.downsamplers.0.resample_conv.bias",
                from_pattern=["encoder.down_blocks.8.resample.1.bias"],
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.2.downsamplers.0.time_conv.conv3d.weight",
                from_pattern=["encoder.down_blocks.8.time_conv.weight"],
                transform=WeightTransforms.transpose_conv3d_weight,
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.2.downsamplers.0.time_conv.conv3d.bias",
                from_pattern=["encoder.down_blocks.8.time_conv.bias"],
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.3.resnets.0.conv1.conv3d.weight",
                from_pattern=["encoder.down_blocks.9.conv1.weight"],
                transform=WeightTransforms.transpose_conv3d_weight,
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.3.resnets.0.conv1.conv3d.bias",
                from_pattern=["encoder.down_blocks.9.conv1.bias"],
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.3.resnets.0.conv2.conv3d.weight",
                from_pattern=["encoder.down_blocks.9.conv2.weight"],
                transform=WeightTransforms.transpose_conv3d_weight,
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.3.resnets.0.conv2.conv3d.bias",
                from_pattern=["encoder.down_blocks.9.conv2.bias"],
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.3.resnets.0.norm1.weight",
                from_pattern=["encoder.down_blocks.9.norm1.gamma"],
                transform=WeightTransforms.reshape_gamma_to_1d,
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.3.resnets.0.norm2.weight",
                from_pattern=["encoder.down_blocks.9.norm2.gamma"],
                transform=WeightTransforms.reshape_gamma_to_1d,
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.3.resnets.1.conv1.conv3d.weight",
                from_pattern=["encoder.down_blocks.10.conv1.weight"],
                transform=WeightTransforms.transpose_conv3d_weight,
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.3.resnets.1.conv1.conv3d.bias",
                from_pattern=["encoder.down_blocks.10.conv1.bias"],
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.3.resnets.1.conv2.conv3d.weight",
                from_pattern=["encoder.down_blocks.10.conv2.weight"],
                transform=WeightTransforms.transpose_conv3d_weight,
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.3.resnets.1.conv2.conv3d.bias",
                from_pattern=["encoder.down_blocks.10.conv2.bias"],
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.3.resnets.1.norm1.weight",
                from_pattern=["encoder.down_blocks.10.norm1.gamma"],
                transform=WeightTransforms.reshape_gamma_to_1d,
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.3.resnets.1.norm2.weight",
                from_pattern=["encoder.down_blocks.10.norm2.gamma"],
                transform=WeightTransforms.reshape_gamma_to_1d,
            ),
            WeightTarget(
                to_pattern="quant_conv.conv3d.weight",
                from_pattern=["quant_conv.weight"],
                transform=WeightTransforms.transpose_conv3d_weight,
            ),
            WeightTarget(
                to_pattern="quant_conv.conv3d.bias",
                from_pattern=["quant_conv.bias"],
            ),
        ]

    @staticmethod
    def get_text_encoder_mapping() -> List[WeightTarget]:
        return [
            WeightTarget(
                to_pattern="encoder.embed_tokens.weight",
                from_pattern=["model.embed_tokens.weight"],
            ),
            WeightTarget(
                to_pattern="encoder.norm.weight",
                from_pattern=["model.norm.weight"],
            ),
            WeightTarget(
                to_pattern="encoder.layers.{layer}.input_layernorm.weight",
                from_pattern=["model.layers.{layer}.input_layernorm.weight"],
            ),
            WeightTarget(
                to_pattern="encoder.layers.{layer}.post_attention_layernorm.weight",
                from_pattern=["model.layers.{layer}.post_attention_layernorm.weight"],
            ),
            WeightTarget(
                to_pattern="encoder.layers.{layer}.self_attn.q_proj.weight",
                from_pattern=["model.layers.{layer}.self_attn.q_proj.weight"],
            ),
            WeightTarget(
                to_pattern="encoder.layers.{layer}.self_attn.q_proj.bias",
                from_pattern=["model.layers.{layer}.self_attn.q_proj.bias"],
            ),
            WeightTarget(
                to_pattern="encoder.layers.{layer}.self_attn.k_proj.weight",
                from_pattern=["model.layers.{layer}.self_attn.k_proj.weight"],
            ),
            WeightTarget(
                to_pattern="encoder.layers.{layer}.self_attn.k_proj.bias",
                from_pattern=["model.layers.{layer}.self_attn.k_proj.bias"],
            ),
            WeightTarget(
                to_pattern="encoder.layers.{layer}.self_attn.v_proj.weight",
                from_pattern=["model.layers.{layer}.self_attn.v_proj.weight"],
            ),
            WeightTarget(
                to_pattern="encoder.layers.{layer}.self_attn.v_proj.bias",
                from_pattern=["model.layers.{layer}.self_attn.v_proj.bias"],
            ),
            WeightTarget(
                to_pattern="encoder.layers.{layer}.self_attn.o_proj.weight",
                from_pattern=["model.layers.{layer}.self_attn.o_proj.weight"],
            ),
            # MLP
            WeightTarget(
                to_pattern="encoder.layers.{layer}.mlp.gate_proj.weight",
                from_pattern=["model.layers.{layer}.mlp.gate_proj.weight"],
            ),
            WeightTarget(
                to_pattern="encoder.layers.{layer}.mlp.up_proj.weight",
                from_pattern=["model.layers.{layer}.mlp.up_proj.weight"],
            ),
            WeightTarget(
                to_pattern="encoder.layers.{layer}.mlp.down_proj.weight",
                from_pattern=["model.layers.{layer}.mlp.down_proj.weight"],
            ),
            WeightTarget(
                to_pattern="encoder.visual.patch_embed.proj.weight",
                from_pattern=["visual.patch_embed.proj.weight"],
                transform=WeightTransforms.transpose_patch_embed,
                required=False,
            ),
            WeightTarget(
                to_pattern="encoder.visual.blocks.{block}.attn.qkv.weight",
                from_pattern=["visual.blocks.{block}.attn.qkv.weight"],
                required=False,
            ),
            WeightTarget(
                to_pattern="encoder.visual.blocks.{block}.attn.qkv.bias",
                from_pattern=["visual.blocks.{block}.attn.qkv.bias"],
                required=False,
            ),
            WeightTarget(
                to_pattern="encoder.visual.blocks.{block}.attn.proj.weight",
                from_pattern=["visual.blocks.{block}.attn.proj.weight"],
                required=False,
            ),
            WeightTarget(
                to_pattern="encoder.visual.blocks.{block}.attn.proj.bias",
                from_pattern=["visual.blocks.{block}.attn.proj.bias"],
                required=False,
            ),
            WeightTarget(
                to_pattern="encoder.visual.blocks.{block}.mlp.gate_proj.weight",
                from_pattern=["visual.blocks.{block}.mlp.gate_proj.weight"],
                required=False,
            ),
            WeightTarget(
                to_pattern="encoder.visual.blocks.{block}.mlp.gate_proj.bias",
                from_pattern=["visual.blocks.{block}.mlp.gate_proj.bias"],
                required=False,
            ),
            WeightTarget(
                to_pattern="encoder.visual.blocks.{block}.mlp.up_proj.weight",
                from_pattern=["visual.blocks.{block}.mlp.up_proj.weight"],
                required=False,
            ),
            WeightTarget(
                to_pattern="encoder.visual.blocks.{block}.mlp.up_proj.bias",
                from_pattern=["visual.blocks.{block}.mlp.up_proj.bias"],
                required=False,
            ),
            WeightTarget(
                to_pattern="encoder.visual.blocks.{block}.mlp.down_proj.weight",
                from_pattern=["visual.blocks.{block}.mlp.down_proj.weight"],
                required=False,
            ),
            WeightTarget(
                to_pattern="encoder.visual.blocks.{block}.mlp.down_proj.bias",
                from_pattern=["visual.blocks.{block}.mlp.down_proj.bias"],
                required=False,
            ),
            WeightTarget(
                to_pattern="encoder.visual.blocks.{block}.norm1.weight",
                from_pattern=["visual.blocks.{block}.norm1.weight"],
                required=False,
            ),
            WeightTarget(
                to_pattern="encoder.visual.blocks.{block}.norm2.weight",
                from_pattern=["visual.blocks.{block}.norm2.weight"],
                required=False,
            ),
            WeightTarget(
                to_pattern="encoder.visual.merger.ln_q.weight",
                from_pattern=["visual.merger.ln_q.weight"],
                required=False,
            ),
            WeightTarget(
                to_pattern="encoder.visual.merger.mlp_0.weight",
                from_pattern=["visual.merger.mlp.0.weight"],
                required=False,
            ),
            WeightTarget(
                to_pattern="encoder.visual.merger.mlp_0.bias",
                from_pattern=["visual.merger.mlp.0.bias"],
                required=False,
            ),
            WeightTarget(
                to_pattern="encoder.visual.merger.mlp_1.weight",
                from_pattern=["visual.merger.mlp.2.weight"],
                required=False,
            ),
            WeightTarget(
                to_pattern="encoder.visual.merger.mlp_1.bias",
                from_pattern=["visual.merger.mlp.2.bias"],
                required=False,
            ),
        ]

    @staticmethod
    def get_mapping() -> List[WeightTarget]:
        return QwenWeightMapping.get_transformer_mapping()
