Update app.py

app.py

@@ -3,12 +3,15 @@ import os
 import sys
 import cv2
 import torch
+import gradio as gr
 import numpy as np
 from PIL import Image, ImageOps
 import io
 import base64
 import traceback
-import gradio as gr
+import tempfile
+from fastapi import FastAPI, File, UploadFile
+from fastapi.middleware.cors import CORSMiddleware
 import spaces
 
 # Import model-specific libraries
@@ -19,7 +22,7 @@ try:
     print("Successfully imported model libraries.")
 except ImportError as e:
     print(f"Error importing model libraries: {e}")
-    print("Please ensure basicsr, gfpgan, realesrgan are installed or in requirements.txt")
+    print("Please ensure basicsr, gfpgan, realesrgan are installed")
     sys.exit(1)
 
 # --- Constants ---
@@ -69,7 +72,7 @@ except Exception as e:
     print(traceback.format_exc())
     print("Warning: GFPGAN will run without background enhancement.")
 
-# --- IMPROVED: Enhanced processing function that handles both file paths and PIL images ---
+# --- Universal processing function ---
 @spaces.GPU(duration=90)
 def process_image(input_image, version, scale):
     """
@@ -265,28 +268,111 @@ def process_image(input_image, version, scale):
             pass
         return error_img, error_msg
 
-# --- Improved API wrapper function to handle both web UI and API calls ---
+# --- Function to handle file upload for API ---
+def handle_file_upload(file_data):
+    """Save uploaded file to temporary directory and return path"""
+    try:
+        print(f"Handling file upload: {type(file_data)}")
+        
+        # Create a temporary file
+        temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.jpg')
+        temp_path = temp_file.name
+        
+        # If it's bytes, write directly
+        if isinstance(file_data, bytes):
+            with open(temp_path, 'wb') as f:
+                f.write(file_data)
+        # If it's a file-like object (from FastAPI/Gradio)
+        elif hasattr(file_data, 'file'):
+            content = file_data.file.read()
+            with open(temp_path, 'wb') as f:
+                f.write(content)
+        # If it's a string path, it's already saved
+        elif isinstance(file_data, str) and os.path.exists(file_data):
+            return file_data
+        else:
+            raise ValueError(f"Unsupported file data type: {type(file_data)}")
+            
+        print(f"File saved to temporary path: {temp_path}")
+        return temp_path
+    
+    except Exception as e:
+        print(f"Error handling file upload: {e}")
+        print(traceback.format_exc())
+        raise
+
+# --- API inference function ---
 @spaces.GPU(duration=90)
 def inference(input_image, version, scale):
     """
     API-friendly wrapper that ensures consistent behavior between web and API interfaces.
     """
-    # Process the image using our universal processor
-    output_pil, base64_or_msg = process_image(input_image, version, scale)
-    
-    # Return the processed results in a format suitable for both UI and API
-    return output_pil, base64_or_msg
+    try:
+        # If input is a file upload (from API), save it to a temporary path
+        if not isinstance(input_image, (str, Image.Image, np.ndarray)) and not (hasattr(input_image, 'name') and os.path.exists(input_image.name)):
+            file_path = handle_file_upload(input_image)
+            input_image = file_path
+        
+        # Process the image
+        output_pil, base64_or_msg = process_image(input_image, version, scale)
+        
+        # Return the processed results
+        return output_pil, base64_or_msg
+    except Exception as e:
+        print(f"Error in inference: {e}")
+        print(traceback.format_exc())
+        # Return a placeholder error image and message
+        error_img = Image.new('RGB', (100, 50), color='red')
+        return error_img, f"Error: {str(e)}"
+
+# --- Get the FastAPI app from Gradio ---
+app = FastAPI()
+
+# Add CORS middleware to allow cross-origin requests
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],  # Allows all origins
+    allow_credentials=True,
+    allow_methods=["*"],  # Allows all methods
+    allow_headers=["*"],  # Allows all headers
+)
+
+# --- Direct API endpoint for file upload ---
+@app.post("/api/direct-process")
+async def direct_process(file: UploadFile = File(...), version: str = "v1.4", scale: float = 2.0):
+    try:
+        # Save the uploaded file
+        temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.jpg')
+        temp_path = temp_file.name
+        with open(temp_path, 'wb') as f:
+            f.write(await file.read())
+        
+        # Process the image
+        _, base64_image = process_image(temp_path, version, scale)
+        
+        # Clean up
+        os.unlink(temp_path)
+        
+        # Return base64 image data
+        if base64_image and base64_image.startswith('data:image'):
+            return {"success": True, "image": base64_image}
+        else:
+            return {"success": False, "error": base64_image or "Unknown error"}
+    except Exception as e:
+        print(f"Error in direct-process API: {e}")
+        print(traceback.format_exc())
+        return {"success": False, "error": str(e)}
 
 # --- Gradio Interface Definition ---
 title = "GFPGAN: Practical Face Restoration"
 description = """Gradio demo for <a href='https://github.com/TencentARC/GFPGAN' target='_blank'><b>GFPGAN: Towards Real-World Blind Face Restoration with Generative Facial Prior</b></a>.
 <br>Restore your <b>old photos</b> or improve <b>AI-generated faces</b>. Upload an image to start.
 <br>If helpful, please ⭐ the <a href='https://github.com/TencentARC/GFPGAN' target='_blank'>Original Github Repo</a>.
-<br>API endpoint available at `/predict`. Expects image file, version string, scale number. Returns Image and Base64 data.
+<br>API endpoint available at `/predict` or `/api/direct-process`. Returns processed image and base64 data.
 """
 article = "Questions? Contact the original creators (see GFPGAN repo)."
 
-# --- IMPROVED: Dual input type handling (both "filepath" for UI and "pil" for API) ---
+# Use upload component for more compatibility
 inputs = [
     gr.Image(type="pil", label="Input Image", sources=["upload", "clipboard"]),
     gr.Radio(
@@ -325,6 +411,10 @@ demo = gr.Interface(
     allow_flagging='never'
 )
 
+# Mount the Gradio app
+app = gr.mount_gradio_app(app, demo, path="/")
+
 # Launch the interface
 if __name__ == "__main__":
-    demo.queue().launch(server_name="0.0.0.0", share=False)
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)