Update app.py

app.py

@@ -1,29 +1,37 @@
 import gradio as gr
 from PIL import Image
 import os
+import base64
+import io
 
 def add_logo(main_image, logo_size_percent=20, margin_percent=2):
-    # Handle API input format
-    if isinstance(main_image, dict) and 'data' in main_image:
-        # Extract base64 from API format
-        import base64
-        import io
-        base64_str = main_image['data']
-        if base64_str.startswith('data:'):
-            base64_str = base64_str.split(',')[1]
-        image_data = base64.b64decode(base64_str)
-        main_image = Image.open(io.BytesIO(image_data))
     """
     Add logo to the top-right corner of the main image
     
     Args:
-        main_image: PIL Image object or file path
+        main_image: PIL Image object, file path, or base64 string (for API)
         logo_size_percent: Size of logo as percentage of main image width (default 20%)
         margin_percent: Margin from edges as percentage of main image width (default 2%)
     """
     if main_image is None:
         return None
     
+    # Handle base64 input from API
+    if isinstance(main_image, str) and main_image.startswith('data:image'):
+        # Remove data URL prefix
+        base64_str = main_image.split(',')[1]
+        # Decode base64 to image
+        image_data = base64.b64decode(base64_str)
+        main_image = Image.open(io.BytesIO(image_data))
+    elif isinstance(main_image, str):
+        # Try to decode as plain base64 (without data URL prefix)
+        try:
+            image_data = base64.b64decode(main_image)
+            main_image = Image.open(io.BytesIO(image_data))
+        except:
+            # If it fails, assume it's a file path
+            main_image = Image.open(main_image)
+    
     # Load the logo
     logo_path = "logo.png"
     if not os.path.exists(logo_path):
@@ -58,7 +66,7 @@ def add_logo(main_image, logo_size_percent=20, margin_percent=2):
     
     return result
 
-# Create Gradio interface
+# Create Gradio interface with API mode
 iface = gr.Interface(
     fn=add_logo,
     inputs=[
@@ -67,7 +75,8 @@ iface = gr.Interface(
     outputs=gr.Image(type="pil", label="Image with Logo"),
     title="Logo Adder",
     description="Upload an image to add logo.png to the top-right corner",
-    examples=None
+    examples=None,
+    api_name="predict"  # Explicit API name
 )
 
 if __name__ == "__main__":