Spaces:
Sleeping
Sleeping
File size: 2,864 Bytes
b98188b |
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 70 |
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, default=1, label="Number of Images"),
outputs=gr.Gallery(label="Generated Artwork").style(grid=[3], height="auto"),
title="Art Generation with GAN",
description="Generate artwork using a trained GAN model."
)
return demo
# The key part: launch the Gradio interface when app.py is run
if __name__ == "__main__":
demo = inference_interface()
demo.launch()
|