"""LLaDA2.2-mini ZeroGPU Space, matching the Ruurd live-diffusion layout.""" from pathlib import Path import hashlib import json import sys import time import gradio as gr import spaces import torch from huggingface_hub import hf_hub_download from transformers import AutoModelForCausalLM, AutoTokenizer sys.path.insert(0, str(Path(__file__).parent / "src")) from diffusion_lm.bigram import BigramScorer, BigramViterbiSelector from diffusion_lm.trigram import TrigramScorer from streaming import MODEL_ID, MODEL_REVISION, MASK, DELETE, INSERT, NativeStreamer, render BASE = Path(__file__).parent torch.set_num_threads(4) CONFIG = json.loads((BASE / "space_model.json").read_text()) TOKENIZER = AutoTokenizer.from_pretrained(MODEL_ID, revision=MODEL_REVISION, trust_remote_code=True) MODEL = AutoModelForCausalLM.from_pretrained( MODEL_ID, revision=MODEL_REVISION, trust_remote_code=True, dtype=torch.bfloat16, attn_implementation="sdpa", ).eval().to("cuda") STREAMER = NativeStreamer(MODEL) def scorer_path(name): revision = CONFIG["ngram_scorer_revision"] if not revision: return BASE / name # local bundle validation before publication path = Path(hf_hub_download(CONFIG["ngram_scorer_repo"], name, revision=revision)) h = hashlib.sha256() with path.open("rb") as f: for chunk in iter(lambda: f.read(8 * 1024 * 1024), b""): h.update(chunk) if h.hexdigest() != CONFIG["assets"][name]["sha256"]: raise RuntimeError(f"Scorer checksum mismatch: {name}") return path BIGRAM = BigramScorer.load(scorer_path("bigram_expanded.pt"), tokenizer=TOKENIZER) TRIGRAM = TrigramScorer.load(scorer_path("trigram_expanded.pt"), bigram=BIGRAM) @spaces.GPU(size='large', duration=30) def generate(question, system_prompt, max_new_tokens, steps, block_length, temperature, top_k, color_mode, method, bigram_weight, threshold, editing_threshold, max_post_steps, early_stopping, max_steps_per_block): if not question.strip(): raise gr.Error("Enter a prompt first.") if int(block_length) not in (32, 64): raise gr.Error("Block length must be 32 or 64 for this model's 32-token MoE routing blocks.") messages = [] if system_prompt.strip(): messages.append(dict(role="system", content=system_prompt)) messages.append(dict(role="user", content=question)) encoded = TOKENIZER.apply_chat_template(messages, add_generation_prompt=True, tokenize=True, return_tensors="pt") inputs = encoded if isinstance(encoded, torch.Tensor) else encoded["input_ids"] if any(token in (MASK, DELETE, INSERT) for token in inputs[0].tolist()): raise gr.Error("Prompt contains a reserved mask/edit token. Remove it before generating.") # The demo bounds dense attention memory; the model itself supports 128K. if inputs.shape[1] + int(max_new_tokens) > 4096: raise gr.Error("This demo supports up to 4096 prompt + answer tokens.") inputs = inputs.to(MODEL.device) selector = None if method == "BG + TG Viterbi (Gumbel)": # CPU scorers use the existing candidate-transfer path and avoid holding # potentially gigabytes of count buffers in a transient GPU allocation. selector = BigramViterbiSelector( BIGRAM, weight=float(bigram_weight), scoring_method="pmi", min_count=1, pmi_clip=2.0, independent_temperature=True, randomized_viterbi=True, trigram_scorer=TRIGRAM, trigram_weight=.5, trigram_min_count=2, trigram_clip=2.0, ) selector = CPUSelector(selector) start = time.monotonic() try: for frame in STREAMER.stream( inputs, selector=selector, gen_length=int(max_new_tokens), steps=int(steps), block_length=int(block_length), temperature=float(temperature), top_k=int(top_k) or None, top_p=None, threshold=float(threshold), editing_threshold=float(editing_threshold), max_post_steps=int(max_post_steps), eos_early_stop=bool(early_stopping), max_steps_per_block=int(max_steps_per_block), ): status = f"{'Complete' if frame.done else 'Denoising'} · {frame.iteration} forwards · {time.monotonic()-start:.1f}s" yield status, render(frame, TOKENIZER, color_mode) except RuntimeError as exc: raise gr.Error(str(exc)) from exc class CPUSelector: """Transfer the small block logits for exact CPU n-gram candidate scoring.""" def __init__(self, selector): self.selector = selector def select_block(self, logits, current_ids, editable, **kwargs): # These blocks are at most 64 x 157184 floats; no extra model forwards. tokens, confidence = self.selector.select_block( logits.cpu(), current_ids.cpu(), editable.cpu(), **kwargs) return tokens.to(logits.device), confidence.to(logits.device) with gr.Blocks(title="LLaDA2.2-mini · masked diffusion") as demo: gr.Markdown("# LLaDA2.2-mini\nFull-precision diffusion with native token rewriting, insertion and deletion.") with gr.Row(): with gr.Column(scale=3): question = gr.Textbox(label="Prompt", value="What do you know about Amsterdam?", lines=4) run = gr.Button("Generate", variant="primary") status = gr.Markdown("Ready") trajectory = gr.HTML(label="Live denoising") with gr.Column(scale=2): system = gr.Textbox(label="System prompt", value="You are a helpful assistant.", lines=2) length = gr.Slider(1, 512, value=64, step=1, label="New tokens") steps = gr.Slider(1, 128, value=32, step=1, label="Denoising steps per block", info="Controls the minimum unmasking schedule. Refinement may add more iterations.") block = gr.Dropdown([32, 64], value=32, label="Block length", info="Must align to the model's 32-token MoE routing blocks.") temperature = gr.Slider(0, 2, value=.7, step=.05, label="Temperature", info="0 is deterministic. Native mode samples masked tokens; BG + TG adds Gumbel noise to ordinary-token paths.") weight = gr.Slider(0, 5, value=1, step=.05, label="Bigram weight") top_k = gr.Slider(0, 100, value=16, step=1, label="Top-k", info="0 means all tokens in native mode; BG + TG uses 16 candidates when set to 0.") color = gr.Radio(["No coloring", "Prediction probability", "Prediction iteration"], value="Prediction iteration", label="Token coloring") method = gr.Radio(["Confidence-guided", "BG + TG Viterbi (Gumbel)"], value="BG + TG Viterbi (Gumbel)", label="Decoding method") early = gr.Checkbox(value=True, label="Early stopping", info="Stop after a completed block contains EOS. Blocks may converge sooner when stable.") with gr.Accordion("Native editing controls", open=False): threshold = gr.Slider(0, 1, value=.5, step=.05, label="Unmasking confidence threshold") editing = gr.Slider(0, 1, value=0, step=.05, label="Rewrite confidence threshold") post = gr.Slider(0, 32, value=16, step=1, label="Maximum post-mask refinement steps") cap = gr.Slider(32, 1000, value=256, step=1, label="Safety cap per block") gr.Markdown( "INSERT opens a mask before the previous token; DELETE shifts tokens left. " "Blocks retain their size by truncating or padding with masks. " "BG + TG guides ordinary text; structural edit predictions remain native.\n\n" f"[Model card](https://huggingface.co/{MODEL_ID}) · " "[Source code](https://github.com/RuurdKuiper/lad-generic)" ) inputs = [question, system, length, steps, block, temperature, top_k, color, method, weight, threshold, editing, post, early, cap] event = run.click(lambda: ("Allocating GPU…", ""), outputs=[status, trajectory], queue=False) event.then(generate, inputs, [status, trajectory], api_name="generate", concurrency_limit=1) method.change(lambda value: gr.update(interactive=value == "BG + TG Viterbi (Gumbel)"), method, weight, queue=False) demo.queue(default_concurrency_limit=1, max_size=8) if __name__ == "__main__": demo.launch()