huggingface deployment

App.py

@@ -0,0 +1,86 @@
+import streamlit as st
+import torch
+import DCGAN
+import SRGAN
+from Utils import color_histogram_mapping, denormalize_images
+import torch.nn as nn
+import random
+
+device = torch.device("cpu")
+
+if torch.cuda.is_available():
+    device = torch.device("cuda")
+
+latent_size = 100
+display_width = 450
+checkpoint_path = "Checkpoints/150epochs.chkpt"
+
+st.title("Generating Abstract Art")
+st.text("start generating (left side bar)")
+st.text("Made by Xingyu B.")
+
+st.sidebar.subheader("Configurations")
+seed = st.sidebar.slider('Seed', -10000, 10000, 0)
+
+num_images = st.sidebar.slider('Number of Images', 1, 10, 1)
+
+use_srgan = st.sidebar.selectbox(
+    'Apply image enhancement',
+    ('Yes', 'No')
+)
+
+generate = st.sidebar.button("Generate")
+
+
+# caching the expensive model loading 
+
+@st.cache(allow_output_mutation=True)
+def load_dcgan():
+    model = torch.jit.load('Checkpoints/dcgan.pt', map_location=device)
+    return model 
+
+@st.cache(allow_output_mutation=True)
+def load_esrgan():
+    model_state_dict = torch.load("Checkpoints/esrgan.pt", map_location=device)
+    return model_state_dict
+
+# if the user wants to generate something new 
+if generate:
+    torch.manual_seed(seed)
+    random.seed(seed)
+    
+    sampled_noise = torch.randn(num_images, latent_size, 1, 1, device=device)
+    generator = load_dcgan()
+    generator.eval()
+
+    with torch.no_grad():
+        fakes = generator(sampled_noise).detach()
+
+    # use srgan for super resolution
+    if use_srgan == "Yes":
+        # restore to the checkpoint
+        st.write("Using DCGAN then ESRGAN upscale...")
+        esrgan_generator = SRGAN.GeneratorRRDB(channels=3, filters=64, num_res_blocks=23).to(device)
+        esrgan_checkpoint = load_esrgan()
+        esrgan_generator.load_state_dict(esrgan_checkpoint)
+
+        esrgan_generator.eval()
+        with torch.no_grad():
+            enhanced_fakes = esrgan_generator(fakes).detach().cpu()
+        color_match = color_histogram_mapping(enhanced_fakes, fakes.cpu())
+
+        for i in range(len(color_match)):
+            # denormalize and permute to correct color channel
+            st.image(denormalize_images(color_match[i]).permute(1, 2, 0).numpy(), width=display_width)
+
+
+    # default setting -> vanilla dcgan generation
+    if use_srgan == "No":
+        fakes = fakes.cpu()
+        st.write("Using DCGAN Model...")
+        for i in range(len(fakes)):
+            st.image(denormalize_images(fakes[i]).permute(1, 2, 0).numpy(), width=display_width)
+
+
+
+

Checkpoints/dcgan.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5dc5b82eab432b52635f67ae7abf6901c36daa58ff71445ff9df01cc6b3193f2
+size 14352101

Checkpoints/esrgan.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4ec7aa5dd51df901367a6ee1d03c2cbbf72acadad01288040ab723860e96ffe4
+size 154489349

DCGAN.py

@@ -0,0 +1,32 @@
+from torch import nn
+import torch
+import torch.nn.functional as F
+    
+# Generator Code
+
+ngf = 64
+num_channels = 3
+
+class Generator(nn.Module):
+    def __init__(self, latent_size):
+        super(Generator, self).__init__()
+
+        self.latent_size = latent_size
+        self.conv1 = nn.ConvTranspose2d(
+            self.latent_size, ngf*8, 4, 1, 0, bias=False)
+        self.bn1 = nn.BatchNorm2d(ngf*8)
+        self.conv2 = nn.ConvTranspose2d(ngf*8, ngf*4, 4, 2, 1, bias=False)
+        self.bn2 = nn.BatchNorm2d(ngf*4)
+        self.conv3 = nn.ConvTranspose2d(ngf*4, ngf*2, 4, 2, 1, bias=False)
+        self.bn3 = nn.BatchNorm2d(ngf*2)
+        self.conv4 = nn.ConvTranspose2d(ngf*2, ngf, 4, 2, 1, bias=False)
+        self.bn4 = nn.BatchNorm2d(ngf)
+
+        self.conv5 = nn.ConvTranspose2d(ngf, num_channels, 4, 2, 1, bias=False)
+
+    def forward(self, x):
+        x = F.relu(self.bn1(self.conv1(x)), inplace=True)
+        x = F.relu(self.bn2(self.conv2(x)), inplace=True)
+        x = F.relu(self.bn3(self.conv3(x)), inplace=True)
+        x = F.relu(self.bn4(self.conv4(x)), inplace=True)
+        return torch.tanh(self.conv5(x))

SRGAN.py

