Update app.py

app.py

@@ -35,57 +35,48 @@ from diffusers import (
 
 device = gr.State("cuda")
 
-
-
-@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", 
+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(
+noise_scheduler = pipe.scheduler
+del pipe
+tokenizer = AutoTokenizer.from_pretrained(
         pretrained_model_name_or_path, subfolder="tokenizer", revision=revision
     )
-    text_encoder = CLIPTextModel.from_pretrained(
+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(
+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.value, dtype=weight_dtype)
-    vae.requires_grad_(False)
+unet.requires_grad_(False)
+unet.to(device.value, 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("")
+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
 
 
+unet = gr.State(unet)
+vae = gr.State(vae)
+text_encoder = gr.State(text_encoder)
+tokenizer = gr.State(tokenizer)
+noise_scheduler = gr.State(noise_scheduler)
 
 
-generator = gr.State()
-unet = gr.State()
-vae = gr.State()
-text_encoder = gr.State()
-tokenizer = gr.State()
-noise_scheduler = gr.State()
 network = gr.State()
-device = gr.State("cuda")
+
 
 
 models_path = snapshot_download(repo_id="Snapchat/w2w")