Update app.py

app.py

@@ -12,7 +12,6 @@ import warnings
 warnings.filterwarnings("ignore")
 from PIL import Image
 import numpy as np
-from utils import load_models
 from editing import get_direction, debias
 from sampling import sample_weights
 from lora_w2w import LoRAw2w
@@ -21,6 +20,43 @@ import spaces
 
 models_path = snapshot_download(repo_id="Snapchat/w2w")
 
+@spaces.GPU
+def load_models(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)
+    noise_scheduler = pipe.scheduler
+    del pipe
+    tokenizer = AutoTokenizer.from_pretrained(
+            pretrained_model_name_or_path, subfolder="tokenizer", revision=revision
+        )
+    text_encoder = CLIPTextModel.from_pretrained(
+            pretrained_model_name_or_path, subfolder="text_encoder", revision=revision
+        )
+    vae = AutoencoderKL.from_pretrained(pretrained_model_name_or_path, subfolder="vae", revision=revision)
+    unet = UNet2DConditionModel.from_pretrained(
+            pretrained_model_name_or_path, subfolder="unet", revision=revision
+        )
+    
+    unet.requires_grad_(False)
+    unet.to(device, dtype=weight_dtype)
+    vae.requires_grad_(False)
+    
+    text_encoder.requires_grad_(False)
+    vae.requires_grad_(False)
+    vae.to(device, dtype=weight_dtype)
+    text_encoder.to(device, dtype=weight_dtype)
+    print("")
+
+    return unet, vae, text_encoder, tokenizer, noise_scheduler
+
 class main():
     def __init__(self):
         super(main, self).__init__()
@@ -43,7 +79,7 @@ class main():
         self.weight_dimensions = weight_dimensions
         self.pinverse = pinverse
         
-        self.unet, self.vae, self.text_encoder, self.tokenizer, self.noise_scheduler = load_models(device)
+        self.unet, self.vae, self.text_encoder, self.tokenizer, self.noise_scheduler = self.load_models(self.device)
         print(self.text_encoder.device)
         self.network = None
     
@@ -89,6 +125,8 @@ class main():
         thick = debias(thick, "Heavy_Makeup", df, pinverse, device)
         self.thick = thick
 
+
+
     def sample_model(self):
         self.unet, _, _, _, _ = load_models(self.device)
         self.network = sample_weights(self.unet, self.proj, self.mean, self.std, self.v[:, :1000], self.device, factor = 1.00)