Update main.py

main.py

@@ -1,33 +1,25 @@
 import flask
 from flask import Flask, request, json, send_file, Response
 import torch
+from diffusers import StableDiffusionPipeline, EulerAncestralDiscreteScheduler
+from random import randrange
 from io import BytesIO
 
-from min_dalle import MinDalle
-
-model = MinDalle(
-    models_root='./pretrained',
-    dtype=torch.float32,
-    device='cuda',
-    is_mega=True, 
-    is_reusable=True
+repo = "Bingsu/my-korean-stable-diffusion-v1-5"
+euler_ancestral_scheduler = EulerAncestralDiscreteScheduler.from_config(repo, subfolder="scheduler")
+pipe = StableDiffusionPipeline.from_pretrained(
+    repo, scheduler=euler_ancestral_scheduler, torch_dtype=torch.float16,
 )
+pipe.to("cuda")
 
 app = Flask(__name__)
 
 @app.post('/sd')
 def generate():
     text = request.json['text']
-    image = model.generate_image(
-        text=text,
-        seed=-1,
-        grid_size=2,
-        is_seamless=False,
-        temperature=1,
-        top_k=256,
-        supercondition_factor=32,
-        is_verbose=False
-    )
+    seed = randrange(1, 9999999999)
+    generator = torch.Generator('cuda').manual_seed(seed)
+    image = pipe(text, num_inference_steps=25, generator=generator).images[0]
     img_io = BytesIO()
     image.save(img_io)
     img_io.seek(0)