| 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() |
|
|
| |
| 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: |
| |
| image = pipe( |
| prompt=item.prompt, |
| negative_prompt=item.negative_prompt, |
| num_inference_steps=item.num_inference_steps, |
| guidance_scale=item.guidance_scale |
| ).images[0] |
|
|
| |
| 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"} |
