File size: 2,255 Bytes
e3a8a3d 819742e e3a8a3d 819742e e3a8a3d 819742e e3a8a3d 819742e e3a8a3d 3e6a450 6feb63c e3a8a3d 6feb63c e3a8a3d 819742e e3a8a3d 25042b6 e3a8a3d 819742e | 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 64 65 66 67 68 69 70 | import io, os
from typing import Optional
from fastapi import FastAPI
from fastapi.responses import StreamingResponse, JSONResponse
from pydantic import BaseModel
from huggingface_hub import InferenceClient
app = FastAPI(title="ImageGen API")
# Choose a model that supports text->image via Inference API.
# Using a model that works well with InferenceClient
MODEL_ID = os.getenv("MODEL", "stabilityai/stable-diffusion-xl-base-1.0")
HF_TOKEN = os.getenv("HF_TOKEN") # set in Space Secrets
# Initialize client with or without token
if HF_TOKEN:
client = InferenceClient(model=MODEL_ID, token=HF_TOKEN)
else:
# Use public inference API (may have rate limits)
client = InferenceClient(model=MODEL_ID)
class GenReq(BaseModel):
prompt: str
width: int = 768
height: int = 768
seed: Optional[int] = None
num_inference_steps: Optional[int] = None # some models ignore this
@app.get("/")
def root():
return {"ok": True, "message": "API is live", "model": MODEL_ID}
@app.post("/generate")
def generate(req: GenReq):
try:
# Different models have different parameter requirements
if "sd-turbo" in MODEL_ID.lower():
# SD Turbo works with minimal parameters
img = client.text_to_image(
prompt=req.prompt,
num_inference_steps=1, # SD Turbo works best with 1 step
guidance_scale=0.0, # SD Turbo doesn't use guidance scale
width=req.width,
height=req.height
)
else:
# Standard SDXL parameters
img = client.text_to_image(
prompt=req.prompt,
width=req.width,
height=req.height,
seed=req.seed,
num_inference_steps=req.num_inference_steps or 20
)
buf = io.BytesIO()
img.save(buf, format="PNG")
buf.seek(0)
return StreamingResponse(buf, media_type="image/png")
except Exception as e:
return JSONResponse({"error": str(e)}, status_code=500)
# Keep your original echo for testing
class Prompt(BaseModel):
text: str
@app.post("/echo")
def echo(p: Prompt):
return {"length": len(p.text), "upper": p.text.upper()}
|