File size: 3,090 Bytes
9dd581e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 | from __future__ import annotations
from pathlib import Path
import gradio as gr
import numpy as np
import torch
from model import ConditionalCodePrior, VectorQuantizedAutoencoder
from PIL import Image
from safetensors.torch import load_file
ARTIFACT_DIR = Path(__file__).resolve().parent / "artifacts" / "vq-pocket"
AUTOENCODER = VectorQuantizedAutoencoder()
AUTOENCODER.load_state_dict(load_file(ARTIFACT_DIR / "vq_vae.safetensors"))
AUTOENCODER.eval()
PRIOR = ConditionalCodePrior(codebook_size=AUTOENCODER.codebook_size)
PRIOR.load_state_dict(load_file(ARTIFACT_DIR / "code_prior.safetensors"))
PRIOR.eval()
@torch.inference_mode()
def generate_tokens(
label: int,
seed: int,
temperature: float,
) -> tuple[Image.Image, Image.Image, dict]:
labels = torch.full((12,), int(label), dtype=torch.long)
codes = PRIOR.generate(labels, seed=int(seed), temperature=float(temperature))
images = AUTOENCODER.decode_indices(codes)
gallery = Image.new("L", (512, 384), color=0)
for index, pixels in enumerate(images):
image = Image.fromarray(
pixels[0].mul(255).clamp(0, 255).to(torch.uint8).numpy(),
mode="L",
).resize((120, 120), Image.Resampling.NEAREST)
gallery.paste(image, ((index % 4) * 128 + 4, (index // 4) * 128 + 4))
palette = np.array(
[
[int((code * 73) % 255), int((code * 151) % 255), int((code * 211) % 255)]
for code in range(AUTOENCODER.codebook_size)
],
dtype=np.uint8,
)
token_map = palette[codes[0].numpy()]
token_image = Image.fromarray(token_map, mode="RGB").resize(
(384, 384),
Image.Resampling.NEAREST,
)
unique_sequences = len({row.numpy().tobytes() for row in codes.flatten(1)})
metadata = {
"digit": int(label),
"tokens_per_image": 16,
"codebook_size": AUTOENCODER.codebook_size,
"unique_sequences_in_batch": unique_sequences,
"temperature": float(temperature),
}
return gallery, token_image, metadata
with gr.Blocks(title="VQ-Pocket") as demo:
gr.Markdown(
"# VQ-Pocket\n"
"Generate digits as discrete 4x4 visual-token sequences, then decode the "
"tokens through a compact VQ-VAE."
)
with gr.Row():
label = gr.Slider(0, 9, value=4, step=1, label="Digit class")
seed = gr.Slider(0, 100_000, value=2053, step=1, label="Sampling seed")
temperature = gr.Slider(
0.1, 1.5, value=0.85, step=0.05, label="Token temperature"
)
initial = generate_tokens(4, 2053, 0.85)
with gr.Row():
gallery = gr.Image(value=initial[0], label="Decoded generations")
tokens = gr.Image(value=initial[1], label="First sample's 4x4 token map")
metrics = gr.JSON(value=initial[2], label="Discrete-latent readout")
button = gr.Button("Sample token sequences", variant="primary")
button.click(
generate_tokens,
inputs=[label, seed, temperature],
outputs=[gallery, tokens, metrics],
)
if __name__ == "__main__":
demo.launch()
|