Add error handling to expose exceptions in API responses

app.py

@@ -6,6 +6,7 @@ import numpy as np
 import random
 import spaces
 import torch
+import traceback
 from diffusers import Flux2KleinPipeline
 from PIL import Image
 
@@ -30,27 +31,45 @@ def infer(
     num_inference_steps: int = 4,
     guidance_scale: float = 1.0,
 ):
-    if isinstance(seed, str): seed = int(seed)
-    if isinstance(randomize_seed, str): randomize_seed = randomize_seed.lower() == "true"
-    if isinstance(width, str): width = int(width)
-    if isinstance(height, str): height = int(height)
-    if isinstance(num_inference_steps, str): num_inference_steps = int(num_inference_steps)
-    if isinstance(guidance_scale, str): guidance_scale = float(guidance_scale)
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-    generator = torch.Generator(device="cuda").manual_seed(seed)
-    pipe_kwargs = {
-        "prompt": prompt,
-        "height": height,
-        "width": width,
-        "num_inference_steps": num_inference_steps,
-        "guidance_scale": guidance_scale,
-        "generator": generator,
-    }
-    if input_images is not None and len(input_images) > 0:
-        pipe_kwargs["image"] = [item[0] if isinstance(item, tuple) else item for item in input_images]
-    image = pipe(**pipe_kwargs).images[0]
-    return image, seed
+    try:
+        if isinstance(seed, str): seed = int(seed)
+        if isinstance(randomize_seed, str): randomize_seed = randomize_seed.lower() == "true"
+        width = int(float(width))
+        height = int(float(height))
+        num_inference_steps = int(float(num_inference_steps))
+        if isinstance(guidance_scale, str): guidance_scale = float(guidance_scale)
+        if randomize_seed:
+            seed = random.randint(0, MAX_SEED)
+        generator = torch.Generator(device="cuda").manual_seed(seed)
+        pipe_kwargs = {
+            "prompt": prompt,
+            "height": height,
+            "width": width,
+            "num_inference_steps": num_inference_steps,
+            "guidance_scale": guidance_scale,
+            "generator": generator,
+        }
+        print(f"input_images type: {type(input_images)}, value: {input_images}")
+        if input_images is not None and len(input_images) > 0:
+            imgs = []
+            for item in input_images:
+                if isinstance(item, tuple):
+                    imgs.append(item[0])
+                elif isinstance(item, Image.Image):
+                    imgs.append(item)
+                else:
+                    print(f"Unknown item type: {type(item)}, value: {item}")
+                    imgs.append(item)
+            pipe_kwargs["image"] = imgs
+        print(f"pipe_kwargs keys: {list(pipe_kwargs.keys())}")
+        print(f"image count: {len(pipe_kwargs.get('image', []))}")
+        result_image = pipe(**pipe_kwargs).images[0]
+        return result_image, seed
+    except Exception as e:
+        tb = traceback.format_exc()
+        print(f"ERROR: {e}")
+        print(tb)
+        raise gr.Error(f"{type(e).__name__}: {e}")
 
 
 with gr.Blocks() as demo: