Update README.md

README.md

@@ -96,11 +96,13 @@ from transformers import AutoTokenizer, AutoModel
 import torch
 
 prompt_encoder_model_name_or_path = "sentence-transformers/all-mpnet-base-v2"
-prompt_encoder_tokenizer = AutoTokenizer.from_pretrained(prompt_encoder_model_name_or_path)
-prompt_encoder = AutoModel.from_pretrained(prompt_encoder_model_name_or_path)
-
 aptp_model_name_or_path = f"rezashkv/APTP"
 aptp_variant = "APTP-Base-CC3M"
+sd_model_name_or_path = "stabilityai/stable-diffusion-2-1"
+
+prompt_encoder = AutoModel.from_pretrained(prompt_encoder_model_name_or_path)
+prompt_encoder_tokenizer = AutoTokenizer.from_pretrained(prompt_encoder_model_name_or_path)
+
 hyper_net = HyperStructure.from_pretrained(aptp_model_name_or_path, subfolder=f"{aptp_variant}/hypernet")
 quantizer = StructureVectorQuantizer.from_pretrained(aptp_model_name_or_path, subfolder=f"{aptp_variant}/quantizer")
 
@@ -110,15 +112,14 @@ prompt_embedding = get_mpnet_embeddings(prompts, prompt_encoder, prompt_encoder_
 arch_embedding = hyper_net(prompt_embedding)
 expert_id = quantizer.get_cosine_sim_min_encoding_indices(arch_embedding)[0].item()
 
-sd_model_name_or_path = "stabilityai/stable-diffusion-2-1"
-
 unet = UNet2DConditionModelPruned.from_pretrained(aptp_model_name_or_path,
                                                   subfolder=f"{aptp_variant}/arch{expert_id}/checkpoint-30000/unet")
-
 noise_scheduler = PNDMScheduler.from_pretrained(sd_model_name_or_path, subfolder="scheduler")
+
 pipeline = StableDiffusionPipeline.from_pretrained(sd_model_name_or_path, unet=unet, scheduler=noise_scheduler)
 
 pipeline.to('cuda')
+
 generator = torch.Generator(device='cuda').manual_seed(43)
 
 image = pipeline(