Upload 7 files

.gitattributes

@@ -1,43 +1,43 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
-examples/demo19.jpg filter=lfs diff=lfs merge=lfs -text
-assets/DA-2K.png filter=lfs diff=lfs merge=lfs -text
-assets/examples_video/basketball.mp4 filter=lfs diff=lfs merge=lfs -text
-assets/examples_video/ferris_wheel.mp4 filter=lfs diff=lfs merge=lfs -text
-assets/examples/demo19.jpg filter=lfs diff=lfs merge=lfs -text
-assets/teaser.png filter=lfs diff=lfs merge=lfs -text
-metric_depth/assets/compare_zoedepth.png filter=lfs diff=lfs merge=lfs -text
-metric_depth/dataset/splits/hypersim/train.txt filter=lfs diff=lfs merge=lfs -text
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+examples/demo19.jpg filter=lfs diff=lfs merge=lfs -text
+assets/DA-2K.png filter=lfs diff=lfs merge=lfs -text
+assets/examples_video/basketball.mp4 filter=lfs diff=lfs merge=lfs -text
+assets/examples_video/ferris_wheel.mp4 filter=lfs diff=lfs merge=lfs -text
+assets/examples/demo19.jpg filter=lfs diff=lfs merge=lfs -text
+assets/teaser.png filter=lfs diff=lfs merge=lfs -text
+metric_depth/assets/compare_zoedepth.png filter=lfs diff=lfs merge=lfs -text
+metric_depth/dataset/splits/hypersim/train.txt filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,14 +1,14 @@
----
-title: BenchLabs Models
-emoji: 🏛️
-colorFrom: red
-colorTo: blue
-sdk: gradio
-sdk_version: 5.50.0
-app_file: app.py
-pinned: false
-license: apache-2.0
----
-
-# BenchLabs Inference Space
-This space runs **PixelModel-v5** (Text-to-Image) and **VoxelModel-v1** (Text-to-3D) dynamically on CPU.
+---
+title: BenchLabs Models
+emoji: 🏛️
+colorFrom: red
+colorTo: blue
+sdk: gradio
+sdk_version: 5.50.0
+app_file: app.py
+pinned: false
+license: apache-2.0
+---
+
+# BenchLabs Inference Space
+This space runs **PixelModel-v5** (Text-to-Image) and **VoxelModel-v1** (Text-to-3D) dynamically on CPU.

app.py

@@ -1,158 +1,219 @@
-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")
+from __future__ import annotations
+
+import tempfile
+import os
+
+os.environ.setdefault("HF_HOME", r"D:\hf-cache")
+os.environ.setdefault("HUGGINGFACE_HUB_CACHE", r"D:\hf-cache\hub")
+
+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, T5EncoderModel, T5TokenizerFast
+
+from pixel_dit import DiT
+from voxel_dit import VoxelDiT
+from mmdit import MMDiT
+
+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] loading PixelModel v6...")
+v6_weights = hf_hub_download("bench-labs/PixelModel-v6", "model.safetensors")
+pixel_model_v6 = MMDiT().to(DEV).eval()
+pixel_model_v6.load_state_dict(load_file(v6_weights), strict=False)
+t5_tokenizer = T5TokenizerFast.from_pretrained("google/flan-t5-base")
+t5_encoder = T5EncoderModel.from_pretrained("google/flan-t5-base").to(DEV).eval()
+vae_v6 = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix").to(DEV).float().eval()
+null_v6_seq, null_v6_mask, null_v6_pool = None, None, None
+
+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)
+
+
+@torch.no_grad()
+def sample_image_v6(prompt: str, steps: int, cfg: float, seed: int, progress=gr.Progress()):
+    global null_v6_seq, null_v6_mask, null_v6_pool
+    if not prompt.strip():
+        raise gr.Error("Type a prompt first.")
+    steps = int(steps)
+
+    def encode_v6(strings):
+        t = t5_tokenizer(strings, padding="max_length", max_length=32, truncation=True, return_tensors="pt").to(DEV)
+        seq = t5_encoder(**t).last_hidden_state.float()
+        _, pool = encode(strings)
+        return seq, t["attention_mask"].float(), pool
+
+    seq, mask, pool = encode_v6([prompt])
+    if null_v6_seq is None:
+        null_v6_seq, null_v6_mask, null_v6_pool = encode_v6([""])
+    g = torch.Generator(device=DEV).manual_seed(int(seed))
+    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_v6(x, t, seq, mask, pool)
+        vu = pixel_model_v6(x, t, null_v6_seq, null_v6_mask, null_v6_pool)
+        x = x + (vu + cfg * (vc - vu)) * dt
+    img = vae_v6.decode(x / vae_v6.config.scaling_factor).sample
+    return (((img.clamp(-1, 1) + 1) / 2).permute(0, 2, 3, 1).numpy()[0] * 255).round().astype(np.uint8)
+
+
+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 v6)"):
+        gr.Markdown("A larger MMDiT model conditioned by T5 and CLIP. CPU generation is slower; 256x256 output.")
+        with gr.Row():
+            with gr.Column():
+                v6_prompt = gr.Textbox(label="Prompt", placeholder="a red fox sitting in a snowy forest")
+                v6_steps = gr.Slider(10, 100, value=50, step=5, label="Detail (sampling steps)")
+                v6_cfg = gr.Slider(1.0, 10.0, value=3.0, step=0.5, label="Prompt strength (CFG)")
+                v6_seed = gr.Number(value=0, precision=0, label="Seed")
+                v6_btn = gr.Button("Generate image", variant="primary")
+            with gr.Column():
+                v6_out = gr.Image(label="Result", type="numpy")
+        v6_btn.click(sample_image_v6, [v6_prompt, v6_steps, v6_cfg, v6_seed], v6_out)
+        gr.Examples(
+            [["a red fox sitting in a snowy forest", 50, 3.0, 0],
+             ["a lighthouse on a cliff at sunset", 50, 3.0, 0],
+             ["a city street at night with neon signs", 50, 3.0, 0]],
+            [v6_prompt, v6_steps, v6_cfg, v6_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],
+        )
+
+    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],
+        )
+
+    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")

