generate_app.py

#!/usr/bin/env python3
"""CategoricalVAE Latent Explorer - Interactive Gradio app for exploring the
latent space of a Categorical VAE model.

Usage:
    python generate_app.py [--share] [--port PORT]

The app loads a checkpoint folder containing:
  - config.yaml            (model hyperparameters)
  - checkpoints/last.ckpt  (PyTorch Lightning checkpoint)

Sliders are generated dynamically based on the model's latent_dim and
categorical_dim.  Each slider picks a category index (0-based) for one
latent dimension.  A temperature slider controls how sharply the one-hot
vector is defined:
  - temperature 0  -> pure one-hot (deterministic)
  - temperature >0 -> softmax-scaled logits (softer blending across categories)
"""

import argparse
import os
import sys

import gradio as gr
import torch
import yaml
from PIL import Image
import numpy as np

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from torchvae import vae_models

STORE = {
    "model": None,
    "latent_dim": None,
    "categorical_dim": None,
    "device": None,
}

LOGIT_VALUE = 10.0
UPSCALE_FACTOR = 4


def load_model(checkpoint_dir: str):
    checkpoint_dir = checkpoint_dir.strip()
    config_path = os.path.join(checkpoint_dir, "config.yaml")
    ckpt_path = os.path.join(checkpoint_dir, "checkpoints", "last.ckpt")

    if not os.path.isfile(config_path):
        raise gr.Error(f"Config not found: {config_path}")
    if not os.path.isfile(ckpt_path):
        raise gr.Error(f"Checkpoint not found: {ckpt_path}")

    with open(config_path, "r") as f:
        config = yaml.safe_load(f)

    model_params = config["model_params"]
    model_name = model_params["name"]

    if model_name not in vae_models:
        raise gr.Error(
            f"Unknown model '{model_name}'. "
            f"Available: {list(vae_models.keys())}"
        )

    model = vae_models[model_name](**model_params)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    try:
        checkpoint = torch.load(ckpt_path, map_location=device, weights_only=False)
    except TypeError:
        checkpoint = torch.load(ckpt_path, map_location=device)

    state_dict = checkpoint["state_dict"]

    cleaned = {}
    for k, v in state_dict.items():
        new_key = k[len("model."):] if k.startswith("model.") else k
        cleaned[new_key] = v

    missing, unexpected = model.load_state_dict(cleaned, strict=False)
    model.to(device)
    model.eval()

    latent_dim = model.latent_dim
    categorical_dim = model.categorical_dim

    STORE["model"] = model
    STORE["latent_dim"] = latent_dim
    STORE["categorical_dim"] = categorical_dim
    STORE["device"] = device

    config_str = f"{latent_dim},{categorical_dim}"
    status = (
        f"Loaded {model_name} from {checkpoint_dir}\n"
        f"latent_dim={latent_dim}, categorical_dim={categorical_dim}\n"
        f"Device: {device}"
    )
    if missing:
        status += f"\nMissing keys ({len(missing)}): {missing[:5]}..."
    if unexpected:
        status += f"\nUnexpected keys ({len(unexpected)}): {unexpected[:5]}..."

    return config_str, status


def generate(*slider_values):
    model = STORE["model"]
    if model is None:
        raise gr.Error("No model loaded. Please load a model first.")

    latent_dim = STORE["latent_dim"]
    categorical_dim = STORE["categorical_dim"]
    device = STORE["device"]

    temperature = slider_values[-1]
    category_indices = slider_values[:latent_dim]

    with torch.no_grad():
        if temperature < 0.01:
            z = torch.zeros(1, latent_dim * categorical_dim, device=device)
            for d, idx in enumerate(category_indices):
                z[0, d * categorical_dim + int(idx)] = 1.0
        else:
            logits = torch.zeros(1, latent_dim, categorical_dim, device=device)
            for d, idx in enumerate(category_indices):
                logits[0, d, int(idx)] = LOGIT_VALUE
            z = torch.softmax(logits / temperature, dim=-1)
            z = z.view(1, -1)

        output = model.decode(z)

    img = output.squeeze(0).cpu().numpy()
    img = ((img + 1.0) / 2.0 * 255.0).clip(0, 255).astype(np.uint8)
    img = img.transpose(1, 2, 0)

    if img.shape[2] == 1:
        img = img.squeeze(2)
        mode = "L"
    else:
        mode = "RGB"

    pil_img = Image.fromarray(img, mode)
    new_size = (pil_img.width * UPSCALE_FACTOR, pil_img.height * UPSCALE_FACTOR)
    pil_img = pil_img.resize(new_size, Image.NEAREST)

    return pil_img


with gr.Blocks(title="CategoricalVAE Latent Explorer") as demo:
    gr.Markdown("# CategoricalVAE Latent Explorer")
    gr.Markdown(
        "Load a checkpoint folder containing `config.yaml` and "
        "`checkpoints/last.ckpt`, then explore the latent space with sliders."
    )

    with gr.Row():
        path_input = gr.Textbox(
            label="Checkpoint Folder Path",
            value="logs/CategoricalVAE/version_5",
            scale=4,
        )
        load_btn = gr.Button("Load Model", variant="primary", scale=1)

    status = gr.Textbox(label="Status", interactive=False, lines=4)
    config_text = gr.Textbox(
        label="Model Config (latent_dim,categorical_dim)",
        interactive=False,
    )

    load_btn.click(
        load_model,
        inputs=[path_input],
        outputs=[config_text, status],
    )

    @gr.render(inputs=[config_text])
    def render_controls(config_str):
        if not config_str:
            return

        parts = config_str.split(",")
        latent_dim = int(parts[0])
        categorical_dim = int(parts[1])

        gr.Markdown(
            f"### Latent space: **{latent_dim}** dimensions × "
            f"**{categorical_dim}** categories"
        )

        sliders = []
        for d in range(latent_dim):
            s = gr.Slider(
                minimum=0,
                maximum=categorical_dim - 1,
                step=1,
                value=0,
                label=f"Dimension {d}",
            )
            sliders.append(s)

        temp_slider = gr.Slider(
            minimum=0.0,
            maximum=5.0,
            step=0.05,
            value=0.0,
            label="Temperature (0 = one-hot, higher = softer)",
        )

        generate_btn = gr.Button("Generate", variant="primary")
        output_img = gr.Image(label="Generated Image", height=300)

        all_inputs = sliders + [temp_slider]
        generate_btn.click(generate, inputs=all_inputs, outputs=[output_img])

        for s in sliders:
            s.change(generate, inputs=all_inputs, outputs=[output_img])
        temp_slider.change(generate, inputs=all_inputs, outputs=[output_img])


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="CategoricalVAE Latent Explorer"
    )
    parser.add_argument(
        "--share", action="store_true", help="Create a public share link"
    )
    parser.add_argument(
        "--port", type=int, default=7860, help="Port to run on"
    )
    args = parser.parse_args()

    demo.launch(share=args.share, server_port=args.port)