import gradio as gr
import torch
from torch import nn
import torchvision
from torchvision.transforms import ToTensor
from types import SimpleNamespace
import matplotlib.pyplot as plt


unet = torch.load("unet_checkpoint13.pt", map_location=torch.device('cpu')).to("cpu")
unet.eval()

@torch.no_grad
def linear_sched(betamin=0.0001,betamax=0.02,n_steps=1000):
    beta = torch.linspace(betamin, betamax, n_steps)
    return SimpleNamespace(a=1.-beta, abar=(1.-beta).cumprod(dim=0), sig=beta.sqrt())
n_steps = 1000
lin_abar = linear_sched(betamax=0.01)
alphabar = lin_abar.abar
alpha = lin_abar.a
sigma = lin_abar.sig

@torch.no_grad()  
def generate():
    model = unet
    sz = (1, 1, 32, 32)

    ps = next(model.parameters())
    x_t = torch.randn(sz).to(ps)
    sample_at = {t for t in range(n_steps) if (t+101)%((t+101)//100)==0}
    preds = []
    img_final = img_799 = img_399 = (x_t[0].float().cpu()+0.5).squeeze().clamp(-1,1).detach().numpy()
    for t in reversed(range(n_steps)):
        t_batch = torch.full((x_t.shape[0],), t, device=ps.device, dtype=torch.long)
        z = (torch.randn(x_t.shape) if t > 0 else torch.zeros(x_t.shape)).to(ps)
        ᾱ_t1 = alphabar[t-1]  if t > 0 else torch.tensor(1)
        b̄_t = 1-alphabar[t]
        b̄_t1 = 1-ᾱ_t1
        if t in sample_at: noise = model(x_t, t_batch).sample
        x_0_hat = ((x_t - b̄_t.sqrt() * noise)/alphabar[t].sqrt())
        x_t = x_0_hat * ᾱ_t1.sqrt()*(1-alpha[t])/b̄_t + x_t * alpha[t].sqrt()*b̄_t1/b̄_t + sigma[t]*z
        if t in sample_at: 
            preds.append(x_t.float().cpu())
            img = (x_t[0].float().cpu()+0.5).squeeze().clamp(-1,1).detach().numpy()

            if t >= 799: 
                img_final = img_799 = img_399 = img
            elif t >= 50: 
                img_final = img_399 = img
            else:
                img_final = img

            yield(img_799,img_399,img_final)
    
with gr.Blocks() as demo:
    gr.HTML("""<h1 align="center">UNet with DPPM</h1>""")
    gr.HTML("""<h1 align="center">trained with FashionMNIST</h1>""")
    session_data = gr.State([])

    sampling_button = gr.Button("Unconditional image generation")

    with gr.Row():
        with gr.Column(scale=2):
            gr.HTML("""<h3 align="left">image at step 800</h1>""")
            step_800_image = gr.Image(height=250,width=200)             
        with gr.Column(scale=2):
            gr.HTML("""<h3 align="left">image at step 50</h1>""")
            step_50_image = gr.Image(height=250,width=200) 
        with gr.Column(scale=2):
            gr.HTML("""<h3 align="left">final image</h1>""")
            step_final_image = gr.Image(height=250,width=200) 
     

    sampling_button.click(
        generate,
        [],
        [step_800_image, step_50_image, step_final_image],
    )

demo.queue().launch(share=False, inbrowser=True)