Upload model with FastAPI app

Dockerfile

@@ -0,0 +1,10 @@
+FROM python:3.10-slim
+
+WORKDIR /app
+
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+COPY . .
+
+CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "8000"]

app.py

@@ -0,0 +1,71 @@
+from typing import Dict, List, Any, Optional, Union
+import torch
+import base64
+import io
+from PIL import Image
+import os
+from diffusers import DiffusionPipeline
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+
+app = FastAPI()
+
+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))

requirements.txt

@@ -2,4 +2,9 @@ diffusers>=0.26.3
 transformers>=4.36.2
 torch>=2.0.0
 numpy>=1.24.0
-pillow>=9.0.0
+pillow>=9.0.0
+huggingface_hub
+requests
+fastapi
+uvicorn
+pydantic