File size: 1,325 Bytes
a31ed9b 46e68f0 a31ed9b 05e0399 a31ed9b 05e0399 46e68f0 a31ed9b 05e0399 a31ed9b 05e0399 a31ed9b 46e68f0 a31ed9b 46e68f0 a31ed9b 05e0399 a31ed9b |
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 |
from typing import Dict, List
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from diffusers import StableDiffusionXLPipeline
import torch
import base64
from io import BytesIO
app = FastAPI()
# Load the model
model_name = "colt12/maxcushion"
pipe = StableDiffusionXLPipeline.from_pretrained(model_name, torch_dtype=torch.float16, use_safetensors=True, variant="fp16")
pipe = pipe.to("cuda" if torch.cuda.is_available() else "cpu")
class Item(BaseModel):
prompt: str
negative_prompt: str = ""
num_inference_steps: int = 30
guidance_scale: float = 7.5
@app.post("/generate")
async def generate(item: Item) -> Dict[str, str]:
try:
# Generate the image
image = pipe(
prompt=item.prompt,
negative_prompt=item.negative_prompt,
num_inference_steps=item.num_inference_steps,
guidance_scale=item.guidance_scale
).images[0]
# Convert to base64
buffered = BytesIO()
image.save(buffered, format="PNG")
image_base64 = base64.b64encode(buffered.getvalue()).decode("utf-8")
return {"image": image_base64}
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@app.get("/")
async def root():
return {"message": "SDXL Image Generation API"} |