Update app.py

app.py

@@ -308,14 +308,15 @@ class CustomImageDataset(Dataset):
         return image
         
 @spaces.GPU
-def invert(self, image, mask, pcs=10000, epochs=400, weight_decay = 1e-10, lr=1e-1):
-            
-    del unet
-    del network
-    unet, _, _, _, _ = load_models(device)
+def invert(image, mask, pcs=10000, epochs=400, weight_decay = 1e-10, lr=1e-1):
+    device = "cuda"
+    mean.to(device)
+    std.to(device)
+    v.to(device)
+    
             
-    proj = torch.zeros(1,pcs).bfloat16().to(device)
-    network = LoRAw2w( proj, mean, std, v[:, :pcs], 
+    weights = torch.zeros(1,pcs).bfloat16().to(device)
+    network = LoRAw2w( weights, mean, std, v[:, :pcs], 
                                 unet,
                                 rank=1,
                                 multiplier=1.0,
@@ -367,18 +368,27 @@ def invert(self, image, mask, pcs=10000, epochs=400, weight_decay = 1e-10, lr=1e
                 optim.zero_grad()
                 loss.backward()
                 optim.step()
-        
-    ### return optimized network
-    return network
+
+
+    #pad to 10000 PCs
+    pcs_original = weights.shape[1]
+    padding =  torch.zeros((1,10000-pcs_original)).to(device)
+    weights = network.proj.detach()
+    weights = torch.cat((weights, padding), 1)
+
+    net = "model_"+str(uuid.uuid4())[:4]+".pt"
+    torch.save(weights, net)
+    
+    return net
 
 
 @spaces.GPU
-def run_inversion(dict, pcs, epochs, weight_decay,lr):
-    print(dict)
-    print(dict.keys())
-    init_image = dict["image"].convert("RGB").resize((512, 512))
-    mask = dict["mask"].convert("RGB").resize((512, 512))
-    network = invert([init_image], mask, pcs, epochs, weight_decay,lr)
+def run_inversion(net, dict, pcs, epochs, weight_decay,lr):
+    init_image = dict["background"].convert("RGB").resize((512, 512))
+    mask = dict["layers"][0].convert("RGB").resize((512, 512))
+    
+    net = invert(init_image, mask, pcs, epochs, weight_decay,lr)
+    
     #sample an image
     prompt = "sks person"
     negative_prompt = "low quality, blurry, unfinished, nudity"
@@ -387,7 +397,7 @@ def run_inversion(dict, pcs, epochs, weight_decay,lr):
     steps = 25
     image = inference( prompt, negative_prompt, cfg, steps, seed)
     torch.save(network.proj, "model.pt" )
-    return image, "model.pt"
+    return net, net, image
     
     
 @spaces.GPU
@@ -408,7 +418,7 @@ def file_upload(file, net):
     cfg = 3.0
     steps = 25
     image = inference(net, prompt, negative_prompt, cfg, steps, seed)
-    return net,image
+    return net, image
     
     
 
@@ -504,8 +514,8 @@ with gr.Blocks(css="style.css") as demo:
     
     
         invert_button.click(fn=run_inversion,
-                        inputs=[input_image, pcs, epochs, weight_decay,lr], 
-                        outputs = [input_image, file_output])
+                        inputs=[net, input_image, pcs, epochs, weight_decay,lr], 
+                        outputs = [net, file_output, input_image])
             
                 
         sample.click(fn=sample_then_run,inputs = [net], outputs=[net, file_output, input_image])