Add learned VQ tokenizer option

README.md

@@ -15,7 +15,12 @@ license: mit
 
 A CPU-friendly Gradio app for teaching image tokenization and image generation as sequence problems.
 
-The first tab takes a real image, splits it into patches, learns a tiny vector-quantized codebook from those patches, and shows the resulting token ID grid and reconstruction. It can use MoMA collection images or an uploaded image.
+The first tab takes a real image and shows image tokenization two ways:
+
+- a pretrained learned VQ tokenizer from `CompVis/ldm-celebahq-256/vqvae`
+- a transparent k-means patch tokenizer for comparison
+
+The learned tokenizer shows real learned codebook IDs, reconstruction from the VQ decoder, token usage, and representative image regions for the most-used codes. The k-means option can learn a tiny codebook from one image or from all loaded MoMA images.
 
 The app also includes a deliberately small, transparent image-token sampler. It does not call a proprietary image model. Instead, it shows the mechanics that matter for a workshop:
 

app.py

@@ -7,10 +7,13 @@ from functools import lru_cache
 from io import BytesIO
 
 os.environ.setdefault("MPLCONFIGDIR", "/tmp/matplotlib")
+os.environ.setdefault("USE_TF", "0")
+os.environ.setdefault("TRANSFORMERS_NO_TF", "1")
 
 import gradio as gr
 import numpy as np
 import requests
+import torch
 from PIL import Image, ImageDraw, ImageFont
 
 try:
@@ -42,6 +45,8 @@ HEADERS = {
 }
 MOMA_SAMPLE_SIZE = 24
 MOMA_CANDIDATES = 120
