Pacicap/stable_diffusion_api

1from fastapi import FastAPI, Request2from pydantic import BaseModel3from fastapi.middleware.cors import CORSMiddleware4from diffusers import DiffusionPipeline5import torch6import uuid7import os8from PIL import Image9from fastapi.staticfiles import StaticFiles10 11app = FastAPI()12 13app.add_middleware(14    CORSMiddleware,15    allow_origins=["*"],  # Accept from all for now16    allow_credentials=True,17    allow_methods=["*"],18    allow_headers=["*"],19)20 21hf_model_ids = {22    "model1": "Pacicap/FineTuned_claude_StableDiffussion_2_1",23    "model2": "Pacicap/FineTuned_gpt4o_StableDiffussion_2_1"24}25 26loaded_models = {}27 28class PromptInput(BaseModel):29    prompt: str30    model: str31 32@app.post("/generate")33def generate(data: PromptInput, request: Request):34    model_key = data.model35 36    if model_key not in hf_model_ids:37        return {"error": "Invalid model selected"}38 39    model_id = hf_model_ids[model_key]40 41    if model_key not in loaded_models:42        pipe = DiffusionPipeline.from_pretrained(43            model_id,44            torch_dtype=torch.float3245        ).to("cpu")  # CPU-safe for Spaces46        loaded_models[model_key] = pipe47    else:48        pipe = loaded_models[model_key]49 50    image = pipe(data.prompt).images[0]51 52    os.makedirs("generated", exist_ok=True)53    filename = f"{uuid.uuid4().hex}.png"54    filepath = os.path.join("generated", filename)55    image.save(filepath)56 57    return {58        "url": f"{request.base_url}generated/{filename}"59    }60 61app.mount("/generated", StaticFiles(directory="generated"), name="generated")62