Upload 3 files

.gitattributes

@@ -33,3 +33,4 @@ 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
+Style_Embedder_v3.safetensors_ filter=lfs diff=lfs merge=lfs -text

Style_Embedder_v3.safetensors_

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

gallery_review.py

@@ -7,6 +7,7 @@ import imageio
 import random
 import matplotlib.pyplot as plt
 import cv2
+import skdim
 
 from torch.utils.data import Dataset, DataLoader
 
@@ -16,6 +17,7 @@ from safetensors.torch import save_file, load_file
 from sklearn.cluster import AgglomerativeClustering
 from sklearn.manifold import TSNE
 from sklearn.neighbors import KDTree
+from sklearn.preprocessing import StandardScaler
 
 from minimal_script import EmbeddingNetwork, closest_interval, adj_size, PLModule
 
@@ -126,9 +128,9 @@ def explore_embedding_space(embeddings, image_paths, model):
     def get_overlay_image(image_path):
         """Get image with gradient overlay"""
         img = Image.open(image_path).convert('RGB')
-        heatmap = compute_gradient_heatmap(image_path)
-        return overlay_heatmap(img, heatmap)
-        #return img
+        #heatmap = compute_gradient_heatmap(image_path)
+        #return overlay_heatmap(img, heatmap)
+        return img
 
     def add_caption_to_image(image, caption):
         """Add text caption to the bottom of an image"""
@@ -217,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_dataset = PredictDataset(image_folder, 5000)
     predict_loader = DataLoader(predict_dataset, batch_size=1, num_workers=5, pin_memory=True)
     trainer = pl.Trainer(accelerator="gpu", logger=False, enable_checkpointing=False, precision="16-mixed")
     predictions_0 = trainer.predict(model, predict_loader)
@@ -227,6 +229,20 @@ def generate_embeddings(image_folder, mode, model):
         for i in pred[1]:
             paths.append(i)
     if mode == 'Grouping':
+        #estimate global intrinsic dimension
+        #scaler = StandardScaler()
+        #normalised_predictions = scaler.fit_transform(predictions)
+        # Initialize estimators
+        estimators = [skdim.id.TwoNN(), skdim.id.CorrInt(), skdim.id.DANCo()]
+        results = {}
+        
+        for est in estimators:
+            est.fit(predictions)
+            results[type(est).__name__] = est.dimension_
+
+        print("Intrinsic Dimension Estimates:")
+        for name, dim in results.items():
+            print(f"{name}: {dim:.2f}")
         labels = cluster_embeddings(predictions)
 
         row_norms = np.linalg.norm(predictions, axis=1)
@@ -297,11 +313,11 @@ def cluster_embeddings(predictions, distance_threshold=32.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()
     state_dict = load_file("Style_Embedder_v2.safetensors")
     model.network.load_state_dict(state_dict)
     # 'Grouping' or 'Explore'
-    generate_embeddings(folder, 'Explore', model)
+    generate_embeddings(folder, 'Grouping', model)

minimal_script.py

@@ -18,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.GroupNorm(1, channels[i+1]))
+            layers.append(nn.InstanceNorm2d(channels[i+1]))
             layers.append(nn.LeakyReLU(inplace=True))
         if dropout > 0.0:
             layers.append(nn.Dropout2d(dropout))
@@ -33,7 +33,7 @@ class ResBlock(nn.Module):
         super().__init__()
         layers = []
         for i in range(num_conv):
-            layers.append(nn.GroupNorm(1, channels))
+            layers.append(nn.InstanceNorm2d(channels))
             if i == num_conv-1 and dropout > 0.0:
                 layers.append(nn.Dropout2d(dropout))
             layers.append(nn.LeakyReLU(inplace=True))
@@ -50,7 +50,7 @@ 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.GroupNorm(1, out_channels))
+        layers.append(nn.InstanceNorm2d(out_channels))
         #layers.append(nn.LeakyReLU(inplace=True))
         self.operations = nn.Sequential(*layers)
 
@@ -95,13 +95,13 @@ class EmbeddingNetwork(nn.Module):
         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.conpactnorm = nn.LayerNorm(1024, elementwise_affine=False)
         self.fc1 = nn.Linear(1024, 1024, bias=False)
-        self.fc1norm = nn.LayerNorm(1024, elementwise_affine=True)
+        self.fc1norm = nn.LayerNorm(1024, elementwise_affine=False)
         self.act = nn.LeakyReLU(inplace=True)
         self.fc2 = nn.Linear(1024, 1024, bias=False)
-        self.fc2norm = nn.LayerNorm(1024, elementwise_affine=True)
-        self.fc3 = nn.Linear(1024, 4)
+        self.fc2norm = nn.LayerNorm(1024, elementwise_affine=False)
+        self.fc3 = nn.Linear(1024, 6)
 
     def forward(self, x):
         x = self.input_conv(x)
@@ -125,6 +125,12 @@ class PLModule(pl.LightningModule):
         super().__init__()
         self.save_hyperparameters()
         self.network = EmbeddingNetwork()
+        self.register_buffer("val_pos_sum", torch.tensor(0.0))
+        self.register_buffer("val_neg_sum", torch.tensor(0.0))
+        self.register_buffer("val_count", torch.tensor(0))
+        self.register_buffer("train_pos_sum", torch.tensor(0.0))
+        self.register_buffer("train_neg_sum", torch.tensor(0.0))
+        self.register_buffer("train_count", torch.tensor(0))
 
     def forward(self, x):
         return self.network(x)
@@ -159,7 +165,7 @@ if __name__ == '__main__':
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     
     model = EmbeddingNetwork()
-    state_dict = load_file("Style_Embedder_v2.safetensors")
+    state_dict = load_file("Style_Embedder_v3.safetensors")
     model.load_state_dict(state_dict)
     
     model.to(device).to(torch.float16)