Update main.py

main.py

@@ -3,6 +3,7 @@ from flask import Flask, request, json, send_file, Response
 import torch
 from diffusers import StableDiffusionPipeline, EulerAncestralDiscreteScheduler
 from random import randrange
+from io import StringIO
 
 repo = "Bingsu/my-korean-stable-diffusion-v1-5"
 euler_ancestral_scheduler = EulerAncestralDiscreteScheduler.from_config(repo, subfolder="scheduler")
@@ -19,7 +20,10 @@ def generate():
     seed = randrange(1, 9999999999)
     generator = torch.Generator('cuda').manual_seed(seed)
     image = pipe(text, num_inference_steps=25, generator=generator).images[0]
-    return send_file(image, mimetype='image/png')
+    img_io = StringIO()
+    image.save(img_io)
+    img_io.seek(0)
+    return send_file(img_io, mimetype='image/png')
 
 
 if __name__ == '__main__':