from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from diffusers import StableDiffusionPipeline, EulerDiscreteScheduler
import torch

app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
    allow_credentials=True,
)

model_id = "runwayml/stable-diffusion-v1-5"
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
pipe = pipe.to("cuda")

def dummy(images, **kwargs):
    return images, False

pipe.safety_checker = dummy

@app.get('/')
def generate_image():
    prompt = request.args.get('prompt')
    image = pipe(prompt).images[0]
    # do something with the generated image
    return image