Update app.py

app.py

@@ -3,6 +3,9 @@ from gradio_client import Client
 import time
 import concurrent.futures
 import os
+from PIL import Image
+import io
+import requests
 
 # Get token from environment variable
 HF_TOKEN = os.getenv('ArtToken')
@@ -26,7 +29,25 @@ class ModelGenerator:
                 randomize_seed=True,
                 api_name="/run"
             )
-            return ("Midjourney", result)
+            
+            # 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
+                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)
+                else:
+                    # If it's already image data
+                    image = Image.open(io.BytesIO(image_data))
+                return ("Midjourney", image)
+            else:
+                return ("Midjourney", f"Error: Unexpected result format: {type(result)}")
         except Exception as e:
             return ("Midjourney", f"Error: {str(e)}")
 
@@ -251,7 +272,8 @@ def main():
                         if isinstance(result, str) and result.startswith("Error"):
                             st.error(result)
                         else:
-                            st.image(result, use_column_width=True)
+                            # Updated to use use_container_width instead of use_column_width
+                            st.image(result, use_container_width=True)
 
 if __name__ == "__main__":
     main()