add increase counts of image and reduce gpu usage

main.py

@@ -4,6 +4,7 @@ import sys
 import re
 import random
 import torch
+from PIL import Image
 from fastapi import FastAPI
 from fastapi.staticfiles import StaticFiles
 from pydantic import BaseModel
@@ -49,13 +50,16 @@ def draw(data: Data):
 
 		#pipe = StableDiffusionPipeline.from_pretrained(model_id)
 		pipe = StableDiffusionPipeline.from_pretrained(model_id, revision='fp16', torch_dtype=torch.float16)
+		pipe.enable_attention_slicing()	# reduce gpu usage
 		pipe = pipe.to('cuda')
 
 		generator = torch.Generator("cuda").manual_seed(seed)
-		image = pipe(prompt, negative_prompt=n_prompt, guidance_scale=7.5, generator=generator).images[0]
+		images = pipe(prompt, negative_prompt=n_prompt, guidance_scale=7.5, generator=generator, num_images_per_prompt=3).images
+
+		grid = image_grid(images, rows=2, cols=2)
 
 		fileName = "sd_" + str(time.time()) + '.png'
-		image.save("/code/tmpdir/" + fileName)
+		grid.save("/code/tmpdir/" + fileName)
 
 		print(fileName)
 
@@ -64,3 +68,15 @@ def draw(data: Data):
 		return {"status": "SORRY! This file is member only.", "file": ""}
 
 app.mount("/static", StaticFiles(directory="/code/tmpdir"), name="/static")
+
+# helper function taken from: https://huggingface.co/blog/stable_diffusion
+def image_grid(imgs, rows, cols):
+    assert len(imgs) == rows*cols
+
+    w, h = imgs[0].size
+    grid = Image.new('RGB', size=(cols*w, rows*h))
+    grid_w, grid_h = grid.size
+
+    for i, img in enumerate(imgs):
+        grid.paste(img, box=(i%cols*w, i//cols*h))
+    return grid