Upload 3 files

Style_Embedder_v2.safetensors

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

gallery_review.py

@@ -12,25 +12,12 @@ from torch.utils.data import Dataset, DataLoader
 
 from PIL import Image
 from matplotlib import cm
-
-from minimal_script import EmbeddingNetworkSmall, closest_interval, adj_size
+from safetensors.torch import save_file, load_file
 from sklearn.cluster import AgglomerativeClustering
 from sklearn.manifold import TSNE
 from sklearn.neighbors import KDTree
 
-
-class PLModule(pl.LightningModule):
-    def __init__(self):
-        super().__init__()
-        self.save_hyperparameters()
-        self.network = EmbeddingNetworkSmall()
-
-    def forward(self, x):
-        return self.network(x)
-
-    def predict_step(self, batch, batch_idx, dataloader_idx=0):
-        outputs = self.forward(batch[0])
-        return outputs, batch[1]
+from minimal_script import EmbeddingNetwork, closest_interval, adj_size, PLModule
 
 
 class PredictDataset(Dataset):
@@ -101,7 +88,7 @@ def explore_embedding_space(embeddings, image_paths, model):
 
         # Move to GPU if available
         device = 'cuda' if torch.cuda.is_available() else 'cpu'
-        img_tensor = img_tensor.to(device)
+        img_tensor = img_tensor.to(device).to(torch.float16)
 
         # Compute embedding and gradient
         with torch.enable_grad():
@@ -122,7 +109,7 @@ def explore_embedding_space(embeddings, image_paths, model):
         heatmap = cm.jet(grad_norm)[..., :3]  # Use jet colormap
         return heatmap
 
-    def overlay_heatmap(original_img, heatmap, alpha=0.6):
+    def overlay_heatmap(original_img, heatmap, alpha=0.4):
         """Overlay heatmap on original image"""
         # Resize heatmap to match original image
         heatmap_img = Image.fromarray((heatmap * 255).astype(np.uint8))
@@ -232,7 +219,7 @@ def explore_embedding_space(embeddings, image_paths, model):
 def generate_embeddings(image_folder, mode, model):
     predict_dataset = PredictDataset(image_folder, 1000)
     predict_loader = DataLoader(predict_dataset, batch_size=1, num_workers=5, pin_memory=True)
-    trainer = pl.Trainer(accelerator="gpu", logger=False, enable_checkpointing=False)
+    trainer = pl.Trainer(accelerator="gpu", logger=False, enable_checkpointing=False, precision="16-mixed")
     predictions_0 = trainer.predict(model, predict_loader)
     predictions = torch.cat([pred[0] for pred in predictions_0], dim=0).numpy()
     paths = []
@@ -250,7 +237,8 @@ def generate_embeddings(image_folder, mode, model):
         plt.ylabel('Average Norm')
         plt.title(f'Average Norm for Each Feature (Column)')
         plt.xticks(range(predictions.shape[1]))
-        plt.show()
+        #plt.show()
+        plt.savefig('Norms.png')
 
         plt.figure(figsize=(8, 6))
         tsne = TSNE(n_components=2, random_state=42)
@@ -263,7 +251,8 @@ def generate_embeddings(image_folder, mode, model):
         plt.legend()
         plt.grid(True)
         plt.axis('equal')
-        plt.show()
+        #plt.show()
+        plt.savefig('Groups.png')
 
         # List unique clusters
         unique_clusters = np.unique(labels)
@@ -291,12 +280,12 @@ def generate_embeddings(image_folder, mode, model):
 
         demo.launch()
     elif mode == 'Explore':
-        demo = explore_embedding_space(predictions, paths, model.to('cuda'))
+        demo = explore_embedding_space(predictions, paths, model.to('cuda').to(torch.float16))
         demo.launch()
 
 
 # Apply Agglomerative Clustering
-def cluster_embeddings(predictions, distance_threshold=6.0):
+def cluster_embeddings(predictions, distance_threshold=32.0):
     agg_clustering = AgglomerativeClustering(
         n_clusters=None,
         distance_threshold=distance_threshold,
@@ -308,9 +297,11 @@ def cluster_embeddings(predictions, distance_threshold=6.0):
 
 
 if __name__ == '__main__':
-    folder = 'Enter Images folder name here'
-    #folder = 'images_for_style_embedding'
+    #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('Style Embedder v1.ckpt')
+    model = PLModule()
+    state_dict = load_file("Style_Embedder_v2.safetensors")
+    model.network.load_state_dict(state_dict)
     # 'Grouping' or 'Explore'
-    generate_embeddings(folder, 'Grouping', model)
+    generate_embeddings(folder, 'Explore', model)

minimal_script.py

@@ -5,7 +5,9 @@ import numpy as np
 import torch.nn as nn
 import lightning.pytorch as pl
 import imageio
+import safetensors
 from torchvision.transforms import v2
+from safetensors.torch import save_file, load_file
 
 
 class BasicBlock(nn.Module):
@@ -16,7 +18,7 @@ class BasicBlock(nn.Module):
         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]))
