Add shared MoMA codebook option

app.py

@@ -189,7 +189,13 @@ TOKENIZER_CODE = """# This tab builds a tiny image tokenizer from one image.
 patches = split_image_into_grid(image, grid_size, patch_size)
 vectors = flatten_each_patch_to_rgb_numbers(patches)
 
-codebook = k_means(vectors, codebook_size)
+# The codebook can be trained on this image only,
+# or on patches from all loaded MoMA images.
+training_vectors = vectors
+if codebook_source == "All loaded MoMA images":
+    training_vectors = flatten_patches_from_many_images(moma_images)
+
+codebook = k_means(training_vectors, codebook_size)
 token_ids = nearest_code_for_each_patch(vectors, codebook)
 
 # The compressed image is now a grid of integers.
@@ -770,14 +776,32 @@ def learned_codebook_rows(centroids, assignments, errors):
 
 
 @spaces.GPU(duration=30)
-def tokenize_image(image, grid_size, patch_size, codebook_size, iterations, seed):
+def tokenize_image(image, grid_size, patch_size, codebook_size, iterations, seed, codebook_source):
     grid_size = max(4, min(32, int(grid_size)))
     patch_size = max(4, min(32, int(patch_size)))
     codebook_size = max(2, min(64, int(codebook_size)))
     prepared = prepare_tokenizer_image(image, grid_size, patch_size)
     patches = image_to_patches(prepared, grid_size, patch_size)
     vectors = patches.reshape(len(patches), -1)
-    centroids, assignments, errors = learn_codebook(vectors, codebook_size, iterations, seed)
+
+    if codebook_source == "All loaded MoMA images":
+        items, _ = load_moma_items()
+        training_vectors = []
+        for item in items:
+            train_image = prepare_tokenizer_image(item["img"], grid_size, patch_size)
+            train_patches = image_to_patches(train_image, grid_size, patch_size)
+            training_vectors.append(train_patches.reshape(len(train_patches), -1))
+        training_vectors = np.concatenate(training_vectors, axis=0)
+        centroids, _, _ = learn_codebook(training_vectors, codebook_size, iterations, seed)
+        source_text = f"all {len(items)} loaded MoMA images ({len(training_vectors)} training patches)"
+    else:
+        training_vectors = vectors
+        centroids, _, _ = learn_codebook(training_vectors, codebook_size, iterations, seed)
+        source_text = f"this image only ({len(training_vectors)} training patches)"
+
+    distances = squared_distances(vectors, centroids)
+    assignments = np.argmin(distances, axis=1).astype(np.int32)
+    errors = distances[np.arange(len(vectors)), assignments]
     reconstruction_vectors = centroids[assignments]
     reconstruction = patches_to_image(reconstruction_vectors.reshape(len(patches), patch_size, patch_size, 3), grid_size, patch_size)
     token_grid = draw_token_id_grid(assignments, grid_size)
@@ -787,25 +811,25 @@ def tokenize_image(image, grid_size, patch_size, codebook_size, iterations, seed
     summary = (
         f"Encoded the image into a {grid_size}x{grid_size} grid: {grid_size * grid_size} discrete token IDs.\n"
         f"Each token represents one {patch_size}x{patch_size} RGB patch. "
-        f"The codebook has {len(centroids)} learned entries from this image.\n"
+        f"The codebook has {len(centroids)} learned entries trained from {source_text}.\n"
         f"Mean reconstruction error per patch: {float(errors.mean()):.6f}. "
         "Lower error means the learned codebook patch is closer to the original patch."
     )
     return prepared, token_grid, reconstruction, error_map, gallery, rows, summary
 
 
-def initial_tokenizer(grid_size, patch_size, codebook_size, iterations, seed):
+def initial_tokenizer(grid_size, patch_size, codebook_size, iterations, seed, codebook_source):
     items, message = load_moma_items()
-    outputs = tokenize_image(items[0]["img"], grid_size, patch_size, codebook_size, iterations, seed)
+    outputs = tokenize_image(items[0]["img"], grid_size, patch_size, codebook_size, iterations, seed, codebook_source)
     summary = f"{message}\n\n{outputs[-1]}"
     return (moma_gallery_items(),) + outputs[:-1] + (summary,)
 
 
-def tokenize_moma_selection(grid_size, patch_size, codebook_size, iterations, seed, evt: gr.SelectData):
+def tokenize_moma_selection(grid_size, patch_size, codebook_size, iterations, seed, codebook_source, evt: gr.SelectData):
     items, _ = load_moma_items()
     index = evt.index if isinstance(evt.index, int) else 0
     index = max(0, min(index, len(items) - 1))
-    return tokenize_image(items[index]["img"], grid_size, patch_size, codebook_size, iterations, seed)
+    return tokenize_image(items[index]["img"], grid_size, patch_size, codebook_size, iterations, seed, codebook_source)
 
 
 def probability_rows(record):
@@ -1055,6 +1079,11 @@ def build_app():
                         with gr.Row():
                             tokenizer_codebook_size = gr.Slider(4, 48, value=16, step=1, label="Codebook size")
                             tokenizer_iterations = gr.Slider(2, 30, value=10, step=1, label="K-means iterations")
+                        tokenizer_codebook_source = gr.Radio(
+                            ["This image only", "All loaded MoMA images"],
+                            value="This image only",
+                            label="Learn codebook from",
+                        )
                         tokenizer_seed = gr.Number(value=7, precision=0, label="Codebook seed")
                         tokenize_button = gr.Button("Tokenize uploaded image", variant="primary")
 
@@ -1235,6 +1264,7 @@ def build_app():
             tokenizer_codebook_size,
             tokenizer_iterations,
             tokenizer_seed,
+            tokenizer_codebook_source,
         ]
         tokenizer_outputs = [
             tokenizer_original,
@@ -1261,6 +1291,7 @@ def build_app():
                 tokenizer_codebook_size,
                 tokenizer_iterations,
                 tokenizer_seed,
+                tokenizer_codebook_source,
             ],
             outputs=tokenizer_outputs,
             show_progress="minimal",
@@ -1274,6 +1305,7 @@ def build_app():
                 tokenizer_codebook_size,
                 tokenizer_iterations,
                 tokenizer_seed,
+                tokenizer_codebook_source,
             ],
             outputs=tokenizer_outputs,
             show_progress="minimal",