import torch
from PIL import Image
from diffusers import StableDiffusionPipeline
from torch import autocast

pipe = StableDiffusionPipeline.from_pretrained("lambdalabs/sd-pokemon-diffusers", torch_dtype=torch.float16)
device = "cuda:0" if torch.cuda.is_available() else "cpu"
pipe = pipe.to(device)


class Service():
    def __init__(self):
        self.scale = 10

    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))
        for i, img in enumerate(imgs):
            grid.paste(img, box=(i % cols * w, i // cols * h))
        return grid

    def generate(self,
                 prompt,
                 n_samples,
                 rows,
                 cols):
        with autocast("device"):
            images = pipe(n_samples * [prompt], guidance_scale=self.scale).images
        return self.image_grid(images, rows=rows, cols=cols)