Update app.py

app.py

@@ -30,22 +30,27 @@ class ModelGenerator:
                 api_name="/run"
             )
             
-            # Handle the result based on its type
-            if isinstance(result, list) and len(result) > 0:
-                # If result is a list of file paths or URLs
+            # Handle different types of results
+            if isinstance(result, tuple):
+                # If it's a tuple, the first element might be the image gallery
+                if len(result) > 0 and isinstance(result[0], list):
+                    image_data = result[0][0]  # Get first image from gallery
+                    if isinstance(image_data, dict) and 'image' in image_data:
+                        return ("Midjourney", image_data['image'])
+                    elif isinstance(image_data, (str, bytes)):
+                        return ("Midjourney", image_data)
+                else:
+                    return ("Midjourney", result[0])  # Try first element of tuple
+            elif isinstance(result, list) and len(result) > 0:
+                # If it's a list, get the first element
                 image_data = result[0]
-                if isinstance(image_data, str):
-                    if image_data.startswith('http'):
-                        # If it's a URL, download the image
-                        response = requests.get(image_data)
-                        image = Image.open(io.BytesIO(response.content))
-                    else:
-                        # If it's a file path
-                        image = Image.open(image_data)
+                if isinstance(image_data, dict) and 'image' in image_data:
+                    return ("Midjourney", image_data['image'])
                 else:
-                    # If it's already image data
-                    image = Image.open(io.BytesIO(image_data))
-                return ("Midjourney", image)
+                    return ("Midjourney", image_data)
+            elif isinstance(result, str):
+                # If it's a direct string (URL or path)
+                return ("Midjourney", result)
             else:
                 return ("Midjourney", f"Error: Unexpected result format: {type(result)}")
         except Exception as e: