import torch
from stylegan2 import legacy
import numpy as np
from PIL import Image
import gradio as gr

torch.autograd.set_grad_enabled(False)
torch.backends.cudnn.benchmark = True

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

with open("network-snapshot-000560.pkl", "rb") as f:
    G = legacy.load_network_pkl(f)['G_ema'].to(device)

def generate_image():
    seed = np.random.randint(0, 2**32)
    z = torch.from_numpy(np.random.RandomState(seed).randn(1, G.z_dim)).to(device)
    label = torch.zeros([1, G.c_dim], device=device)
    img = G(z, label, truncation_psi=1.0, noise_mode='const')
    img = (img.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to(torch.uint8)

    pil_img = Image.fromarray(img[0].cpu().numpy(), 'RGB')
    resized = pil_img.resize((512, 512), Image.LANCZOS)
    return resized

gr.Interface(
    fn=generate_image,
    inputs=[],
    outputs="image",
    title="StyleGAN2 Batik Generator",
    description="Click to generate a random image",
    allow_flagging="never"
).launch()