added files

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.ipynb filter=lfs diff=lfs merge=lfs -text
+*.jpg filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,6 @@
+__pycache__/
+.DS_Store
+.vscode/
+
+test.*
+test/

.python-version

@@ -0,0 +1 @@
+ai-env

README.md

@@ -1,12 +1,16 @@
 ---
 title: Style Transfer
-emoji: 📈
-colorFrom: purple
-colorTo: purple
+emoji: 👨‍🎨
+colorFrom: pink
+colorTo: yellow
 sdk: gradio
 sdk_version: 4.36.1
 app_file: app.py
 pinned: false
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Style Transfer
+
+## References
+
+[1] Huang, Xun, and Serge Belongie. "Arbitrary style transfer in real-time with adaptive instance normalization." *Proceedings of the IEEE international conference on computer vision*. 2017.

app.py

@@ -0,0 +1,86 @@
+import gradio as gr
+import torch
+import torchvision.transforms.functional as TF
+import torchvision.transforms as transforms
+from src.model import Model
+import os
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+
+def denorm_img(img: torch.Tensor):
+    std = torch.Tensor([0.229, 0.224, 0.225]).reshape(-1, 1, 1)
+    mean = torch.Tensor([0.485, 0.456, 0.406]).reshape(-1, 1, 1)
+    return torch.clip(img * std + mean, min=0, max=1)
+
+
+def main(inp1, inp2, alph, out_size=256):
+    model = Model()
+    model.load_state_dict(torch.load("models/model_puddle.pt", map_location=torch.device(device)))
+    model.eval()
+
+    model.alpha = alph
+
+    style = TF.to_tensor(inp1)
+    content = TF.to_tensor(inp2)
+
+    norm = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+    transform = transforms.Compose(
+        [transforms.Resize(out_size, antialias=True), transforms.CenterCrop(out_size)]
+    )
+
+    style, content = norm(style), norm(content)
+    style, content = transform(style), transform(content)
+
+    style, content = style.unsqueeze(0).to(device), content.unsqueeze(0).to(device)
+
+    out = model(content, style)
+
+    return denorm_img(out[0].detach()).permute(1, 2, 0).numpy()
+
+
+with gr.Blocks() as demo:
+    gr.Markdown("# Style Transfer with AdaIN")
+    with gr.Row(variant="compact"):
+        inp1 = gr.Image(type="pil", sources=["upload", "clipboard"], label="Style")
+        inp2 = gr.Image(type="pil", sources=["upload", "clipboard"], label="Content")
+        out = gr.Image(type="numpy", label="Output")
+    with gr.Row():
+        out_size = (
+            gr.Dropdown(
+                choices=[256, 512],
+                value=256,
+                multiselect=False,
+                interactive=True,
+                allow_custom_value=True,
+                label="Output size",
+                info="Size of the output image",
+            ),
+        )
+        alph = gr.Slider(0, 1, value=1, label="Alpha", info="How much to change the original image", interactive=True, scale=3)
+
+    with gr.Row():
+        with gr.Column():
+            gr.Markdown("## Style Examples")
+            gr.Examples(
+                examples=[
+                    os.path.join(os.path.dirname(__file__), "data/styles/25.jpg"),
+                    os.path.join(os.path.dirname(__file__), "data/styles/2272.jpg"),
+                    os.path.join(os.path.dirname(__file__), "data/styles/2314.jpg"),
+                ],
+                inputs=inp1,
+            )
+        with gr.Column():
+            gr.Markdown("## Content Examples")
+            gr.Examples(
+                examples=[
+                    os.path.join(os.path.dirname(__file__), "data/content/bear.jpg"),
+                    os.path.join(os.path.dirname(__file__), "data/content/cow.jpg"),
+                    os.path.join(os.path.dirname(__file__), "data/content/ducks.jpg"),
+                ],
+                inputs=inp2,
+            )
+    btn = gr.Button("Run")
+    btn.click(fn=main, inputs=[inp1, inp2, alph, out_size[0]], outputs=out)
+
+demo.launch()

models/checkpoint_puddle_70k.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:090b259580c0a6a7fbece489de7177aa5decebeea7dc6f26dab4d9e9aeb6f700
+size 41942833

models/checkpoint_puddle_79k.pt

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

models/model_puddle.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8a215c24cf629cf5f2b2cf96014347fda704c06acd68e312720c5783eebde5ca
+size 23333701

nb.ipynb

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:53f653d323480c9ccce8283220a78f5f68fba868ca181923d41da2869408bd46
+size 60328933

requirements.txt

@@ -0,0 +1,4 @@
+torch
+torchvision
+gradio
+torchinfo

src/adain.py

