import gradio as gr
import torch, torchvision
from torchvision import transforms
import torch.nn.functional as F
import numpy as np
from time import time, ctime
from PIL import Image, ImageColor
from diffusers import DDPMPipeline
from diffusers import DDIMScheduler
from tqdm import tqdm

device = (
    "mps"
    if torch.backends.mps.is_available()
    else "cuda"
    if torch.cuda.is_available()
    else "cpu"
)

pipeline_name = 'WiNE-iNEFF/Minecraft-Skin-Diffusion'
image_pipe = DDPMPipeline.from_pretrained(pipeline_name).to(device)

# Set up the scheduler
scheduler = DDIMScheduler.from_pretrained(pipeline_name)
scheduler.set_timesteps(num_inference_steps=40)

def show_images_save(x):
    """Given a batch of images x, make a grid and convert to PIL"""
    x = x * 0.5 + 0.5  # Map from (-1, 1) back to (0, 1)
    grid = torchvision.utils.make_grid(x, nrow=4)
    grid_im = grid.detach().cpu().permute(1, 2, 0).clip(0, 1) * 255
    grid_im = Image.fromarray(np.array(grid_im).astype(np.uint8))
    grid_im.save(f"test.png")
    return grid_im

def generate():
    x = torch.randn(1, 4, 64, 64).to(device)
    # Minimal sampling loop
    for i, t in enumerate(scheduler.timesteps):
        model_input = scheduler.scale_model_input(x, t)
        with torch.no_grad():
            noise_pred = image_pipe.unet(model_input, t)["sample"]
        x = scheduler.step(noise_pred, t, x).prev_sample
    # View the results
    return show_images_save(x)

def ex():
    t = time()
    print(ctime(t))
    return generate(), generate(), generate(), generate()

demo = gr.Blocks(css="#img_size {max-height: 128px} .container {max-width: 730px; margin: auto;} .min-h-\[15rem\]{min-height: 5rem !important;}")

with demo:
  gr.HTML(
        """
            <div style="text-align: center; margin: 0 auto;">
              <div style="display: inline-flex;align-items: center;gap: 0.8rem;font-size: 1.75rem;">
                <h1 style="font-weight: 900; margin-bottom: 7px;margin-top:5px">
                  Minecraft Skin Diffusion
                </h1>
              </div>
              <p style="margin-bottom: 10px; font-size: 94%; line-height: 23px;">
                Gradio demo for Minecraft Skin Diffusion. This is simple Unconditional Diffusion Model that will help you generate skins for game Minecraft.
              </p>
            </div>
        """
    )
  with gr.Column():
    with gr.Row().style(equal_height=True):
        out = gr.Image(shape=(64,64), image_mode='RGBA', type='pil', elem_id='img_size')
        out2 = gr.Image(shape=(64,64), image_mode='RGBA', type='pil', elem_id='img_size')
    with gr.Row().style(equal_height=True):
        out3 = gr.Image(shape=(64,64), image_mode='RGBA', type='pil', elem_id='img_size')
        out4 = gr.Image(shape=(64,64), image_mode='RGBA', type='pil', elem_id='img_size')
  greet_btn = gr.Button("Generate")
  greet_btn.click(fn=ex, inputs=None, outputs=[out, out2, out3, out4])
  gr.HTML(
            """
                <div class="footer">
                    <div style='text-align: center;'>Minecraft Skin Diffusion by <a href='https://twitter.com/wine_ineff' target='_blank'>Artsem Holub (WiNE-iNEFF)</a> | 
                      <center>
                        <img src='https://visitor-badge.glitch.me/badge?page_id=WiNE-iNEFF_MinecraftSkin-Diffusion' alt='visitor badge'>
                      </center>
                    </div>
               </div>
           """
        )

demo.launch()