Update app.py

app.py

@@ -92,8 +92,39 @@ class main():
         self.weight_dimensions = weight_dimensions
         self.pinverse = pinverse
         
-        self.unet, self.vae, self.text_encoder, self.tokenizer, self.noise_scheduler = load_models(self.device)
-        print(self.text_encoder.device)
+        pretrained_model_name_or_path = "stablediffusionapi/realistic-vision-v51" 
+    
+        revision = None
+        rank = 1
+        weight_dtype = torch.bfloat16
+        
+        # Load scheduler, tokenizer and models.
+        pipe = StableDiffusionPipeline.from_pretrained("stablediffusionapi/realistic-vision-v51", 
+                                                    torch_dtype=torch.float16,safety_checker = None,
+                                                    requires_safety_checker = False).to(device)
+        self.noise_scheduler = pipe.scheduler
+        del pipe
+        self.tokenizer = AutoTokenizer.from_pretrained(
+                pretrained_model_name_or_path, subfolder="tokenizer", revision=revision
+            )
+        self.text_encoder = CLIPTextModel.from_pretrained(
+                pretrained_model_name_or_path, subfolder="text_encoder", revision=revision
+            )
+        self.vae = AutoencoderKL.from_pretrained(pretrained_model_name_or_path, subfolder="vae", revision=revision)
+        self.unet = UNet2DConditionModel.from_pretrained(
+                pretrained_model_name_or_path, subfolder="unet", revision=revision
+            )
+        
+        self.unet.requires_grad_(False)
+        self.unet.to(device, dtype=weight_dtype)
+        self.vae.requires_grad_(False)
+        
+        self.text_encoder.requires_grad_(False)
+        self.vae.requires_grad_(False)
+        self.vae.to(device, dtype=weight_dtype)
+        self.text_encoder.to(device, dtype=weight_dtype)
+        print("")
+        
         self.network = None
     
         young = get_direction(df, "Young", pinverse, 1000, device)