File size: 1,940 Bytes
3eda9ab |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 |
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() |