Create app.py

app.py

@@ -0,0 +1,113 @@
+import os 
+os.system("git clone https://github.com/mchong6/SOAT.git")
+
+
+import os
+import torch
+import torchvision
+from torch import nn
+import numpy as np
+import torch.backends.cudnn as cudnn
+cudnn.benchmark = True
+
+import math
+import matplotlib.pyplot as plt
+import torch.nn.functional as F
+from model import *
+from tqdm import tqdm as tqdm
+import pickle
+from copy import deepcopy
+import warnings
+warnings.filterwarnings("ignore", category=UserWarning) # get rid of interpolation warning
+import kornia.filters as k
+from torchvision.utils import save_image
+from util import *
+import scipy
+
+import gradio as gr
+
+import PIL
+
+from torchvision import transforms
+
+device = 'cpu' #@param ['cuda', 'cpu']
+
+generator = Generator(256, 512, 8, channel_multiplier=2).eval().to(device)
+truncation = 0.7
+
+def display_image(image, size=None, mode='nearest', unnorm=False, title=''):
+    # image is [3,h,w] or [1,3,h,w] tensor [0,1]
+    if image.is_cuda:
+        image = image.cpu()
+    if size is not None and image.size(-1) != size:
+        image = F.interpolate(image, size=(size,size), mode=mode)
+    if image.dim() == 4:
+        image = image[0]
+    image = ((image.clamp(-1,1)+1)/2).permute(1, 2, 0).detach().numpy()
+    return image
+    
+
+def inferece(num, seed):
+    model_type = 'landscape' 
+    num_im =  int(num)
+    random_seed =  int(seed)
+
+    plt.rcParams['figure.dpi'] = 300
+
+    mean_latent = load_model(generator, f'{model_type}.pt')
+
+    # pad determines how much of an image is involve in the blending
+    pad = 512//4
+
+    all_im = []
+
+    random_state = np.random.RandomState(random_seed)
+
+    # latent smoothing
+    with torch.no_grad():
+        z = random_state.randn(num_im, 512).astype(np.float32)
+        z = scipy.ndimage.gaussian_filter(z, [.7, 0], mode='wrap')
+        z /= np.sqrt(np.mean(np.square(z)))
+        z = torch.from_numpy(z).cuda()
+
+        source = generator.get_latent(z, truncation=truncation, mean_latent=mean_latent)
+            
+        # merge images 2 at a time
+        for i in range(num_im-1):
+            source1 = index_layers(source, i)
+            source2 = index_layers(source, i+1)
+            all_im.append(generator.merge_extension(source1, source2))
+
+    # display intermediate generations    
+    # for i in all_im:
+    #     display_image(i)
+        
+
+    b,c,h,w = all_im[0].shape
+    panorama_im = torch.zeros(b,c,h,512+(num_im-2)*256)
+
+    # We created a series of 2-blended images which we can overlay to form a large panorama
+    # add first image
+    coord = 256+pad
+    panorama_im[..., :coord] = all_im[0][..., :coord]
+
+    for im in all_im[1:]:
+        panorama_im[..., coord:coord+512-2*pad] = im[..., pad:-pad]
+        coord += 512-2*pad
+    panorama_im[..., coord:] = all_im[-1][..., 512-pad:]
+
+    img = display_image(panorama_im)
+    return img
+
+title = "SOAT"
+description = "Gradio demo for SOAT Panorama Generaton for landscapes. Generate a panorama using a pretrained stylegan by stitching intermediate activations. To use it, simply add the number of images and random seed number . Read more at the links below."
+article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2111.01619' target='_blank'>StyleGAN of All Trades: Image Manipulation with Only Pretrained StyleGAN</a> | <a href='https://github.com/mchong6/SOAT' target='_blank'>Github Repo</a></p>"
+
+gr.Interface(
+    inferece, 
+    [gr.inputs.Number(default=5, label="Number of Images")
+,gr.inputs.Number(default=90, label="Random Seed")],
+    gr.outputs.Image(type="numpy", label="Output"),
+    title=title,
+    description=description,
+    article=article, theme="huggingface",enable_queue=True).launch(debug=True)