Update app.py

app.py

@@ -196,7 +196,7 @@ def edit_inference(input_image, prompt, negative_prompt, guidance_scale, ddim_st
     network.proj = torch.nn.Parameter(original_weights)
     network.reset()
 
-    return (image, input_image['composite'])
+    return (image, input_image["background"])
 
 
 def sample_then_run():
@@ -344,10 +344,11 @@ def invert(image, mask, pcs=10000, epochs=400, weight_decay = 1e-10, lr=1e-1):
 
 
 
-def run_inversion(dict, pcs, epochs, weight_decay,lr):
+def run_inversion(input_image, pcs, epochs, weight_decay,lr):
     global network
-    init_image = dict["image"].convert("RGB").resize((512, 512))
-    mask = dict["mask"].convert("RGB").resize((512, 512))
+    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)
 
 
@@ -516,7 +517,7 @@ 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])
+                    outputs = [image_slider, file_output])
         
             
         sample.click(fn=sample_then_run, outputs=[input_image, file_output])