Update app.py

app.py

@@ -32,19 +32,24 @@ class ModelGenerator:
                 api_name="/run"
             )
             
-            if isinstance(result, list) and len(result) > 0:
-                image_data = result[0]
+            if isinstance(result, tuple):
+                image_data = result[0] if len(result) > 0 else None
+            elif isinstance(result, list):
+                image_data = result[0] if len(result) > 0 else None
+            else:
+                image_data = result
+
+            if image_data:
                 if isinstance(image_data, str):
                     if image_data.startswith('http'):
                         response = requests.get(image_data)
-                        image = Image.open(io.BytesIO(response.content))
-                    else:
-                        image = Image.open(image_data)
-                else:
-                    image = Image.open(io.BytesIO(image_data))
-                return ("Midjourney", image)
-            else:
-                return ("Midjourney", f"Error: Unexpected result format: {type(result)}")
+                        return ("Midjourney", Image.open(io.BytesIO(response.content)))
+                    return ("Midjourney", Image.open(image_data))
+                elif isinstance(image_data, bytes):
+                    return ("Midjourney", Image.open(io.BytesIO(image_data)))
+                elif hasattr(image_data, 'read'):  # File-like object
+                    return ("Midjourney", Image.open(image_data))
+            return ("Midjourney", "Error: No valid image data found")
         except Exception as e:
             return ("Midjourney", f"Error: {str(e)}")
 
@@ -65,7 +70,11 @@ class ModelGenerator:
                 num_images_per_prompt=1,
                 api_name="/run"
             )
-            return ("Stable Cascade", result)
+            if isinstance(result, (str, bytes)):
+                return ("Stable Cascade", Image.open(io.BytesIO(result) if isinstance(result, bytes) else result))
+            elif isinstance(result, list) and len(result) > 0:
+                return ("Stable Cascade", Image.open(io.BytesIO(result[0]) if isinstance(result[0], bytes) else result[0]))
+            return ("Stable Cascade", "Error: No valid image data found")
         except Exception as e:
             return ("Stable Cascade", f"Error: {str(e)}")
 
@@ -84,7 +93,23 @@ class ModelGenerator:
                 num_inference_steps=28,
                 api_name="/infer"
             )
-            return ("SD 3 Medium", result)
+            if isinstance(result, bytes):
+                return ("SD 3 Medium", Image.open(io.BytesIO(result)))
+            elif isinstance(result, str):
+                if result.startswith('http'):
+                    response = requests.get(result)
+                    return ("SD 3 Medium", Image.open(io.BytesIO(response.content)))
+                return ("SD 3 Medium", Image.open(result))
+            elif isinstance(result, list) and len(result) > 0:
+                image_data = result[0]
+                if isinstance(image_data, bytes):
+                    return ("SD 3 Medium", Image.open(io.BytesIO(image_data)))
+                elif isinstance(image_data, str):
+                    if image_data.startswith('http'):
+                        response = requests.get(image_data)
+                        return ("SD 3 Medium", Image.open(io.BytesIO(response.content)))
+                    return ("SD 3 Medium", Image.open(image_data))
+            return ("SD 3 Medium", "Error: No valid image data found")
         except Exception as e:
             return ("SD 3 Medium", f"Error: {str(e)}")
 
@@ -103,7 +128,23 @@ class ModelGenerator:
                 num_inference_steps=40,
                 api_name="/infer"
             )
-            return ("SD 3.5 Large", result)
+            if isinstance(result, bytes):
+                return ("SD 3.5 Large", Image.open(io.BytesIO(result)))
+            elif isinstance(result, str):
+                if result.startswith('http'):
+                    response = requests.get(result)
+                    return ("SD 3.5 Large", Image.open(io.BytesIO(response.content)))
+                return ("SD 3.5 Large", Image.open(result))
+            elif isinstance(result, list) and len(result) > 0:
+                image_data = result[0]
+                if isinstance(image_data, bytes):
+                    return ("SD 3.5 Large", Image.open(io.BytesIO(image_data)))
+                elif isinstance(image_data, str):
+                    if image_data.startswith('http'):
+                        response = requests.get(image_data)
+                        return ("SD 3.5 Large", Image.open(io.BytesIO(response.content)))
+                    return ("SD 3.5 Large", Image.open(image_data))
+            return ("SD 3.5 Large", "Error: No valid image data found")
         except Exception as e:
             return ("SD 3.5 Large", f"Error: {str(e)}")
 
@@ -123,8 +164,15 @@ class ModelGenerator:
                 True,   # randomize seed
                 api_name="/run"
             )
-            if result and isinstance(result, tuple) and result[0]:
-                return ("Playground v2.5", result[0][0]['image'])
+            if isinstance(result, tuple) and result[0] and len(result[0]) > 0:
+                image_data = result[0][0].get('image')
+                if image_data:
+                    if isinstance(image_data, str):
+                        if image_data.startswith('http'):
+                            response = requests.get(image_data)
+                            return ("Playground v2.5", Image.open(io.BytesIO(response.content)))
+                        return ("Playground v2.5", Image.open(image_data))
+                    return ("Playground v2.5", Image.open(io.BytesIO(image_data)))
             return ("Playground v2.5", "Error: No image generated")
         except Exception as e:
             return ("Playground v2.5", f"Error: {str(e)}")
@@ -150,7 +198,12 @@ def generate_images(prompt, selected_models):
                 futures.append(executor.submit(model_map[model], prompt))
         
         for future in concurrent.futures.as_completed(futures):
-            results.append(future.result())
+            try:
+                result = future.result()
+                if result:
+                    results.append(result)
+            except Exception as e:
+                st.error(f"Error during image generation: {str(e)}")
     
     return results