Update app.py

app.py

@@ -19,7 +19,7 @@ from editing import get_direction, debias
 from sampling import sample_weights
 from lora_w2w import LoRAw2w
 from huggingface_hub import snapshot_download
-
+import spaces
 
 
 
@@ -32,9 +32,7 @@ global tokenizer
 global noise_scheduler
 global network
 device = "cuda:0"
-generator = torch.Generator(device=device)
-
-
+#generator = torch.Generator(device=device)
 
 models_path = snapshot_download(repo_id="Snapchat/w2w")
 
@@ -49,7 +47,6 @@ pinverse = torch.load(f"{models_path}/files/pinverse_1000pc.pt").bfloat16().to(d
 
 unet, vae, text_encoder, tokenizer, noise_scheduler = load_models(device)
 
-
 def sample_model():
     global unet
     del unet
@@ -59,6 +56,7 @@ def sample_model():
     network = sample_weights(unet, proj, mean, std, v[:, :1000], device, factor = 1.00)
  
 @torch.no_grad()
+@spaces.GPU
 def inference( prompt, negative_prompt, guidance_scale, ddim_steps, seed):
     global device
     global generator 
@@ -109,6 +107,7 @@ def inference( prompt, negative_prompt, guidance_scale, ddim_steps, seed):
 
 
 @torch.no_grad()
+@spaces.GPU
 def edit_inference(prompt, negative_prompt, guidance_scale, ddim_steps, seed, start_noise, a1, a2, a3, a4):
 
     global device
@@ -193,7 +192,8 @@ def edit_inference(prompt, negative_prompt, guidance_scale, ddim_steps, seed, st
     network.reset()
 
     return image
-    
+
+@spaces.GPU
 def sample_then_run():
     sample_model()    
     prompt = "sks person"
@@ -267,7 +267,8 @@ class CustomImageDataset(Dataset):
         if self.transform:
             image = self.transform(image)
         return image
-    
+
+@spaces.GPU
 def invert(image, mask, pcs=10000, epochs=400, weight_decay = 1e-10, lr=1e-1):
     global unet
     del unet
@@ -332,7 +333,7 @@ def invert(image, mask, pcs=10000, epochs=400, weight_decay = 1e-10, lr=1e-1):
     return network
 
 
-
+@spaces.GPU
 def run_inversion(dict, pcs, epochs, weight_decay,lr):
     global network
     init_image = dict["image"].convert("RGB").resize((512, 512))
@@ -351,7 +352,7 @@ def run_inversion(dict, pcs, epochs, weight_decay,lr):
     return image, "model.pt"
     
     
-
+@spaces.GPU
 def file_upload(file):
     global unet
     del unet