import torch
import gradio as gr
from torchvision.utils import save_image
from torchvision.transforms import ToPILImage
from model import Generator

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
NOISE_DIM = 256

G = Generator().to(DEVICE)
G.load_state_dict(torch.load("generator.pth", map_location=DEVICE))
G.eval()

to_pil = ToPILImage()

def generate_image():
    noise = torch.randn(1, NOISE_DIM).to(DEVICE)

    with torch.no_grad():
        image = G(noise)

    image = (image + 1) / 2
    return to_pil(image.squeeze(0))

demo = gr.Interface(
    fn=generate_image,
    inputs=None,
    outputs=gr.Image(),
    title="GAN Image Generator API"
)

demo.launch()