Update app.py

app.py

@@ -40,7 +40,7 @@ v = torch.load(f"{models_path}/files/V.pt", map_location=torch.device('cpu')).bf
 proj = torch.load(f"{models_path}/files/proj_1000pc.pt", map_location=torch.device('cpu')).bfloat16().to(device)
 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)
+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)
 
@@ -109,7 +109,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):
-
+    start_items()
     global device
     #global generator 
     global unet
@@ -207,7 +207,6 @@ def sample_then_run():
 
 #@spaces.GPU
 def start_items():
-    device = "cpu"
     print("Starting items")
     global young
     global pointy
@@ -221,7 +220,6 @@ def start_items():
     young = debias(young, "No_Beard", df, pinverse, device)
     young = debias(young, "Mustache", df, pinverse, device)
     
-    
     pointy = get_direction(df, "Pointy_Nose", pinverse, 1000, device)
     pointy = debias(pointy, "Young", df, pinverse, device)
     pointy = debias(pointy, "Male", df, pinverse, device)
@@ -229,16 +227,13 @@ def start_items():
     pointy = debias(pointy, "Chubby", df, pinverse, device)
     pointy = debias(pointy, "Heavy_Makeup", df, pinverse, device)
     
-    
-    
     wavy = get_direction(df, "Wavy_Hair", pinverse, 1000, device)
     wavy = debias(wavy, "Young", df, pinverse, device)
     wavy = debias(wavy, "Male", df, pinverse, device)
     wavy = debias(wavy, "Pointy_Nose", df, pinverse, device)
     wavy = debias(wavy, "Chubby", df, pinverse, device)
     wavy = debias(wavy, "Heavy_Makeup", df, pinverse, device)
-    
-    
+
     thick = get_direction(df, "Bushy_Eyebrows", pinverse, 1000, device)
     thick = debias(thick, "Male", df, pinverse, device)
     thick = debias(thick, "Young", df, pinverse, device)
@@ -254,8 +249,6 @@ def start_items():
     thick = debias(thick, "Pale_Skin", df, pinverse, device)
     thick = debias(thick, "Heavy_Makeup", df, pinverse, device)
 
-start_items()
-
 class CustomImageDataset(Dataset):
     def __init__(self, images, transform=None):
         self.images = images
@@ -488,11 +481,10 @@ with gr.Blocks(css="style.css") as demo:
         
             
         sample.click(fn=sample_then_run, outputs=[input_image, file_output])
-
         
         submit.click(
-        fn=edit_inference, inputs=[prompt, negative_prompt, cfg, steps, seed, injection_step, a1, a2, a3, a4], outputs=[gallery]
-    )
+            fn=edit_inference, inputs=[prompt, negative_prompt, cfg, steps, seed, injection_step, a1, a2, a3, a4], outputs=[gallery]
+        )
         file_input.change(fn=file_upload, inputs=file_input, outputs = gallery)