Update app.py

app.py

@@ -28,38 +28,52 @@ def load_fn(models):
 
 load_fn(models)
 
-def save_image_to_tmp(image_data):
+def process_image_output(image_data):
+    """Handle various image output formats and save to temp directory"""
     if image_data is None:
         return None
         
-    # Handle case where image_data is already a file path
-    if isinstance(image_data, str):
-        if image_data.startswith(('http://', 'https://')):
+    # If it's already a file path that exists, use it
+    if isinstance(image_data, str) and os.path.isfile(image_data):
+        return image_data
+        
+    # If it's a URL, download it
+    if isinstance(image_data, str) and image_data.startswith(('http://', 'https://')):
+        try:
             response = requests.get(image_data)
             img = Image.open(BytesIO(response.content))
-        elif os.path.isfile(image_data):
-            img = Image.open(image_data)
-        else:
+        except:
             return None
     else:
-        # Assume it's a PIL Image or compatible object
+        # Assume it's a PIL Image or numpy array
         img = image_data
     
+    # Save to temp file
     filename = f"{uuid.uuid4()}.jpg"
     filepath = os.path.join(IMAGE_DIR, filename)
-    img.save(filepath)
-    return filepath  # Return full path for Gradio to handle
+    
+    try:
+        if isinstance(img, Image.Image):
+            img.save(filepath)
+        else:
+            # Handle numpy arrays or other formats
+            Image.fromarray(img).save(filepath)
+        return filepath
+    except Exception as e:
+        print(f"Error saving image: {e}")
+        return None
 
 def gen_fn(model_str, prompt):
     if model_str == 'NA':
         return None
+        
     noise = str(randint(0, 4294967296))
     klir = '| ultra detail, ultra elaboration, ultra quality, perfect'
     result = models_load[model_str](f'{prompt} {klir} {noise}')
     
     if isinstance(result, list):
-        return [save_image_to_tmp(img) for img in result if img is not None]
-    return save_image_to_tmp(result)
+        return [process_image_output(img) for img in result if img is not None]
+    return process_image_output(result)
 
 def create_interface():
     with gr.Blocks(css=css_code) as demo: