File size: 1,112 Bytes
8332ab6 |
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 |
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from gradio_client import Client
# Init FastAPI app
app = FastAPI()
# Initialize the Gradio Client
client = Client("Efficient-Large-Model/SanaSprint")
# Request body schema
class GenerationRequest(BaseModel):
prompt: str
model_size: str = "1.6B"
seed: int = 0
randomize_seed: bool = True
width: int = 1024
height: int = 1024
guidance_scale: float = 4.5
num_inference_steps: int = 2
@app.post("/generate")
async def generate_image(request: GenerationRequest):
try:
result = client.predict(
prompt=request.prompt,
model_size=request.model_size,
seed=request.seed,
randomize_seed=request.randomize_seed,
width=request.width,
height=request.height,
guidance_scale=request.guidance_scale,
num_inference_steps=request.num_inference_steps,
api_name="/infer"
)
return {"result": result}
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
|