from fastapi import FastAPI
import torch
from diffusers import StableDiffusionPipeline

app = FastAPI()

# Load model
model_path = "/app/models/Stable-diffusion/model.safetensors"
pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16)
pipe.to("cuda")  # Gunakan GPU jika tersedia

@app.post("/generate")
async def generate_image(prompt: str):
    image = pipe(prompt).images[0]
    image_path = "/app/generated_image.png"
    image.save(image_path)
    return {"image_url": f"/generated_image.png"}

@app.get("/generated_image.png")
async def get_image():
    return FileResponse("/app/generated_image.png")