@@ -0,0 +1,79 @@
+import torch.nn as nn 
+import torch
+
+class DenseResidualBlock(nn.Module):
+    """
+    The core module of paper: (Residual Dense Network for Image Super-Resolution, CVPR 18)
+    """
+
+    def __init__(self, filters, res_scale=0.2):
+        super(DenseResidualBlock, self).__init__()
+        self.res_scale = res_scale
+
+        def block(in_features, non_linearity=True):
+            layers = [nn.Conv2d(in_features, filters, 3, 1, 1, bias=True)]
+            if non_linearity:
+                layers += [nn.LeakyReLU()]
+            return nn.Sequential(*layers)
+
+        self.b1 = block(in_features=1 * filters)
+        self.b2 = block(in_features=2 * filters)
+        self.b3 = block(in_features=3 * filters)
+        self.b4 = block(in_features=4 * filters)
+        self.b5 = block(in_features=5 * filters, non_linearity=False)
+        self.blocks = [self.b1, self.b2, self.b3, self.b4, self.b5]
+
+    def forward(self, x):
+        inputs = x
+        for block in self.blocks:
+            out = block(inputs)
+            inputs = torch.cat([inputs, out], 1)
+        return out.mul(self.res_scale) + x
+
+
+class ResidualInResidualDenseBlock(nn.Module):
+    def __init__(self, filters, res_scale=0.2):
+        super(ResidualInResidualDenseBlock, self).__init__()
+        self.res_scale = res_scale
+        self.dense_blocks = nn.Sequential(
+            DenseResidualBlock(filters), DenseResidualBlock(filters), DenseResidualBlock(filters)
+        )
+
+    def forward(self, x):
+        return self.dense_blocks(x).mul(self.res_scale) + x
+
+
+class GeneratorRRDB(nn.Module):
+    def __init__(self, channels, filters=64, num_res_blocks=16, num_upsample=2):
+        super(GeneratorRRDB, self).__init__()
+
+        # First layer
+        self.conv1 = nn.Conv2d(channels, filters, kernel_size=3, stride=1, padding=1)
+        # Residual blocks
+        self.res_blocks = nn.Sequential(*[ResidualInResidualDenseBlock(filters) for _ in range(num_res_blocks)])
+        # Second conv layer post residual blocks
+        self.conv2 = nn.Conv2d(filters, filters, kernel_size=3, stride=1, padding=1)
+        # Upsampling layers
+        upsample_layers = []
+        for _ in range(num_upsample):
+            upsample_layers += [
+                nn.Conv2d(filters, filters * 4, kernel_size=3, stride=1, padding=1),
+                nn.LeakyReLU(),
+                nn.PixelShuffle(upscale_factor=2),
+            ]
+        self.upsampling = nn.Sequential(*upsample_layers)
+        # Final output block
+        self.conv3 = nn.Sequential(
+            nn.Conv2d(filters, filters, kernel_size=3, stride=1, padding=1),
+            nn.LeakyReLU(),
+            nn.Conv2d(filters, channels, kernel_size=3, stride=1, padding=1),
+        )
+
+    def forward(self, x):
+        out1 = self.conv1(x)
+        out = self.res_blocks(out1)
+        out2 = self.conv2(out)
+        out = torch.add(out1, out2)
+        out = self.upsampling(out)
+        out = self.conv3(out)
+        return out

Utils.py

@@ -0,0 +1,40 @@
+import matplotlib.pyplot as plt
+import numpy as np
+import torchvision.utils as vutils
+import torchvision.transforms as transforms
+from skimage.exposure import match_histograms
+import torch
+
+# contains utility functions that we need in the main program
+
+# matches the color histogram of original and the super resolution output
+def color_histogram_mapping(images, references):
+    matched_list = []
+    for i in range(len(images)):
+        matched = match_histograms(images[i].permute(1, 2, 0).numpy(), references[i].permute(1, 2, 0).numpy(),
+                                   channel_axis=-1)
+        matched_list.append(matched)
+    return torch.tensor(np.array(matched_list)).permute(0, 3, 1, 2)
+
+
+def visualize_generations(seed, images):
+    plt.figure(figsize=(16, 16))
+    plt.title(f"Seed: {seed}")
+    plt.axis("off")
+    plt.imshow(np.transpose(vutils.make_grid(images, padding=2, nrow=5, normalize=True), (2, 1, 0)))
+    plt.show()
+
+
+# denormalize the images for proper display
+def denormalize_images(images):
+    mean= [0.5, 0.5, 0.5]
+    std= [0.5, 0.5, 0.5]
+    inv_normalize = transforms.Normalize(
+        mean=[-m / s for m, s in zip(mean, std)],
+        std=[1 / s for s in std]
+    )
+    return inv_normalize(images)
+
+
+
+

requirements.txt

@@ -0,0 +1,6 @@
+matplotlib==3.5.2
+numpy==1.23.0
+torch==1.12.0
+torchvision==0.13.0
+scikit-image~=0.19.3
+streamlit==1.11.0