Update app.py

app.py

@@ -32,6 +32,44 @@ from diffusers import (
     StableDiffusionPipeline
 )
 
+@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
+
 
 
 device = gr.State()
@@ -101,43 +139,6 @@ 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()