DocuGAN / main.py
chainyo's picture
Update main.py
fb680d6
import gradio as gr
import torch
import torchvision.transforms as T
from model import DocuGAN
chk_path = "best_model.ckpt"
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)
pred = pred.mul(0.5).add(0.5)
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 - This document doesn't exist",
description="Select your seed and click on `Submit` to generate a new document",
article="<p>The SN-GAN model has been trained on the `invoice` part of RVL-CDIP dataset, available <a href='https://huggingface.co/datasets/ChainYo/rvl-cdip-invoice' target='_blank'>here</a>.<br> You can see the full implementation on the dedicated <a href='https://colab.research.google.com/drive/1u6Ct3KnNl7rcgla0268cp-XGTMmVUuJL?usp=sharing' target='_blank'>Colab notebook</a>. <br> Made with ❀️ by <a href='https://huggingface.co/ChainYo' target='_blank'>@ChainYo</a></p>",
css=".panel { padding: 5px } .moflo-link { color: #999 }"
).launch()