Update app.py

app.py

@@ -2,13 +2,18 @@ import os
 import sys
 
 import torch
+import torch.nn.functional as F
 import gradio as gr
 from tokenizers import Tokenizer
 from huggingface_hub import hf_hub_download, snapshot_download
+from transformers import AutoModel, AutoModelForCausalLM, AutoTokenizer
 
 # ── Repo IDs ───────────────────────────────────────────────────────────────────
 REPO_V1 = "IvmeLabs/Ivme-Conversate-v1-Base"
 REPO_V2 = "IvmeLabs/Ivme-Conversate-v2-Base"
+REPO_CODER = "IvmeLabs/Ivme-Coder-v1"
+REPO_DIFF_BASE = "IvmeLabs/ExpIvme-DiffusionConversate-v1"
+REPO_DIFF_INSTRUCT = "IvmeLabs/ExpIvme-DiffusionConversate-v1-Instruct"
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
@@ -41,7 +46,7 @@ def load_v1():
         or 1024
     )
     eos_id = tok.token_to_id("<|eos|>")
-    return {"tokenizer": tok, "model": model, "max_ctx": max_ctx, "eos_id": eos_id}
+    return {"kind": "ar-raw", "tokenizer": tok, "model": model, "max_ctx": max_ctx, "eos_id": eos_id}
 
 
 # ── Load v2 ────────────────────────────────────────────────────────────────────
@@ -75,17 +80,54 @@ def load_v2():
 
     max_ctx = getattr(cfg, "context_len", 1024)
     eos_id = tok.token_to_id("<|endoftext|>")
-    return {"tokenizer": tok, "model": model, "max_ctx": max_ctx, "eos_id": eos_id}
+    return {"kind": "ar-raw", "tokenizer": tok, "model": model, "max_ctx": max_ctx, "eos_id": eos_id}
+
+
+# ── Load Coder-v1 (standard transformers AutoModelForCausalLM) ────────────────
+def load_coder():
+    tokenizer = AutoTokenizer.from_pretrained(REPO_CODER, trust_remote_code=True)
+    model = AutoModelForCausalLM.from_pretrained(
+        REPO_CODER, trust_remote_code=True, dtype=torch.float32,
+    ).to(device).eval()
+    return {"kind": "ar-hf", "tokenizer": tokenizer, "model": model}
+
+
+# ── Load diffusion base + instruct (custom masked-diffusion sampler) ──────────
+def load_diffusion(repo_id, instruct):
+    tokenizer = AutoTokenizer.from_pretrained(repo_id, trust_remote_code=True)
+    model = AutoModel.from_pretrained(
+        repo_id, trust_remote_code=True,
+    ).to(device).eval()
+    bundle = {
+        "kind": "diffusion-instruct" if instruct else "diffusion-base",
+        "tokenizer": tokenizer,
+        "model": model,
+        "mask_token_id": model.config.mask_token_id,
+    }
+    if instruct:
+        bundle["user_token_id"] = model.config.user_token_id
+        bundle["assistant_token_id"] = model.config.assistant_token_id
+        bundle["endturn_token_id"] = model.config.endturn_token_id
+    return bundle
 
 
 print("Loading İvme-Conversate-v1-Base...")
 V1 = load_v1()
 print("Loading İvme-Conversate-v2-Base...")
 V2 = load_v2()