+DEFAULT_VQ_MODEL = "CompVis/ldm-celebahq-256"
+DEFAULT_VQ_SUBFOLDER = "vqvae"
 
 TOKENS = [
     {"id": "sky", "label": "Sky", "base": (91, 169, 230), "initial": "S"},
@@ -255,6 +260,27 @@ def randomize_seed():
     return int(rng.integers(0, MAX_SEED))
 
 
+def current_device():
+    if torch.cuda.is_available():
+        return torch.device("cuda")
+    if getattr(torch.backends, "mps", None) and torch.backends.mps.is_available():
+        return torch.device("mps")
+    return torch.device("cpu")
+
+
+@lru_cache(maxsize=1)
+def load_vq_tokenizer(model_id=DEFAULT_VQ_MODEL, subfolder=DEFAULT_VQ_SUBFOLDER):
+    from diffusers import VQModel
+
+    device = current_device()
+    model = VQModel.from_pretrained(model_id, subfolder=subfolder, use_safetensors=False)
+    model.to(device)
+    model.eval()
+    for param in model.parameters():
+        param.requires_grad_(False)
+    return model, device
+
+
 def prompt_bias(prompt):
     prompt = (prompt or "").lower()
     bias = np.zeros(TOKEN_COUNT, dtype=np.float32)
@@ -775,8 +801,140 @@ def learned_codebook_rows(centroids, assignments, errors):
     return rows
 
 
-@spaces.GPU(duration=30)
-def tokenize_image(image, grid_size, patch_size, codebook_size, iterations, seed, codebook_source):
+def prepare_vq_image(image, size=256):
+    return center_crop_square(image).resize((size, size), Image.Resampling.BICUBIC)
+
+
+def pil_to_vq_tensor(image, device):
+    arr = np.asarray(image).astype(np.float32) / 255.0
+    arr = arr * 2.0 - 1.0
+    tensor = torch.from_numpy(arr).permute(2, 0, 1).unsqueeze(0).to(device=device, dtype=torch.float32)
+    return tensor
+
+
+def vq_tensor_to_pil(tensor):
+    arr = tensor.detach().float().cpu().clamp(-1, 1).squeeze(0).permute(1, 2, 0).numpy()
+    arr = ((arr + 1.0) * 127.5).round().clip(0, 255).astype(np.uint8)
+    return Image.fromarray(arr, mode="RGB")
+
+
+def draw_learned_token_grid(indices, height, width, cell=None):
+    if cell is None:
+        cell = max(8, min(18, 720 // max(1, width)))
+    grid = indices.reshape(height, width)
+    unique = sorted(int(value) for value in np.unique(grid))
+    palette = {}
+    for order, token_id_value in enumerate(unique):
+        hue = order / max(1, len(unique))
+        r = int(80 + 135 * abs(math.sin(math.tau * hue)))
+        g = int(80 + 135 * abs(math.sin(math.tau * (hue + 0.33))))
+        b = int(80 + 135 * abs(math.sin(math.tau * (hue + 0.66))))
+        palette[token_id_value] = (r, g, b)
+
+    image = Image.new("RGB", (width * cell, height * cell), (31, 36, 44))
+    draw = ImageDraw.Draw(image)
+    font = ImageFont.load_default()
+    for row in range(height):
+        for col in range(width):
+            token_id_value = int(grid[row, col])
+            x0 = col * cell
+            y0 = row * cell
+            draw.rectangle((x0, y0, x0 + cell, y0 + cell), fill=palette[token_id_value], outline=(24, 28, 34))
+            if cell >= 18:
+                text = str(token_id_value % 100)
+                bbox = draw.textbbox((0, 0), text, font=font)
+                draw.text(
+                    (x0 + (cell - (bbox[2] - bbox[0])) / 2, y0 + (cell - (bbox[3] - bbox[1])) / 2),
+                    text,
+                    fill=(10, 14, 20),
+                    font=font,
+                )
+    return image
+
+
+def learned_vq_gallery(image, indices, latent_height, latent_width, distances, max_items=48):
+    grid = indices.reshape(latent_height, latent_width)
+    counts = np.bincount(indices)
+    used_ids = [int(index) for index in np.argsort(counts)[::-1] if counts[index] > 0][:max_items]
+    cell_w = image.width / latent_width
+    cell_h = image.height / latent_height
+    gallery = []
+
+    for token_id_value in used_ids:
+        positions = np.argwhere(grid == token_id_value)
+        flat_positions = positions[:, 0] * latent_width + positions[:, 1]
+        best_flat = int(flat_positions[np.argmin(distances[flat_positions])])
+        row = best_flat // latent_width
+        col = best_flat % latent_width
+        left = int(round(col * cell_w))
+        top = int(round(row * cell_h))
+        right = int(round((col + 1) * cell_w))
+        bottom = int(round((row + 1) * cell_h))
+        crop = image.crop((left, top, right, bottom)).resize((96, 96), Image.Resampling.NEAREST)
+        gallery.append((crop, f"Token {token_id_value}\nused {int(counts[token_id_value])} positions"))
+    return gallery
+
+
+def learned_vq_rows(indices, distances):
+    counts = np.bincount(indices)
+    total = max(1, int(counts.sum()))
+    rows = []
+    for token_id_value in np.argsort(counts)[::-1]:
+        count = int(counts[token_id_value])
+        if count == 0:
+            continue
+        members = distances[indices == token_id_value]
+        rows.append(
+            [
+                f"Token {int(token_id_value)}",
+                count,
+                round(count / total, 3),
+                "learned VQ embedding",
+                round(float(members.mean()) if len(members) else 0.0, 6),
+            ]
+        )
+    return rows
+
+
+def learned_vq_tokenize(image, model_id=DEFAULT_VQ_MODEL, subfolder=DEFAULT_VQ_SUBFOLDER):
+    model, device = load_vq_tokenizer(model_id, subfolder)
+    sample_size = int(getattr(model.config, "sample_size", 256) or 256)
+    prepared = prepare_vq_image(image, sample_size)
+    tensor = pil_to_vq_tensor(prepared, device)
+
+    with torch.inference_mode():
+        latents = model.encode(tensor).latents
+        quantized, _, info = model.quantize(latents)
+        indices = info[2].detach().cpu().numpy().astype(np.int64).reshape(-1)
+        reconstruction = model.decode(latents).sample
+
+        z = latents.permute(0, 2, 3, 1).contiguous().view(-1, model.quantize.vq_embed_dim)
+        embeddings = model.quantize.embedding.weight
+        chosen = embeddings[torch.from_numpy(indices).to(device)]
+        distances = torch.mean((z - chosen) ** 2, dim=1).detach().cpu().numpy()
+
+    latent_height, latent_width = int(latents.shape[2]), int(latents.shape[3])
+    learned_cell = max(8, min(18, 720 // max(1, latent_width)))
+    token_grid = draw_learned_token_grid(indices, latent_height, latent_width, cell=learned_cell)
+    reconstruction_image = vq_tensor_to_pil(reconstruction)
+    error_map = draw_error_heatmap(distances, latent_height, cell=learned_cell)
+    gallery = learned_vq_gallery(prepared, indices, latent_height, latent_width, distances)
+    rows = learned_vq_rows(indices, distances)
+    summary = (
+        f"Encoded the image with the pretrained learned tokenizer `{model_id}/{subfolder}`.\n"
+        f"The VQ model produced a {latent_height}x{latent_width} grid: {latent_height * latent_width} learned token IDs. "
+        f"It used {len(rows)} unique codebook entries from a vocabulary of {int(model.config.num_vq_embeddings)}.\n"
+        f"Mean latent-to-codebook distance: {float(distances.mean()):.6f}. "
+        "The gallery shows representative image regions for the most-used learned token IDs."
+    )
+    return prepared, token_grid, reconstruction_image, error_map, gallery, rows, summary
+
+
+@spaces.GPU(duration=120)
+def tokenize_image(image, grid_size, patch_size, codebook_size, iterations, seed, tokenizer_method, codebook_source):
+    if tokenizer_method == "Learned VQ tokenizer":
+        return learned_vq_tokenize(image)
+
     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)))
@@ -818,18 +976,18 @@ def tokenize_image(image, grid_size, patch_size, codebook_size, iterations, seed
     return prepared, token_grid, reconstruction, error_map, gallery, rows, summary
 
 
-def initial_tokenizer(grid_size, patch_size, codebook_size, iterations, seed, codebook_source):
+def initial_tokenizer(grid_size, patch_size, codebook_size, iterations, seed, tokenizer_method, codebook_source):
     items, message = load_moma_items()
-    outputs = tokenize_image(items[0]["img"], grid_size, patch_size, codebook_size, iterations, seed, codebook_source)
+    outputs = tokenize_image(items[0]["img"], grid_size, patch_size, codebook_size, iterations, seed, tokenizer_method, 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, codebook_source, evt: gr.SelectData):
+def tokenize_moma_selection(grid_size, patch_size, codebook_size, iterations, seed, tokenizer_method, 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, codebook_source)
+    return tokenize_image(items[index]["img"], grid_size, patch_size, codebook_size, iterations, seed, tokenizer_method, codebook_source)
 
 
 def probability_rows(record):
@@ -1073,6 +1231,11 @@ def build_app():
                     )
                     tokenizer_upload = gr.Image(label="Upload image", type="pil", sources=["upload", "clipboard"])
                     with gr.Accordion("Tokenizer settings", open=True):
+                        tokenizer_method = gr.Radio(
+                            ["K-means patch tokenizer", "Learned VQ tokenizer"],
+                            value="Learned VQ tokenizer",
+                            label="Tokenizer",
+                        )
                         with gr.Row():
                             tokenizer_grid_size = gr.Slider(6, 28, value=16, step=1, label="Token grid")
                             tokenizer_patch_size = gr.Slider(6, 24, value=14, step=1, label="Patch pixels")
@@ -1264,6 +1427,7 @@ def build_app():
             tokenizer_codebook_size,
             tokenizer_iterations,
             tokenizer_seed,
+            tokenizer_method,
             tokenizer_codebook_source,
         ]
         tokenizer_outputs = [
@@ -1291,6 +1455,7 @@ def build_app():
                 tokenizer_codebook_size,
                 tokenizer_iterations,
                 tokenizer_seed,
+                tokenizer_method,
                 tokenizer_codebook_source,
             ],
             outputs=tokenizer_outputs,
@@ -1305,6 +1470,7 @@ def build_app():
                 tokenizer_codebook_size,
                 tokenizer_iterations,
                 tokenizer_seed,
+                tokenizer_method,
                 tokenizer_codebook_source,
             ],
             outputs=tokenizer_outputs,

requirements.txt

@@ -1,5 +1,10 @@
 gradio>=5.22.0
 spaces
-numpy
+numpy<2
 pillow
 requests
+torch>=2.8.0
+diffusers>=0.35.0
+accelerate
+safetensors
+huggingface_hub>=0.34.0,<1.0