Rename main.py to app.py

main.py → app.py

@@ -1,17 +1,8 @@
-from fastapi import FastAPI
-from fastapi.middleware.cors import CORSMiddleware
-from diffusers import StableDiffusionPipeline, EulerDiscreteScheduler
+from flask import Flask, request
+from diffusers import StableDiffusionPipeline
 import torch
 
-app = FastAPI()
-
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_methods=["*"],
-    allow_headers=["*"],
-    allow_credentials=True,
-)
+app = Flask(__name__)
 
 model_id = "runwayml/stable-diffusion-v1-5"
 pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
@@ -22,9 +13,12 @@ def dummy(images, **kwargs):
 
 pipe.safety_checker = dummy
 
-@app.get('/')
+@app.route('/')
 def generate_image():
     prompt = request.args.get('prompt')
     image = pipe(prompt).images[0]
     # do something with the generated image
-    return image
+    return image
+
+if __name__ == '__main__':
+    app.run()