from diffusers import StableDiffusionPipeline, DDIMScheduler
from time import time
from PIL import Image
from einops import rearrange
import numpy as np
import torch
from torch import autocast
from torchvision.utils import make_grid

torch.manual_seed(42)

prompts = ["a photograph of an astronaut riding a horse"]
prompts.append("a photograph of the eiffel tower on the moon")
prompts.append("an oil painting of a futuristic forest gives")

#pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-3-diffusers", use_auth_token=True)  # make sure you're logged in with `huggingface-cli login`
pipe = StableDiffusionPipeline.from_pretrained("fusing/sd-v1-3", use_auth_token=True)  # make sure you're logged in with `huggingface-cli login`

all_images = []
num_rows = 1
num_columns = 4
for prompt in prompts:
    with autocast("cuda"):
        images = pipe(num_columns * [prompt], guidance_scale=7.5, output_type="np")["sample"]  # image here is in [PIL format](https://pillow.readthedocs.io/en/stable/)
        all_images.append(torch.from_numpy(images))

# additionally, save as grid
grid = torch.stack(all_images, 0)
grid = rearrange(grid, 'n b h w c -> (n b) h w c')
grid = rearrange(grid, 'n h w c -> n c h w')
grid = make_grid(grid, nrow=num_rows)

# to image
grid = 255. * rearrange(grid, 'c h w -> h w c').cpu().numpy()
image = Image.fromarray(grid.astype(np.uint8))

image.save(f"./images/diffusers/final_{'_'.join(prompt.split())}_{round(time())}.png")