mmdit.py

@@ -0,0 +1,141 @@
+from __future__ import annotations
+
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def modulate(x, shift, scale):
+    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+
+
+def timestep_embedding(t, dim, max_period=10000):
+    half = dim // 2
+    freqs = torch.exp(-math.log(max_period) * torch.arange(half, device=t.device) / half)
+    args = t[:, None].float() * freqs[None]
+    emb = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+    return torch.cat([emb, torch.zeros_like(emb[:, :1])], dim=-1) if dim % 2 else emb
+
+
+def rope_freqs(positions, dim, base=10000.0):
+    inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
+    return torch.outer(positions.float(), inv_freq)
+
+
+def rope_cos_sin(freqs):
+    emb = torch.cat([freqs, freqs], dim=-1)
+    return emb.cos(), emb.sin()
+
+
+def rotate_half(x):
+    x1, x2 = x.chunk(2, dim=-1)
+    return torch.cat([-x2, x1], dim=-1)
+
+
+def apply_rope(x, cos, sin):
+    return x * cos + rotate_half(x) * sin
+
+
+def apply_rope_2d(x, row_cos, row_sin, col_cos, col_sin):
+    x1, x2 = x.chunk(2, dim=-1)
+    return torch.cat([apply_rope(x1, row_cos, row_sin), apply_rope(x2, col_cos, col_sin)], dim=-1)
+
+
+class RMSNormHead(nn.Module):
+    def __init__(self, head_dim, eps=1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(head_dim))
+        self.eps = eps
+
+    def forward(self, x):
+        return x * x.pow(2).mean(-1, keepdim=True).add(self.eps).rsqrt() * self.weight
+
+
+class SwiGLU(nn.Module):
+    def __init__(self, dim, hidden):
+        super().__init__()
+        self.gate = nn.Linear(dim, hidden)
+        self.up = nn.Linear(dim, hidden)
+        self.down = nn.Linear(hidden, dim)
+
+    def forward(self, x):
+        return self.down(F.silu(self.gate(x)) * self.up(x))
+
+
+class JointBlock(nn.Module):
+    def __init__(self, dim, heads, mlp_hidden):
+        super().__init__()
+        self.heads, self.head_dim = heads, dim // heads
+        self.norm1_img = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+        self.norm1_txt = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+        self.qkv_img = nn.Linear(dim, 3 * dim)
+        self.qkv_txt = nn.Linear(dim, 3 * dim)
+        self.qn_img, self.kn_img = RMSNormHead(self.head_dim), RMSNormHead(self.head_dim)
+        self.qn_txt, self.kn_txt = RMSNormHead(self.head_dim), RMSNormHead(self.head_dim)
+        self.proj_img, self.proj_txt = nn.Linear(dim, dim), nn.Linear(dim, dim)
+        self.norm2_img = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+        self.norm2_txt = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+        self.mlp_img, self.mlp_txt = SwiGLU(dim, mlp_hidden), SwiGLU(dim, mlp_hidden)
+        self.ada_img = nn.Sequential(nn.SiLU(), nn.Linear(dim, 6 * dim))
+        self.ada_txt = nn.Sequential(nn.SiLU(), nn.Linear(dim, 6 * dim))
+
+    def forward(self, img, txt, c, rope_img, rope_txt, key_valid):
+        s1i, sc1i, g1i, s2i, sc2i, g2i = self.ada_img(c).chunk(6, dim=-1)
+        s1t, sc1t, g1t, s2t, sc2t, g2t = self.ada_txt(c).chunk(6, dim=-1)
+        xi, xt = modulate(self.norm1_img(img), s1i, sc1i), modulate(self.norm1_txt(txt), s1t, sc1t)
+        b, ni, cdim = xi.shape; nt = xt.shape[1]; h, d = self.heads, self.head_dim
+        qi, ki, vi = self.qkv_img(xi).reshape(b, ni, 3, h, d).permute(2, 0, 3, 1, 4)
+        qt, kt, vt = self.qkv_txt(xt).reshape(b, nt, 3, h, d).permute(2, 0, 3, 1, 4)
+        qi, ki, qt, kt = self.qn_img(qi), self.kn_img(ki), self.qn_txt(qt), self.kn_txt(kt)
+        rc, rs, cc, cs = rope_img; tc, ts = rope_txt
+        qi, ki = apply_rope_2d(qi, rc, rs, cc, cs), apply_rope_2d(ki, rc, rs, cc, cs)
+        qt, kt = apply_rope(qt, tc, ts), apply_rope(kt, tc, ts)
+        q, k, v = torch.cat([qi, qt], 2), torch.cat([ki, kt], 2), torch.cat([vi, vt], 2)
+        o = F.scaled_dot_product_attention(q, k, v, attn_mask=key_valid[:, None, None, :])
+        o = o.transpose(1, 2).reshape(b, ni + nt, cdim)
+        oi, ot = o[:, :ni], o[:, ni:]
+        img = img + g1i[:, None] * self.proj_img(oi)
+        txt = txt + g1t[:, None] * self.proj_txt(ot)
+        img = img + g2i[:, None] * self.mlp_img(modulate(self.norm2_img(img), s2i, sc2i))
+        txt = txt + g2t[:, None] * self.mlp_txt(modulate(self.norm2_txt(txt), s2t, sc2t))
+        return img, txt
+
+
+class MMDiT(nn.Module):
+    def __init__(self, latent_ch=4, latent_size=32, patch=2, dim=512, depth=16, heads=8,
+                 t5_dim=768, clip_dim=512, t5_len=32, mlp_hidden=1408, repa_dim=384, repa_layer=8):
+        super().__init__()
+        self.latent_ch, self.latent_size, self.patch = latent_ch, latent_size, patch
+        self.grid, self.patch_dim, self.dim = latent_size // patch, latent_ch * patch * patch, dim
+        self.t5_len, self.repa_layer = t5_len, repa_layer
+        self.x_embed = nn.Linear(self.patch_dim, dim)
+        self.t_mlp = nn.Sequential(nn.Linear(dim, dim), nn.SiLU(), nn.Linear(dim, dim))
+        self.clip_proj, self.t5_proj = nn.Linear(clip_dim, dim), nn.Linear(t5_dim, dim)
+        self.blocks = nn.ModuleList([JointBlock(dim, heads, mlp_hidden) for _ in range(depth)])
+        self.norm_out = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+        self.ada_out, self.head = nn.Sequential(nn.SiLU(), nn.Linear(dim, 2 * dim)), nn.Linear(dim, self.patch_dim)
+        self.repa_head = nn.Sequential(nn.Linear(dim, dim), nn.GELU(approximate="tanh"), nn.Linear(dim, repa_dim))
+        hd2 = (dim // heads) // 2
+        rows = torch.arange(self.grid).repeat_interleave(self.grid); cols = torch.arange(self.grid).repeat(self.grid)
+        for name, value in zip(("row_cos", "row_sin", "col_cos", "col_sin"), (*rope_cos_sin(rope_freqs(rows, hd2)), *rope_cos_sin(rope_freqs(cols, hd2)))):
+            self.register_buffer(name, value, persistent=False)
+        self.register_buffer("t_cos", rope_cos_sin(rope_freqs(torch.arange(t5_len), dim // heads))[0], persistent=False)
+        self.register_buffer("t_sin", rope_cos_sin(rope_freqs(torch.arange(t5_len), dim // heads))[1], persistent=False)
+
+    def patchify(self, x):
+        b, c, h, w = x.shape; p = self.patch
+        return x.reshape(b, c, h // p, p, w // p, p).permute(0, 2, 4, 1, 3, 5).reshape(b, (h // p) * (w // p), c * p * p)
+
+    def unpatchify(self, x):
+        b, _, _ = x.shape; p, g, c = self.patch, self.grid, self.latent_ch
+        return x.reshape(b, g, g, c, p, p).permute(0, 3, 1, 4, 2, 5).reshape(b, c, g * p, g * p)
+
+    def forward(self, x, t, t5_seq, t5_mask, clip_pool):
+        b = x.shape[0]; img = self.x_embed(self.patchify(x)); txt = self.t5_proj(t5_seq)
+        c = self.t_mlp(timestep_embedding(t, self.dim)) + self.clip_proj(clip_pool)
+        valid = torch.cat([torch.ones(b, img.shape[1], dtype=torch.bool, device=x.device), t5_mask.bool()], dim=1)
+        ri = (self.row_cos, self.row_sin, self.col_cos, self.col_sin); rt = (self.t_cos, self.t_sin)
+        for block in self.blocks: img, txt = block(img, txt, c, ri, rt, valid)
+        shift, scale = self.ada_out(c).chunk(2, dim=-1)
+        return self.unpatchify(self.head(modulate(self.norm_out(img), shift, scale)))

pixel_dit.py

@@ -1,133 +1,133 @@
-from __future__ import annotations
-
-import math
-import numpy as np
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-def modulate(x, shift, scale):
-    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
-
-def timestep_embedding(t, dim, max_period=10000):
-    half = dim // 2
-    freqs = torch.exp(-math.log(max_period) * torch.arange(half, device=t.device) / half)
-    args = t[:, None].float() * freqs[None]
-    emb = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
-    if dim % 2:
-        emb = torch.cat([emb, torch.zeros_like(emb[:, :1])], dim=-1)
-    return emb
-
-def sincos_2d(dim, grid_size):
-    g = np.arange(grid_size, dtype=np.float32)
-    gx, gy = np.meshgrid(g, g, indexing="xy")
-    assert dim % 4 == 0
-    d4 = dim // 4
-    omega = 1.0 / (10000 ** (np.arange(d4, dtype=np.float32) / d4))
-    def emb1(p):
-        out = p.reshape(-1)[:, None] * omega[None]
-        return np.concatenate([np.sin(out), np.cos(out)], axis=1)
-    pe = np.concatenate([emb1(gx), emb1(gy)], axis=1)
-    return torch.from_numpy(pe).float()
-
-class Attention(nn.Module):
-    def __init__(self, dim, heads):
-        super().__init__()
-        self.heads = heads
-        self.q = nn.Linear(dim, dim)
-        self.kv = nn.Linear(dim, dim * 2)
-        self.proj = nn.Linear(dim, dim)
-
-    def forward(self, x, ctx=None):
-        ctx = x if ctx is None else ctx
-        B, N, C = x.shape
-        M = ctx.shape[1]
-        h = self.heads
-        q = self.q(x).reshape(B, N, h, C // h).transpose(1, 2)
-        kv = self.kv(ctx).reshape(B, M, 2, h, C // h).permute(2, 0, 3, 1, 4)
-        k, v = kv[0], kv[1]
-        o = F.scaled_dot_product_attention(q, k, v)
-        o = o.transpose(1, 2).reshape(B, N, C)
-        return self.proj(o)
-
-class Block(nn.Module):
-    def __init__(self, dim, heads, mlp_ratio=4.0):
-        super().__init__()
-        self.norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
-        self.attn = Attention(dim, heads)
-        self.norm_ca = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
-        self.cross = Attention(dim, heads)
-        self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
-        hidden = int(dim * mlp_ratio)
-        self.mlp = nn.Sequential(nn.Linear(dim, hidden), nn.GELU(approximate="tanh"),
-                                 nn.Linear(hidden, dim))
-        self.ada = nn.Sequential(nn.SiLU(), nn.Linear(dim, 6 * dim))
-        self.cross_gate = nn.Parameter(torch.zeros(1))
-
-    def forward(self, x, c, text):
-        shift1, scale1, gate1, shift2, scale2, gate2 = self.ada(c).chunk(6, dim=1)
-        x = x + gate1.unsqueeze(1) * self.attn(modulate(self.norm1(x), shift1, scale1))
-        x = x + self.cross_gate * self.cross(self.norm_ca(x), text)
-        x = x + gate2.unsqueeze(1) * self.mlp(modulate(self.norm2(x), shift2, scale2))
-        return x
-
-class DiT(nn.Module):
-    def __init__(self, latent_ch=4, latent_size=32, patch=2, dim=384, depth=12,
-                 heads=6, text_dim=512, mlp_ratio=4.0):
-        super().__init__()
-        self.latent_ch = latent_ch
-        self.latent_size = latent_size
-        self.patch = patch
-        self.grid = latent_size // patch
-        self.patch_dim = latent_ch * patch * patch
-        self.x_embed = nn.Linear(self.patch_dim, dim)
-        self.register_buffer("pos", sincos_2d(dim, self.grid).unsqueeze(0))
-        self.t_mlp = nn.Sequential(nn.Linear(dim, dim), nn.SiLU(), nn.Linear(dim, dim))
-        self.text_proj = nn.Linear(text_dim, dim)
-        self.text_pool = nn.Linear(text_dim, dim)
-        self.blocks = nn.ModuleList([Block(dim, heads, mlp_ratio) for _ in range(depth)])
-        self.norm_out = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
-        self.ada_out = nn.Sequential(nn.SiLU(), nn.Linear(dim, 2 * dim))
-        self.head = nn.Linear(dim, self.patch_dim)
-        self.dim = dim
-        self._init()
-
-    def _init(self):
-        for m in self.modules():
-            if isinstance(m, nn.Linear):
-                nn.init.xavier_uniform_(m.weight)
-                if m.bias is not None:
-                    nn.init.zeros_(m.bias)
-        for b in self.blocks:
-            nn.init.zeros_(b.ada[-1].weight); nn.init.zeros_(b.ada[-1].bias)
-        nn.init.zeros_(self.ada_out[-1].weight); nn.init.zeros_(self.ada_out[-1].bias)
-        nn.init.zeros_(self.head.weight); nn.init.zeros_(self.head.bias)
-
-    def patchify(self, x):
-        B, C, H, W = x.shape
-        p = self.patch
-        x = x.reshape(B, C, H // p, p, W // p, p)
-        x = x.permute(0, 2, 4, 1, 3, 5).reshape(B, (H // p) * (W // p), C * p * p)
-        return x
-
-    def unpatchify(self, x):
-        B, N, _ = x.shape
-        p = self.patch
-        g = self.grid
-        C = self.latent_ch
-        x = x.reshape(B, g, g, C, p, p).permute(0, 3, 1, 4, 2, 5)
-        return x.reshape(B, C, g * p, g * p)
-
-    def forward(self, x, t, text_seq, text_pool):
-        x = self.x_embed(self.patchify(x)) + self.pos
-        c = self.t_mlp(timestep_embedding(t, self.dim)) + self.text_pool(text_pool)
-        text = self.text_proj(text_seq)
-        for blk in self.blocks:
-            x = blk(x, c, text)
-        shift, scale = self.ada_out(c).chunk(2, dim=1)
-        x = modulate(self.norm_out(x), shift, scale)
-        x = self.head(x)
-        return self.unpatchify(x)
-
-    def num_params(self):
-        return sum(p.numel() for p in self.parameters())
+from __future__ import annotations
+
+import math
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+def modulate(x, shift, scale):
+    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+
+def timestep_embedding(t, dim, max_period=10000):
+    half = dim // 2
+    freqs = torch.exp(-math.log(max_period) * torch.arange(half, device=t.device) / half)
+    args = t[:, None].float() * freqs[None]
+    emb = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+    if dim % 2:
+        emb = torch.cat([emb, torch.zeros_like(emb[:, :1])], dim=-1)
+    return emb
+
+def sincos_2d(dim, grid_size):
+    g = np.arange(grid_size, dtype=np.float32)
+    gx, gy = np.meshgrid(g, g, indexing="xy")
+    assert dim % 4 == 0
+    d4 = dim // 4
+    omega = 1.0 / (10000 ** (np.arange(d4, dtype=np.float32) / d4))
+    def emb1(p):
+        out = p.reshape(-1)[:, None] * omega[None]
+        return np.concatenate([np.sin(out), np.cos(out)], axis=1)
+    pe = np.concatenate([emb1(gx), emb1(gy)], axis=1)
+    return torch.from_numpy(pe).float()
+
+class Attention(nn.Module):
+    def __init__(self, dim, heads):
+        super().__init__()
+        self.heads = heads
+        self.q = nn.Linear(dim, dim)
+        self.kv = nn.Linear(dim, dim * 2)
+        self.proj = nn.Linear(dim, dim)
+
+    def forward(self, x, ctx=None):
+        ctx = x if ctx is None else ctx
+        B, N, C = x.shape
+        M = ctx.shape[1]
+        h = self.heads
+        q = self.q(x).reshape(B, N, h, C // h).transpose(1, 2)
+        kv = self.kv(ctx).reshape(B, M, 2, h, C // h).permute(2, 0, 3, 1, 4)
+        k, v = kv[0], kv[1]
+        o = F.scaled_dot_product_attention(q, k, v)
+        o = o.transpose(1, 2).reshape(B, N, C)
+        return self.proj(o)
+
+class Block(nn.Module):
+    def __init__(self, dim, heads, mlp_ratio=4.0):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+        self.attn = Attention(dim, heads)
+        self.norm_ca = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+        self.cross = Attention(dim, heads)
+        self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+        hidden = int(dim * mlp_ratio)
+        self.mlp = nn.Sequential(nn.Linear(dim, hidden), nn.GELU(approximate="tanh"),
+                                 nn.Linear(hidden, dim))
+        self.ada = nn.Sequential(nn.SiLU(), nn.Linear(dim, 6 * dim))
+        self.cross_gate = nn.Parameter(torch.zeros(1))
+
+    def forward(self, x, c, text):
+        shift1, scale1, gate1, shift2, scale2, gate2 = self.ada(c).chunk(6, dim=1)
+        x = x + gate1.unsqueeze(1) * self.attn(modulate(self.norm1(x), shift1, scale1))
+        x = x + self.cross_gate * self.cross(self.norm_ca(x), text)
+        x = x + gate2.unsqueeze(1) * self.mlp(modulate(self.norm2(x), shift2, scale2))
+        return x
+
+class DiT(nn.Module):
+    def __init__(self, latent_ch=4, latent_size=32, patch=2, dim=384, depth=12,
+                 heads=6, text_dim=512, mlp_ratio=4.0):
+        super().__init__()
+        self.latent_ch = latent_ch
+        self.latent_size = latent_size
+        self.patch = patch
+        self.grid = latent_size // patch
+        self.patch_dim = latent_ch * patch * patch
+        self.x_embed = nn.Linear(self.patch_dim, dim)
+        self.register_buffer("pos", sincos_2d(dim, self.grid).unsqueeze(0))
+        self.t_mlp = nn.Sequential(nn.Linear(dim, dim), nn.SiLU(), nn.Linear(dim, dim))
+        self.text_proj = nn.Linear(text_dim, dim)
+        self.text_pool = nn.Linear(text_dim, dim)
+        self.blocks = nn.ModuleList([Block(dim, heads, mlp_ratio) for _ in range(depth)])
+        self.norm_out = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+        self.ada_out = nn.Sequential(nn.SiLU(), nn.Linear(dim, 2 * dim))
+        self.head = nn.Linear(dim, self.patch_dim)
+        self.dim = dim
+        self._init()
+
+    def _init(self):
+        for m in self.modules():
+            if isinstance(m, nn.Linear):
+                nn.init.xavier_uniform_(m.weight)
+                if m.bias is not None:
+                    nn.init.zeros_(m.bias)
+        for b in self.blocks:
+            nn.init.zeros_(b.ada[-1].weight); nn.init.zeros_(b.ada[-1].bias)
+        nn.init.zeros_(self.ada_out[-1].weight); nn.init.zeros_(self.ada_out[-1].bias)
+        nn.init.zeros_(self.head.weight); nn.init.zeros_(self.head.bias)
+
+    def patchify(self, x):
+        B, C, H, W = x.shape
+        p = self.patch
+        x = x.reshape(B, C, H // p, p, W // p, p)
+        x = x.permute(0, 2, 4, 1, 3, 5).reshape(B, (H // p) * (W // p), C * p * p)
+        return x
+
+    def unpatchify(self, x):
+        B, N, _ = x.shape
+        p = self.patch
+        g = self.grid
+        C = self.latent_ch
+        x = x.reshape(B, g, g, C, p, p).permute(0, 3, 1, 4, 2, 5)
+        return x.reshape(B, C, g * p, g * p)
+
+    def forward(self, x, t, text_seq, text_pool):
+        x = self.x_embed(self.patchify(x)) + self.pos
+        c = self.t_mlp(timestep_embedding(t, self.dim)) + self.text_pool(text_pool)
+        text = self.text_proj(text_seq)
+        for blk in self.blocks:
+            x = blk(x, c, text)
+        shift, scale = self.ada_out(c).chunk(2, dim=1)
+        x = modulate(self.norm_out(x), shift, scale)
+        x = self.head(x)
+        return self.unpatchify(x)
+
+    def num_params(self):
+        return sum(p.numel() for p in self.parameters())

requirements.txt

@@ -1,11 +1,13 @@
-gradio
-torch
-diffusers==0.31.0
-transformers==4.49.0
-safetensors
-numpy
-Pillow
-huggingface_hub
-trimesh
-scipy
-audioop-lts
+gradio
+torch
+diffusers==0.31.0
+transformers==4.49.0
+safetensors
+numpy
+Pillow
+huggingface_hub
+trimesh
+scipy
+audioop-lts
+sentencepiece
+accelerate

voxel_dit.py

@@ -1,134 +1,134 @@
-from __future__ import annotations
-
-import math
-import numpy as np
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-def modulate(x, shift, scale):
-    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
-
-def timestep_embedding(t, dim, max_period=10000):
-    half = dim // 2
-    freqs = torch.exp(-math.log(max_period) * torch.arange(half, device=t.device) / half)
-    args = t[:, None].float() * freqs[None]
-    emb = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
-    if dim % 2:
-        emb = torch.cat([emb, torch.zeros_like(emb[:, :1])], dim=-1)
-    return emb
-
-def sincos_3d(dim, grid_size):
-    assert dim % 6 == 0
-    g = np.arange(grid_size, dtype=np.float32)
-    gz, gy, gx = np.meshgrid(g, g, g, indexing="ij")
-    d6 = dim // 6
-    omega = 1.0 / (10000 ** (np.arange(d6, dtype=np.float32) / d6))
-    def emb1(p):
-        out = p.reshape(-1)[:, None] * omega[None]
-        return np.concatenate([np.sin(out), np.cos(out)], axis=1)
-    pe = np.concatenate([emb1(gz), emb1(gy), emb1(gx)], axis=1)
-    return torch.from_numpy(pe).float()
-
-class Attention(nn.Module):
-    def __init__(self, dim, heads):
-        super().__init__()
-        self.heads = heads
-        self.q = nn.Linear(dim, dim)
-        self.kv = nn.Linear(dim, dim * 2)
-        self.proj = nn.Linear(dim, dim)
-
-    def forward(self, x, ctx=None):
-        ctx = x if ctx is None else ctx
-        B, N, C = x.shape
-        M = ctx.shape[1]
-        h = self.heads
-        q = self.q(x).reshape(B, N, h, C // h).transpose(1, 2)
-        kv = self.kv(ctx).reshape(B, M, 2, h, C // h).permute(2, 0, 3, 1, 4)
-        k, v = kv[0], kv[1]
-        o = F.scaled_dot_product_attention(q, k, v)
-        o = o.transpose(1, 2).reshape(B, N, C)
-        return self.proj(o)
-
-class Block(nn.Module):
-    def __init__(self, dim, heads, mlp_ratio=4.0):
-        super().__init__()
-        self.norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
-        self.attn = Attention(dim, heads)
-        self.norm_ca = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
-        self.cross = Attention(dim, heads)
-        self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
-        hidden = int(dim * mlp_ratio)
-        self.mlp = nn.Sequential(nn.Linear(dim, hidden), nn.GELU(approximate="tanh"),
-                                 nn.Linear(hidden, dim))
-        self.ada = nn.Sequential(nn.SiLU(), nn.Linear(dim, 6 * dim))
-        self.cross_gate = nn.Parameter(torch.zeros(1))
-
-    def forward(self, x, c, text):
-        shift1, scale1, gate1, shift2, scale2, gate2 = self.ada(c).chunk(6, dim=1)
-        x = x + gate1.unsqueeze(1) * self.attn(modulate(self.norm1(x), shift1, scale1))
-        x = x + self.cross_gate * self.cross(self.norm_ca(x), text)
-        x = x + gate2.unsqueeze(1) * self.mlp(modulate(self.norm2(x), shift2, scale2))
-        return x
-
-class VoxelDiT(nn.Module):
-    def __init__(self, vox_ch=1, vox_size=32, patch=4, dim=384, depth=12,
-                 heads=6, text_dim=512, mlp_ratio=4.0):
-        super().__init__()
-        self.vox_ch = vox_ch
-        self.vox_size = vox_size
-        self.patch = patch
-        self.grid = vox_size // patch
-        self.patch_dim = vox_ch * patch * patch * patch
-        self.x_embed = nn.Linear(self.patch_dim, dim)
-        self.register_buffer("pos", sincos_3d(dim, self.grid).unsqueeze(0))
-        self.t_mlp = nn.Sequential(nn.Linear(dim, dim), nn.SiLU(), nn.Linear(dim, dim))
-        self.text_proj = nn.Linear(text_dim, dim)
-        self.text_pool = nn.Linear(text_dim, dim)
-        self.blocks = nn.ModuleList([Block(dim, heads, mlp_ratio) for _ in range(depth)])
-        self.norm_out = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
-        self.ada_out = nn.Sequential(nn.SiLU(), nn.Linear(dim, 2 * dim))
-        self.head = nn.Linear(dim, self.patch_dim)
-        self.dim = dim
-        self._init()
-
-    def _init(self):
-        for m in self.modules():
-            if isinstance(m, nn.Linear):
-                nn.init.xavier_uniform_(m.weight)
-                if m.bias is not None:
-                    nn.init.zeros_(m.bias)
-        for b in self.blocks:
-            nn.init.zeros_(b.ada[-1].weight); nn.init.zeros_(b.ada[-1].bias)
-        nn.init.zeros_(self.ada_out[-1].weight); nn.init.zeros_(self.ada_out[-1].bias)
-        nn.init.zeros_(self.head.weight); nn.init.zeros_(self.head.bias)
-
-    def patchify(self, x):
-        B, C, D, H, W = x.shape
-        p = self.patch
-        x = x.reshape(B, C, D // p, p, H // p, p, W // p, p)
-        x = x.permute(0, 2, 4, 6, 1, 3, 5, 7)
-        return x.reshape(B, (D // p) * (H // p) * (W // p), C * p * p * p)
-
-    def unpatchify(self, x):
-        B, N, _ = x.shape
-        p = self.patch
-        g = self.grid
-        C = self.vox_ch
-        x = x.reshape(B, g, g, g, C, p, p, p)
-        x = x.permute(0, 4, 1, 5, 2, 6, 3, 7)
-        return x.reshape(B, C, g * p, g * p, g * p)
-
-    def forward(self, x, t, text_seq, text_pool):
-        x = self.x_embed(self.patchify(x)) + self.pos
-        c = self.t_mlp(timestep_embedding(t, self.dim)) + self.text_pool(text_pool)
-        text = self.text_proj(text_seq)
-        for blk in self.blocks:
-            x = blk(x, c, text)
-        shift, scale = self.ada_out(c).chunk(2, dim=1)
-        x = modulate(self.norm_out(x), shift, scale)
-        x = self.head(x)
-        return self.unpatchify(x)
-
-    def num_params(self):
-        return sum(p.numel() for p in self.parameters())
+from __future__ import annotations
+
+import math
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+def modulate(x, shift, scale):
+    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+
+def timestep_embedding(t, dim, max_period=10000):
+    half = dim // 2
+    freqs = torch.exp(-math.log(max_period) * torch.arange(half, device=t.device) / half)
+    args = t[:, None].float() * freqs[None]
+    emb = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+    if dim % 2:
+        emb = torch.cat([emb, torch.zeros_like(emb[:, :1])], dim=-1)
+    return emb
+
+def sincos_3d(dim, grid_size):
+    assert dim % 6 == 0
+    g = np.arange(grid_size, dtype=np.float32)
+    gz, gy, gx = np.meshgrid(g, g, g, indexing="ij")
+    d6 = dim // 6
+    omega = 1.0 / (10000 ** (np.arange(d6, dtype=np.float32) / d6))
+    def emb1(p):
+        out = p.reshape(-1)[:, None] * omega[None]
+        return np.concatenate([np.sin(out), np.cos(out)], axis=1)
+    pe = np.concatenate([emb1(gz), emb1(gy), emb1(gx)], axis=1)
+    return torch.from_numpy(pe).float()
+
+class Attention(nn.Module):
+    def __init__(self, dim, heads):
+        super().__init__()
+        self.heads = heads
+        self.q = nn.Linear(dim, dim)
+        self.kv = nn.Linear(dim, dim * 2)
+        self.proj = nn.Linear(dim, dim)
+
+    def forward(self, x, ctx=None):
+        ctx = x if ctx is None else ctx
+        B, N, C = x.shape
+        M = ctx.shape[1]
+        h = self.heads
+        q = self.q(x).reshape(B, N, h, C // h).transpose(1, 2)
+        kv = self.kv(ctx).reshape(B, M, 2, h, C // h).permute(2, 0, 3, 1, 4)
+        k, v = kv[0], kv[1]
+        o = F.scaled_dot_product_attention(q, k, v)
+        o = o.transpose(1, 2).reshape(B, N, C)
+        return self.proj(o)
+
+class Block(nn.Module):
+    def __init__(self, dim, heads, mlp_ratio=4.0):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+        self.attn = Attention(dim, heads)
+        self.norm_ca = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+        self.cross = Attention(dim, heads)
+        self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+        hidden = int(dim * mlp_ratio)
+        self.mlp = nn.Sequential(nn.Linear(dim, hidden), nn.GELU(approximate="tanh"),
+                                 nn.Linear(hidden, dim))
+        self.ada = nn.Sequential(nn.SiLU(), nn.Linear(dim, 6 * dim))
+        self.cross_gate = nn.Parameter(torch.zeros(1))
+
+    def forward(self, x, c, text):
+        shift1, scale1, gate1, shift2, scale2, gate2 = self.ada(c).chunk(6, dim=1)
+        x = x + gate1.unsqueeze(1) * self.attn(modulate(self.norm1(x), shift1, scale1))
+        x = x + self.cross_gate * self.cross(self.norm_ca(x), text)
+        x = x + gate2.unsqueeze(1) * self.mlp(modulate(self.norm2(x), shift2, scale2))
+        return x
+
+class VoxelDiT(nn.Module):
+    def __init__(self, vox_ch=1, vox_size=32, patch=4, dim=384, depth=12,
+                 heads=6, text_dim=512, mlp_ratio=4.0):
+        super().__init__()
+        self.vox_ch = vox_ch
+        self.vox_size = vox_size
+        self.patch = patch
+        self.grid = vox_size // patch
+        self.patch_dim = vox_ch * patch * patch * patch
+        self.x_embed = nn.Linear(self.patch_dim, dim)
+        self.register_buffer("pos", sincos_3d(dim, self.grid).unsqueeze(0))
+        self.t_mlp = nn.Sequential(nn.Linear(dim, dim), nn.SiLU(), nn.Linear(dim, dim))
+        self.text_proj = nn.Linear(text_dim, dim)
+        self.text_pool = nn.Linear(text_dim, dim)
+        self.blocks = nn.ModuleList([Block(dim, heads, mlp_ratio) for _ in range(depth)])
+        self.norm_out = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+        self.ada_out = nn.Sequential(nn.SiLU(), nn.Linear(dim, 2 * dim))
+        self.head = nn.Linear(dim, self.patch_dim)
+        self.dim = dim
+        self._init()
+
+    def _init(self):
+        for m in self.modules():
+            if isinstance(m, nn.Linear):
+                nn.init.xavier_uniform_(m.weight)
+                if m.bias is not None:
+                    nn.init.zeros_(m.bias)
+        for b in self.blocks:
+            nn.init.zeros_(b.ada[-1].weight); nn.init.zeros_(b.ada[-1].bias)
+        nn.init.zeros_(self.ada_out[-1].weight); nn.init.zeros_(self.ada_out[-1].bias)
+        nn.init.zeros_(self.head.weight); nn.init.zeros_(self.head.bias)
+
+    def patchify(self, x):
+        B, C, D, H, W = x.shape
+        p = self.patch
+        x = x.reshape(B, C, D // p, p, H // p, p, W // p, p)
+        x = x.permute(0, 2, 4, 6, 1, 3, 5, 7)
+        return x.reshape(B, (D // p) * (H // p) * (W // p), C * p * p * p)
+
+    def unpatchify(self, x):
+        B, N, _ = x.shape
+        p = self.patch
+        g = self.grid
+        C = self.vox_ch
+        x = x.reshape(B, g, g, g, C, p, p, p)
+        x = x.permute(0, 4, 1, 5, 2, 6, 3, 7)
+        return x.reshape(B, C, g * p, g * p, g * p)
+
+    def forward(self, x, t, text_seq, text_pool):
+        x = self.x_embed(self.patchify(x)) + self.pos
+        c = self.t_mlp(timestep_embedding(t, self.dim)) + self.text_pool(text_pool)
+        text = self.text_proj(text_seq)
+        for blk in self.blocks:
+            x = blk(x, c, text)
+        shift, scale = self.ada_out(c).chunk(2, dim=1)
+        x = modulate(self.norm_out(x), shift, scale)
+        x = self.head(x)
+        return self.unpatchify(x)
+
+    def num_params(self):
+        return sum(p.numel() for p in self.parameters())