Upload 2 files

Final_8.ckpt

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

minimal_script.py

@@ -0,0 +1,143 @@
+import os
+import torch
+import math
+import numpy as np
+import torch.nn as nn
+import lightning.pytorch as pl
+import imageio
+from torchvision.transforms import v2
+
+
+class BasicBlock(nn.Module):
+    def __init__(self, channels, kernel_size=(3,3)):
+        super().__init__()
+        layers = []
+        num_conv = len(channels)-1
+        for i in range(num_conv):
+            layers.append(nn.Conv2d(channels[i], channels[i+1],
+                                    kernel_size=kernel_size, padding='same', padding_mode='reflect', bias=False))
+            layers.append(nn.InstanceNorm2d(channels[i+1], affine=False))
+            layers.append(nn.ReLU())
+        self.operations = nn.Sequential(*layers)
+
+    def forward(self, x):
+        return self.operations(x)
+
+
+class ResBlock(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size=(3,3), num_conv=2):
+        super().__init__()
+        layers = []
+        if in_channels == out_channels:
+            self.mapping = nn.Identity()
+        else:
+            self.mapping = nn.Conv2d(in_channels, out_channels, 1)
+        for i in range(num_conv):
+            layers.append(nn.Conv2d(in_channels if i == 0 else out_channels, out_channels,
+                                    kernel_size=kernel_size, padding='same', padding_mode='reflect', bias=False))
+            layers.append(nn.InstanceNorm2d(out_channels, affine=False))
+            layers.append(nn.ReLU())
+        self.operations = nn.Sequential(*layers)
+
+    def forward(self, x):
+        return (self.mapping(x) + self.operations(x)) / math.sqrt(2)
+
+
+class ConvPool(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+        layers = []
+        layers.append(nn.Conv2d(in_channels, out_channels, 4, 2, 1, bias=False, padding_mode='reflect'))
+        layers.append(nn.InstanceNorm2d(out_channels, affine=False))
+        layers.append(nn.ReLU(inplace=True))
+        self.operations = nn.Sequential(*layers)
+
+    def forward(self, x):
+        return self.operations(x)
+
+
+class EmbeddingNetworkSmall(nn.Module):
+    def __init__(self):
+        super(EmbeddingNetworkSmall, self).__init__()
+        self.conv1 = BasicBlock((3, 8, 16), (3, 3))
+        self.pool1 = ConvPool(16, 32) # 2
+        self.conv2 = ResBlock(32, 32, (3, 3), 3)
+        self.pool2 = ConvPool(32, 64) # 4
+        self.conv3 = ResBlock(64, 64, (3, 3), 3)
+        self.drop1 = nn.Dropout2d(p=0.25)
+        self.pool3 = ConvPool(64, 128) # 8
+        self.conv4 = ResBlock(128, 128, (3, 3), 3)
+        self.adpool = nn.AdaptiveAvgPool2d(1)
+        self.poolnorm = nn.LayerNorm(128, elementwise_affine=False)
+        self.flatten = nn.Flatten()
+        self.drop2 = nn.Dropout(p=0.33)
+        self.fc1 = nn.Linear(128, 128, bias=False)
+        self.fc1norm = nn.LayerNorm(128, elementwise_affine=False)
+        self.act = nn.ReLU()
+        self.fc2 = nn.Linear(128, 128, bias=False)
+        self.fc2norm = nn.LayerNorm(128, elementwise_affine=False)
+        self.fc3 = nn.Linear(128, 8)
+
+        self.use_checkpoint = False
+
+    def forward(self, x):
+        x = self.pool1(self.conv1(x))
+        x = self.pool2(self.conv2(x))
+        x = self.pool3(self.drop1(self.conv3(x)))
+        x = self.conv4(x)
+
+        x = self.adpool(x)
+        x = self.poolnorm(self.flatten(x))
+        x = self.act(self.drop2(x))
+        x = self.act(self.fc1norm(self.fc1(x)))
+        x = self.act(self.fc2norm(self.fc2(x)))
+        x = self.fc3(x)
+        return x
+
+
+
+class PLModule(pl.LightningModule):
+    def __init__(self):
+        super().__init__()
+        self.save_hyperparameters()
+        self.network = EmbeddingNetworkSmall()
+
+    def forward(self, x):
+        return self.network(x)
+
+
+def down_to_1k(img, size=1024):
+    h, w = img.shape[1], img.shape[2]
+    area = h * w
+    if area > size ** 2:
+        scale_factor = (size ** 2 / area) ** 0.5
+        new_h = math.floor(h * scale_factor)
+        new_w = math.floor(w * scale_factor)
+        img = v2.functional.resize(img, (new_w, new_h))
+    return img
+
+
+def closest_interval(img, interval=8):
+    c, h, w = img.shape
+    new_h = h - (h % interval) if h % interval != 0 else h
+    new_w = w - (w % interval) if w % interval != 0 else w
+    h_start = (h - new_h) // 2
+    w_start = (w - new_w) // 2
+    new_h, new_w = max(new_h, interval), max(new_w, interval)
+    return img[:, h_start:h_start + new_h, w_start:w_start + new_w]
+
+
+if __name__ == '__main__':
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model = PLModule.load_from_checkpoint('Final_8.ckpt')
+    model.to(device)
+    model.eval()
+
+    img = imageio.v3.imread('images_for_style_embedding/6857740.webp').copy()
+    img = torch.from_numpy(img).permute(2, 0, 1)
+    img = closest_interval(down_to_1k(img))
+    img = 2*(img/255)-1
+    img = img.unsqueeze(0).to(device)
+
+    pred = model(img)
+    print(pred)