| import argparse |
| import os |
|
|
| import torch |
| from torchvision.datasets.utils import download_url |
|
|
| from diffusers import AutoencoderKL, DDIMScheduler, DiTPipeline, Transformer2DModel |
|
|
|
|
| pretrained_models = {512: "DiT-XL-2-512x512.pt", 256: "DiT-XL-2-256x256.pt"} |
|
|
|
|
| def download_model(model_name): |
| """ |
| Downloads a pre-trained DiT model from the web. |
| """ |
| local_path = f"pretrained_models/{model_name}" |
| if not os.path.isfile(local_path): |
| os.makedirs("pretrained_models", exist_ok=True) |
| web_path = f"https://dl.fbaipublicfiles.com/DiT/models/{model_name}" |
| download_url(web_path, "pretrained_models") |
| model = torch.load(local_path, map_location=lambda storage, loc: storage) |
| return model |
|
|
|
|
| def main(args): |
| state_dict = download_model(pretrained_models[args.image_size]) |
|
|
| state_dict["pos_embed.proj.weight"] = state_dict["x_embedder.proj.weight"] |
| state_dict["pos_embed.proj.bias"] = state_dict["x_embedder.proj.bias"] |
| state_dict.pop("x_embedder.proj.weight") |
| state_dict.pop("x_embedder.proj.bias") |
|
|
| for depth in range(28): |
| state_dict[f"transformer_blocks.{depth}.norm1.emb.timestep_embedder.linear_1.weight"] = state_dict[ |
| "t_embedder.mlp.0.weight" |
| ] |
| state_dict[f"transformer_blocks.{depth}.norm1.emb.timestep_embedder.linear_1.bias"] = state_dict[ |
| "t_embedder.mlp.0.bias" |
| ] |
| state_dict[f"transformer_blocks.{depth}.norm1.emb.timestep_embedder.linear_2.weight"] = state_dict[ |
| "t_embedder.mlp.2.weight" |
| ] |
| state_dict[f"transformer_blocks.{depth}.norm1.emb.timestep_embedder.linear_2.bias"] = state_dict[ |
| "t_embedder.mlp.2.bias" |
| ] |
| state_dict[f"transformer_blocks.{depth}.norm1.emb.class_embedder.embedding_table.weight"] = state_dict[ |
| "y_embedder.embedding_table.weight" |
| ] |
|
|
| state_dict[f"transformer_blocks.{depth}.norm1.linear.weight"] = state_dict[ |
| f"blocks.{depth}.adaLN_modulation.1.weight" |
| ] |
| state_dict[f"transformer_blocks.{depth}.norm1.linear.bias"] = state_dict[ |
| f"blocks.{depth}.adaLN_modulation.1.bias" |
| ] |
|
|
| q, k, v = torch.chunk(state_dict[f"blocks.{depth}.attn.qkv.weight"], 3, dim=0) |
| q_bias, k_bias, v_bias = torch.chunk(state_dict[f"blocks.{depth}.attn.qkv.bias"], 3, dim=0) |
|
|
| state_dict[f"transformer_blocks.{depth}.attn1.to_q.weight"] = q |
| state_dict[f"transformer_blocks.{depth}.attn1.to_q.bias"] = q_bias |
| state_dict[f"transformer_blocks.{depth}.attn1.to_k.weight"] = k |
| state_dict[f"transformer_blocks.{depth}.attn1.to_k.bias"] = k_bias |
| state_dict[f"transformer_blocks.{depth}.attn1.to_v.weight"] = v |
| state_dict[f"transformer_blocks.{depth}.attn1.to_v.bias"] = v_bias |
|
|
| state_dict[f"transformer_blocks.{depth}.attn1.to_out.0.weight"] = state_dict[ |
| f"blocks.{depth}.attn.proj.weight" |
| ] |
| state_dict[f"transformer_blocks.{depth}.attn1.to_out.0.bias"] = state_dict[f"blocks.{depth}.attn.proj.bias"] |
|
|
| state_dict[f"transformer_blocks.{depth}.ff.net.0.proj.weight"] = state_dict[f"blocks.{depth}.mlp.fc1.weight"] |
| state_dict[f"transformer_blocks.{depth}.ff.net.0.proj.bias"] = state_dict[f"blocks.{depth}.mlp.fc1.bias"] |
| state_dict[f"transformer_blocks.{depth}.ff.net.2.weight"] = state_dict[f"blocks.{depth}.mlp.fc2.weight"] |
| state_dict[f"transformer_blocks.{depth}.ff.net.2.bias"] = state_dict[f"blocks.{depth}.mlp.fc2.bias"] |
|
|
| state_dict.pop(f"blocks.{depth}.attn.qkv.weight") |
| state_dict.pop(f"blocks.{depth}.attn.qkv.bias") |
| state_dict.pop(f"blocks.{depth}.attn.proj.weight") |
| state_dict.pop(f"blocks.{depth}.attn.proj.bias") |
| state_dict.pop(f"blocks.{depth}.mlp.fc1.weight") |
| state_dict.pop(f"blocks.{depth}.mlp.fc1.bias") |
| state_dict.pop(f"blocks.{depth}.mlp.fc2.weight") |
| state_dict.pop(f"blocks.{depth}.mlp.fc2.bias") |
| state_dict.pop(f"blocks.{depth}.adaLN_modulation.1.weight") |
| state_dict.pop(f"blocks.{depth}.adaLN_modulation.1.bias") |
|
|
| state_dict.pop("t_embedder.mlp.0.weight") |
| state_dict.pop("t_embedder.mlp.0.bias") |
| state_dict.pop("t_embedder.mlp.2.weight") |
| state_dict.pop("t_embedder.mlp.2.bias") |
| state_dict.pop("y_embedder.embedding_table.weight") |
|
|
| state_dict["proj_out_1.weight"] = state_dict["final_layer.adaLN_modulation.1.weight"] |
| state_dict["proj_out_1.bias"] = state_dict["final_layer.adaLN_modulation.1.bias"] |
| state_dict["proj_out_2.weight"] = state_dict["final_layer.linear.weight"] |
| state_dict["proj_out_2.bias"] = state_dict["final_layer.linear.bias"] |
|
|
| state_dict.pop("final_layer.linear.weight") |
| state_dict.pop("final_layer.linear.bias") |
| state_dict.pop("final_layer.adaLN_modulation.1.weight") |
| state_dict.pop("final_layer.adaLN_modulation.1.bias") |
|
|
| |
| transformer = Transformer2DModel( |
| sample_size=args.image_size // 8, |
| num_layers=28, |
| attention_head_dim=72, |
| in_channels=4, |
| out_channels=8, |
| patch_size=2, |
| attention_bias=True, |
| num_attention_heads=16, |
| activation_fn="gelu-approximate", |
| num_embeds_ada_norm=1000, |
| norm_type="ada_norm_zero", |
| norm_elementwise_affine=False, |
| ) |
| transformer.load_state_dict(state_dict, strict=True) |
|
|
| scheduler = DDIMScheduler( |
| num_train_timesteps=1000, |
| beta_schedule="linear", |
| prediction_type="epsilon", |
| clip_sample=False, |
| ) |
|
|
| vae = AutoencoderKL.from_pretrained(args.vae_model) |
|
|
| pipeline = DiTPipeline(transformer=transformer, vae=vae, scheduler=scheduler) |
|
|
| if args.save: |
| pipeline.save_pretrained(args.checkpoint_path) |
|
|
|
|
| if __name__ == "__main__": |
| parser = argparse.ArgumentParser() |
|
|
| parser.add_argument( |
| "--image_size", |
| default=256, |
| type=int, |
| required=False, |
| help="Image size of pretrained model, either 256 or 512.", |
| ) |
| parser.add_argument( |
| "--vae_model", |
| default="stabilityai/sd-vae-ft-ema", |
| type=str, |
| required=False, |
| help="Path to pretrained VAE model, either stabilityai/sd-vae-ft-mse or stabilityai/sd-vae-ft-ema.", |
| ) |
| parser.add_argument( |
| "--save", default=True, type=bool, required=False, help="Whether to save the converted pipeline or not." |
| ) |
| parser.add_argument( |
| "--checkpoint_path", default=None, type=str, required=True, help="Path to the output pipeline." |
| ) |
|
|
| args = parser.parse_args() |
| main(args) |
|
|
