import gradio as gr
import torch
import torchvision.transforms as T

from model import DocuGAN


chk_path = "best_model.ckpt"
model = DocuGAN(hidden_size=64, num_channel=1, latent_size=128, batch_size=128)
# model = DocuGAN.load_from_checkpoint(chk_path, strict=False)
model.eval()
transform = T.ToPILImage()


def fn(seed: int = 42):
    torch.manual_seed(seed)
    noise = torch.randn(1, 128, 1, 1)
    with torch.no_grad():
        pred = model(noise)
        img = transform(pred.squeeze(1))
    return img


gr.Interface(
    fn,
    inputs=[
        gr.inputs.Slider(minimum=0, maximum=999999999, step=1, default=298422436, label='Random Seed')
    ],
    outputs='image',
    examples=[],
    enable_queue=True,
    title="DocuGAN",
    description="Select random seed and click on Submit to generate a new Document",
    article="DocuGAN, Document Generator by ChainYo",
    css=".panel { padding: 5px } .moflo-link { color: #999 }"
).launch()