+print("Loading İvme-Coder-v1...")
+CODER = load_coder()
+print("Loading ExpİvmeDiffusionConversate-v1 (base)...")
+DIFF_BASE = load_diffusion(REPO_DIFF_BASE, instruct=False)
+print("Loading ExpİvmeDiffusionConversate-v1-Instruct...")
+DIFF_INSTRUCT = load_diffusion(REPO_DIFF_INSTRUCT, instruct=True)
 
 REGISTRY = {
     "İvme-Conversate-v2-Base (recommended)": V2,
     "İvme-Conversate-v1-Base": V1,
+    "İvme-Coder-v1 (Python code)": CODER,
+    "Expİvme-DiffusionConversate-v1 (experimental)": DIFF_BASE,
+    "Expİvme-DiffusionConversate-v1-Instruct (experimental)": DIFF_INSTRUCT,
 }
 
 BENCH = {
@@ -96,9 +138,9 @@ BENCH = {
 }
 
 
-# ── Shared generation core ────────────────────────────────────────────────────
+# ── Generation core: raw checkpoint AR models (v1/v2) ─────────────────────────
 @torch.no_grad()
-def _generate(bundle, prompt, max_new_tokens, temperature, top_k, repetition_penalty):
+def _generate_ar_raw(bundle, prompt, max_new_tokens, temperature, top_k, repetition_penalty):
     tokenizer = bundle["tokenizer"]
     model = bundle["model"]
     max_ctx = bundle["max_ctx"]
@@ -165,15 +207,210 @@ def _generate(bundle, prompt, max_new_tokens, temperature, top_k, repetition_pen
         yield prompt
 
 
-def continue_text(model_choice, prompt, max_new_tokens, temperature, top_k, repetition_penalty):
+# ── Generation core: HF transformers AR models (Coder-v1) ─────────────────────
+@torch.no_grad()
+def _generate_ar_hf(bundle, prompt, max_new_tokens, temperature, top_k, repetition_penalty):
+    tokenizer = bundle["tokenizer"]
+    model = bundle["model"]
+
+    prompt = prompt or ""
+    inputs = tokenizer(prompt, return_tensors="pt").to(device)
+    if inputs["input_ids"].shape[1] == 0:
+        yield prompt
+        return
+
+    generated = inputs["input_ids"]
+    response_tokens: list[int] = []
+    temperature = max(float(temperature), 1e-6)
+    eos_id = tokenizer.eos_token_id
+
+    for _ in range(int(max_new_tokens)):
+        out = model(generated)
+        logits = out.logits[:, -1, :].float()
+        vocab_size = logits.size(-1)
+
+        if repetition_penalty and repetition_penalty != 1.0:
+            seen = torch.unique(generated[0])
+            scores = logits[0, seen]
+            scores = torch.where(
+                scores > 0, scores / repetition_penalty, scores * repetition_penalty
+            )
+            logits[0, seen] = scores
+
+        logits = logits / temperature
+
+        k = int(top_k)
+        if k > 0:
+            k = min(k, vocab_size)
+            topk_vals, _ = torch.topk(logits, k)
+            logits[logits < topk_vals[:, -1:]] = float("-inf")
+
+        probs = torch.softmax(logits, dim=-1)
+        if not torch.isfinite(probs).all() or probs.sum() <= 0:
+            next_tok = torch.argmax(logits, dim=-1, keepdim=True)
+        else:
+            next_tok = torch.multinomial(probs, num_samples=1)
+
+        tok_id = next_tok.item()
+        if eos_id is not None and tok_id == eos_id:
+            break
+
+        response_tokens.append(tok_id)
+        generated = torch.cat([generated, next_tok], dim=1)
+
+        yield prompt + tokenizer.decode(response_tokens)
+
+    if not response_tokens:
+        yield prompt
+
+
+# ── Generation core: masked-diffusion base model (unconditional/continuation) ─
+@torch.no_grad()
+def _generate_diffusion_base(bundle, prompt, length, steps, temperature, gumbel_temp):
+    tokenizer = bundle["tokenizer"]
+    model = bundle["model"]
+    mask_token_id = bundle["mask_token_id"]
+
+    length = int(length)
+    steps = max(int(steps), 1)
+
+    prefix_ids = tokenizer.encode(prompt) if prompt else []
+    prefix_len = len(prefix_ids)
+    total_len = prefix_len + length
+
+    input_ids = torch.full((1, total_len), mask_token_id, dtype=torch.long, device=device)
+    if prefix_len > 0:
+        input_ids[0, :prefix_len] = torch.tensor(prefix_ids, dtype=torch.long, device=device)
+
+    response_start = prefix_len
+
+    for step in range(steps):
+        logits = model(input_ids=input_ids).logits
+        probs = F.softmax(logits / max(float(temperature), 1e-6), dim=-1)
+        sampled = torch.multinomial(probs.view(-1, probs.size(-1)), 1).view(1, total_len)
+
+        # Never allow the fixed prefix to be resampled.
+        is_masked = input_ids == mask_token_id
+        if prefix_len > 0:
+            is_masked[0, :prefix_len] = False
+
+        n_masked = is_masked.sum().item()
+        if n_masked == 0:
+            break
+
+        frac_remaining = 1.0 - (step + 1) / steps
+        denom = max(1 - step / steps, 1e-6)
+        n_to_unmask = min(max(1, int(n_masked * (1 - frac_remaining / denom))), n_masked)
+
+        conf = probs.gather(-1, sampled.unsqueeze(-1)).squeeze(-1)
+        log_conf = torch.log(conf.clamp(min=1e-9))
+        u = torch.rand_like(conf).clamp(min=1e-9, max=1 - 1e-9)
+        gumbel_noise = -torch.log(-torch.log(u))
+        score = (log_conf + gumbel_temp * gumbel_noise).masked_fill(~is_masked, float("-inf"))
+
+        topk = torch.topk(score, k=n_to_unmask, dim=-1).indices
+        update_mask = torch.zeros_like(is_masked).scatter_(1, topk, True)
+        input_ids = torch.where(update_mask, sampled, input_ids)
+
+        partial = input_ids[0, response_start:].tolist()
+        yield (prompt or "") + tokenizer.decode(partial)
+
+    final = input_ids[0, response_start:].tolist()
+    yield (prompt or "") + tokenizer.decode(final)
+
+
+# ── Generation core: masked-diffusion instruct model (chat) ───────────────────
+@torch.no_grad()
+def _generate_diffusion_instruct(bundle, user_message, max_response_len, steps, temperature,
+                                  gumbel_temp, presence_penalty):
+    tokenizer = bundle["tokenizer"]
+    model = bundle["model"]
+    mask_id = bundle["mask_token_id"]
+    user_id = bundle["user_token_id"]
+    assistant_id = bundle["assistant_token_id"]
+    endturn_id = bundle["endturn_token_id"]
+
+    max_response_len = int(max_response_len)
+    steps = max(int(steps), 1)
+
+    prefix_ids = [user_id] + tokenizer.encode(user_message or "") + [endturn_id, assistant_id]
+    input_ids = torch.tensor(
+        [prefix_ids + [mask_id] * max_response_len], dtype=torch.long, device=device,
+    )
+    prefix_len = len(prefix_ids)
+    vocab_size = model.config.vocab_size
+
+    for step in range(steps):
+        logits = model(input_ids=input_ids).logits
+
+        if presence_penalty > 0:
+            response_span = input_ids[:, prefix_len:]
+            visible = response_span.masked_fill(response_span == mask_id, -1)
+            counts = torch.zeros(1, vocab_size, device=device)
+            valid = visible[0][visible[0] >= 0]
+            if len(valid) > 0:
+                counts[0].scatter_add_(0, valid, torch.ones_like(valid, dtype=torch.float))
+            logits = logits - presence_penalty * counts.unsqueeze(1)
+
+        probs = F.softmax(logits / max(float(temperature), 1e-6), dim=-1)
+        sampled = torch.multinomial(probs.view(-1, probs.size(-1)), 1).view(input_ids.shape)
+
+        is_masked = input_ids == mask_id
+        n_masked = is_masked.sum().item()
+        if n_masked == 0:
+            break
+
+        frac_remaining = 1.0 - (step + 1) / steps
+        denom = max(1 - step / steps, 1e-6)
+        n_to_unmask = min(max(1, int(n_masked * (1 - frac_remaining / denom))), n_masked)
+
+        conf = probs.gather(-1, sampled.unsqueeze(-1)).squeeze(-1)
+        log_conf = torch.log(conf.clamp(min=1e-9))
+        u = torch.rand_like(conf).clamp(min=1e-9, max=1 - 1e-9)
+        gumbel_noise = -torch.log(-torch.log(u))
+        score = (log_conf + gumbel_temp * gumbel_noise).masked_fill(~is_masked, float("-inf"))
+
+        topk = torch.topk(score.view(1, -1), k=n_to_unmask, dim=-1).indices
+        update_mask = torch.zeros_like(is_masked).view(1, -1).scatter_(1, topk, True).view(is_masked.shape)
+        input_ids = torch.where(update_mask, sampled, input_ids)
+
+        response_tokens = input_ids[0, prefix_len:].tolist()
+        if endturn_id in response_tokens:
+            response_tokens = response_tokens[:response_tokens.index(endturn_id)]
+        yield tokenizer.decode(response_tokens)
+
+    response_tokens = input_ids[0, prefix_len:].tolist()
+    if endturn_id in response_tokens:
+        response_tokens = response_tokens[:response_tokens.index(endturn_id)]
+    yield tokenizer.decode(response_tokens)
+
+
+# ── Unified dispatcher used by the Playground tab ──────────────────────────────
+def continue_text(model_choice, prompt, max_new_tokens, temperature, top_k, repetition_penalty,
+                   diff_steps, gumbel_temp, presence_penalty):
     bundle = REGISTRY[model_choice]
-    yield from _generate(bundle, prompt, max_new_tokens, temperature, top_k, repetition_penalty)
+    kind = bundle["kind"]
+
+    if kind == "ar-raw":
+        yield from _generate_ar_raw(bundle, prompt, max_new_tokens, temperature, top_k, repetition_penalty)
+    elif kind == "ar-hf":
+        yield from _generate_ar_hf(bundle, prompt, max_new_tokens, temperature, top_k, repetition_penalty)
+    elif kind == "diffusion-base":
+        yield from _generate_diffusion_base(
+            bundle, prompt, max_new_tokens, diff_steps, temperature, gumbel_temp
+        )
+    elif kind == "diffusion-instruct":
+        yield from _generate_diffusion_instruct(
+            bundle, prompt, max_new_tokens, diff_steps, temperature, gumbel_temp, presence_penalty
+        )
+    else:
+        yield prompt
 
 
 def compare_generate(prompt, max_new_tokens, temperature, top_k, repetition_penalty):
     """Run v1 and v2 on the same prompt/settings, streaming both in parallel steps."""
-    gen_v1 = _generate(V1, prompt, max_new_tokens, temperature, top_k, repetition_penalty)
-    gen_v2 = _generate(V2, prompt, max_new_tokens, temperature, top_k, repetition_penalty)
+    gen_v1 = _generate_ar_raw(V1, prompt, max_new_tokens, temperature, top_k, repetition_penalty)
+    gen_v2 = _generate_ar_raw(V2, prompt, max_new_tokens, temperature, top_k, repetition_penalty)
 
     last_v1, last_v2 = prompt, prompt
     done_v1 = done_v2 = False
@@ -219,11 +456,85 @@ EXAMPLES = [
     "Python is a programming language that",
 ]
 
+CODE_EXAMPLES = [
+    "def fibonacci(n):",
+    "class BinaryTree:",
+    "import numpy as np\n\ndef normalize(",
+    "# Sort a list using quicksort\ndef quicksort(arr):",
+]
+
+CHAT_EXAMPLES = [
+    "Hi there, how are you?",
+    "What's your favorite color?",
+    "Can you help me plan my day?",
+    "Tell me something interesting.",
+]
+
+DIFFUSION_KEYS = {
+    "Expİvme-DiffusionConversate-v1 (experimental)",
+    "Expİvme-DiffusionConversate-v1-Instruct (experimental)",
+}
+INSTRUCT_KEY = "Expİvme-DiffusionConversate-v1-Instruct (experimental)"
+CODER_KEY = "İvme-Coder-v1 (Python code)"
+
+MODEL_NOTES = {
+    "İvme-Conversate-v2-Base (recommended)": (
+        "Autoregressive base model, general text. Not instruction-tuned — continues text, doesn't chat."
+    ),
+    "İvme-Conversate-v1-Base": (
+        "Autoregressive base model, general text (earlier version). Not instruction-tuned."
+    ),
+    CODER_KEY: (
+        "Autoregressive base model trained only on Python source. Writes code-*shaped* text reliably; "
+        "does not reliably write *correct* code. Not instruction-tuned — give it a code prefix to continue."
+    ),
+    "Expİvme-DiffusionConversate-v1 (experimental)": (
+        "🧪 Experimental masked-diffusion model (not autoregressive). Generates a fixed-length span via "
+        "iterative denoising instead of left-to-right decoding. Not instruction-tuned, no chat behavior. "
+        "Weak general capability (near-chance on ARC-Easy) — expect local fluency, not coherent long-form text."
+    ),
+    INSTRUCT_KEY: (
+        "🧪 Experimental masked-diffusion model, SFT'd for basic chat. Enter a single user message (not a "
+        "free-form prompt). Known limitation per the model card: output is not reliably grammatical — "
+        "locally plausible words that often don't compose into coherent sentences."
+    ),
+}
+
+
+def on_model_change(model_choice):
+    """Toggle which settings are relevant/visible and swap in the right examples + notes."""
+    is_diffusion = model_choice in DIFFUSION_KEYS
+    is_instruct = model_choice == INSTRUCT_KEY
+    is_coder = model_choice == CODER_KEY
+
+    if is_instruct:
+        examples = CHAT_EXAMPLES
+        prompt_label = "User message"
+        prompt_placeholder = "Hi there, how are you?"
+    elif is_coder:
+        examples = CODE_EXAMPLES
+        prompt_label = "Prompt (Python)"
+        prompt_placeholder = "def fibonacci(n):"
+    else:
+        examples = EXAMPLES
+        prompt_label = "Prompt"
+        prompt_placeholder = "The theory of relativity states that…"
+
+    return (
+        gr.update(visible=not is_diffusion),  # AR-only settings group
+        gr.update(visible=is_diffusion),  # diffusion-only settings group
+        gr.update(visible=is_instruct),  # presence penalty (instruct diffusion only)
+        gr.update(label=prompt_label, placeholder=prompt_placeholder),
+        gr.Dataset(samples=[[e] for e in examples]),
+        gr.update(value=MODEL_NOTES.get(model_choice, "")),
+    )
+
+
 with gr.Blocks(css=CSS, title="İvme-Conversate") as demo:
     gr.Markdown(
-        "## İvme-Conversate — Tiny Language Models, Text Continuation\n"
-        "Sub-25M-parameter decoder-only **base** models · not instruction-tuned · "
-        "these models **continue** text, they do not chat or answer questions."
+        "## İvme-Conversate — Tiny Language Models\n"
+        "A family of sub-130M-parameter language models from IvmeLabs: autoregressive base models, "
+        "a Python-only coder model, and experimental masked-diffusion models."
     )
 
     with gr.Tabs():
@@ -235,6 +546,8 @@ with gr.Blocks(css=CSS, title="İvme-Conversate") as demo:
                 label="Model",
             )
 
+            model_note = gr.Markdown(MODEL_NOTES["İvme-Conversate-v2-Base (recommended)"])
+
             prompt_box = gr.Textbox(
                 label="Prompt",
                 placeholder="The theory of relativity states that…",
@@ -247,33 +560,84 @@ with gr.Blocks(css=CSS, title="İvme-Conversate") as demo:
                 clear_btn = gr.Button("Clear", scale=1)
 
             output_box = gr.Textbox(
-                label="Continuation",
+                label="Output",
                 lines=12,
                 show_copy_button=True,
                 elem_classes="ivme-output",
                 interactive=False,
             )
 
-            gr.Examples(examples=[[e] for e in EXAMPLES], inputs=prompt_box, label="Try a prompt")
+            example_set = gr.Examples(examples=[[e] for e in EXAMPLES], inputs=prompt_box, label="Try a prompt")
 
             with gr.Accordion("Settings", open=False):
-                with gr.Row():
-                    max_tokens = gr.Slider(16, 512, value=200, step=8, label="Max new tokens")
-                    temperature = gr.Slider(0.1, 2.0, value=0.7, step=0.05, label="Temperature")
-                with gr.Row():
-                    top_k = gr.Slider(0, 200, value=40, step=1, label="Top-k (0 = disabled)")
-                    rep_penalty = gr.Slider(1.0, 2.0, value=1.15, step=0.05, label="Repetition penalty")
-
-            gen_inputs = [model_picker, prompt_box, max_tokens, temperature, top_k, rep_penalty]
-            gen_btn.click(continue_text, gen_inputs, output_box)
-            prompt_box.submit(continue_text, gen_inputs, output_box)
+                with gr.Group(visible=True) as ar_settings:
+                    with gr.Row():
+                        max_tokens = gr.Slider(16, 512, value=200, step=8, label="Max new tokens")
+                        temperature = gr.Slider(0.1, 2.0, value=0.7, step=0.05, label="Temperature")
+                    with gr.Row():
+                        top_k = gr.Slider(0, 200, value=40, step=1, label="Top-k (0 = disabled)")
+                        rep_penalty = gr.Slider(1.0, 2.0, value=1.15, step=0.05, label="Repetition penalty")
+
+                with gr.Group(visible=False) as diff_settings:
+                    gr.Markdown(
+                        "Masked-diffusion sampling: the model denoises a fully-masked span over a fixed "
+                        "number of steps rather than decoding left-to-right."
+                    )
+                    with gr.Row():
+                        diff_length = gr.Slider(16, 256, value=96, step=8, label="Response length (tokens)")
+                        diff_steps = gr.Slider(4, 64, value=32, step=2, label="Diffusion steps")
+                    with gr.Row():
+                        diff_temperature = gr.Slider(0.1, 2.0, value=1.0, step=0.05, label="Temperature")
+                        diff_gumbel_temp = gr.Slider(0.0, 2.0, value=1.0, step=0.1, label="Gumbel temp (unmask noise)")
+                    diff_presence_penalty = gr.Slider(
+                        0.0, 3.0, value=1.2, step=0.1,
+                        label="Presence penalty (Instruct only — suppresses repetition)",
+                        visible=False,
+                    )
+
+            gen_inputs = [
+                model_picker, prompt_box, max_tokens, temperature, top_k, rep_penalty,
+                diff_steps, diff_gumbel_temp, diff_presence_penalty,
+            ]
+            # Note: for diffusion models, `max_tokens` slider doubles as response length via diff_length
+            # binding below; wire diff_length into the same "max_new_tokens" slot dynamically:
+
+            def route_generate(model_choice, prompt, max_new_tokens, temperature, top_k, repetition_penalty,
+                                length, steps, d_temperature, gumbel_temp, presence_penalty):
+                bundle = REGISTRY[model_choice]
+                kind = bundle["kind"]
+                if kind == "diffusion-base":
+                    yield from _generate_diffusion_base(bundle, prompt, length, steps, d_temperature, gumbel_temp)
+                elif kind == "diffusion-instruct":
+                    yield from _generate_diffusion_instruct(
+                        bundle, prompt, length, steps, d_temperature, gumbel_temp, presence_penalty
+                    )
+                elif kind == "ar-hf":
+                    yield from _generate_ar_hf(bundle, prompt, max_new_tokens, temperature, top_k, repetition_penalty)
+                else:
+                    yield from _generate_ar_raw(bundle, prompt, max_new_tokens, temperature, top_k, repetition_penalty)
+
+            full_inputs = [
+                model_picker, prompt_box, max_tokens, temperature, top_k, rep_penalty,
+                diff_length, diff_steps, diff_temperature, diff_gumbel_temp, diff_presence_penalty,
+            ]
+            gen_btn.click(route_generate, full_inputs, output_box)
+            prompt_box.submit(route_generate, full_inputs, output_box)
             clear_btn.click(lambda: ("", ""), None, [prompt_box, output_box], queue=False)
 
+            model_picker.change(
+                on_model_change,
+                inputs=model_picker,
+                outputs=[ar_settings, diff_settings, diff_presence_penalty, prompt_box, example_set.dataset, model_note],
+            )
+
         # ── Tab 2: side-by-side compare ──────────────────────────────────────
         with gr.Tab("Compare v1 vs v2"):
             gr.Markdown(
-                "Run the **same prompt and settings** through both models at once "
-                "to see the difference training data made, plus the benchmark deltas below."
+                "Run the **same prompt and settings** through both autoregressive base models at once "
+                "to see the difference training data made, plus the benchmark deltas below. "
+                "(Coder-v1 and the diffusion models aren't included here since they use different "
+                "generation mechanics — try them individually in the Playground tab.)"
             )
 
             cmp_prompt = gr.Textbox(
@@ -325,4 +689,53 @@ with gr.Blocks(css=CSS, title="İvme-Conversate") as demo:
             cmp_prompt.submit(compare_generate, cmp_inputs, [cmp_out_v1, cmp_out_v2])
             cmp_clear_btn.click(lambda: ("", "", ""), None, [cmp_prompt, cmp_out_v1, cmp_out_v2], queue=False)
 
-demo.queue().launch()
+        # ── Tab 3: diffusion chat (Instruct model, dedicated chat-style UI) ──
+        with gr.Tab("Diffusion Chat (experimental)"):
+            gr.Markdown(
+                "### Expİvme-DiffusionConversate-v1-Instruct\n"
+                "🧪 **Experimental.** A 130M-parameter masked-diffusion model, SFT'd for basic chat. "
+                "Per the model card: output is **not reliably grammatical** — expect locally plausible "
+                "word choice that often doesn't compose into coherent sentences. Included here in the "
+                "spirit of the model card's own honesty about its limitations, not as a working assistant."
+            )
+
+            chat_input = gr.Textbox(
+                label="Your message",
+                placeholder="Hi there, how are you?",
+                lines=2,
+            )
+
+            with gr.Row():
+                chat_btn = gr.Button("Send", variant="primary", scale=3)
+                chat_clear_btn = gr.Button("Clear", scale=1)
+
+            chat_output = gr.Textbox(
+                label="Assistant (diffusion-sampled)",
+                lines=6,
+                show_copy_button=True,
+                elem_classes="ivme-output",
+                interactive=False,
+            )
+
+            gr.Examples(examples=[[e] for e in CHAT_EXAMPLES], inputs=chat_input, label="Try a message")
+
+            with gr.Accordion("Settings", open=False):
+                with gr.Row():
+                    chat_len = gr.Slider(16, 128, value=64, step=8, label="Max response length")
+                    chat_steps = gr.Slider(4, 64, value=32, step=2, label="Diffusion steps")
+                with gr.Row():
+                    chat_temperature = gr.Slider(0.1, 2.0, value=1.0, step=0.05, label="Temperature")
+                    chat_gumbel = gr.Slider(0.0, 2.0, value=1.0, step=0.1, label="Gumbel temp")
+                chat_presence = gr.Slider(0.0, 3.0, value=1.2, step=0.1, label="Presence penalty")
+
+            def diffusion_chat(user_message, length, steps, temperature, gumbel_temp, presence_penalty):
+                yield from _generate_diffusion_instruct(
+                    DIFF_INSTRUCT, user_message, length, steps, temperature, gumbel_temp, presence_penalty
+                )
+
+            chat_inputs = [chat_input, chat_len, chat_steps, chat_temperature, chat_gumbel, chat_presence]
+            chat_btn.click(diffusion_chat, chat_inputs, chat_output)
+            chat_input.submit(diffusion_chat, chat_inputs, chat_output)
+            chat_clear_btn.click(lambda: ("", ""), None, [chat_input, chat_output], queue=False)
+
+demo.queue().launch()