app.py

from __future__ import annotations

import tempfile

import gradio as gr
import numpy as np
import torch
import trimesh
from diffusers import AutoencoderKL
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
from transformers import CLIPTextModel, CLIPTokenizer

from pixel_dit import DiT
from voxel_dit import VoxelDiT

DEV = "cpu"
SCALE = 0.18215
CLIP_ID = "openai/clip-vit-base-patch32"
MAX_TOKENS = 40

print("[boot] loading shared CLIP text encoder...")
tokenizer = CLIPTokenizer.from_pretrained(CLIP_ID)
text_encoder = CLIPTextModel.from_pretrained(CLIP_ID).to(DEV).eval()


@torch.no_grad()
def encode(strings: list[str]):
    t = tokenizer(strings, padding="max_length", max_length=MAX_TOKENS, truncation=True, return_tensors="pt").to(DEV)
    o = text_encoder(**t)
    return o.last_hidden_state.float(), o.pooler_output.float()


null_seq, null_pool = encode([""])

print("[boot] loading PixelModel v5...")
pm5_weights = hf_hub_download("bench-labs/PixelModel-v5", "model.safetensors")
pm5_state = load_file(pm5_weights)
pixel_model = DiT(dim=384, depth=12, heads=6).to(DEV).eval()
pixel_model.load_state_dict({k[len("dit."):]: v for k, v in pm5_state.items() if k.startswith("dit.")})
vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse").to(DEV).eval()

print("[boot] loading VoxelModel v1...")
vm1_weights = hf_hub_download("bench-labs/VoxelModel-v1", "model.safetensors")
voxel_model = VoxelDiT().to(DEV).eval()
voxel_model.load_state_dict(load_file(vm1_weights))

print("[boot] ready.")


@torch.no_grad()
def sample_image(prompt: str, steps: int, cfg: float, seed: int, progress=gr.Progress()):
    if not prompt.strip():
        raise gr.Error("Type a prompt first.")
    steps = int(steps)
    g = torch.Generator(device=DEV).manual_seed(int(seed))
    seq, pool = encode([prompt])
    x = torch.randn(1, 4, 32, 32, device=DEV, generator=g)
    dt = 1.0 / steps
    for i in progress.tqdm(range(steps), desc="sampling"):
        t = torch.full((1,), i * dt, device=DEV)
        vc = pixel_model(x, t, seq, pool)
        vu = pixel_model(x, t, null_seq, null_pool)
        x = x + (vu + cfg * (vc - vu)) * dt
    img = vae.decode((x / SCALE)).sample
    img = ((img.clamp(-1, 1) + 1) / 2).permute(0, 2, 3, 1).numpy()[0]
    return (img * 255).round().astype(np.uint8)


@torch.no_grad()
def sample_voxel(prompt: str, steps: int, cfg: float, threshold: float, seed: int, progress=gr.Progress()):
    if not prompt.strip():
        raise gr.Error("Type a prompt first.")
    steps = int(steps)
    g = torch.Generator(device=DEV).manual_seed(int(seed))
    seq, pool = encode([prompt])
    x = torch.randn(1, 1, 32, 32, 32, device=DEV, generator=g)
    dt = 1.0 / steps
    for i in progress.tqdm(range(steps), desc="sampling"):
        t = torch.full((1,), i * dt, device=DEV)
        vc = voxel_model(x, t, seq, pool)
        vu = voxel_model(x, t, null_seq, null_pool)
        x = x + (vu + cfg * (vc - vu)) * dt
    grid = (x[0, 0] > threshold).numpy()
    if not grid.any():
        raise gr.Error("Nothing came back above the occupancy threshold — try lowering it or re-rolling the seed.")
    return grid_to_glb(grid)


def grid_to_glb(grid: np.ndarray) -> str:
    voxel = trimesh.voxel.VoxelGrid(encoding=grid)
    mesh = voxel.as_boxes()
    mesh.visual.face_colors = [180, 180, 190, 255]
    path = tempfile.NamedTemporaryFile(suffix=".glb", delete=False).name
    mesh.export(path)
    return path


with gr.Blocks(title="BenchLabs Models") as demo:
    gr.Markdown(
        "# BenchLabs Models\n"
        "Two tiny diffusion models, running live on CPU, no GPU behind this Space. "
        "Both are under 45M trained parameters, so generation is slower than a hosted API "
        "but the whole model fits in a PNG image if you're curious — see the model pages linked below."
    )
    with gr.Tab("Text → Image (PixelModel v5)"):
        gr.Markdown(
            "Good at material and light: food, landscapes, skies, interiors. "
            "Weak on faces, hands, and anything needing precise structure or text."
        )
        with gr.Row():
            with gr.Column():
                img_prompt = gr.Textbox(label="Prompt", placeholder="a bowl of ramen with a soft boiled egg")
                img_steps = gr.Slider(10, 50, value=25, step=1, label="Detail (sampling steps)")
                img_cfg = gr.Slider(1.0, 10.0, value=5.0, step=0.5, label="Prompt strength (CFG)")
                img_seed = gr.Number(value=0, precision=0, label="Seed")
                img_btn = gr.Button("Generate image", variant="primary")
            with gr.Column():
                img_out = gr.Image(label="Result", type="numpy")
        img_btn.click(sample_image, [img_prompt, img_steps, img_cfg, img_seed], img_out)
        gr.Examples(
            [["a bowl of ramen with a soft boiled egg", 25, 5.0, 0],
             ["a wet cobblestone street at night", 25, 5.0, 0],
             ["a library of wooden shelves", 25, 5.0, 0]],
            [img_prompt, img_steps, img_cfg, img_seed],
        )

    with gr.Tab("Text → 3D (VoxelModel v1)"):
        gr.Markdown(
            "Good at bulky objects: chairs, tables, cars, mushrooms. "
            "Thin objects (swords, keys) don't survive 32³ voxelization, in the training "
            "data or the model, so expect a blob rather than a blade."
        )
        with gr.Row():
            with gr.Column():
                vox_prompt = gr.Textbox(label="Prompt", placeholder="a wooden chair")
                vox_steps = gr.Slider(10, 50, value=25, step=1, label="Detail (sampling steps)")
                vox_cfg = gr.Slider(1.0, 10.0, value=5.0, step=0.5, label="Prompt strength (CFG)")
                vox_thresh = gr.Slider(-1.0, 1.0, value=0.0, step=0.05, label="Occupancy threshold")
                vox_seed = gr.Number(value=0, precision=0, label="Seed")
                vox_btn = gr.Button("Generate 3D model", variant="primary")
            with gr.Column():
                vox_out = gr.Model3D(label="Result")
        vox_btn.click(sample_voxel, [vox_prompt, vox_steps, vox_cfg, vox_thresh, vox_seed], vox_out)
        gr.Examples(
            [["a wooden chair", 25, 5.0, 0.0, 0],
             ["a purple mushroom", 25, 5.0, 0.0, 0],
             ["a small boat", 25, 5.0, 0.0, 0]],
            [vox_prompt, vox_steps, vox_cfg, vox_thresh, vox_seed],
        )

    gr.Markdown(
        "Models: [PixelModel v5](https://huggingface.co/bench-labs/PixelModel-v5) · "
        "[VoxelModel v1](https://huggingface.co/bench-labs/VoxelModel-v1)"
    )

if __name__ == "__main__":
    demo.queue(max_size=20).launch(server_name="0.0.0.0")