@@ -0,0 +1,17 @@
+import torch
+from torch import nn
+
+
+def mi(x: torch.Tensor) -> torch.Tensor:
+    return torch.sum(x, dim=(2, 3), keepdim=True) / (x.shape[2] * x.shape[3])
+
+def sigma(x: torch.Tensor, epsilon=1e-5) -> torch.Tensor:
+    return torch.sqrt(torch.sum(((x - mi(x))**2 + epsilon), dim=(2, 3), keepdim=True) / (x.shape[2] * x.shape[3]))
+
+class AdaIN(nn.Module):
+    def __init__(self, epsilon=1e-5):
+        super().__init__()
+        self.epsilon = epsilon
+
+    def forward(self, content: torch.Tensor, style: torch.Tensor) -> torch.Tensor:
+        return (torch.mul(sigma(style, self.epsilon), ((content - mi(content)) / sigma(content, self.epsilon))) + mi(style))

src/loss.py

@@ -0,0 +1,32 @@
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+from adain import mi, sigma
+
+
+class Loss(nn.Module):
+    def __init__(self, lamb=8):
+        super().__init__()
+        self.lamb = lamb
+
+    def content_loss(self, enc_out: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
+        return F.mse_loss(enc_out, t)
+    
+    def style_loss(self, out_activations: dict, style_activations: dict) -> torch.Tensor:
+        means, sds = 0, 0
+        for out_act, style_act in zip(out_activations.values(), style_activations.values()):
+            means += F.mse_loss(mi(out_act), mi(style_act))
+            sds += F.mse_loss(sigma(out_act), sigma(style_act))
+            
+        return means + sds
+
+    def forward(self, enc_out: torch.Tensor, t: torch.Tensor, out_activations: dict, style_activations: dict) -> torch.Tensor:
+        self.loss_c = self.content_loss(enc_out, t)
+        self.loss_s = self.style_loss(out_activations, style_activations)
+        
+        return (self.loss_c + self.lamb * self.loss_s)
+    
+
+
+    

src/main.py

@@ -0,0 +1,200 @@
+import os
+import glob
+import numpy as np
+import wandb
+import copy
+import argparse
+import matplotlib.pyplot as plt
+
+import torch
+import torchvision.transforms as transforms
+from torchinfo import summary
+
+from utils import StyleContentDataset, DataStore, denorm_img
+from loss import Loss
+from model import Model
+
+
+config = {
+    "lr": 1e-4,
+    "max_iter": 80000,
+    "logging_interval": 100,
+    "preview_interval": 1000,
+    "batch_size": 4,
+    "activations": "ReLU",
+    "optimizer": "Adam",
+    "lambda": 7
+}
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+print(f"Using {device} device")
+
+def prepare_data(style_dir, content_dir, preview_dir):
+    norm = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+    
+    # Training images
+    transform = transforms.Compose([transforms.Resize(512), transforms.RandomCrop(256)])
+    style_imgs = glob.glob(os.path.join(style_dir, '*.jpg'))
+    content_imgs = glob.glob(os.path.join(content_dir, '*.jpg'))
+
+    train_dataset = StyleContentDataset(style_imgs, content_imgs, transform=transform, normalize=norm)
+    datastore = DataStore(train_dataset, batch_size=config['batch_size'], shuffle=True)
+    
+    # Preview images
+    transform = transforms.Compose([transforms.Resize(256), transforms.CenterCrop(256)])
+    preview_style_imgs = glob.glob(os.path.join(preview_dir, 'style/*.jpg'))
+    preview_content_imgs = glob.glob(os.path.join(preview_dir, 'content/*.jpg'))
+
+    # preview_dataset = StyleContentDataset(preview_style_imgs, preview_content_imgs, transform=transform, normalize=norm)
+    preview_dataset = StyleContentDataset(preview_style_imgs, [preview_content_imgs[8]] * len(preview_style_imgs), transform=transform, normalize=norm)
+    preview_datastore = DataStore(preview_dataset, batch_size=len(preview_dataset), shuffle=False)
+    
+    return datastore, preview_datastore
+
+
+def preview(model: Model, datastore: DataStore, iteration, save=False, use_wandb=False):
+    model.eval()
+    with torch.no_grad():
+        # np.random.shuffle(datastore.dataset.style_imgs)
+        # np.random.shuffle(datastore.dataset.content_imgs)
+        
+        style, content = datastore.get()
+        style, content = style.to(device), content.to(device)
+        out = model(content, style)
+        
+        fig, axs = plt.subplots(8, 6, figsize=(20, 26))
+        axs = axs.flatten()
+        i = 0
+        for (s, c, o) in zip(style, content, out): # style, content, out
+            axs[i].imshow(denorm_img(s.cpu()).permute(1, 2, 0))
+            axs[i].axis('off')
+            axs[i].set_title('style')
+            axs[i+1].imshow(denorm_img(c.cpu()).permute(1, 2, 0))
+            axs[i+1].axis('off')
+            axs[i+1].set_title('content')
+            axs[i+2].imshow(denorm_img(o.cpu()).permute(1, 2, 0))
+            axs[i+2].axis('off')
+            axs[i+2].set_title('output')
+            i += 3
+         
+        if save:
+            fig.savefig(f'outputs/{iteration}_preview.png')
+            plt.close(fig)
+        
+        if use_wandb:
+            wandb.log({'preview': wandb.Image(f'outputs/{iteration}_preview.png')}, step=iteration)    
+
+
+def train_one_iter(datastore: DataStore, model: Model, optimizer: torch.optim.Adam, loss_fn: Loss):
+    model.train()
+    
+    style, content = datastore.get()
+    style, content = style.to(device), content.to(device)
+    
+    optimizer.zero_grad()
+
+    # Forward
+    out = model(content, style)
+
+    # Save activations
+    style_activations = copy.deepcopy(model.activations)
+    
+    enc_out = model.encoder(out)
+    out_activations = model.activations
+
+    # Compute loss
+    loss = loss_fn(enc_out, model.t, out_activations, style_activations)
+
+    # Update parameters
+    loss.backward()
+    optimizer.step()
+
+    return loss.item(), loss_fn.loss_c.item(), loss_fn.loss_s.item()
+
+
+def train(datastore, preview_datastore, model: Model, optimizer: torch.optim.Adam, use_wandb=False):
+    train_history = {'style_loss': [], 'content_loss': [], 'loss': []}
+    
+    # optimizer = torch.optim.Adam(model.decoder.parameters(), lr=config['lr'])
+    loss_fn = Loss(lamb=config['lambda'])
+
+    for i in range(config['max_iter']):
+        loss, content_loss, style_loss = train_one_iter(datastore, model, optimizer, loss_fn)
+        train_history['loss'].append(loss)
+        train_history['style_loss'].append(style_loss)
+        train_history['content_loss'].append(content_loss)
+
+        if i%config['logging_interval'] == 0:
+            print(f'iter: {i}')
+            print(f'loss: {loss:>5f}, style loss: {style_loss:>5f}, content loss: {content_loss:>5f}')
+            print('-------------------------------')
+
+            if use_wandb:
+                wandb.log({
+                    'iter': i, 'loss': loss, 'style_loss': style_loss, 'content_loss': content_loss
+                })
+
+        if i%config['preview_interval'] == 0:
+            torch.save({
+                'iter': i, 'model_state': model.state_dict(), 'optimizer_state': optimizer.state_dict()
+            }, 'outputs/checkpoint.pt')
+            preview(model, preview_datastore, i, save=True, use_wandb=use_wandb)
+
+    return train_history
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--content_path', type=str, help='path to content dataset')
+    parser.add_argument('--style_path', type=str, help='path to content dataset')
+    parser.add_argument('--preview_path', type=str, help='path to preview dataset')
+    parser.add_argument('--wandb', type=str, help='wandb id')
+    parser.add_argument('--model_path', type=str, help='path to model')
+    args = parser.parse_args()
+
+    use_wandb = False
+    wandb_key = args.wandb
+    if wandb_key:
+        wandb.login(key=wandb_key)
+        wandb.init(project="assignment-3", name="", reinit=True, config=config)    
+        use_wandb = True
+
+    if args.content_path and args.style_path and args.preview_path:
+        content_dir = args.content_path
+        style_dir = args.style_path
+        preview_dir = args.preview_path
+    else:
+        print('You didnt specify the data path >:(')
+        return
+    
+    if not os.path.isdir('outputs'):
+        os.mkdir('outputs')
+
+    datastore, preview_datastore = prepare_data(style_dir, content_dir, preview_dir)
+    
+    model = Model()
+    optimizer = torch.optim.Adam(model.decoder.parameters(), lr=config['lr'])
+    if args.model_path:
+        # From checkpoint
+        checkpoint = torch.load('outputs/checkpoint.pt')
+        model.load_state_dict(checkpoint['model_state'])
+        optimizer.load_state_dict(checkpoint['optimizer_state'])
+        config['max_iter'] -= checkpoint['iter']
+
+        # From final model
+        # model.load_state_dict(torch.load(args.model_path, map_location=torch.device(device)))
+    # print(summary(model))
+    model.to(device)
+    
+    train(datastore, preview_datastore, model, optimizer, use_wandb)
+    
+    torch.save(model.state_dict(), 'outputs/model.pt')
+    if use_wandb:
+        artifact = wandb.Artifact('model', type='model')
+        artifact.add_file('outputs/model.pt')
+        wandb.log_artifact(artifact)
+        wandb.finish()
+
+
+if __name__ == '__main__':
+    main() 

src/model.py

@@ -0,0 +1,66 @@
+from torch import nn
+from torchvision.models import vgg19
+import torchvision
+from src.adain import AdaIN
+
+class Model(nn.Module):
+    def __init__(self, alpha=1.0):
+        super().__init__()
+        self.alpha = alpha
+        
+        self.encoder = nn.Sequential(*list(vgg19(weights=torchvision.models.VGG19_Weights.DEFAULT).features)[:21])
+
+        for param in self.encoder.parameters():
+            param.requires_grad = False
+        
+        # set padding in conv layers to reflect
+        # create dict for saving activations used in the style loss
+        self.activations = {}
+        for i, module in enumerate(self.encoder.children()):
+            if isinstance(module, nn.Conv2d):
+                module.padding_mode = 'reflect'
+
+            if i in [1, 6, 11, 20]:
+                module.register_forward_hook(self._save_activations(i))
+        
+        self.AdaIN = AdaIN()
+        
+        self.decoder = nn.Sequential(
+            nn.Upsample(scale_factor=2.0, mode='nearest'),
+            nn.Conv2d(512, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), padding_mode='reflect'),
+            nn.ReLU(),
+
+            nn.Upsample(scale_factor=2.0, mode='nearest'),
+            nn.Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), padding_mode='reflect'),
+            nn.ReLU(),
+            nn.Conv2d(256, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), padding_mode='reflect'),
+            nn.ReLU(),
+            
+            nn.Upsample(scale_factor=2.0, mode='nearest'),
+            nn.Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), padding_mode='reflect'),
+            nn.ReLU(),
+            nn.Conv2d(128, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), padding_mode='reflect'),
+            nn.ReLU(),
+            
+            nn.Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), padding_mode='reflect'),
+            nn.ReLU(),
+            nn.Conv2d(64, 3, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), padding_mode='reflect'),
+            nn.Tanh()
+        )
+
+    # https://stackoverflow.com/a/68854535
+    def _save_activations(self, name):
+        def hook(module, input, output):
+            self.activations[name] = output
+        return hook
+
+    def forward(self, content, style):
+        enc_content = self.encoder(content)
+        enc_style = self.encoder(style)
+        
+        self.t = self.AdaIN(enc_content, enc_style)
+        self.t = (1.0 - self.alpha) * enc_content + self.alpha * self.t
+        out = self.decoder(self.t)
+
+        return out
+    

