Create app.py

app.py

@@ -0,0 +1,72 @@
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+import requests
+import io
+import random
+import os
+from PIL import Image
+import base64
+
+app = FastAPI()
+
+API_URL = f"https://api-inference.huggingface.co/models/{os.getenv("HF_MODEL")}"
+headers = {"Authorization": f"Bearer {os.getenv("HF_TOKEN")}"}
+timeout = 100
+
+class ImageRequest(BaseModel):
+    prompt: str
+    negative_prompt: str = "(deformed, distorted, disfigured), poorly drawn, bad anatomy"
+    steps: int = 4
+    cfg_scale: float = 7.0
+    sampler: str = "DPM++ SDE Karras"
+    seed: int = -1
+    strength: float = 0.7
+
+def query(prompt: str, negative_prompt: str, steps: int, cfg_scale: float, 
+          sampler: str, seed: int, strength: float):
+    if not prompt:
+        raise HTTPException(status_code=400, detail="Prompt is required")
+
+    payload = {
+        "inputs": prompt,
+        "is_negative": bool(negative_prompt),
+        "steps": steps,
+        "cfg_scale": cfg_scale,
+        "seed": seed if seed != -1 else random.randint(1, 1000000000),
+        "strength": strength
+    }
+
+    if negative_prompt:
+        payload["negative_prompt"] = negative_prompt
+
+    response = requests.post(API_URL, headers=headers, json=payload, timeout=timeout)
+    
+    if response.status_code != 200:
+        raise HTTPException(status_code=response.status_code, detail=response.text)
+    
+    image_bytes = response.content
+    image = Image.open(io.BytesIO(image_bytes))
+    
+    buffered = io.BytesIO()
+    image.save(buffered, format="PNG")
+    img_str = base64.b64encode(buffered.getvalue()).decode()
+    
+    return {"image": f"data:image/png;base64,{img_str}"}
+
+@app.post("/generate")
+async def generate_image(request: ImageRequest):
+    try:
+        result = query(
+            prompt=request.prompt,
+            negative_prompt=request.negative_prompt,
+            steps=request.steps,
+            cfg_scale=request.cfg_scale,
+            sampler=request.sampler,
+            seed=request.seed,
+            strength=request.strength
+        )
+        return result
+    except HTTPException as e:
+        raise e
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))