Upload server.py

server.py

@@ -185,16 +185,30 @@ async def edit_image(
                 fn_index=2
             )
 
+        print(f"DEBUG: Result from {model}: {result}")
         output_image_data = result[0]
         
         # Some models return a list of images as the first element
         if isinstance(output_image_data, list) and len(output_image_data) > 0:
             output_image_data = output_image_data[0]
             
+        gradio_temp_path = None
         if isinstance(output_image_data, dict):
-            gradio_temp_path = output_image_data['path']
-        else:
+            # Try common Gradio keys for image paths
+            gradio_temp_path = output_image_data.get('path') or output_image_data.get('name') or output_image_data.get('url')
+            
+            # Special case: nested 'image' key (found in some models)
+            if not gradio_temp_path and 'image' in output_image_data:
+                img_val = output_image_data['image']
+                if isinstance(img_val, str):
+                    gradio_temp_path = img_val
+                elif isinstance(img_val, dict):
+                    gradio_temp_path = img_val.get('path') or img_val.get('name')
+        elif isinstance(output_image_data, str):
             gradio_temp_path = output_image_data
+
+        if not gradio_temp_path:
+            raise Exception(f"Could not extract image path from result: {output_image_data}")
             
         output_filename = f"{model}_edited_{uuid.uuid4()}.webp"
         final_output_path = os.path.join(OUTPUT_DIR, output_filename)