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 FluxWeightMapping(WeightMapping):
    @staticmethod
    def get_transformer_mapping() -> List[WeightTarget]:
        return [
            WeightTarget(
                to_pattern="x_embedder.weight",
                from_pattern=["x_embedder.weight"],
            ),
            WeightTarget(
                to_pattern="x_embedder.bias",
                from_pattern=["x_embedder.bias"],
            ),
            WeightTarget(
                to_pattern="context_embedder.weight",
                from_pattern=["context_embedder.weight"],
            ),
            WeightTarget(
                to_pattern="context_embedder.bias",
                from_pattern=["context_embedder.bias"],
            ),
            WeightTarget(
                to_pattern="proj_out.weight",
                from_pattern=["proj_out.weight"],
            ),
            WeightTarget(
                to_pattern="proj_out.bias",
                from_pattern=["proj_out.bias"],
            ),
            # TimeTextEmbed
            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="time_text_embed.text_embedder.linear_1.weight",
                from_pattern=["time_text_embed.text_embedder.linear_1.weight"],
            ),
            WeightTarget(
                to_pattern="time_text_embed.text_embedder.linear_1.bias",
                from_pattern=["time_text_embed.text_embedder.linear_1.bias"],
            ),
            WeightTarget(
                to_pattern="time_text_embed.text_embedder.linear_2.weight",
                from_pattern=["time_text_embed.text_embedder.linear_2.weight"],
            ),
            WeightTarget(
                to_pattern="time_text_embed.text_embedder.linear_2.bias",
                from_pattern=["time_text_embed.text_embedder.linear_2.bias"],
            ),
            WeightTarget(
                to_pattern="time_text_embed.guidance_embedder.linear_1.weight",
                from_pattern=["time_text_embed.guidance_embedder.linear_1.weight"],
                required=False,
            ),
            WeightTarget(
                to_pattern="time_text_embed.guidance_embedder.linear_1.bias",
                from_pattern=["time_text_embed.guidance_embedder.linear_1.bias"],
                required=False,
            ),
            WeightTarget(
                to_pattern="time_text_embed.guidance_embedder.linear_2.weight",
                from_pattern=["time_text_embed.guidance_embedder.linear_2.weight"],
                required=False,
            ),
            WeightTarget(
                to_pattern="time_text_embed.guidance_embedder.linear_2.bias",
                from_pattern=["time_text_embed.guidance_embedder.linear_2.bias"],
                required=False,
            ),
            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="transformer_blocks.{block}.norm1.linear.weight",
                from_pattern=["transformer_blocks.{block}.norm1.linear.weight"],
            ),
            WeightTarget(
                to_pattern="transformer_blocks.{block}.norm1.linear.bias",
                from_pattern=["transformer_blocks.{block}.norm1.linear.bias"],
            ),
            WeightTarget(
                to_pattern="transformer_blocks.{block}.norm1_context.linear.weight",
                from_pattern=["transformer_blocks.{block}.norm1_context.linear.weight"],
            ),
            WeightTarget(
                to_pattern="transformer_blocks.{block}.norm1_context.linear.bias",
                from_pattern=["transformer_blocks.{block}.norm1_context.linear.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.to_out.0.weight",
                from_pattern=["transformer_blocks.{block}.attn.to_out.0.weight"],
            ),
            WeightTarget(
                to_pattern="transformer_blocks.{block}.attn.to_out.0.bias",
                from_pattern=["transformer_blocks.{block}.attn.to_out.0.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.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}.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}.ff.linear1.weight",
                from_pattern=["transformer_blocks.{block}.ff.net.0.proj.weight"],
            ),
            WeightTarget(
                to_pattern="transformer_blocks.{block}.ff.linear1.bias",
                from_pattern=["transformer_blocks.{block}.ff.net.0.proj.bias"],
            ),
            WeightTarget(
                to_pattern="transformer_blocks.{block}.ff.linear2.weight",
                from_pattern=["transformer_blocks.{block}.ff.net.2.weight"],
            ),
            WeightTarget(
                to_pattern="transformer_blocks.{block}.ff.linear2.bias",
                from_pattern=["transformer_blocks.{block}.ff.net.2.bias"],
            ),
            WeightTarget(
                to_pattern="transformer_blocks.{block}.ff_context.linear1.weight",
                from_pattern=["transformer_blocks.{block}.ff_context.net.0.proj.weight"],
            ),
            WeightTarget(
                to_pattern="transformer_blocks.{block}.ff_context.linear1.bias",
                from_pattern=["transformer_blocks.{block}.ff_context.net.0.proj.bias"],
            ),
            WeightTarget(
                to_pattern="transformer_blocks.{block}.ff_context.linear2.weight",
                from_pattern=["transformer_blocks.{block}.ff_context.net.2.weight"],
            ),
            WeightTarget(
                to_pattern="transformer_blocks.{block}.ff_context.linear2.bias",
                from_pattern=["transformer_blocks.{block}.ff_context.net.2.bias"],
            ),
            WeightTarget(
                to_pattern="single_transformer_blocks.{block}.norm.linear.weight",
                from_pattern=["single_transformer_blocks.{block}.norm.linear.weight"],
                max_blocks=38,
            ),
            WeightTarget(
                to_pattern="single_transformer_blocks.{block}.norm.linear.bias",
                from_pattern=["single_transformer_blocks.{block}.norm.linear.bias"],
                max_blocks=38,
            ),
            WeightTarget(
                to_pattern="single_transformer_blocks.{block}.attn.to_q.weight",
                from_pattern=["single_transformer_blocks.{block}.attn.to_q.weight"],
                max_blocks=38,
            ),
            WeightTarget(
                to_pattern="single_transformer_blocks.{block}.attn.to_q.bias",
                from_pattern=["single_transformer_blocks.{block}.attn.to_q.bias"],
                max_blocks=38,
            ),
            WeightTarget(
                to_pattern="single_transformer_blocks.{block}.attn.to_k.weight",
                from_pattern=["single_transformer_blocks.{block}.attn.to_k.weight"],
                max_blocks=38,
            ),
            WeightTarget(
                to_pattern="single_transformer_blocks.{block}.attn.to_k.bias",
                from_pattern=["single_transformer_blocks.{block}.attn.to_k.bias"],
                max_blocks=38,
            ),
            WeightTarget(
                to_pattern="single_transformer_blocks.{block}.attn.to_v.weight",
                from_pattern=["single_transformer_blocks.{block}.attn.to_v.weight"],
                max_blocks=38,
            ),
            WeightTarget(
                to_pattern="single_transformer_blocks.{block}.attn.to_v.bias",
                from_pattern=["single_transformer_blocks.{block}.attn.to_v.bias"],
                max_blocks=38,
            ),
            WeightTarget(
                to_pattern="single_transformer_blocks.{block}.attn.norm_q.weight",
                from_pattern=["single_transformer_blocks.{block}.attn.norm_q.weight"],
                max_blocks=38,
            ),
            WeightTarget(
                to_pattern="single_transformer_blocks.{block}.attn.norm_k.weight",
                from_pattern=["single_transformer_blocks.{block}.attn.norm_k.weight"],
                max_blocks=38,
            ),
            WeightTarget(
                to_pattern="single_transformer_blocks.{block}.proj_mlp.weight",
                from_pattern=["single_transformer_blocks.{block}.proj_mlp.weight"],
                max_blocks=38,
            ),
            WeightTarget(
                to_pattern="single_transformer_blocks.{block}.proj_mlp.bias",
                from_pattern=["single_transformer_blocks.{block}.proj_mlp.bias"],
                max_blocks=38,
            ),
            WeightTarget(
                to_pattern="single_transformer_blocks.{block}.proj_out.weight",
                from_pattern=["single_transformer_blocks.{block}.proj_out.weight"],
                max_blocks=38,
            ),
            WeightTarget(
                to_pattern="single_transformer_blocks.{block}.proj_out.bias",
                from_pattern=["single_transformer_blocks.{block}.proj_out.bias"],
                max_blocks=38,
            ),
        ]

    @staticmethod
    def get_controlnet_transformer_mapping() -> List[WeightTarget]:
        return FluxWeightMapping.get_transformer_mapping() + [
            WeightTarget(
                to_pattern="controlnet_x_embedder.weight",
                from_pattern=["controlnet_x_embedder.weight"],
            ),
            WeightTarget(
                to_pattern="controlnet_x_embedder.bias",
                from_pattern=["controlnet_x_embedder.bias"],
            ),
            WeightTarget(
                to_pattern="controlnet_blocks.{block}.weight",
                from_pattern=["controlnet_blocks.{block}.weight"],
                max_blocks=19,
            ),
            WeightTarget(
                to_pattern="controlnet_blocks.{block}.bias",
                from_pattern=["controlnet_blocks.{block}.bias"],
                max_blocks=19,
            ),
            WeightTarget(
                to_pattern="controlnet_single_blocks.{block}.weight",
                from_pattern=["controlnet_single_blocks.{block}.weight"],
                max_blocks=38,
            ),
            WeightTarget(
                to_pattern="controlnet_single_blocks.{block}.bias",
                from_pattern=["controlnet_single_blocks.{block}.bias"],
                max_blocks=38,
            ),
        ]

    @staticmethod
    def get_vae_mapping() -> List[WeightTarget]:
        return [
            WeightTarget(
                to_pattern="decoder.conv_in.conv2d.weight",
                from_pattern=["decoder.conv_in.weight"],
                transform=WeightTransforms.transpose_conv2d_weight,
            ),
            WeightTarget(
                to_pattern="decoder.conv_in.conv2d.bias",
                from_pattern=["decoder.conv_in.bias"],
            ),
            WeightTarget(
                to_pattern="decoder.conv_out.conv2d.weight",
                from_pattern=["decoder.conv_out.weight"],
                transform=WeightTransforms.transpose_conv2d_weight,
            ),
            WeightTarget(
                to_pattern="decoder.conv_out.conv2d.bias",
                from_pattern=["decoder.conv_out.bias"],
            ),
            WeightTarget(
                to_pattern="decoder.conv_norm_out.norm.weight",
                from_pattern=["decoder.conv_norm_out.weight"],
            ),
            WeightTarget(
                to_pattern="decoder.conv_norm_out.norm.bias",
                from_pattern=["decoder.conv_norm_out.bias"],
            ),
            WeightTarget(
                to_pattern="encoder.conv_in.conv2d.weight",
                from_pattern=["encoder.conv_in.weight"],
                transform=WeightTransforms.transpose_conv2d_weight,
            ),
            WeightTarget(
                to_pattern="encoder.conv_in.conv2d.bias",
                from_pattern=["encoder.conv_in.bias"],
            ),
            WeightTarget(
                to_pattern="encoder.conv_out.conv2d.weight",
                from_pattern=["encoder.conv_out.weight"],
                transform=WeightTransforms.transpose_conv2d_weight,
            ),
            WeightTarget(
                to_pattern="encoder.conv_out.conv2d.bias",
                from_pattern=["encoder.conv_out.bias"],
            ),
            WeightTarget(
                to_pattern="encoder.conv_norm_out.norm.weight",
                from_pattern=["encoder.conv_norm_out.weight"],
            ),
            WeightTarget(
                to_pattern="encoder.conv_norm_out.norm.bias",
                from_pattern=["encoder.conv_norm_out.bias"],
            ),
            WeightTarget(
                to_pattern="decoder.mid_block.resnets.{i}.norm1.weight",
                from_pattern=["decoder.mid_block.resnets.{i}.norm1.weight"],
            ),
            WeightTarget(
                to_pattern="decoder.mid_block.resnets.{i}.norm1.bias",
                from_pattern=["decoder.mid_block.resnets.{i}.norm1.bias"],
            ),
            WeightTarget(
                to_pattern="decoder.mid_block.resnets.{i}.conv1.weight",
                from_pattern=["decoder.mid_block.resnets.{i}.conv1.weight"],
                transform=WeightTransforms.transpose_conv2d_weight,
            ),
            WeightTarget(
                to_pattern="decoder.mid_block.resnets.{i}.conv1.bias",
                from_pattern=["decoder.mid_block.resnets.{i}.conv1.bias"],
            ),
            WeightTarget(
                to_pattern="decoder.mid_block.resnets.{i}.norm2.weight",
                from_pattern=["decoder.mid_block.resnets.{i}.norm2.weight"],
            ),
            WeightTarget(
                to_pattern="decoder.mid_block.resnets.{i}.norm2.bias",
                from_pattern=["decoder.mid_block.resnets.{i}.norm2.bias"],
            ),
            WeightTarget(
                to_pattern="decoder.mid_block.resnets.{i}.conv2.weight",
                from_pattern=["decoder.mid_block.resnets.{i}.conv2.weight"],
                transform=WeightTransforms.transpose_conv2d_weight,
            ),
            WeightTarget(
                to_pattern="decoder.mid_block.resnets.{i}.conv2.bias",
                from_pattern=["decoder.mid_block.resnets.{i}.conv2.bias"],
            ),
            WeightTarget(
                to_pattern="encoder.mid_block.resnets.{i}.norm1.weight",
                from_pattern=["encoder.mid_block.resnets.{i}.norm1.weight"],
            ),
            WeightTarget(
                to_pattern="encoder.mid_block.resnets.{i}.norm1.bias",
                from_pattern=["encoder.mid_block.resnets.{i}.norm1.bias"],
            ),
            WeightTarget(
                to_pattern="encoder.mid_block.resnets.{i}.conv1.weight",
                from_pattern=["encoder.mid_block.resnets.{i}.conv1.weight"],
                transform=WeightTransforms.transpose_conv2d_weight,
            ),
            WeightTarget(
                to_pattern="encoder.mid_block.resnets.{i}.conv1.bias",
                from_pattern=["encoder.mid_block.resnets.{i}.conv1.bias"],
            ),
            WeightTarget(
                to_pattern="encoder.mid_block.resnets.{i}.norm2.weight",
                from_pattern=["encoder.mid_block.resnets.{i}.norm2.weight"],
            ),
            WeightTarget(
                to_pattern="encoder.mid_block.resnets.{i}.norm2.bias",
                from_pattern=["encoder.mid_block.resnets.{i}.norm2.bias"],
            ),
            WeightTarget(
                to_pattern="encoder.mid_block.resnets.{i}.conv2.weight",
                from_pattern=["encoder.mid_block.resnets.{i}.conv2.weight"],
                transform=WeightTransforms.transpose_conv2d_weight,
            ),
            WeightTarget(
                to_pattern="encoder.mid_block.resnets.{i}.conv2.bias",
                from_pattern=["encoder.mid_block.resnets.{i}.conv2.bias"],
            ),
            WeightTarget(
                to_pattern="decoder.mid_block.attentions.0.group_norm.weight",
                from_pattern=["decoder.mid_block.attentions.0.group_norm.weight"],
            ),
            WeightTarget(
                to_pattern="decoder.mid_block.attentions.0.group_norm.bias",
                from_pattern=["decoder.mid_block.attentions.0.group_norm.bias"],
            ),
            WeightTarget(
                to_pattern="decoder.mid_block.attentions.0.to_q.weight",
                from_pattern=["decoder.mid_block.attentions.0.to_q.weight"],
            ),
            WeightTarget(
                to_pattern="decoder.mid_block.attentions.0.to_q.bias",
                from_pattern=["decoder.mid_block.attentions.0.to_q.bias"],
            ),
            WeightTarget(
                to_pattern="decoder.mid_block.attentions.0.to_k.weight",
                from_pattern=["decoder.mid_block.attentions.0.to_k.weight"],
            ),
            WeightTarget(
                to_pattern="decoder.mid_block.attentions.0.to_k.bias",
                from_pattern=["decoder.mid_block.attentions.0.to_k.bias"],
            ),
            WeightTarget(
                to_pattern="decoder.mid_block.attentions.0.to_v.weight",
                from_pattern=["decoder.mid_block.attentions.0.to_v.weight"],
            ),
            WeightTarget(
                to_pattern="decoder.mid_block.attentions.0.to_v.bias",
                from_pattern=["decoder.mid_block.attentions.0.to_v.bias"],
            ),
            WeightTarget(
                to_pattern="decoder.mid_block.attentions.0.to_out.0.weight",
                from_pattern=["decoder.mid_block.attentions.0.to_out.0.weight"],
            ),
            WeightTarget(
                to_pattern="decoder.mid_block.attentions.0.to_out.0.bias",
                from_pattern=["decoder.mid_block.attentions.0.to_out.0.bias"],
            ),
            WeightTarget(
                to_pattern="encoder.mid_block.attentions.0.group_norm.weight",
                from_pattern=["encoder.mid_block.attentions.0.group_norm.weight"],
            ),
            WeightTarget(
                to_pattern="encoder.mid_block.attentions.0.group_norm.bias",
                from_pattern=["encoder.mid_block.attentions.0.group_norm.bias"],
            ),
            WeightTarget(
                to_pattern="encoder.mid_block.attentions.0.to_q.weight",
                from_pattern=["encoder.mid_block.attentions.0.to_q.weight"],
            ),
            WeightTarget(
                to_pattern="encoder.mid_block.attentions.0.to_q.bias",
                from_pattern=["encoder.mid_block.attentions.0.to_q.bias"],
            ),
            WeightTarget(
                to_pattern="encoder.mid_block.attentions.0.to_k.weight",
                from_pattern=["encoder.mid_block.attentions.0.to_k.weight"],
            ),
            WeightTarget(
                to_pattern="encoder.mid_block.attentions.0.to_k.bias",
                from_pattern=["encoder.mid_block.attentions.0.to_k.bias"],
            ),
            WeightTarget(
                to_pattern="encoder.mid_block.attentions.0.to_v.weight",
                from_pattern=["encoder.mid_block.attentions.0.to_v.weight"],
            ),
            WeightTarget(
                to_pattern="encoder.mid_block.attentions.0.to_v.bias",
                from_pattern=["encoder.mid_block.attentions.0.to_v.bias"],
            ),
            WeightTarget(
                to_pattern="encoder.mid_block.attentions.0.to_out.0.weight",
                from_pattern=["encoder.mid_block.attentions.0.to_out.0.weight"],
            ),
            WeightTarget(
                to_pattern="encoder.mid_block.attentions.0.to_out.0.bias",
                from_pattern=["encoder.mid_block.attentions.0.to_out.0.bias"],
            ),
            WeightTarget(
                to_pattern="decoder.up_blocks.{block}.resnets.{res}.norm1.weight",
                from_pattern=["decoder.up_blocks.{block}.resnets.{res}.norm1.weight"],
            ),
            WeightTarget(
                to_pattern="decoder.up_blocks.{block}.resnets.{res}.norm1.bias",
                from_pattern=["decoder.up_blocks.{block}.resnets.{res}.norm1.bias"],
            ),
            WeightTarget(
                to_pattern="decoder.up_blocks.{block}.resnets.{res}.conv1.weight",
                from_pattern=["decoder.up_blocks.{block}.resnets.{res}.conv1.weight"],
                transform=WeightTransforms.transpose_conv2d_weight,
            ),
            WeightTarget(
                to_pattern="decoder.up_blocks.{block}.resnets.{res}.conv1.bias",
                from_pattern=["decoder.up_blocks.{block}.resnets.{res}.conv1.bias"],
            ),
            WeightTarget(
                to_pattern="decoder.up_blocks.{block}.resnets.{res}.norm2.weight",
                from_pattern=["decoder.up_blocks.{block}.resnets.{res}.norm2.weight"],
            ),
            WeightTarget(
                to_pattern="decoder.up_blocks.{block}.resnets.{res}.norm2.bias",
                from_pattern=["decoder.up_blocks.{block}.resnets.{res}.norm2.bias"],
            ),
            WeightTarget(
                to_pattern="decoder.up_blocks.{block}.resnets.{res}.conv2.weight",
                from_pattern=["decoder.up_blocks.{block}.resnets.{res}.conv2.weight"],
                transform=WeightTransforms.transpose_conv2d_weight,
            ),
            WeightTarget(
                to_pattern="decoder.up_blocks.{block}.resnets.{res}.conv2.bias",
                from_pattern=["decoder.up_blocks.{block}.resnets.{res}.conv2.bias"],
            ),
            WeightTarget(
                to_pattern="decoder.up_blocks.{block}.resnets.{res}.conv_shortcut.weight",
                from_pattern=["decoder.up_blocks.{block}.resnets.{res}.conv_shortcut.weight"],
                transform=WeightTransforms.transpose_conv2d_weight,
                required=False,
            ),
            WeightTarget(
                to_pattern="decoder.up_blocks.{block}.resnets.{res}.conv_shortcut.bias",
                from_pattern=["decoder.up_blocks.{block}.resnets.{res}.conv_shortcut.bias"],
                required=False,
            ),
            WeightTarget(
                to_pattern="decoder.up_blocks.{block}.upsamplers.0.conv.weight",
                from_pattern=["decoder.up_blocks.{block}.upsamplers.0.conv.weight"],
                transform=WeightTransforms.transpose_conv2d_weight,
                required=False,
            ),
            WeightTarget(
                to_pattern="decoder.up_blocks.{block}.upsamplers.0.conv.bias",
                from_pattern=["decoder.up_blocks.{block}.upsamplers.0.conv.bias"],
                required=False,
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.{block}.resnets.{res}.norm1.weight",
                from_pattern=["encoder.down_blocks.{block}.resnets.{res}.norm1.weight"],
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.{block}.resnets.{res}.norm1.bias",
                from_pattern=["encoder.down_blocks.{block}.resnets.{res}.norm1.bias"],
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.{block}.resnets.{res}.conv1.weight",
                from_pattern=["encoder.down_blocks.{block}.resnets.{res}.conv1.weight"],
                transform=WeightTransforms.transpose_conv2d_weight,
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.{block}.resnets.{res}.conv1.bias",
                from_pattern=["encoder.down_blocks.{block}.resnets.{res}.conv1.bias"],
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.{block}.resnets.{res}.norm2.weight",
                from_pattern=["encoder.down_blocks.{block}.resnets.{res}.norm2.weight"],
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.{block}.resnets.{res}.norm2.bias",
                from_pattern=["encoder.down_blocks.{block}.resnets.{res}.norm2.bias"],
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.{block}.resnets.{res}.conv2.weight",
                from_pattern=["encoder.down_blocks.{block}.resnets.{res}.conv2.weight"],
                transform=WeightTransforms.transpose_conv2d_weight,
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.{block}.resnets.{res}.conv2.bias",
                from_pattern=["encoder.down_blocks.{block}.resnets.{res}.conv2.bias"],
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.{block}.resnets.{res}.conv_shortcut.weight",
                from_pattern=["encoder.down_blocks.{block}.resnets.{res}.conv_shortcut.weight"],
                transform=WeightTransforms.transpose_conv2d_weight,
                required=False,
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.{block}.resnets.{res}.conv_shortcut.bias",
                from_pattern=["encoder.down_blocks.{block}.resnets.{res}.conv_shortcut.bias"],
                required=False,
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.{block}.downsamplers.0.conv.weight",
                from_pattern=["encoder.down_blocks.{block}.downsamplers.0.conv.weight"],
                transform=WeightTransforms.transpose_conv2d_weight,
                required=False,
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.{block}.downsamplers.0.conv.bias",
                from_pattern=["encoder.down_blocks.{block}.downsamplers.0.conv.bias"],
                required=False,
            ),
            WeightTarget(
                to_pattern="quant_conv.weight",
                from_pattern=["quant_conv.weight"],
                transform=WeightTransforms.transpose_conv2d_weight,
            ),
            WeightTarget(
                to_pattern="quant_conv.bias",
                from_pattern=["quant_conv.bias"],
            ),
            WeightTarget(
                to_pattern="post_quant_conv.weight",
                from_pattern=["post_quant_conv.weight"],
                transform=WeightTransforms.transpose_conv2d_weight,
            ),
            WeightTarget(
                to_pattern="post_quant_conv.bias",
                from_pattern=["post_quant_conv.bias"],
            ),
        ]

    @staticmethod
    def get_t5_encoder_mapping() -> List[WeightTarget]:
        return [
            WeightTarget(
                to_pattern="shared.weight",
                from_pattern=["shared.weight"],
            ),
            WeightTarget(
                to_pattern="final_layer_norm.weight",
                from_pattern=["encoder.final_layer_norm.weight"],
            ),
            WeightTarget(
                to_pattern="t5_blocks.{block}.attention.layer_norm.weight",
                from_pattern=["encoder.block.{block}.layer.0.layer_norm.weight"],
            ),
            WeightTarget(
                to_pattern="t5_blocks.{block}.attention.SelfAttention.q.weight",
                from_pattern=["encoder.block.{block}.layer.0.SelfAttention.q.weight"],
            ),
            WeightTarget(
                to_pattern="t5_blocks.{block}.attention.SelfAttention.k.weight",
                from_pattern=["encoder.block.{block}.layer.0.SelfAttention.k.weight"],
            ),
            WeightTarget(
                to_pattern="t5_blocks.{block}.attention.SelfAttention.v.weight",
                from_pattern=["encoder.block.{block}.layer.0.SelfAttention.v.weight"],
            ),
            WeightTarget(
                to_pattern="t5_blocks.{block}.attention.SelfAttention.o.weight",
                from_pattern=["encoder.block.{block}.layer.0.SelfAttention.o.weight"],
            ),
            WeightTarget(
                to_pattern="t5_blocks.{block}.attention.SelfAttention.relative_attention_bias.weight",
                from_pattern=["encoder.block.0.layer.0.SelfAttention.relative_attention_bias.weight"],
                max_blocks=24,
            ),
            WeightTarget(
                to_pattern="t5_blocks.{block}.ff.layer_norm.weight",
                from_pattern=["encoder.block.{block}.layer.1.layer_norm.weight"],
            ),
            WeightTarget(
                to_pattern="t5_blocks.{block}.ff.DenseReluDense.wi_0.weight",
                from_pattern=["encoder.block.{block}.layer.1.DenseReluDense.wi_0.weight"],
            ),
            WeightTarget(
                to_pattern="t5_blocks.{block}.ff.DenseReluDense.wi_1.weight",
                from_pattern=["encoder.block.{block}.layer.1.DenseReluDense.wi_1.weight"],
            ),
            WeightTarget(
                to_pattern="t5_blocks.{block}.ff.DenseReluDense.wo.weight",
                from_pattern=["encoder.block.{block}.layer.1.DenseReluDense.wo.weight"],
            ),
        ]

    @staticmethod
    def get_clip_encoder_mapping() -> List[WeightTarget]:
        return [
            WeightTarget(
                to_pattern="text_model.embeddings.token_embedding.weight",
                from_pattern=["text_model.embeddings.token_embedding.weight"],
            ),
            WeightTarget(
                to_pattern="text_model.embeddings.position_embedding.weight",
                from_pattern=["text_model.embeddings.position_embedding.weight"],
            ),
            WeightTarget(
                to_pattern="text_model.encoder.layers.{block}.self_attn.q_proj.weight",
                from_pattern=["text_model.encoder.layers.{block}.self_attn.q_proj.weight"],
                max_blocks=12,
            ),
            WeightTarget(
                to_pattern="text_model.encoder.layers.{block}.self_attn.q_proj.bias",
                from_pattern=["text_model.encoder.layers.{block}.self_attn.q_proj.bias"],
                max_blocks=12,
            ),
            WeightTarget(
                to_pattern="text_model.encoder.layers.{block}.self_attn.k_proj.weight",
                from_pattern=["text_model.encoder.layers.{block}.self_attn.k_proj.weight"],
                max_blocks=12,
            ),
            WeightTarget(
                to_pattern="text_model.encoder.layers.{block}.self_attn.k_proj.bias",
                from_pattern=["text_model.encoder.layers.{block}.self_attn.k_proj.bias"],
                max_blocks=12,
            ),
            WeightTarget(
                to_pattern="text_model.encoder.layers.{block}.self_attn.v_proj.weight",
                from_pattern=["text_model.encoder.layers.{block}.self_attn.v_proj.weight"],
                max_blocks=12,
            ),
            WeightTarget(
                to_pattern="text_model.encoder.layers.{block}.self_attn.v_proj.bias",
                from_pattern=["text_model.encoder.layers.{block}.self_attn.v_proj.bias"],
                max_blocks=12,
            ),
            WeightTarget(
                to_pattern="text_model.encoder.layers.{block}.self_attn.out_proj.weight",
                from_pattern=["text_model.encoder.layers.{block}.self_attn.out_proj.weight"],
                max_blocks=12,
            ),
            WeightTarget(
                to_pattern="text_model.encoder.layers.{block}.self_attn.out_proj.bias",
                from_pattern=["text_model.encoder.layers.{block}.self_attn.out_proj.bias"],
                max_blocks=12,
            ),
            WeightTarget(
                to_pattern="text_model.encoder.layers.{block}.layer_norm1.weight",
                from_pattern=["text_model.encoder.layers.{block}.layer_norm1.weight"],
                max_blocks=12,
            ),
            WeightTarget(
                to_pattern="text_model.encoder.layers.{block}.layer_norm1.bias",
                from_pattern=["text_model.encoder.layers.{block}.layer_norm1.bias"],
                max_blocks=12,
            ),
            WeightTarget(
                to_pattern="text_model.encoder.layers.{block}.mlp.fc1.weight",
                from_pattern=["text_model.encoder.layers.{block}.mlp.fc1.weight"],
                max_blocks=12,
            ),
            WeightTarget(
                to_pattern="text_model.encoder.layers.{block}.mlp.fc1.bias",
                from_pattern=["text_model.encoder.layers.{block}.mlp.fc1.bias"],
                max_blocks=12,
            ),
            WeightTarget(
                to_pattern="text_model.encoder.layers.{block}.mlp.fc2.weight",
                from_pattern=["text_model.encoder.layers.{block}.mlp.fc2.weight"],
                max_blocks=12,
            ),
            WeightTarget(
                to_pattern="text_model.encoder.layers.{block}.mlp.fc2.bias",
                from_pattern=["text_model.encoder.layers.{block}.mlp.fc2.bias"],
                max_blocks=12,
            ),
            WeightTarget(
                to_pattern="text_model.encoder.layers.{block}.layer_norm2.weight",
                from_pattern=["text_model.encoder.layers.{block}.layer_norm2.weight"],
                max_blocks=12,
            ),
            WeightTarget(
                to_pattern="text_model.encoder.layers.{block}.layer_norm2.bias",
                from_pattern=["text_model.encoder.layers.{block}.layer_norm2.bias"],
                max_blocks=12,
            ),
            WeightTarget(
                to_pattern="text_model.final_layer_norm.weight",
                from_pattern=["text_model.final_layer_norm.weight"],
            ),
            WeightTarget(
                to_pattern="text_model.final_layer_norm.bias",
                from_pattern=["text_model.final_layer_norm.bias"],
            ),
        ]
