Upload app.py with huggingface_hub

app.py

@@ -1,71 +1,13 @@
-from typing import Dict, List, Any, Optional, Union
-import torch
-import base64
-import io
-from PIL import Image
+
+from fastapi import FastAPI, Request
+from handler import EndpointHandler
 import os
-from diffusers import DiffusionPipeline
-from fastapi import FastAPI, HTTPException
-from pydantic import BaseModel
 
 app = FastAPI()
+handler = EndpointHandler(os.getcwd())
 
-class DiffSketcherInput(BaseModel):
-    prompt: str
-    negative_prompt: Optional[str] = None
-    num_paths: Optional[int] = 96
-    num_iter: Optional[int] = 800
-    guidance_scale: Optional[float] = 7.5
-    width: Optional[float] = 2.0
-    seed: Optional[int] = None
-
-class DiffSketcherOutput(BaseModel):
-    svg: str
-    image: str
-
-# Load the model
-model = None
-
-def load_model():
-    global model
-    if model is None:
-        model = DiffusionPipeline.from_pretrained(".")
-        model.to(torch.device("cuda" if torch.cuda.is_available() else "cpu"))
-    return model
-
-@app.post("/", response_model=DiffSketcherOutput)
-async def generate(input_data: DiffSketcherInput):
-    # Load the model
-    model = load_model()
-    
-    # Set the seed if provided
-    if input_data.seed is not None:
-        torch.manual_seed(input_data.seed)
-    
-    try:
-        # Generate the SVG
-        output = model(
-            prompt=input_data.prompt,
-            negative_prompt=input_data.negative_prompt,
-            num_paths=input_data.num_paths,
-            num_iter=input_data.num_iter,
-            guidance_scale=input_data.guidance_scale,
-            width=input_data.width
-        )
-        
-        # Get the SVG and image
-        svg = output.svg
-        image = output.images[0]
-        
-        # Convert the image to base64
-        buffered = io.BytesIO()
-        image.save(buffered, format="PNG")
-        img_str = base64.b64encode(buffered.getvalue()).decode()
-        
-        # Return the results
-        return {
-            "svg": svg,
-            "image": img_str
-        }
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
+@app.post("/")
+async def process_request(request: Request):
+    json_data = await request.json()
+    return handler(json_data)
+