Fix: Update API and config for text-to-image task

api.py

@@ -1,27 +1,36 @@
-#!/usr/bin/env python
-# -*- coding: utf-8 -*-
-
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+from typing import Dict, Any, Optional, List, Union
+import base64
+import io
+from PIL import Image
+import torch
 import os
-import json
-from fastapi import FastAPI, Request, HTTPException
-import uvicorn
+import sys
 
-# Import the handler directly
-from handler import handler
+# Import the handler
+from handler import EndpointHandler
 
+# Initialize the app
 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 process_request(request: Request):
+async def text_to_image(request: TextToImageRequest):
     try:
-        json_data = await request.json()
-        return handler(json_data)
+        # Process the request
+        result = model(request.dict())
+        return result
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
+# Add a health check endpoint
 @app.get("/health")
-async def health_check():
-    return {"status": "ok"}
-
-if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+async def health():
+    return {"status": "ok"}

config.json

@@ -1,34 +1,22 @@
 {
-  "model_type": "vector_graphics",
-  "name": "DiffSketcher",
   "task": "text-to-image",
-  "parameters": {
-    "prompt": {
-      "type": "string",
-      "description": "Text description of the desired output",
-      "required": true
-    },
-    "negative_prompt": {
-      "type": "string",
-      "description": "Text to avoid in the generation",
-      "required": false
-    },
-    "num_paths": {
-      "type": "integer",
-      "description": "Number of paths in the SVG",
-      "required": false,
-      "default": 100
-    },
-    "guidance_scale": {
-      "type": "number",
-      "description": "Guidance scale for the diffusion model",
-      "required": false,
-      "default": 7.5
-    },
-    "seed": {
-      "type": "integer",
-      "description": "Random seed for reproducibility",
-      "required": false
+  "framework": "custom",
+  "model-type": "text-to-image",
+  "inference": {
+    "task_type": "text-to-image",
+    "text_to_image": {
+      "inputs": {
+        "text": {
+          "type": "string",
+          "description": "Input text prompt"
+        }
+      },
+      "outputs": {
+        "image": {
+          "type": "image",
+          "description": "Generated image"
+        }
+      }
     }
   }
 }