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


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

    @staticmethod
    def get_vae_mapping() -> list[WeightTarget]:
        return [
            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.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="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="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="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="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="decoder.conv_in.conv.weight",
                from_pattern=["decoder.conv_in.weight"],
                transform=WeightTransforms.transpose_conv2d_weight,
            ),
            WeightTarget(
                to_pattern="decoder.conv_in.conv.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.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="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="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="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="decoder.conv_out.conv.weight",
                from_pattern=["decoder.conv_out.weight"],
                transform=WeightTransforms.transpose_conv2d_weight,
            ),
            WeightTarget(
                to_pattern="decoder.conv_out.conv.bias",
                from_pattern=["decoder.conv_out.bias"],
            ),
        ]

    @staticmethod
    def get_transformer_mapping() -> list[WeightTarget]:
        return [
            WeightTarget(
                to_pattern="t_embedder.linear1.weight",
                from_pattern=["t_embedder.mlp.0.weight"],
            ),
            WeightTarget(
                to_pattern="t_embedder.linear1.bias",
                from_pattern=["t_embedder.mlp.0.bias"],
            ),
            WeightTarget(
                to_pattern="t_embedder.linear2.weight",
                from_pattern=["t_embedder.mlp.2.weight"],
            ),
            WeightTarget(
                to_pattern="t_embedder.linear2.bias",
                from_pattern=["t_embedder.mlp.2.bias"],
            ),
            WeightTarget(
                to_pattern="cap_embedder.0.weight",
                from_pattern=["cap_embedder.0.weight"],
            ),
            WeightTarget(
                to_pattern="cap_embedder.1.weight",
                from_pattern=["cap_embedder.1.weight"],
            ),
            WeightTarget(
                to_pattern="cap_embedder.1.bias",
                from_pattern=["cap_embedder.1.bias"],
            ),
            WeightTarget(
                to_pattern="x_pad_token",
                from_pattern=["x_pad_token"],
            ),
            WeightTarget(
                to_pattern="cap_pad_token",
                from_pattern=["cap_pad_token"],
            ),
            WeightTarget(
                to_pattern="all_x_embedder.2-1.weight",
                from_pattern=["all_x_embedder.2-1.weight"],
            ),
            WeightTarget(
                to_pattern="all_x_embedder.2-1.bias",
                from_pattern=["all_x_embedder.2-1.bias"],
            ),
            WeightTarget(
                to_pattern="all_final_layer.2-1.linear.weight",
                from_pattern=["all_final_layer.2-1.linear.weight"],
            ),
            WeightTarget(
                to_pattern="all_final_layer.2-1.linear.bias",
                from_pattern=["all_final_layer.2-1.linear.bias"],
            ),
            WeightTarget(
                to_pattern="all_final_layer.2-1.adaLN_modulation.0.weight",
                from_pattern=["all_final_layer.2-1.adaLN_modulation.1.weight"],
            ),
            WeightTarget(
                to_pattern="all_final_layer.2-1.adaLN_modulation.0.bias",
                from_pattern=["all_final_layer.2-1.adaLN_modulation.1.bias"],
            ),
            WeightTarget(
                to_pattern="noise_refiner.{layer}.adaLN_modulation.0.weight",
                from_pattern=["noise_refiner.{layer}.adaLN_modulation.0.weight"],
            ),
            WeightTarget(
                to_pattern="noise_refiner.{layer}.adaLN_modulation.0.bias",
                from_pattern=["noise_refiner.{layer}.adaLN_modulation.0.bias"],
            ),
            WeightTarget(
                to_pattern="noise_refiner.{layer}.attention.to_q.weight",
                from_pattern=["noise_refiner.{layer}.attention.to_q.weight"],
            ),
            WeightTarget(
                to_pattern="noise_refiner.{layer}.attention.to_k.weight",
                from_pattern=["noise_refiner.{layer}.attention.to_k.weight"],
            ),
            WeightTarget(
                to_pattern="noise_refiner.{layer}.attention.to_v.weight",
                from_pattern=["noise_refiner.{layer}.attention.to_v.weight"],
            ),
            WeightTarget(
                to_pattern="noise_refiner.{layer}.attention.to_out.0.weight",
                from_pattern=["noise_refiner.{layer}.attention.to_out.0.weight"],
            ),
            WeightTarget(
                to_pattern="noise_refiner.{layer}.attention.norm_q.weight",
                from_pattern=["noise_refiner.{layer}.attention.norm_q.weight"],
            ),
            WeightTarget(
                to_pattern="noise_refiner.{layer}.attention.norm_k.weight",
                from_pattern=["noise_refiner.{layer}.attention.norm_k.weight"],
            ),
            WeightTarget(
                to_pattern="noise_refiner.{layer}.attention_norm1.weight",
                from_pattern=["noise_refiner.{layer}.attention_norm1.weight"],
            ),
            WeightTarget(
                to_pattern="noise_refiner.{layer}.attention_norm2.weight",
                from_pattern=["noise_refiner.{layer}.attention_norm2.weight"],
            ),
            WeightTarget(
                to_pattern="noise_refiner.{layer}.ffn_norm1.weight",
                from_pattern=["noise_refiner.{layer}.ffn_norm1.weight"],
            ),
            WeightTarget(
                to_pattern="noise_refiner.{layer}.ffn_norm2.weight",
                from_pattern=["noise_refiner.{layer}.ffn_norm2.weight"],
            ),
            WeightTarget(
                to_pattern="noise_refiner.{layer}.feed_forward.w1.weight",
                from_pattern=["noise_refiner.{layer}.feed_forward.w1.weight"],
            ),
            WeightTarget(
                to_pattern="noise_refiner.{layer}.feed_forward.w2.weight",
                from_pattern=["noise_refiner.{layer}.feed_forward.w2.weight"],
            ),
            WeightTarget(
                to_pattern="noise_refiner.{layer}.feed_forward.w3.weight",
                from_pattern=["noise_refiner.{layer}.feed_forward.w3.weight"],
            ),
            WeightTarget(
                to_pattern="context_refiner.{layer}.attention.to_q.weight",
                from_pattern=["context_refiner.{layer}.attention.to_q.weight"],
            ),
            WeightTarget(
                to_pattern="context_refiner.{layer}.attention.to_k.weight",
                from_pattern=["context_refiner.{layer}.attention.to_k.weight"],
            ),
            WeightTarget(
                to_pattern="context_refiner.{layer}.attention.to_v.weight",
                from_pattern=["context_refiner.{layer}.attention.to_v.weight"],
            ),
            WeightTarget(
                to_pattern="context_refiner.{layer}.attention.to_out.0.weight",
                from_pattern=["context_refiner.{layer}.attention.to_out.0.weight"],
            ),
            WeightTarget(
                to_pattern="context_refiner.{layer}.attention.norm_q.weight",
                from_pattern=["context_refiner.{layer}.attention.norm_q.weight"],
            ),
            WeightTarget(
                to_pattern="context_refiner.{layer}.attention.norm_k.weight",
                from_pattern=["context_refiner.{layer}.attention.norm_k.weight"],
            ),
            WeightTarget(
                to_pattern="context_refiner.{layer}.attention_norm1.weight",
                from_pattern=["context_refiner.{layer}.attention_norm1.weight"],
            ),
            WeightTarget(
                to_pattern="context_refiner.{layer}.attention_norm2.weight",
                from_pattern=["context_refiner.{layer}.attention_norm2.weight"],
            ),
            WeightTarget(
                to_pattern="context_refiner.{layer}.ffn_norm1.weight",
                from_pattern=["context_refiner.{layer}.ffn_norm1.weight"],
            ),
            WeightTarget(
                to_pattern="context_refiner.{layer}.ffn_norm2.weight",
                from_pattern=["context_refiner.{layer}.ffn_norm2.weight"],
            ),
            WeightTarget(
                to_pattern="context_refiner.{layer}.feed_forward.w1.weight",
                from_pattern=["context_refiner.{layer}.feed_forward.w1.weight"],
            ),
            WeightTarget(
                to_pattern="context_refiner.{layer}.feed_forward.w2.weight",
                from_pattern=["context_refiner.{layer}.feed_forward.w2.weight"],
            ),
            WeightTarget(
                to_pattern="context_refiner.{layer}.feed_forward.w3.weight",
                from_pattern=["context_refiner.{layer}.feed_forward.w3.weight"],
            ),
            WeightTarget(
                to_pattern="layers.{layer}.adaLN_modulation.0.weight",
                from_pattern=["layers.{layer}.adaLN_modulation.0.weight"],
            ),
            WeightTarget(
                to_pattern="layers.{layer}.adaLN_modulation.0.bias",
                from_pattern=["layers.{layer}.adaLN_modulation.0.bias"],
            ),
            WeightTarget(
                to_pattern="layers.{layer}.attention.to_q.weight",
                from_pattern=["layers.{layer}.attention.to_q.weight"],
            ),
            WeightTarget(
                to_pattern="layers.{layer}.attention.to_k.weight",
                from_pattern=["layers.{layer}.attention.to_k.weight"],
            ),
            WeightTarget(
                to_pattern="layers.{layer}.attention.to_v.weight",
                from_pattern=["layers.{layer}.attention.to_v.weight"],
            ),
            WeightTarget(
                to_pattern="layers.{layer}.attention.to_out.0.weight",
                from_pattern=["layers.{layer}.attention.to_out.0.weight"],
            ),
            WeightTarget(
                to_pattern="layers.{layer}.attention.norm_q.weight",
                from_pattern=["layers.{layer}.attention.norm_q.weight"],
            ),
            WeightTarget(
                to_pattern="layers.{layer}.attention.norm_k.weight",
                from_pattern=["layers.{layer}.attention.norm_k.weight"],
            ),
            WeightTarget(
                to_pattern="layers.{layer}.attention_norm1.weight",
                from_pattern=["layers.{layer}.attention_norm1.weight"],
            ),
            WeightTarget(
                to_pattern="layers.{layer}.attention_norm2.weight",
                from_pattern=["layers.{layer}.attention_norm2.weight"],
            ),
            WeightTarget(
                to_pattern="layers.{layer}.ffn_norm1.weight",
                from_pattern=["layers.{layer}.ffn_norm1.weight"],
            ),
            WeightTarget(
                to_pattern="layers.{layer}.ffn_norm2.weight",
                from_pattern=["layers.{layer}.ffn_norm2.weight"],
            ),
            WeightTarget(
                to_pattern="layers.{layer}.feed_forward.w1.weight",
                from_pattern=["layers.{layer}.feed_forward.w1.weight"],
            ),
            WeightTarget(
                to_pattern="layers.{layer}.feed_forward.w2.weight",
                from_pattern=["layers.{layer}.feed_forward.w2.weight"],
            ),
            WeightTarget(
                to_pattern="layers.{layer}.feed_forward.w3.weight",
                from_pattern=["layers.{layer}.feed_forward.w3.weight"],
            ),
        ]
