Update app.py

app.py

@@ -2,31 +2,63 @@ import gradio as gr
 import threading
 import os
 import torch
+from PIL import Image
+import numpy as np
 
 os.environ["OMP_NUM_THREADS"] = str(os.cpu_count())
 torch.set_num_threads(os.cpu_count())
 
+stop_event = threading.Event()
+
 # Load models via gradio's built-in client
 model1 = gr.load("models/prithivMLmods/SD3.5-Turbo-Realism-2.0-LoRA")
 model2 = gr.load("models/Purz/face-projection")
 
-stop_event = threading.Event()
-
-def ensure_image_type(possible_img):
+def unwrap_image_output(raw_output):
     """
-    Unwraps the various container types that might be returned
-    (tuple, list, dict) until we end up with an actual PIL or
-    numpy image object that gr.Image() can handle.
+    Unwraps raw_output until we (hopefully) get a PIL Image or NumPy array.
+    Gradio's Image component can display:
+      - A PIL image,
+      - A NumPy array (H x W x C),
+      - A file path (string) pointing to an image on disk, or
+      - A base64-encoded image string.
+    This function attempts to handle the first two cases.
+
+    If the model returns a tuple/list/dict that does not
+    contain an actual image, we raise a ValueError so that
+    it's clear how to proceed.
     """
-    # If it's a list or tuple, assume the first element is the image
-    while isinstance(possible_img, (list, tuple)):
-        possible_img = possible_img[0]
-    
-    # If it's a dict with an 'image' key, get that
-    if isinstance(possible_img, dict) and "image" in possible_img:
-        possible_img = possible_img["image"]
-    
-    return possible_img
+    # 1) Keep unwrapping if it's a tuple/list
+    #    (e.g. sometimes pipelines return (image, text))
+    while isinstance(raw_output, (list, tuple)):
+        if not raw_output:
+            raise ValueError("Received an empty list/tuple from the model.")
+        raw_output = raw_output[0]  # Take the first element
+
+    # 2) Check if it's a dict containing an "image" key
+    if isinstance(raw_output, dict):
+        # Common pattern: raw_output.get("image") might be the real image
+        if "image" in raw_output:
+            raw_output = raw_output["image"]
+        else:
+            raise ValueError(
+                f"Received a dictionary but could not find an 'image' key. "
+                f"Keys present: {list(raw_output.keys())}"
+            )
+
+    # 3) At this point, raw_output should be either a PIL image, NumPy array, or a path
+    if isinstance(raw_output, Image.Image):
+        return raw_output
+    elif isinstance(raw_output, np.ndarray):
+        return raw_output
+    elif isinstance(raw_output, str):
+        # If it's a file path or base64 string, we can return it as-is
+        return raw_output
+
+    # If we get here, it means we still have some unexpected type
+    raise ValueError(
+        f"Could not convert output to an image. It is of type: {type(raw_output)}"
+    )
 
 def generate_images(text, selected_model):
     stop_event.clear()
@@ -44,11 +76,13 @@ def generate_images(text, selected_model):
             return ["Image generation stopped by user."] * 3
 
         modified_text = f"{text} variation {i+1}"
+
+        # Call the model
         raw_output = model(modified_text)
 
-        # Unwrap raw_output to get a proper image object
-        image = ensure_image_type(raw_output)
-        results.append(image)
+        # Unwrap it to get a proper single image
+        img = unwrap_image_output(raw_output)
+        results.append(img)
 
     return results
 
@@ -59,9 +93,9 @@ def stop_generation():
 
 with gr.Blocks() as interface:
     gr.Markdown(
-        "### ⚠ Sorry for the inconvenience. The Space is currently running on the CPU, which might affect performance."
+        "### ⚠ Sorry for the inconvenience. The Space is currently running on CPU, which might affect performance."
     )
-    
+
     text_input = gr.Textbox(
         label="Welcome to EpicFrame. Set free your imagination!",
         placeholder="Type your prompt. Example: A mob boss smoking a cigar outside a tiny cafe."
@@ -71,16 +105,16 @@ with gr.Blocks() as interface:
         label="Select Model",
         value="Model 1 (Turbo Realism)"
     )
-    
+
     with gr.Row():
         generate_button = gr.Button("Generate 3 Images 🎨")
         stop_button = gr.Button("Stop Image Generation")
-    
+
     with gr.Row():
         output1 = gr.Image(label="Generated Image 1")
         output2 = gr.Image(label="Generated Image 2")
         output3 = gr.Image(label="Generated Image 3")
-    
+
     generate_button.click(
         fn=generate_images,
         inputs=[text_input, model_selector],
@@ -92,4 +126,4 @@ with gr.Blocks() as interface:
         outputs=[output1, output2, output3]
     )
 
-interface.launch()
+interface.launch()