import gradio as gr
import torch
import numpy as np
from PIL import Image
import random
from huggingface_hub import hf_hub_download
from generator import Generator  # Import your generator class

# Import your generator class
# from generator import Generator  # Uncomment and adjust to your file

wts = ['trial_0_G (1).pth' , 'trial_0_G (2).pth' , 'trial_0_G (3).pth' , 'trial_0_G (4).pth' , 'trial_0_G (5).pth' , 'trial_0_G.pth' ]
random_wt = random.choice(wts)

# Load trained model weights from Hugging Face Hub
weights_path = hf_hub_download(
    repo_id="keysun89/image_generation",  # Replace with your repo
    filename= random_wt        # Replace with your weights file
)


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

# Configure your generator parameters
z_dim = 512
w_dim = 512
img_resolution = 256  # Adjust to your training resolution
img_channels = 3

model = Generator(
    z_dim=z_dim,
    w_dim=w_dim,
    img_resolution=img_resolution,
    img_channels=img_channels
)

# Load weights
model.load_state_dict(torch.load(weights_path, map_location=device))
model.to(device)
model.eval()

def generate():
    """Generate a random image"""
    with torch.no_grad():
        # Generate random latent vector
        z = torch.randn(1, z_dim, device=device)
        
        # Generate image
        img = model(z, use_truncation=True, truncation_psi=0.7)
        
        # Convert to PIL Image
        img = img.squeeze(0).cpu().numpy()
        img = np.transpose(img, (1, 2, 0))  # CHW to HWC
        img = (img * 127.5 + 128).clip(0, 255).astype(np.uint8)
    
    return Image.fromarray(img)

# Gradio interface
demo = gr.Interface(
    fn=generate,
    inputs=None,
    outputs=gr.Image(type="pil"),
    title="StyleGAN2 Image Generator",
    description="Click 'Submit' or refresh the page to generate a new random image",
    allow_flagging="never"
)

if __name__ == "__main__":
    demo.launch()