from flask import url_for from diffusers import StableDiffusionPipeline import torch from fastapi import FastAPI, Response from fastapi.middleware.cors import CORSMiddleware app = FastAPI() app.add_middleware( # add the middleware CORSMiddleware, allow_credentials=True, # allow credentials allow_origins=["*"], # allow all origins allow_methods=["*"], # allow all methods allow_headers=["*"], # allow all headers ) model_id = "runwayml/stable-diffusion-v1-5" pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16) pipe = pipe.to("cpu") pipe.enable_attention_slicing() def dummy(images, **kwargs): return images, False pipe.safety_checker = dummy @app.get("/") def hello(): return "Hello, I'm Linlada" @app.get("/gen/{prompt}") def generate_image(prompt: str): image = pipe(prompt, guidance_scale=8.5 # how strict to follow the prompt ).images[0] # Save the image image.save('static/image.png') # do something with the generated image image_data = image.tobytes().hex() image_url = url_for('static', filename='image.png') return {'image_data': image_data, 'image_url': image_url}