Update: Improve API implementation

api.py

@@ -1,58 +1,56 @@
-from fastapi import FastAPI, HTTPException, Request
+from fastapi import FastAPI, Response, HTTPException
 from pydantic import BaseModel
-from typing import Dict, Any, Optional, List, Union
-import base64
-import io
-from PIL import Image
-import torch
+from typing import Union, Dict, Any, Optional
 import os
+import io
 import sys
-import json
-
-# Import the handler
 from handler import EndpointHandler
 
-# Initialize the app
+# Add debug logging
+def debug_log(message):
+    print(f"DEBUG API: {message}")
+    sys.stdout.flush()
+
+debug_log("Initializing API")
+
 app = FastAPI()
 
-# Initialize the model
-model = EndpointHandler(model_dir="/code")
+# Initialize the handler with the model directory
+model_dir = os.environ.get("MODEL_DIR", "/code")
+debug_log(f"Using model_dir: {model_dir}")
+handler = EndpointHandler(model_dir)
+
+class TextRequest(BaseModel):
+    inputs: Union[str, Dict[str, Any]]
+
+@app.get("/")
+def read_root():
+    return {"message": "Vector Graphics Generation API"}
 
 @app.post("/")
-async def process_request(request: Request):
+async def generate(request: TextRequest):
     try:
-        # Get the raw request body
-        body = await request.body()
+        debug_log(f"Received request: {request}")
         
-        # Try to parse as JSON
-        try:
-            data = json.loads(body)
-        except:
-            # If not JSON, treat as plain text
-            data = {"inputs": body.decode("utf-8")}
+        # Call the handler
+        result = handler(request.dict())
+        debug_log(f"Handler returned result of type: {type(result)}")
         
-        # 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)}
+        # If the result is a PIL Image, convert it to bytes
+        if hasattr(result, "save"):
+            debug_log("Converting PIL Image to bytes")
+            img_byte_arr = io.BytesIO()
+            result.save(img_byte_arr, format="PNG")
+            img_byte_arr.seek(0)
+            
+            # Return the image as a response
+            return Response(content=img_byte_arr.getvalue(), media_type="image/png")
         else:
-            # Not a dict, use as is
-            data = {"inputs": str(data)}
-        
-        # Process the request
-        result = model(data)
-        return result
+            # Return the result as JSON
+            debug_log("Returning result as JSON")
+            return result
     except Exception as e:
+        debug_log(f"Error in generate endpoint: {e}")
+        import traceback
+        debug_log(traceback.format_exc())
         raise HTTPException(status_code=500, detail=str(e))
-
-# Add a health check endpoint
-@app.get("/health")
-async def health():
-    return {"status": "ok"}