Fix: Update API to better handle text-to-image requests

api.py

@@ -1,4 +1,4 @@
-from fastapi import FastAPI, HTTPException
+from fastapi import FastAPI, HTTPException, Request
 from pydantic import BaseModel
 from typing import Dict, Any, Optional, List, Union
 import base64
@@ -7,6 +7,7 @@ from PIL import Image
 import torch
 import os
 import sys
+import json
 
 # Import the handler
 from handler import EndpointHandler
@@ -17,15 +18,36 @@ app = FastAPI()
 # Initialize the model
 model = EndpointHandler(model_dir="/code")
 
-class TextToImageRequest(BaseModel):
-    inputs: Union[str, Dict[str, Any]]
-    parameters: Optional[Dict[str, Any]] = None
-
 @app.post("/")
-async def text_to_image(request: TextToImageRequest):
+async def process_request(request: Request):
     try:
+        # Get the raw request body
+        body = await request.body()
+        
+        # Try to parse as JSON
+        try:
+            data = json.loads(body)
+        except:
+            # If not JSON, treat as plain text
+            data = {"inputs": body.decode("utf-8")}
+        
+        # Handle different input formats
+        if isinstance(data, dict):
+            if "inputs" in data:
+                # Standard format
+                pass
+            elif "text" in data:
+                # Text field directly
+                data = {"inputs": data["text"]}
+            else:
+                # No recognized fields, use the whole dict as input
+                data = {"inputs": str(data)}
+        else:
+            # Not a dict, use as is
+            data = {"inputs": str(data)}
+        
         # Process the request
-        result = model(request.dict())
+        result = model(data)
         return result
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))