Update app.py

app.py

@@ -17,6 +17,23 @@ from editing import get_direction, debias
 from lora_w2w import LoRAw2w
 from huggingface_hub import snapshot_download
 import spaces
+from transformers import CLIPTextModel
+from lora_w2w import LoRAw2w
+from diffusers import AutoencoderKL, DDPMScheduler, DiffusionPipeline, UNet2DConditionModel, LMSDiscreteScheduler
+from transformers import AutoTokenizer, PretrainedConfig
+import warnings
+warnings.filterwarnings("ignore")
+from diffusers import (
+    AutoencoderKL,
+    DDPMScheduler,
+    DiffusionPipeline,
+    DPMSolverMultistepScheduler,
+    UNet2DConditionModel,
+    PNDMScheduler, 
+    StableDiffusionPipeline
+)
+
+
 
 device = gr.State()
 generator = gr.State()
@@ -83,6 +100,47 @@ thick.value = debias(thick.value, "Pale_Skin", df, pinverse, device.value)
 thick.value = debias(thick.value, "Heavy_Makeup", df, pinverse, device.value)
 
 
+
+
+@torch.no_grad()
+@spaces.GPU
+def load_models():
+    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.value)
+    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.value, dtype=weight_dtype)
+    text_encoder.to(device.value, dtype=weight_dtype)
+    print("")
+
+    return unet, vae, text_encoder, tokenizer, noise_scheduler
+
+
 @torch.no_grad()
 @spaces.GPU
 def sample_weights(unet, proj, mean, std, v, device, factor = 1.0):