change to zerogpu

app.py

@@ -33,7 +33,7 @@ global network
 device = "cuda:0"
 generator = torch.Generator(device=device)
 from gradio_imageslider import ImageSlider
-
+import spaces
 
 
 models_path = snapshot_download(repo_id="Snapchat/w2w")
@@ -46,9 +46,10 @@ 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").bfloat16().to(device)
 
+@spaces.GPU()
 unet, vae, text_encoder, tokenizer, noise_scheduler = load_models(device)
 
-
+@spaces.GPU()
 def sample_model():
     global unet
     del unet
@@ -57,9 +58,7 @@ def sample_model():
     unet, _, _, _, _ = load_models(device)
     network = sample_weights(unet, proj, mean, std, v[:, :1000], device, factor = 1.00)
  
- 
 
- 
 @torch.no_grad()
 def inference(prompt, negative_prompt, guidance_scale, ddim_steps, seed):
     global device
@@ -111,7 +110,7 @@ def inference(prompt, negative_prompt, guidance_scale, ddim_steps, seed):
 
 
 
-
+@spaces.GPU()
 @torch.no_grad()
 def edit_inference(input_image, prompt, negative_prompt, guidance_scale, ddim_steps, seed, start_noise, a1, a2, a3, a4):
 
@@ -198,7 +197,7 @@ def edit_inference(input_image, prompt, negative_prompt, guidance_scale, ddim_st
 
     return (image, input_image["background"])
 
-
+@spaces.GPU()
 def sample_then_run():
     sample_model()    
     prompt = "sks person"
@@ -270,7 +269,7 @@ class CustomImageDataset(Dataset):
         if self.transform:
             image = self.transform(image)
         return image
-    
+
 def invert(image, mask, pcs=10000, epochs=400, weight_decay = 1e-10, lr=1e-1):
     global unet
     del unet
@@ -343,10 +342,9 @@ def invert(image, mask, pcs=10000, epochs=400, weight_decay = 1e-10, lr=1e-1):
     return network
 
 
-
+@spaces.GPU(duration=200)
 def run_inversion(input_image, pcs, epochs, weight_decay,lr):
     global network
-    print(len(input_image["layers"]))
     init_image = input_image["background"].convert("RGB").resize((512, 512))
     mask = input_image["layers"][0].convert("RGB").resize((512, 512))
     network = invert([init_image], mask, pcs, epochs, weight_decay,lr)
@@ -366,7 +364,7 @@ def run_inversion(input_image, pcs, epochs, weight_decay,lr):
 
 
 
-
+@spaces.GPU()
 def file_upload(file):
     global unet
     del unet