Upload 3 files

Style Embedder v1.ckpt

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

gallery_review.py

@@ -13,7 +13,7 @@ from torch.utils.data import Dataset, DataLoader
 from PIL import Image
 from matplotlib import cm
 
-from minimal_script import EmbeddingNetworkSmall, closest_interval, down_to_1k
+from minimal_script import EmbeddingNetworkSmall, closest_interval, adj_size
 from sklearn.cluster import AgglomerativeClustering
 from sklearn.manifold import TSNE
 from sklearn.neighbors import KDTree
@@ -50,7 +50,7 @@ class PredictDataset(Dataset):
         path = self.image_paths[idx]
         image = imageio.v3.imread(path).copy()
         image = torch.from_numpy(image).permute(2, 0, 1)
-        processed = closest_interval(down_to_1k(image, 1024))
+        processed = closest_interval(adj_size(image, 1024))
         processed = 2*(processed/255)-1
         return processed.detach(), path
 
@@ -95,7 +95,7 @@ def explore_embedding_space(embeddings, image_paths, model):
         # Load and preprocess image
         img = imageio.v3.imread(image_path).copy()
         img = torch.from_numpy(img).permute(2, 0, 1)
-        img_tensor = closest_interval(down_to_1k(img, 1024)).unsqueeze(0)
+        img_tensor = closest_interval(adj_size(img, 1024)).unsqueeze(0)
         img_tensor = 2*(img_tensor/255)-1
         img_tensor.requires_grad_(True)
 
@@ -311,6 +311,6 @@ if __name__ == '__main__':
     folder = 'Enter Images folder name here'
     #folder = 'images_for_style_embedding'
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    model = PLModule.load_from_checkpoint('Final_8.ckpt')
+    model = PLModule.load_from_checkpoint('Style Embedder v1.ckpt')
     # 'Grouping' or 'Explore'
     generate_embeddings(folder, 'Grouping', model)

minimal_script.py

@@ -9,15 +9,17 @@ from torchvision.transforms import v2
 
 
 class BasicBlock(nn.Module):
-    def __init__(self, channels, kernel_size=(3,3)):
+    def __init__(self, channels, kernel_size=(3,3), dropout=0.0):
         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())
+            layers.append(nn.InstanceNorm2d(channels[i+1]))
+            layers.append(nn.LeakyReLU(inplace=True))
+        if dropout > 0.0:
+            layers.append(nn.Dropout2d(dropout))
         self.operations = nn.Sequential(*layers)
 
     def forward(self, x):
@@ -25,70 +27,88 @@ class BasicBlock(nn.Module):
 
 
 class ResBlock(nn.Module):
-    def __init__(self, in_channels, out_channels, kernel_size=(3,3), num_conv=2):
+    def __init__(self, channels, kernel_size=(3,3), num_conv=2, dropout=0.0):
         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,
+            layers.append(nn.Conv2d(channels, channels,
                                     kernel_size=kernel_size, padding='same', padding_mode='reflect', bias=False))
-            layers.append(nn.InstanceNorm2d(out_channels, affine=False))
-            layers.append(nn.ReLU())
+            layers.append(nn.InstanceNorm2d(channels))
+            layers.append(nn.LeakyReLU(inplace=True))
+        self.norm = nn.InstanceNorm2d(channels)
+        self.dropout = nn.Dropout2d(dropout) if dropout > 0 else nn.Identity()
         self.operations = nn.Sequential(*layers)
 
     def forward(self, x):
-        return (self.mapping(x) + self.operations(x)) / math.sqrt(2)
+        return self.dropout(self.norm(x + self.operations(x)))
 
 
 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.Conv2d(in_channels, out_channels, 4, 2, 1, padding_mode='reflect', bias=False))
         layers.append(nn.InstanceNorm2d(out_channels, affine=False))
-        layers.append(nn.ReLU(inplace=True))
+        layers.append(nn.LeakyReLU(inplace=True))
         self.operations = nn.Sequential(*layers)
 
     def forward(self, x):
         return self.operations(x)
 
 
+class CompactGramMatrix(nn.Module):
+    def __init__(self, in_channels):
+        super().__init__()
+        self.in_channels = in_channels
+        # Precompute indices for lower triangle (including diagonal)
+        self.register_buffer('tril_indices',
+                             torch.tril_indices(in_channels, in_channels, offset=0, dtype=torch.int32))
+
+    def forward(self, x):
+        """
+        Input: (B, C, H, W)
+        Output: (B, C*(C+1)//2) compact Gram features
+        """
+        b, c, h, w = x.size()
+        x = x.view(b, c, -1) / ((h * w) ** 0.5)  # Flatten spatial dimensions -> (B, C, H*W), then normalise
+
+        # Compute full Gram matrix (still needed temporarily)
+        gram = torch.bmm(x, x.transpose(1, 2))  # (B, C, C)
+
+        # Extract lower triangle including diagonal
+        compact_gram = gram[:, self.tril_indices[0], self.tril_indices[1]]  # (B, n_unique)
+        return compact_gram
+
+
 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
+        self.conv1 = BasicBlock((3, 8, 16, 24), (3, 3))
+        self.pool1 = ConvPool(24, 48) # 2
+        self.conv2 = ResBlock(48, (3, 3), 3, 0.2)
+        self.pool2 = ConvPool(48, 96) # 4
+        self.conv3 = ResBlock(96, (3, 3), 2, 0.25)
+        self.pool3 = ConvPool(96, 192) # 8
+        self.conv4 = ResBlock(192, (3, 3), 2, 0.3)
+        self.gram = CompactGramMatrix(192)
+        self.compact = nn.Linear(192*(192+1)//2, 192, bias=False)
+        self.conpactnorm = nn.LayerNorm(192, elementwise_affine=False)
+        self.fc1 = nn.Linear(192, 192, bias=False)
+        self.fc1norm = nn.LayerNorm(192, elementwise_affine=False)
+        self.act = nn.LeakyReLU(inplace=True)
+        self.fc2 = nn.Linear(192, 192, bias=False)
+        self.fc2norm = nn.LayerNorm(192, elementwise_affine=False)
+        self.fc3 = nn.Linear(192, 8)
 
     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.pool3(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.gram(x)
+        x = self.compact(x)
+        x = self.conpactnorm(x)
         x = self.act(self.fc1norm(self.fc1(x)))
         x = self.act(self.fc2norm(self.fc2(x)))
         x = self.fc3(x)
@@ -106,7 +126,7 @@ class PLModule(pl.LightningModule):
         return self.network(x)
 
 
-def down_to_1k(img, size=1024):
+def adj_size(img, size=512):
     h, w = img.shape[1], img.shape[2]
     area = h * w
     if area > size ** 2:
@@ -129,13 +149,13 @@ def closest_interval(img, interval=8):
 
 if __name__ == '__main__':
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    model = PLModule.load_from_checkpoint('Final_8.ckpt')
+    model = PLModule.load_from_checkpoint('Style Embedder v1.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 = closest_interval(adj_size(img))
     img = 2*(img/255)-1
     img = img.unsqueeze(0).to(device)