Spaces:
Paused
Paused
File size: 1,992 Bytes
da72ab7 c6bb57f da72ab7 a1df811 d885fee da72ab7 c6bb57f da72ab7 c6bb57f da72ab7 c6bb57f da72ab7 c6bb57f da72ab7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 | from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import requests
import random
import os
from fastapi.responses import Response
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++ 2M 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)
return response.content
@app.post("/generate")
async def generate_image(request: ImageRequest):
try:
raw_data = 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 Response(content=raw_data, media_type="application/octet-stream")
except HTTPException as e:
raise e
except Exception as e:
raise HTTPException(status_code=500, detail=str(e)) |