import torch
import torchvision
import gradio as gr
from PIL import Image
from cli import iterative_refinement
from viz import grid_of_images_default
models = {
    "convae": torch.load("convae.th", map_location="cpu"),
    "deep_convae": torch.load("deep_convae.th", map_location="cpu"),
}
def gen(md, model, seed, nb_iter, nb_samples, width, height):
    torch.manual_seed(int(seed))
    bs = 64
    model = models[model]
    samples = iterative_refinement(
        model, 
        nb_iter=int(nb_iter),
        nb_examples=int(nb_samples), 
        w=int(width), h=int(height), c=1, 
        batch_size=bs,
    )
    grid = grid_of_images_default(samples.reshape((samples.shape[0]*samples.shape[1], int(height), int(width), 1)).numpy(), shape=(samples.shape[0], samples.shape[1])) 
    grid = (grid*255).astype("uint8")
    return Image.fromarray(grid)

text = """
Interface with ConvAE model (from [here](https://arxiv.org/pdf/1606.04345.pdf)) and DeepConvAE model (from [here](https://tel.archives-ouvertes.fr/tel-01838272/file/75406_CHERTI_2018_diffusion.pdf), Section 10.1 with `L=3`)

These models were trained on MNIST only (digits), but were found to generate new kinds of symbols, see the references for more details.
"""
iface = gr.Interface(
    fn=gen,
    inputs=[
        gr.Markdown(text),
        gr.Dropdown(list(models.keys()), value="deep_convae"), gr.Number(value=0), gr.Number(value=20), gr.Number(value=1), gr.Number(value=28), gr.Number(value=28)
    ],
    outputs="image"
)
iface.launch()