Update app.py

app.py

@@ -35,44 +35,37 @@ 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
-    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
+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(
+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, dtype=weight_dtype)
-    vae.requires_grad_(False)
+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
-
-
-
+text_encoder.requires_grad_(False)
+vae.requires_grad_(False)
+vae.to(device, dtype=weight_dtype)
+text_encoder.to(device, dtype=weight_dtype)
+print("")
 
+    
 
 device="cuda"
 mean = torch.load(f"{models_path}/files/mean.pt", map_location=torch.device('cpu')).bfloat16().to(device)
@@ -83,7 +76,7 @@ df = torch.load(f"{models_path}/files/identity_df.pt")
 weight_dimensions = torch.load(f"{models_path}/files/weight_dimensions.pt")
 pinverse = torch.load(f"{models_path}/files/pinverse_1000pc.pt", map_location=torch.device('cpu')).bfloat16().to(device)
 
-unet, vae, text_encoder, tokenizer, noise_scheduler = load_models(device)
+
 
 
 @spaces.GPU