Spaces:
Runtime error
Runtime error
File size: 4,085 Bytes
6b8aec5 |
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 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 |
import torch
from torchvision.utils import make_grid
from torchvision import transforms
import torchvision.transforms.functional as TF
from torch import nn, optim
from torch.optim.lr_scheduler import CosineAnnealingLR
from torch.utils.data import DataLoader, Dataset
from huggingface_hub import hf_hub_download
import requests
import gradio as gr
class Upsample(nn.Module):
def __init__(self, in_channels, out_channels, kernel_size=4, stride=2, padding=1, dropout=True):
super(Upsample, self).__init__()
self.dropout = dropout
self.block = nn.Sequential(
nn.ConvTranspose2d(in_channels, out_channels, kernel_size, stride, padding, bias=nn.InstanceNorm2d),
nn.InstanceNorm2d(out_channels),
nn.ReLU(inplace=True)
)
self.dropout_layer = nn.Dropout2d(0.5)
def forward(self, x, shortcut=None):
x = self.block(x)
if self.dropout:
x = self.dropout_layer(x)
if shortcut is not None:
x = torch.cat([x, shortcut], dim=1)
return x
class Downsample(nn.Module):
def __init__(self, in_channels, out_channels, kernel_size=4, stride=2, padding=1, apply_instancenorm=True):
super(Downsample, self).__init__()
self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, bias=nn.InstanceNorm2d)
self.norm = nn.InstanceNorm2d(out_channels)
self.relu = nn.LeakyReLU(0.2, inplace=True)
self.apply_norm = apply_instancenorm
def forward(self, x):
x = self.conv(x)
if self.apply_norm:
x = self.norm(x)
x = self.relu(x)
return x
class CycleGAN_Unet_Generator(nn.Module):
def __init__(self, filter=64):
super(CycleGAN_Unet_Generator, self).__init__()
self.downsamples = nn.ModuleList([
Downsample(3, filter, kernel_size=4, apply_instancenorm=False), # (b, filter, 128, 128)
Downsample(filter, filter * 2), # (b, filter * 2, 64, 64)
Downsample(filter * 2, filter * 4), # (b, filter * 4, 32, 32)
Downsample(filter * 4, filter * 8), # (b, filter * 8, 16, 16)
Downsample(filter * 8, filter * 8), # (b, filter * 8, 8, 8)
Downsample(filter * 8, filter * 8), # (b, filter * 8, 4, 4)
Downsample(filter * 8, filter * 8), # (b, filter * 8, 2, 2)
])
self.upsamples = nn.ModuleList([
Upsample(filter * 8, filter * 8),
Upsample(filter * 16, filter * 8),
Upsample(filter * 16, filter * 8),
Upsample(filter * 16, filter * 4, dropout=False),
Upsample(filter * 8, filter * 2, dropout=False),
Upsample(filter * 4, filter, dropout=False)
])
self.last = nn.Sequential(
nn.ConvTranspose2d(filter * 2, 3, kernel_size=4, stride=2, padding=1),
nn.Tanh()
)
def forward(self, x):
skips = []
for l in self.downsamples:
x = l(x)
skips.append(x)
skips = reversed(skips[:-1])
for l, s in zip(self.upsamples, skips):
x = l(x, s)
out = self.last(x)
return out
class ImageTransform:
def __init__(self, img_size=256):
self.transform = {
'train': transforms.Compose([
transforms.Resize((img_size, img_size)),
transforms.RandomHorizontalFlip(),
transforms.RandomVerticalFlip(),
transforms.ToTensor(),
transforms.Normalize(mean=[0.5], std=[0.5])
]),
'test': transforms.Compose([
transforms.Resize((img_size, img_size)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.5], std=[0.5])
})}
def __call__(self, img, phase='train'):
img = self.transform[phase](img)
return img
path = hf_hub_download('huggan/NeonGAN', 'model.bin')
model_gen_n = torch.load(path, map_location=torch.device('cpu'))
|