+            layers.append(nn.GroupNorm(1, channels[i+1]))
             layers.append(nn.LeakyReLU(inplace=True))
         if dropout > 0.0:
             layers.append(nn.Dropout2d(dropout))
@@ -27,20 +29,20 @@ class BasicBlock(nn.Module):
 
 
 class ResBlock(nn.Module):
-    def __init__(self, channels, kernel_size=(3,3), num_conv=2, dropout=0.0):
+    def __init__(self, channels, kernel_size=3, num_conv=2, dropout=0.0):
         super().__init__()
         layers = []
         for i in range(num_conv):
+            layers.append(nn.GroupNorm(1, channels))
+            if i == num_conv-1 and dropout > 0.0:
+                layers.append(nn.Dropout2d(dropout))
+            layers.append(nn.LeakyReLU(inplace=True))
             layers.append(nn.Conv2d(channels, channels,
                                     kernel_size=kernel_size, padding='same', padding_mode='reflect', bias=False))
-            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.dropout(self.norm(x + self.operations(x)))
+        return x + self.operations(x)
 
 
 class ConvPool(nn.Module):
@@ -48,8 +50,8 @@ class ConvPool(nn.Module):
         super().__init__()
         layers = []
         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.LeakyReLU(inplace=True))
+        layers.append(nn.GroupNorm(1, out_channels))
+        #layers.append(nn.LeakyReLU(inplace=True))
         self.operations = nn.Sequential(*layers)
 
     def forward(self, x):
@@ -80,27 +82,29 @@ class CompactGramMatrix(nn.Module):
         return compact_gram
 
 
-class EmbeddingNetworkSmall(nn.Module):
+class EmbeddingNetwork(nn.Module):
     def __init__(self):
-        super(EmbeddingNetworkSmall, self).__init__()
-        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)
+        super(EmbeddingNetwork, self).__init__()
+        self.input_conv = nn.Conv2d(3, 32, 5, padding='same', padding_mode='reflect', bias=False)
+        self.conv1 = ResBlock(32, 3, 3)
+        self.pool1 = ConvPool(32, 64) # 2
+        self.conv2 = ResBlock(64, 3, 3)
+        self.pool2 = ConvPool(64, 128) # 4
+        self.conv3 = ResBlock(128, 3, 3)
+        self.pool3 = ConvPool(128, 256) # 8
+        self.conv4 = ResBlock(256, 3, 3)
+        self.gram = CompactGramMatrix(256)
+        self.compact = nn.Linear(256*(256+1)//2, 1024, bias=False)
+        self.conpactnorm = nn.LayerNorm(1024, elementwise_affine=True)
+        self.fc1 = nn.Linear(1024, 1024, bias=False)
+        self.fc1norm = nn.LayerNorm(1024, elementwise_affine=True)
         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)
+        self.fc2 = nn.Linear(1024, 1024, bias=False)
+        self.fc2norm = nn.LayerNorm(1024, elementwise_affine=True)
+        self.fc3 = nn.Linear(1024, 4)
 
     def forward(self, x):
+        x = self.input_conv(x)
         x = self.pool1(self.conv1(x))
         x = self.pool2(self.conv2(x))
         x = self.pool3(self.conv3(x))
@@ -120,13 +124,17 @@ class PLModule(pl.LightningModule):
     def __init__(self):
         super().__init__()
         self.save_hyperparameters()
-        self.network = EmbeddingNetworkSmall()
+        self.network = EmbeddingNetwork()
 
     def forward(self, x):
         return self.network(x)
 
+    def predict_step(self, batch, batch_idx, dataloader_idx=0):
+        outputs = self.forward(batch[0])
+        return outputs, batch[1]
+
 
-def adj_size(img, size=512):
+def adj_size(img, size=1536):
     h, w = img.shape[1], img.shape[2]
     area = h * w
     if area > size ** 2:
@@ -149,15 +157,19 @@ 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('Style Embedder v1.ckpt')
-    model.to(device)
+    
+    model = EmbeddingNetwork()
+    state_dict = load_file("Style_Embedder_v2.safetensors")
+    model.load_state_dict(state_dict)
+    
+    model.to(device).to(torch.float16)
     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(adj_size(img))
     img = 2*(img/255)-1
-    img = img.unsqueeze(0).to(device)
+    img = img.unsqueeze(0).to(device).to(torch.float16)
 
     pred = model(img)
-    print(pred)
+    print(pred)