Create app.py

app.py

@@ -0,0 +1,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()