Update app.py

app.py

@@ -7,6 +7,9 @@ from datetime import datetime
 
 # Load the models
 HF_TOKEN = os.getenv("HF_TOKEN", None)
+from PIL import Image
+import io
+
 
 def load_models(models):
     loaded_models = {}
@@ -28,14 +31,27 @@ async def infer(model_str, prompt, seed=-1):
     result = await asyncio.wait_for(task, timeout=300)
     return result
 
+
 def generate_image(model_name, prompt, seed):
     loop = asyncio.new_event_loop()
     asyncio.set_event_loop(loop)
     try:
-        return loop.run_until_complete(infer(model_name, prompt, seed))
+        # Get the result from inference
+        result = loop.run_until_complete(infer(model_name, prompt, seed))
+        if isinstance(result, tuple):
+            # Assuming the first element is the image data
+            result = result[0]
+        if isinstance(result, bytes):
+            # Convert bytes to PIL Image if necessary
+            return Image.open(io.BytesIO(result))
+        elif isinstance(result, Image.Image):
+            return result
+        else:
+            raise ValueError(f"Unexpected output type: {type(result)}")
     finally:
         loop.close()
 
+        
 # Interface
 with gr.Blocks() as demo:
     with gr.Column():