Add Folders

Author: BasharatWali

HF_App/App_Code.py

@@ -0,0 +1,69 @@
+from PIL import Image
+import torchvision.transforms as transforms
+import gradio as gr
+import torch
+import torch.nn as nn
+
+latent_dim = 100
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+class Generator(nn.Module):
+    def __init__(self, latent_dim=100, img_channels=3, feature_map_size=32):
+        super(Generator, self).__init__()
+        self.net = nn.Sequential(
+            nn.ConvTranspose2d(latent_dim, feature_map_size * 8, 4, 1, 0, bias=False),
+            nn.BatchNorm2d(feature_map_size * 8),
+            nn.ReLU(True),
+            nn.ConvTranspose2d(feature_map_size * 8, feature_map_size * 4, 4, 2, 1, bias=False),
+            nn.BatchNorm2d(feature_map_size * 4),
+            nn.ReLU(True),
+            nn.ConvTranspose2d(feature_map_size * 4, feature_map_size * 2, 4, 2, 1, bias=False),
+            nn.BatchNorm2d(feature_map_size * 2),
+            nn.ReLU(True),
+            nn.ConvTranspose2d(feature_map_size * 2, feature_map_size, 4, 2, 1, bias=False),
+            nn.BatchNorm2d(feature_map_size),
+            nn.ReLU(True),
+            nn.ConvTranspose2d(feature_map_size, img_channels, 4, 2, 1, bias=False),
+            nn.Tanh()
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+def generate_artwork(generator, latent_dim=latent_dim, device=device, num_images=1):
+    generator.eval()
+    with torch.no_grad():
+        noise = torch.randn(num_images, latent_dim, 1, 1, device=device)
+        fake_images = generator(noise)
+        fake_images = fake_images * 0.5 + 0.5
+    return fake_images.detach().cpu()
+
+def inference_interface(latent_dim=latent_dim, device=device):
+    # Create model and load weights
+    generator = Generator(latent_dim=latent_dim)
+    generator = nn.DataParallel(generator)
+    generator.load_state_dict(torch.load("generator_final.pth", map_location=device))
+    
+    if isinstance(generator, nn.DataParallel):
+        generator = generator.module
+    generator.to(device)
+
+    def generate(num_images):
+        fake_images = generate_artwork(generator, latent_dim=latent_dim, device=device, num_images=num_images)
+        images = [transforms.ToPILImage()(img) for img in fake_images]
+        upscaled_images = [img.resize((256, 256), resample=Image.LANCZOS) for img in images]
+        return upscaled_images
+
+    demo = gr.Interface(fn=generate,
+      inputs=gr.Slider(minimum=1, maximum=9, step=1, value=1, label="Number of Images"),
+      outputs=gr.Gallery(label="Generated Artwork", columns=3, height="auto"),
+      title="This Artwork Doesn’t Exist",
+      description="Generate artwork using our Wiki_ArtGAN."
+      )
+
+    return demo
+
+# The key part: launch the Gradio interface when app.py is run
+if __name__ == "__main__":
+    demo = inference_interface()
+    demo.launch()