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") # DiT XL/2 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)