update save function

app.py

@@ -129,12 +129,18 @@ def download_image(image):
     image.save("output.png", "PNG")
     return "output.png"
 
-def save_details(image, mask, prompt, lora_model, strength, seed, guidance_scale, num_inference_steps):
+def save_details(result, edit_image, prompt, lora_model, strength, seed, guidance_scale, num_inference_steps):
+    image = edit_image["background"]
+    mask = edit_image["layers"][0]
+
+    if isinstance(result, np.ndarray):
+        result = Image.fromarray(result)
     if isinstance(image, np.ndarray):
         image = Image.fromarray(image)
     if isinstance(mask, np.ndarray):
         mask = Image.fromarray(mask)
 
+    result.save("saved_result.png", "PNG")
     image.save("saved_image.png", "PNG")
     mask.save("saved_mask.png", "PNG")
 
@@ -152,6 +158,7 @@ def save_details(image, mask, prompt, lora_model, strength, seed, guidance_scale
 
     # Create a ZIP file
     with zipfile.ZipFile("output.zip", "w") as zipf:
+        zipf.write("saved_result.png")
         zipf.write("saved_image.png")
         zipf.write("saved_mask.png")
         zipf.write("details.json")