Update app.py

app.py

@@ -1,28 +1,13 @@
-from fastapi import FastAPI, File, UploadFile, Form
-from fastapi.responses import StreamingResponse
+import gradio as gr
+import torch
 from diffusers import StableDiffusionPipeline
-import torch, io
-from PIL import Image
 
-app = FastAPI()
-pipe = StableDiffusionPipeline.from_pretrained(
-    "nitrosocke/Ghibli-Diffusion",
-    torch_dtype=torch.float16,
-    safety_checker=None  # Optional: disable safety checker
-).to("cuda")
+model_id = "runwayml/stable-diffusion-v1-5"
+pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
+pipe = pipe.to("cuda")
 
-# Health check endpoint
-@app.get("/health")
-async def health_check():
-    return {"status": "ok", "message": "FastAPI server is running!"}
+def generate(prompt):
+    image = pipe(prompt).images[0]
+    return image
 
-@app.post("/ghibli")
-async def generate(
-    file: UploadFile = File(...), prompt: str = Form(...)
-):
-    image = Image.open(io.BytesIO(await file.read())).convert("RGB")
-    result = pipe(prompt=prompt, image=image, guidance_scale=7.5).images[0]
-    buf = io.BytesIO()
-    result.save(buf, format="PNG")
-    buf.seek(0)
-    return StreamingResponse(buf, media_type="image/png")
+gr.Interface(fn=generate, inputs="text", outputs="image").launch()