SunilKrishna/image-generative-ai

1from fastapi import FastAPI, HTTPException2from fastapi.responses import Response, FileResponse3from diffusers import StableDiffusionPipeline4import torch5import io6 7app = FastAPI()8 9# Serve frontend10@app.get("/")11def serve_frontend():12    return FileResponse("index.html")13 14# Device detection15device = "cuda" if torch.cuda.is_available() else "cpu"16print(f"Using device: {device}")17 18# Load model19pipe = StableDiffusionPipeline.from_pretrained(20    "CompVis/stable-diffusion-v1-4",21    torch_dtype=torch.float16 if device == "cuda" else torch.float3222)23pipe = pipe.to(device)24 25# Cache to store generated images for repeated prompts26cache = {}27 28# Generate image endpoint29@app.post("/generate")30def generate(request_data: dict):31    prompt = request_data.get("prompt")32    if not prompt:33        raise HTTPException(status_code=400, detail="No prompt provided.")34 35    # Check cache36    if prompt in cache:37        buf = cache[prompt]38        buf.seek(0)39        return Response(buf.getvalue(), media_type="image/png")40 41    try:42        # Generate image43        image = pipe(prompt, num_inference_steps=20).images[0]44 45        # Save to buffer46        buf = io.BytesIO()47        image.save(buf, format="PNG")48        buf.seek(0)49 50        # Store in cache51        cache[prompt] = buf52 53        return Response(buf.getvalue(), media_type="image/png")54    except Exception as e:55        raise HTTPException(status_code=500, detail=str(e))56