| import gradio as gr | |
| from huggingface_hub import hf_hub_download | |
| import torch | |
| import matplotlib.pyplot as plt | |
| import numpy as np | |
| from torch import nn | |
| class Generator(nn.Module): | |
| def __init__(self): | |
| super(Generator, self).__init__() | |
| self.main = nn.Sequential( | |
| nn.ConvTranspose2d(100, 64 * 8, 4, 1, 0, bias=False), | |
| nn.BatchNorm2d(64 * 8), | |
| nn.ReLU(True), | |
| nn.ConvTranspose2d(64 * 8, 64 * 4, 4, 2, 1, bias=False), | |
| nn.BatchNorm2d(64 * 4), | |
| nn.ReLU(True), | |
| nn.ConvTranspose2d(64 * 4, 64 * 2, 4, 2, 1, bias=False), | |
| nn.BatchNorm2d(64 * 2), | |
| nn.ReLU(True), | |
| nn.ConvTranspose2d(64 * 2, 64, 4, 2, 1, bias=False), | |
| nn.BatchNorm2d(64), | |
| nn.ReLU(True), | |
| nn.ConvTranspose2d(64, 3, 4, 2, 1, bias=False), | |
| nn.Tanh() | |
| ) | |
| def forward(self, input): | |
| return self.main(input) | |
| path = hf_hub_download('huggan/ArtGAN', 'ArtGAN.pt') | |
| model = torch.load(path, map_location=torch.device('cpu')) | |
| device = 'cuda' if torch.cuda.is_available() else 'cpu' | |
| def generate(seed): | |
| with torch.no_grad(): | |
| noise = torch.randn(seed, 100, 1, 1, device=device) | |
| with torch.no_grad(): | |
| art = model(noise).detach().cpu() | |
| gen = np.transpose(art[-1], (1, 2, 0)) | |
| fig = plt.figure(figsize=(5, 5)) | |
| plt.imshow(gen) | |
| plt.axis('off') | |
| return fig | |
| gr.Interface( | |
| fn=generate, | |
| inputs=[ | |
| gr.inputs.Slider | |
| ( | |
| label='noise', | |
| minimum=10, | |
| maximum=100, | |
| step=1, | |
| default=25 | |
| ) | |
| ], | |
| outputs=gr.outputs.Image(type='plot'), | |
| title='ArtGAN', | |
| description='Generate A Abstract Art Using ArtGAN', | |
| ).launch() | |