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 FIBOWeightMapping(WeightMapping):
    @staticmethod
    def get_transformer_mapping() -> List[WeightTarget]:
        return [
            WeightTarget(
                to_pattern="time_embed.timestep_embedder.linear_1.weight",
                from_pattern=["time_embed.timestep_embedder.linear_1.weight"],
            ),
            WeightTarget(
                to_pattern="time_embed.timestep_embedder.linear_1.bias",
                from_pattern=["time_embed.timestep_embedder.linear_1.bias"],
            ),
            WeightTarget(
                to_pattern="time_embed.timestep_embedder.linear_2.weight",
                from_pattern=["time_embed.timestep_embedder.linear_2.weight"],
            ),
            WeightTarget(
                to_pattern="time_embed.timestep_embedder.linear_2.bias",
                from_pattern=["time_embed.timestep_embedder.linear_2.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="x_embedder.weight",
                from_pattern=["x_embedder.weight"],
            ),
            WeightTarget(
                to_pattern="x_embedder.bias",
                from_pattern=["x_embedder.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.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.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.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.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}.norm2.weight",
                from_pattern=["transformer_blocks.{block}.norm2.weight"],
            ),
            WeightTarget(
                to_pattern="transformer_blocks.{block}.norm2.bias",
                from_pattern=["transformer_blocks.{block}.norm2.bias"],
            ),
            WeightTarget(
                to_pattern="transformer_blocks.{block}.ff.net.0.proj.weight",
                from_pattern=["transformer_blocks.{block}.ff.net.0.proj.weight"],
            ),
            WeightTarget(
                to_pattern="transformer_blocks.{block}.ff.net.0.proj.bias",
                from_pattern=["transformer_blocks.{block}.ff.net.0.proj.bias"],
            ),
            WeightTarget(
                to_pattern="transformer_blocks.{block}.ff.net.2.weight",
                from_pattern=["transformer_blocks.{block}.ff.net.2.weight"],
            ),
            WeightTarget(
                to_pattern="transformer_blocks.{block}.ff.net.2.bias",
                from_pattern=["transformer_blocks.{block}.ff.net.2.bias"],
            ),
            WeightTarget(
                to_pattern="transformer_blocks.{block}.norm2_context.weight",
                from_pattern=["transformer_blocks.{block}.norm2_context.weight"],
            ),
            WeightTarget(
                to_pattern="transformer_blocks.{block}.norm2_context.bias",
                from_pattern=["transformer_blocks.{block}.norm2_context.bias"],
            ),
            WeightTarget(
                to_pattern="transformer_blocks.{block}.ff_context.net.0.proj.weight",
                from_pattern=["transformer_blocks.{block}.ff_context.net.0.proj.weight"],
            ),
            WeightTarget(
                to_pattern="transformer_blocks.{block}.ff_context.net.0.proj.bias",
                from_pattern=["transformer_blocks.{block}.ff_context.net.0.proj.bias"],
            ),
            WeightTarget(
                to_pattern="transformer_blocks.{block}.ff_context.net.2.weight",
                from_pattern=["transformer_blocks.{block}.ff_context.net.2.weight"],
            ),
            WeightTarget(
                to_pattern="transformer_blocks.{block}.ff_context.net.2.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"],
            ),
            WeightTarget(
                to_pattern="single_transformer_blocks.{block}.norm.linear.bias",
                from_pattern=["single_transformer_blocks.{block}.norm.linear.bias"],
            ),
            WeightTarget(
                to_pattern="single_transformer_blocks.{block}.attn.norm_q.weight",
                from_pattern=["single_transformer_blocks.{block}.attn.norm_q.weight"],
            ),
            WeightTarget(
                to_pattern="single_transformer_blocks.{block}.attn.norm_k.weight",
                from_pattern=["single_transformer_blocks.{block}.attn.norm_k.weight"],
            ),
            WeightTarget(
                to_pattern="single_transformer_blocks.{block}.attn.to_q.weight",
                from_pattern=["single_transformer_blocks.{block}.attn.to_q.weight"],
            ),
            WeightTarget(
                to_pattern="single_transformer_blocks.{block}.attn.to_q.bias",
                from_pattern=["single_transformer_blocks.{block}.attn.to_q.bias"],
            ),
            WeightTarget(
                to_pattern="single_transformer_blocks.{block}.attn.to_k.weight",
                from_pattern=["single_transformer_blocks.{block}.attn.to_k.weight"],
            ),
            WeightTarget(
                to_pattern="single_transformer_blocks.{block}.attn.to_k.bias",
                from_pattern=["single_transformer_blocks.{block}.attn.to_k.bias"],
            ),
            WeightTarget(
                to_pattern="single_transformer_blocks.{block}.attn.to_v.weight",
                from_pattern=["single_transformer_blocks.{block}.attn.to_v.weight"],
            ),
            WeightTarget(
                to_pattern="single_transformer_blocks.{block}.attn.to_v.bias",
                from_pattern=["single_transformer_blocks.{block}.attn.to_v.bias"],
            ),
            WeightTarget(
                to_pattern="single_transformer_blocks.{block}.proj_mlp.weight",
                from_pattern=["single_transformer_blocks.{block}.proj_mlp.weight"],
            ),
            WeightTarget(
                to_pattern="single_transformer_blocks.{block}.proj_mlp.bias",
                from_pattern=["single_transformer_blocks.{block}.proj_mlp.bias"],
            ),
            WeightTarget(
                to_pattern="single_transformer_blocks.{block}.proj_out.weight",
                from_pattern=["single_transformer_blocks.{block}.proj_out.weight"],
            ),
            WeightTarget(
                to_pattern="single_transformer_blocks.{block}.proj_out.bias",
                from_pattern=["single_transformer_blocks.{block}.proj_out.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="caption_projection.{layer}.linear.weight",
                from_pattern=["caption_projection.{layer}.linear.weight"],
            ),
        ]

    @staticmethod
    def get_text_encoder_mapping() -> List[WeightTarget]:
        return [
            WeightTarget(
                to_pattern="embed_tokens.weight",
                from_pattern=["model.embed_tokens.weight"],
            ),
            WeightTarget(
                to_pattern="layers.{block}.self_attn.q_proj.weight",
                from_pattern=["model.layers.{block}.self_attn.q_proj.weight"],
            ),
            WeightTarget(
                to_pattern="layers.{block}.self_attn.k_proj.weight",
                from_pattern=["model.layers.{block}.self_attn.k_proj.weight"],
            ),
            WeightTarget(
                to_pattern="layers.{block}.self_attn.v_proj.weight",
                from_pattern=["model.layers.{block}.self_attn.v_proj.weight"],
            ),
            WeightTarget(
                to_pattern="layers.{block}.self_attn.o_proj.weight",
                from_pattern=["model.layers.{block}.self_attn.o_proj.weight"],
            ),
            WeightTarget(
                to_pattern="layers.{block}.mlp.gate_proj.weight",
                from_pattern=["model.layers.{block}.mlp.gate_proj.weight"],
            ),
            WeightTarget(
                to_pattern="layers.{block}.mlp.up_proj.weight",
                from_pattern=["model.layers.{block}.mlp.up_proj.weight"],
            ),
            WeightTarget(
                to_pattern="layers.{block}.mlp.down_proj.weight",
                from_pattern=["model.layers.{block}.mlp.down_proj.weight"],
            ),
            WeightTarget(
                to_pattern="layers.{block}.input_layernorm.weight",
                from_pattern=["model.layers.{block}.input_layernorm.weight"],
            ),
            WeightTarget(
                to_pattern="layers.{block}.post_attention_layernorm.weight",
                from_pattern=["model.layers.{block}.post_attention_layernorm.weight"],
            ),
            WeightTarget(
                to_pattern="norm.weight",
                from_pattern=["model.norm.weight"],
            ),
        ]

    @staticmethod
    def get_vae_mapping() -> List[WeightTarget]:
        return [
            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.down_blocks.{block}.resnets.{res}.norm1.weight",
                from_pattern=["encoder.down_blocks.{block}.resnets.{res}.norm1.gamma"],
                transform=WeightTransforms.reshape_gamma_to_1d,
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.{block}.resnets.{res}.conv1.conv3d.weight",
                from_pattern=["encoder.down_blocks.{block}.resnets.{res}.conv1.weight"],
                transform=WeightTransforms.transpose_conv3d_weight,
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.{block}.resnets.{res}.conv1.conv3d.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.gamma"],
                transform=WeightTransforms.reshape_gamma_to_1d,
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.{block}.resnets.{res}.conv2.conv3d.weight",
                from_pattern=["encoder.down_blocks.{block}.resnets.{res}.conv2.weight"],
                transform=WeightTransforms.transpose_conv3d_weight,
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.{block}.resnets.{res}.conv2.conv3d.bias",
                from_pattern=["encoder.down_blocks.{block}.resnets.{res}.conv2.bias"],
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.{block}.resnets.{res}.conv_shortcut.conv3d.weight",
                from_pattern=["encoder.down_blocks.{block}.resnets.{res}.conv_shortcut.weight"],
                transform=WeightTransforms.transpose_conv3d_weight,
                required=False,
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.{block}.resnets.{res}.conv_shortcut.conv3d.bias",
                from_pattern=["encoder.down_blocks.{block}.resnets.{res}.conv_shortcut.bias"],
                required=False,
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.{block}.downsampler.resample_conv.weight",
                from_pattern=["encoder.down_blocks.{block}.downsampler.resample.1.weight"],
                transform=WeightTransforms.transpose_conv2d_weight,
                required=False,
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.{block}.downsampler.resample_conv.bias",
                from_pattern=["encoder.down_blocks.{block}.downsampler.resample.1.bias"],
                required=False,
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.{block}.downsampler.time_conv.conv3d.weight",
                from_pattern=["encoder.down_blocks.{block}.downsampler.time_conv.weight"],
                transform=WeightTransforms.transpose_conv3d_weight,
                required=False,
            ),
            WeightTarget(
                to_pattern="encoder.down_blocks.{block}.downsampler.time_conv.conv3d.bias",
                from_pattern=["encoder.down_blocks.{block}.downsampler.time_conv.bias"],
                required=False,
            ),
            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}.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}.norm2.weight",
                from_pattern=["encoder.mid_block.resnets.{i}.norm2.gamma"],
                transform=WeightTransforms.reshape_gamma_to_1d,
            ),
            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.attentions.{i}.norm.weight",
                from_pattern=["encoder.mid_block.attentions.{i}.norm.gamma"],
                transform=WeightTransforms.reshape_gamma_to_1d,
            ),
            WeightTarget(
                to_pattern="encoder.mid_block.attentions.{i}.to_qkv.weight",
                from_pattern=["encoder.mid_block.attentions.{i}.to_qkv.weight"],
                transform=WeightTransforms.transpose_conv2d_weight,
            ),
            WeightTarget(
                to_pattern="encoder.mid_block.attentions.{i}.to_qkv.bias",
                from_pattern=["encoder.mid_block.attentions.{i}.to_qkv.bias"],
            ),
            WeightTarget(
                to_pattern="encoder.mid_block.attentions.{i}.proj.weight",
                from_pattern=["encoder.mid_block.attentions.{i}.proj.weight"],
                transform=WeightTransforms.transpose_conv2d_weight,
            ),
            WeightTarget(
                to_pattern="encoder.mid_block.attentions.{i}.proj.bias",
                from_pattern=["encoder.mid_block.attentions.{i}.proj.bias"],
            ),
            WeightTarget(
                to_pattern="encoder.norm_out.weight",
                from_pattern=["encoder.norm_out.gamma"],
                transform=WeightTransforms.reshape_gamma_to_1d,
            ),
            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="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.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}.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}.norm2.weight",
                from_pattern=["decoder.mid_block.resnets.{i}.norm2.gamma"],
                transform=WeightTransforms.reshape_gamma_to_1d,
            ),
            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.attentions.{i}.norm.weight",
                from_pattern=["decoder.mid_block.attentions.{i}.norm.gamma"],
                transform=WeightTransforms.reshape_gamma_to_1d,
            ),
            WeightTarget(
                to_pattern="decoder.mid_block.attentions.{i}.to_qkv.weight",
                from_pattern=["decoder.mid_block.attentions.{i}.to_qkv.weight"],
                transform=WeightTransforms.transpose_conv2d_weight,
            ),
            WeightTarget(
                to_pattern="decoder.mid_block.attentions.{i}.to_qkv.bias",
                from_pattern=["decoder.mid_block.attentions.{i}.to_qkv.bias"],
            ),
            WeightTarget(
                to_pattern="decoder.mid_block.attentions.{i}.proj.weight",
                from_pattern=["decoder.mid_block.attentions.{i}.proj.weight"],
                transform=WeightTransforms.transpose_conv2d_weight,
            ),
            WeightTarget(
                to_pattern="decoder.mid_block.attentions.{i}.proj.bias",
                from_pattern=["decoder.mid_block.attentions.{i}.proj.bias"],
            ),
            WeightTarget(
                to_pattern="decoder.up_blocks.{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_blocks.{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_blocks.{block}.resnets.{res}.conv1.conv3d.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.gamma"],
                transform=WeightTransforms.reshape_gamma_to_1d,
            ),
            WeightTarget(
                to_pattern="decoder.up_blocks.{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_blocks.{block}.resnets.{res}.conv2.conv3d.bias",
                from_pattern=["decoder.up_blocks.{block}.resnets.{res}.conv2.bias"],
            ),
            WeightTarget(
                to_pattern="decoder.up_blocks.{block}.resnets.{res}.conv_shortcut.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_blocks.{block}.resnets.{res}.conv_shortcut.conv3d.bias",
                from_pattern=["decoder.up_blocks.{block}.resnets.{res}.conv_shortcut.bias"],
                required=False,
            ),
            WeightTarget(
                to_pattern="decoder.up_blocks.{block}.upsampler.time_conv.conv3d.weight",
                from_pattern=["decoder.up_blocks.{block}.upsampler.time_conv.weight"],
                transform=WeightTransforms.transpose_conv3d_weight,
                required=False,
            ),
            WeightTarget(
                to_pattern="decoder.up_blocks.{block}.upsampler.time_conv.conv3d.bias",
                from_pattern=["decoder.up_blocks.{block}.upsampler.time_conv.bias"],
                required=False,
            ),
            WeightTarget(
                to_pattern="decoder.up_blocks.{block}.upsampler.resample_conv.weight",
                from_pattern=["decoder.up_blocks.{block}.upsampler.resample.1.weight"],
                transform=WeightTransforms.transpose_conv2d_weight,
                required=False,
            ),
            WeightTarget(
                to_pattern="decoder.up_blocks.{block}.upsampler.resample_conv.bias",
                from_pattern=["decoder.up_blocks.{block}.upsampler.resample.1.bias"],
                required=False,
            ),
            WeightTarget(
                to_pattern="decoder.norm_out.weight",
                from_pattern=["decoder.norm_out.gamma"],
                transform=WeightTransforms.reshape_gamma_to_1d,
            ),
            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="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"],
            ),
            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"],
            ),
        ]
