import argparse

import torch

from huggingface_hub import hf_hub_download

from diffusers.models.transformers.auraflow_transformer_2d import AuraFlowTransformer2DModel

def load_original_state_dict(args):

model_pt = hf_hub_download(repo_id=args.original_state_dict_repo_id, filename="aura_diffusion_pytorch_model.bin")

state_dict = torch.load(model_pt, map_location="cpu")

return state_dict

def calculate_layers(state_dict_keys, key_prefix):

dit_layers = set()

for k in state_dict_keys:

if key_prefix in k:

dit_layers.add(int(k.split(".")[2]))

print(f"{key_prefix}: {len(dit_layers)}")

return len(dit_layers)

# similar to SD3 but only for the last norm layer

def swap_scale_shift(weight, dim):

shift, scale = weight.chunk(2, dim=0)

new_weight = torch.cat([scale, shift], dim=0)

return new_weight

def convert_transformer(state_dict):

converted_state_dict = {}

state_dict_keys = list(state_dict.keys())

converted_state_dict["register_tokens"] = state_dict.pop("model.register_tokens")

converted_state_dict["pos_embed.pos_embed"] = state_dict.pop("model.positional_encoding")

converted_state_dict["pos_embed.proj.weight"] = state_dict.pop("model.init_x_linear.weight")

converted_state_dict["pos_embed.proj.bias"] = state_dict.pop("model.init_x_linear.bias")

converted_state_dict["time_step_proj.linear_1.weight"] = state_dict.pop("model.t_embedder.mlp.0.weight")

converted_state_dict["time_step_proj.linear_1.bias"] = state_dict.pop("model.t_embedder.mlp.0.bias")

converted_state_dict["time_step_proj.linear_2.weight"] = state_dict.pop("model.t_embedder.mlp.2.weight")

converted_state_dict["time_step_proj.linear_2.bias"] = state_dict.pop("model.t_embedder.mlp.2.bias")

converted_state_dict["context_embedder.weight"] = state_dict.pop("model.cond_seq_linear.weight")

mmdit_layers = calculate_layers(state_dict_keys, key_prefix="double_layers")

single_dit_layers = calculate_layers(state_dict_keys, key_prefix="single_layers")

# MMDiT blocks 🎸.

for i in range(mmdit_layers):

# feed-forward

path_mapping = {"mlpX": "ff", "mlpC": "ff_context"}

weight_mapping = {"c_fc1": "linear_1", "c_fc2": "linear_2", "c_proj": "out_projection"}

for orig_k, diffuser_k in path_mapping.items():

for k, v in weight_mapping.items():

converted_state_dict[f"joint_transformer_blocks.{i}.{diffuser_k}.{v}.weight"] = state_dict.pop(

f"model.double_layers.{i}.{orig_k}.{k}.weight"

)

# norms

path_mapping = {"modX": "norm1", "modC": "norm1_context"}

for orig_k, diffuser_k in path_mapping.items():

converted_state_dict[f"joint_transformer_blocks.{i}.{diffuser_k}.linear.weight"] = state_dict.pop(

f"model.double_layers.{i}.{orig_k}.1.weight"

)

# attns

x_attn_mapping = {"w2q": "to_q", "w2k": "to_k", "w2v": "to_v", "w2o": "to_out.0"}

context_attn_mapping = {"w1q": "add_q_proj", "w1k": "add_k_proj", "w1v": "add_v_proj", "w1o": "to_add_out"}

for attn_mapping in [x_attn_mapping, context_attn_mapping]:

for k, v in attn_mapping.items():

converted_state_dict[f"joint_transformer_blocks.{i}.attn.{v}.weight"] = state_dict.pop(

f"model.double_layers.{i}.attn.{k}.weight"

)

# Single-DiT blocks.

for i in range(single_dit_layers):

# feed-forward

mapping = {"c_fc1": "linear_1", "c_fc2": "linear_2", "c_proj": "out_projection"}

for k, v in mapping.items():

converted_state_dict[f"single_transformer_blocks.{i}.ff.{v}.weight"] = state_dict.pop(

f"model.single_layers.{i}.mlp.{k}.weight"

)

# norms

converted_state_dict[f"single_transformer_blocks.{i}.norm1.linear.weight"] = state_dict.pop(

f"model.single_layers.{i}.modCX.1.weight"

)

# attns

x_attn_mapping = {"w1q": "to_q", "w1k": "to_k", "w1v": "to_v", "w1o": "to_out.0"}

for k, v in x_attn_mapping.items():

converted_state_dict[f"single_transformer_blocks.{i}.attn.{v}.weight"] = state_dict.pop(

f"model.single_layers.{i}.attn.{k}.weight"

)

# Final blocks.

converted_state_dict["proj_out.weight"] = state_dict.pop("model.final_linear.weight")

converted_state_dict["norm_out.linear.weight"] = swap_scale_shift(state_dict.pop("model.modF.1.weight"), dim=None)

return converted_state_dict

@torch.no_grad()

def populate_state_dict(args):

original_state_dict = load_original_state_dict(args)

state_dict_keys = list(original_state_dict.keys())

mmdit_layers = calculate_layers(state_dict_keys, key_prefix="double_layers")

single_dit_layers = calculate_layers(state_dict_keys, key_prefix="single_layers")

converted_state_dict = convert_transformer(original_state_dict)

model_diffusers = AuraFlowTransformer2DModel(

num_mmdit_layers=mmdit_layers, num_single_dit_layers=single_dit_layers

)

model_diffusers.load_state_dict(converted_state_dict, strict=True)

return model_diffusers

if __name__ == "__main__":

parser = argparse.ArgumentParser()

parser.add_argument("--original_state_dict_repo_id", default="AuraDiffusion/auradiffusion-v0.1a0", type=str)

parser.add_argument("--dump_path", default="aura-flow", type=str)

parser.add_argument("--hub_id", default=None, type=str)

args = parser.parse_args()

model_diffusers = populate_state_dict(args)

model_diffusers.save_pretrained(args.dump_path)

if args.hub_id is not None:

model_diffusers.push_to_hub(args.hub_id)