Update main.py

main.py

@@ -1,25 +1,33 @@
 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
 
-repo = "runwayml/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,
+from min_dalle import MinDalle
+
+model = MinDalle(
+    models_root='./pretrained',
+    dtype=torch.float32,
+    device='cuda',
+    is_mega=True, 
+    is_reusable=True
 )
-pipe.to("cuda")
 
 app = Flask(__name__)
 
 @app.post('/sd')
 def generate():
     text = request.json['text']
-    seed = randrange(1, 9999999999)
-    generator = torch.Generator('cuda').manual_seed(seed)
-    image = pipe(text, num_inference_steps=25, generator=generator).images[0]
+    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
+    )
     img_io = BytesIO()
     image.save(img_io)
     img_io.seek(0)