src/utils.py

@@ -0,0 +1,62 @@
+import torch
+from torch.utils.data import Dataset, DataLoader
+from torchvision.io import read_image, ImageReadMode
+import numpy as np
+
+
+def denorm_img(img: torch.Tensor) -> torch.Tensor:
+    std = torch.Tensor([0.229, 0.224, 0.225]).reshape(-1, 1, 1)
+    mean = torch.Tensor([0.485, 0.456, 0.406]).reshape(-1, 1, 1)
+    return torch.clip(img * std + mean, min=0, max=1)
+
+
+class StyleContentDataset(Dataset):
+    def __init__(self, style_imgs, content_imgs, transform=None, normalize=None):
+        self.style_imgs = style_imgs
+        self.content_imgs = content_imgs
+        self.transform = transform
+        self.normalize = normalize
+
+    def __len__(self):
+        if len(self.style_imgs) < len(self.content_imgs):
+            return len(self.style_imgs)
+        else:
+            return len(self.content_imgs)
+    
+    def __getitem__(self, idx):
+        try:
+            style = read_image(self.style_imgs[idx], ImageReadMode.RGB).float() / 255.0
+            content = read_image(self.content_imgs[idx], ImageReadMode.RGB).float() / 255.0
+        except RuntimeError:
+            print(self.style_imgs[idx])
+            print(self.content_imgs[idx])
+            style = read_image(self.style_imgs[0], ImageReadMode.RGB).float() / 255.0
+            content = read_image(self.content_imgs[0], ImageReadMode.RGB).float() / 255.0
+
+        if self.normalize:
+            style = self.normalize(style)
+            content = self.normalize(content)
+        
+        if self.transform:
+            style = self.transform(style)
+            content = self.transform(content)
+        
+        return style, content
+    
+
+class DataStore():
+    def __init__(self, dataset: StyleContentDataset, batch_size, shuffle=False):
+        self.dataset = dataset
+        self.dataloader = DataLoader(self.dataset, batch_size=batch_size, shuffle=shuffle, num_workers=2)
+        self.iterator = iter(self.dataloader)
+
+    def get(self):
+        try:
+           style, content = next(self.iterator)
+        except (StopIteration):
+            # print('| Repeating |')
+            # np.random.shuffle(self.dataset.style_imgs)
+            self.iterator = iter(self.dataloader)
+            style, content = next(self.iterator)
+        
+        return style, content