diffusion-chatbot

Sleeping

File size: 34,824 Bytes

import sys
import types
import os
import math
import json
import copy
import torch
import torch.nn.functional as F
from flask import Flask, request, jsonify, Response
from transformers import AutoTokenizer, AutoModelForMaskedLM, AutoModelForCausalLM, TextIteratorStreamer
from threading import Thread

HTML_UI = """
<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>LLM API Tester</title>
    <style>
        * { box-sizing: border-box; }
        body {
            font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, sans-serif;
            max-width: 900px;
            margin: 0 auto;
            padding: 20px;
            background: #f5f5f5;
            color: #333;
        }
        h1 { margin-top: 0; font-size: 1.5rem; }
        .card {
            background: #fff;
            border-radius: 8px;
            padding: 20px;
            margin-bottom: 16px;
            box-shadow: 0 1px 3px rgba(0,0,0,0.1);
        }
        label {
            display: block;
            font-weight: 600;
            margin-bottom: 6px;
            font-size: 0.9rem;
        }
        textarea, input, select {
            width: 100%;
            padding: 10px;
            border: 1px solid #ddd;
            border-radius: 6px;
            font-size: 0.95rem;
            font-family: inherit;
        }
        textarea { resize: vertical; min-height: 80px; }
        .row {
            display: grid;
            grid-template-columns: repeat(auto-fit, minmax(140px, 1fr));
            gap: 12px;
            margin-bottom: 12px;
        }
        .field { margin-bottom: 12px; }
        .field.inline {
            display: flex;
            align-items: center;
            gap: 8px;
        }
        .field.inline label { margin: 0; }
        .field.inline input, .field.inline select {
            width: auto;
            flex: 1;
        }
        button {
            background: #2563eb;
            color: #fff;
            border: none;
            padding: 10px 20px;
            border-radius: 6px;
            font-size: 1rem;
            cursor: pointer;
            font-weight: 600;
        }
        button:hover { background: #1d4ed8; }
        button:disabled { background: #93c5fd; cursor: not-allowed; }
        .output {
            background: #1e1e1e;
            color: #e4e4e4;
            padding: 16px;
            border-radius: 6px;
            font-family: ui-monospace, SFMono-Regular, "SF Mono", Menlo, Consolas, monospace;
            font-size: 0.9rem;
            white-space: pre-wrap;
            word-break: break-word;
            min-height: 120px;
            max-height: 500px;
            overflow-y: auto;
        }
        .output:empty::before {
            content: "Response will appear here...";
            color: #666;
        }
        .status {
            font-size: 0.85rem;
            color: #666;
            margin-top: 8px;
        }
        .error { color: #dc2626; }
        .success { color: #16a34a; }
        .route-badge {
            display: inline-block;
            background: #e5e7eb;
            padding: 2px 8px;
            border-radius: 4px;
            font-size: 0.8rem;
            font-weight: 600;
            margin-bottom: 8px;
        }
    </style>
</head>
<body>
    <h1>LLM API Tester</h1>

    <div class="card">
        <div class="route-badge" id="routeBadge">/generate</div>
        <div class="field">
            <label for="route">Route</label>
            <select id="route">
                <option value="/generate">/generate (sync JSON)</option>
                <option value="/generate_stream">/generate_stream (sync JSON + intermediates)</option>
                <option value="/generate_sse">/generate_sse (streaming SSE)</option>
            </select>
        </div>

        <div class="field">
            <label for="prompt">Prompt</label>
            <textarea id="prompt" placeholder="Enter your prompt here...">Write a short poem about stars</textarea>
        </div>

        <div class="row">
            <div class="field">
                <label for="max_new_tokens">max_new_tokens</label>
                <input type="number" id="max_new_tokens" value="150" min="1" max="2048">
            </div>
            <div class="field">
                <label for="temperature">temperature</label>
                <input type="number" id="temperature" value="0.0" min="0" max="2" step="0.1">
            </div>
            <div class="field">
                <label for="steps">steps (diffusion)</label>
                <input type="number" id="steps" value="256" min="1">
            </div>
            <div class="field">
                <label for="block_size">block_size</label>
                <input type="number" id="block_size" value="32" min="1">
            </div>
        </div>

        <div class="row">
            <div class="field">
                <label for="cfg_scale">cfg_scale</label>
                <input type="number" id="cfg_scale" value="0.0" min="0" step="0.1">
            </div>
            <div class="field">
                <label for="remasking">remasking</label>
                <select id="remasking">
                    <option value="low_confidence">low_confidence</option>
                    <option value="random">random</option>
                </select>
            </div>
            <div class="field">
                <label for="capture_interval">capture_interval</label>
                <input type="number" id="capture_interval" value="10" min="1">
            </div>
        </div>

        <button id="sendBtn">Send Request</button>
        <div class="status" id="status"></div>
    </div>

    <div class="card">
        <label>Response</label>
        <div class="output" id="output"></div>
    </div>

    <script>
        const $ = id => document.getElementById(id);
        const routeSelect = $('route');
        const routeBadge = $('routeBadge');
        const sendBtn = $('sendBtn');
        const output = $('output');
        const status = $('status');

        routeSelect.addEventListener('change', () => {
            routeBadge.textContent = routeSelect.value;
        });

        function setStatus(msg, isError = false) {
            status.textContent = msg;
            status.className = 'status ' + (isError ? 'error' : 'success');
        }

        function appendOutput(text, clear = false) {
            if (clear) output.textContent = '';
            output.textContent += text;
            output.scrollTop = output.scrollHeight;
        }

        function getPayload() {
            return {
                prompt: $('prompt').value,
                max_new_tokens: parseInt($('max_new_tokens').value),
                temperature: parseFloat($('temperature').value),
                steps: parseInt($('steps').value),
                block_size: parseInt($('block_size').value),
                cfg_scale: parseFloat($('cfg_scale').value),
                remasking: $('remasking').value,
                capture_interval: parseInt($('capture_interval').value)
            };
        }

        async function handleGenerate() {
            const payload = getPayload();
            // Remove diffusion-only fields for non-diffusion if needed, but server ignores extras
            const t0 = performance.now();
            const res = await fetch('/generate', {
                method: 'POST',
                headers: { 'Content-Type': 'application/json' },
                body: JSON.stringify(payload)
            });
            const data = await res.json();
            const ms = Math.round(performance.now() - t0);
            if (res.ok) {
                appendOutput(`[${ms}ms]\\n${data.generated_text || JSON.stringify(data, null, 2)}\\n\\n`, true);
                setStatus(`OK — ${ms}ms`);
            } else {
                appendOutput(`Error ${res.status}:\\n${JSON.stringify(data, null, 2)}\\n\\n`, true);
                setStatus(`HTTP ${res.status}`, true);
            }
        }

        async function handleGenerateStream() {
            const payload = getPayload();
            const t0 = performance.now();
            const res = await fetch('/generate_stream', {
                method: 'POST',
                headers: { 'Content-Type': 'application/json' },
                body: JSON.stringify(payload)
            });
            const data = await res.json();
            const ms = Math.round(performance.now() - t0);
            if (res.ok) {
                let text = `[${ms}ms]\\nGenerated text:\n${data.generated_text}\\n\\n`;
                if (data.intermediate_states && data.intermediate_states.length) {
                    text += `Intermediate states (${data.intermediate_states.length}):\\n`;
                    data.intermediate_states.forEach((s, i) => {
                        text += "  Step " + s.step + ": " + s.text.substring(0,120).replace(/\\n/g, ' ') + "...\\n";
                    });
                }
                appendOutput(text + '\\n', true);
                setStatus(`OK — ${ms}ms, ${data.intermediate_states?.length || 0} intermediates`);
            } else {
                appendOutput(`Error ${res.status}:\\n${JSON.stringify(data, null, 2)}\\n\\n`, true);
                setStatus(`HTTP ${res.status}`, true);
            }
        }

        async function handleGenerateSSE() {
            const payload = getPayload();
            const t0 = performance.now();
            appendOutput('', true);
            setStatus('Connecting SSE...');

            const res = await fetch('/generate_sse', {
                method: 'POST',
                headers: { 'Content-Type': 'application/json' },
                body: JSON.stringify(payload)
            });

            if (!res.ok) {
                const data = await res.json().catch(() => ({}));
                appendOutput(`Error ${res.status}:\\n${JSON.stringify(data, null, 2)}`, true);
                setStatus(`HTTP ${res.status}`, true);
                return;
            }

            const reader = res.body.getReader();
            const decoder = new TextDecoder();
            let buffer = '';
            let finalText = '';
            let eventCount = 0;

            while (true) {
                const { done, value } = await reader.read();
                if (done) break;
                buffer += decoder.decode(value, { stream: true });
                const lines = buffer.split('\\n');
                buffer = lines.pop(); // keep incomplete line in buffer

                for (const line of lines) {
                    if (!line.startsWith('data: ')) continue;
                    const jsonStr = line.slice(6).trim();
                    if (!jsonStr) continue;
                    try {
                        const event = JSON.parse(jsonStr);
                        eventCount++;
                        if (event.type === 'final') {
                            finalText = event.text;
                            const ms = Math.round(performance.now() - t0);
                            appendOutput(`[${ms}ms | ${eventCount} events]\\n${finalText}\\n`, true);
                            setStatus(`Done — ${ms}ms, ${eventCount} events, ${event.total_steps || '?'} steps`);
                        } else if (event.type === 'intermediate' || event.type === 'token') {
                            // Live update: overwrite with latest accumulated text
                            appendOutput(`${event.text}`, true);
                            setStatus(`Streaming... (${eventCount} events)`);
                        }
                    } catch (e) {
                        // ignore malformed lines
                    }
                }
            }

            if (!finalText && eventCount === 0) {
                setStatus('Stream ended with no events', true);
            }
        }

        sendBtn.addEventListener('click', async () => {
            sendBtn.disabled = true;
            setStatus('Sending...');
            try {
                const route = routeSelect.value;
                if (route === '/generate') await handleGenerate();
                else if (route === '/generate_stream') await handleGenerateStream();
                else if (route === '/generate_sse') await handleGenerateSSE();
            } catch (err) {
                appendOutput(`Network/JS Error:\\n${err.message}\\n\\n`, true);
                setStatus(err.message, true);
            } finally {
                sendBtn.disabled = false;
            }
        });
    </script>
</body>
</html>
"""

# 1. Environment Parsing & Architecture Strategy Mapping
MODEL_NAME = os.getenv("MODEL_NAME", "dllm-hub/Qwen3-0.6B-diffusion-bd3lm-v0.1")
IS_DIFFUSION = "diffusion" in MODEL_NAME.lower()

# Dynamic initialization layer targeting Diffusion Language Models
if IS_DIFFUSION:
    try:
        import dllm.utils
        import dllm.pipelines
        import dllm.data
        import dllm.core
    except ImportError:
        pass
    if 'dllm' not in sys.modules:
        dllm_mock = types.ModuleType('dllm')
        dllm_mock.core = sys.modules.get('dllm.core')
        dllm_mock.data = sys.modules.get('dllm.data')
        dllm_mock.pipelines = sys.modules.get('dllm.pipelines')
        dllm_mock.utils = sys.modules.get('dllm.utils')
        sys.modules['dllm'] = dllm_mock

app = Flask(__name__)
model = None
tokenizer = None
device = None

# ==========================================================
# SYSTEM WORKSPACE PIPELINES: CORE DIFFUSION SAMPLING LOOPS
# ==========================================================

def add_gumbel_noise(logits, temperature):
    """Add Gumbel noise using float32 (faster than float64 on most GPUs)."""
    if temperature == 0:
        return logits
    logits = logits.float()
    noise = torch.rand_like(logits)
    g = (-torch.log(noise)) ** temperature
    return logits.exp() / g


def get_num_transfer_tokens(mask_index, steps):
    mask_num = mask_index.sum(dim=1, keepdim=True)
    base = mask_num // steps
    rem = mask_num % steps
    out = torch.zeros(mask_num.size(0), steps, device=mask_index.device, dtype=torch.long) + base
    for i in range(mask_num.size(0)):
        out[i, : rem[i]] += 1
    return out


def build_staircase_attention_mask(x, block_size, pad_id):
    B, T = x.shape
    device = x.device
    valid = x != pad_id
    pos_raw = torch.cumsum(valid.long(), dim=-1)
    position_ids = torch.where(valid, pos_raw - 1, torch.zeros_like(pos_raw)).long()
    col = torch.arange(T, device=device)
    block_ids = (col // block_size).view(1, T).expand(B, T)
    block_ids = torch.where(valid, block_ids, torch.full_like(block_ids, -1))
    q = block_ids.view(B, 1, T, 1)
    k = block_ids.view(B, 1, 1, T)
    attn = (k <= q) & (q >= 0) & (k >= 0)
    return attn, position_ids


def clone_past_key_values(pkv):
    """Clone KV-cache. Fast path for tuples and Cache objects; falls back to deepcopy."""
    if pkv is None:
        return None
    # Fast path: legacy tuple format
    if isinstance(pkv, tuple):
        return tuple(
            (k.clone() if k is not None else None, v.clone() if v is not None else None)
            for k, v in pkv
        )
    # Fast path: transformers Cache objects (DynamicCache, etc.)
    if hasattr(pkv, 'key_cache') and hasattr(pkv, 'value_cache'):
        try:
            new_cache = pkv.__class__()
            new_cache.key_cache = [k.clone() for k in pkv.key_cache]
            new_cache.value_cache = [v.clone() for v in pkv.value_cache]
            for attr in ('_seen_tokens', 'seen_tokens'):
                if hasattr(pkv, attr):
                    setattr(new_cache, attr, getattr(pkv, attr))
            return new_cache
        except Exception:
            pass
    # Fallback
    return copy.deepcopy(pkv)


def diffusion_step_block(logits, x_block, mask_block, num_transfer, temperature, remasking):
    """Vectorized diffusion step — no per-sample Python loops."""
    B, L, _ = logits.shape
    if not mask_block.any():
        return x_block
    noisy = add_gumbel_noise(logits, temperature)
    x0 = noisy.argmax(dim=-1)
    if remasking == "low_confidence":
        p = F.softmax(logits, dim=-1)
        conf = p.gather(-1, x0.unsqueeze(-1)).squeeze(-1)
    elif remasking == "random":
        conf = torch.rand((B, L), device=logits.device)
    else:
        raise ValueError(remasking)
    x0 = torch.where(mask_block, x0, x_block)
    conf = conf.masked_fill(~mask_block, float("-inf"))
    k_max = int(num_transfer.max().item())
    if k_max > 0:
        k = min(k_max, L)
        topk_vals, topk_idx = torch.topk(conf, k=k, dim=-1)
        commit = torch.zeros_like(x_block, dtype=torch.bool)
        valid_mask = torch.arange(k, device=x_block.device).view(1, k) < num_transfer.view(B, 1)
        commit.scatter_(1, topk_idx, valid_mask)
        x_block = torch.where(commit, x0, x_block)
    return x_block


@torch.inference_mode()
def generate(model, tokenizer, prompt, steps=128, max_new_tokens=128, block_size=32,
             temperature=0.0, cfg_scale=0.0, remasking="low_confidence", capture_interval=0):
    device = model.device
    mask_id = tokenizer.mask_token_id
    pad_id = tokenizer.pad_token_id
    if pad_id is None:
        pad_id = tokenizer.eos_token_id if tokenizer.eos_token_id is not None else tokenizer.mask_token_id
    if isinstance(prompt, torch.Tensor):
        x = prompt.to(device).long()
    else:
        if isinstance(prompt[0], (list, tuple)):
            max_len = max(len(p) for p in prompt)
            x = torch.full((len(prompt), max_len), pad_id, device=device, dtype=torch.long)
            for i, p in enumerate(prompt):
                x[i, : len(p)] = torch.tensor(p, device=device)
        else:
            x = torch.tensor(prompt, device=device).long()
    if x.dim() == 1:
        x = x.unsqueeze(0)
    B = x.size(0)
    finished = torch.zeros(B, dtype=torch.bool, device=device)
    num_blocks = math.ceil(max_new_tokens / block_size)
    steps_per_block = math.ceil(steps / num_blocks)
    generated = 0
    intermediates = []
    total_step = 0
    while generated < max_new_tokens:
        if finished.all():
            break
        T_prefix = x.size(1)
        offset = T_prefix % block_size
        room = block_size if offset == 0 else block_size - offset
        cur_len = min(room, max_new_tokens - generated)
        if cur_len <= 0:
            break
        attn_pfx, pos_pfx = build_staircase_attention_mask(x, block_size, pad_id)
        out = model(x, attention_mask=attn_pfx, position_ids=pos_pfx, use_cache=True)
        cond_past = out.past_key_values
        if cfg_scale > 0:
            un_x = x.clone()
            un_x[:] = mask_id
            out_un = model(un_x, attention_mask=attn_pfx, position_ids=pos_pfx, use_cache=True)
            uncond_past = out_un.past_key_values
        else:
            uncond_past = None
        block = torch.full((B, cur_len), mask_id, device=device, dtype=torch.long)
        block[finished] = pad_id
        x = torch.cat([x, block], dim=1)
        T_total = x.size(1)
        block_mask = x[:, -cur_len:] == mask_id
        num_transfer = get_num_transfer_tokens(block_mask, steps_per_block)
        eff_steps = num_transfer.size(1)
        full_attn, full_pos = build_staircase_attention_mask(x, block_size, pad_id)
        attn_blk = full_attn[:, :, T_prefix:T_total, :]
        pos_blk = full_pos[:, T_prefix:T_total]
        for t in range(eff_steps):
            x_blk = x[:, T_prefix:T_total]
            m_blk = x_blk == mask_id
            cond_logits = model(
                x_blk, attention_mask=attn_blk, position_ids=pos_blk,
                past_key_values=clone_past_key_values(cond_past), use_cache=False
            ).logits
            logits = cond_logits
            if cfg_scale > 0:
                un_logits = model(
                    x_blk, attention_mask=attn_blk, position_ids=pos_blk,
                    past_key_values=clone_past_key_values(uncond_past), use_cache=False
                ).logits
                logits = un_logits + (cfg_scale + 1.0) * (cond_logits - un_logits)
            x_blk_new = diffusion_step_block(
                logits, x_blk, m_blk, num_transfer[:, t], temperature, remasking
            )
            x[:, T_prefix:T_total] = x_blk_new
            if capture_interval > 0 and total_step % capture_interval == 0:
                intermediates.append(x.clone())
            total_step += 1
            if tokenizer.eos_token_id is not None:
                finished |= (x_blk_new == tokenizer.eos_token_id).any(dim=1)
        generated += cur_len
        if finished.all():
            break
    if capture_interval > 0:
        return x, intermediates
    return x


@torch.inference_mode()
def generate_stream(model, tokenizer, prompt, steps=128, max_new_tokens=128, block_size=32,
                    temperature=0.0, cfg_scale=0.0, remasking="low_confidence", capture_interval=10):
    device = model.device
    mask_id = tokenizer.mask_token_id
    pad_id = tokenizer.pad_token_id
    if pad_id is None:
        pad_id = tokenizer.eos_token_id if tokenizer.eos_token_id is not None else tokenizer.mask_token_id
    if isinstance(prompt, torch.Tensor):
        x = prompt.to(device).long()
    else:
        if isinstance(prompt[0], (list, tuple)):
            max_len = max(len(p) for p in prompt)
            x = torch.full((len(prompt), max_len), pad_id, device=device, dtype=torch.long)
            for i, p in enumerate(prompt):
                x[i, : len(p)] = torch.tensor(p, device=device)
        else:
            x = torch.tensor(prompt, device=device).long()
    if x.dim() == 1:
        x = x.unsqueeze(0)
    B = x.size(0)
    finished = torch.zeros(B, dtype=torch.bool, device=device)
    num_blocks = math.ceil(max_new_tokens / block_size)
    steps_per_block = math.ceil(steps / num_blocks)
    generated = 0
    total_step = 0
    prompt_len = x.size(1)
    while generated < max_new_tokens:
        if finished.all():
            break
        T_prefix = x.size(1)
        offset = T_prefix % block_size
        room = block_size if offset == 0 else block_size - offset
        cur_len = min(room, max_new_tokens - generated)
        if cur_len <= 0:
            break
        attn_pfx, pos_pfx = build_staircase_attention_mask(x, block_size, pad_id)
        out = model(x, attention_mask=attn_pfx, position_ids=pos_pfx, use_cache=True)
        cond_past = out.past_key_values
        if cfg_scale > 0:
            un_x = x.clone()
            un_x[:] = mask_id
            out_un = model(un_x, attention_mask=attn_pfx, position_ids=pos_pfx, use_cache=True)
            uncond_past = out_un.past_key_values
        else:
            uncond_past = None
        block = torch.full((B, cur_len), mask_id, device=device, dtype=torch.long)
        block[finished] = pad_id
        x = torch.cat([x, block], dim=1)
        T_total = x.size(1)
        block_mask = x[:, -cur_len:] == mask_id
        num_transfer = get_num_transfer_tokens(block_mask, steps_per_block)
        eff_steps = num_transfer.size(1)
        full_attn, full_pos = build_staircase_attention_mask(x, block_size, pad_id)
        attn_blk = full_attn[:, :, T_prefix:T_total, :]
        pos_blk = full_pos[:, T_prefix:T_total]
        for t in range(eff_steps):
            x_blk = x[:, T_prefix:T_total]
            m_blk = x_blk == mask_id
            cond_logits = model(
                x_blk, attention_mask=attn_blk, position_ids=pos_blk,
                past_key_values=clone_past_key_values(cond_past), use_cache=False
            ).logits
            logits = cond_logits
            if cfg_scale > 0:
                un_logits = model(
                    x_blk, attention_mask=attn_blk, position_ids=pos_blk,
                    past_key_values=clone_past_key_values(uncond_past), use_cache=False
                ).logits
                logits = un_logits + (cfg_scale + 1.0) * (cond_logits - un_logits)
            x_blk_new = diffusion_step_block(
                logits, x_blk, m_blk, num_transfer[:, t], temperature, remasking
            )
            x[:, T_prefix:T_total] = x_blk_new
            if total_step % capture_interval == 0:
                new_tokens = x[0, prompt_len:prompt_len + max_new_tokens].tolist()
                text = tokenizer.decode(new_tokens, skip_special_tokens=True)
                yield {
                    "type": "intermediate",
                    "step": total_step,
                    "text": text,
                    "total_steps": steps
                }
            total_step += 1
            if tokenizer.eos_token_id is not None:
                finished |= (x_blk_new == tokenizer.eos_token_id).any(dim=1)
            if finished.all():
                break
        generated += cur_len
        if finished.all():
            break
    new_tokens = x[0, prompt_len:prompt_len + max_new_tokens].tolist()
    final_text = tokenizer.decode(new_tokens, skip_special_tokens=True)
    yield {
        "type": "final",
        "text": final_text,
        "total_steps": total_step
    }


# ==========================================================
# ARCHITECTURE ROUTING LAYERS & TRANSLATION ENGINE CODES
# ==========================================================

def load_model():
    global model, tokenizer, device
    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"Initializing {MODEL_NAME} on {device}... (Diffusion Strategy Flag = {IS_DIFFUSION})")
    if IS_DIFFUSION:
        model = AutoModelForMaskedLM.from_pretrained(
            MODEL_NAME,
            torch_dtype=torch.bfloat16,
            trust_remote_code=True
        ).to(device).eval()
    else:
        model = AutoModelForCausalLM.from_pretrained(
            MODEL_NAME,
            torch_dtype=torch.bfloat16,
            trust_remote_code=False
        ).to(device).eval()
    # Compile model for faster inference — ONLY for standard causal models
    if not IS_DIFFUSION:
        try:
            model = torch.compile(model, mode="reduce-overhead", fullgraph=False)
            print("Model compiled with torch.compile.")
        except Exception as e:
            print(f"torch.compile skipped: {e}")
    else:
        print("Diffusion model loaded without torch.compile (custom FX code incompatible with Dynamo).")
    tokenizer = AutoTokenizer.from_pretrained(
        MODEL_NAME,
        trust_remote_code=IS_DIFFUSION
    )
    print("Model compilation completed and loaded into memory workspace.")


@app.route('/health', methods=['GET'])
def health():
    return jsonify({"status": "healthy", "model_loaded": model is not None, "is_diffusion": IS_DIFFUSION})


@app.route('/generate', methods=['POST'])
def generate_text():
    if model is None or tokenizer is None:
        return jsonify({"error": "Model initialization missing"}), 503
    data = request.get_json() or {}
    if 'prompt' not in data:
        return jsonify({"error": "Missing 'prompt' operational field"}), 400
    prompt = data['prompt']
    max_new_tokens = data.get('max_new_tokens', 256)
    temperature = data.get('temperature', 0.0)
    system_prompt = data.get('system_prompt', 'You are an expert real-time translation assistant.')
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": prompt}
    ]
    # enable_thinking=False for ALL routes to prevent Qwen3 from leaking internal monologue
    encoded = tokenizer.apply_chat_template(
        messages,
        add_generation_prompt=True,
        tokenize=True,
        enable_thinking=False
    )
    if IS_DIFFUSION:
        input_ids = torch.tensor([encoded], dtype=torch.long, device=device)
        steps = data.get('steps', 256)
        block_size = data.get('block_size', 32)
        cfg_scale = data.get('cfg_scale', 0.0)
        remasking = data.get('remasking', 'low_confidence')
        output = generate(
            model, tokenizer, input_ids,
            steps=steps, max_new_tokens=max_new_tokens, block_size=block_size,
            temperature=temperature, cfg_scale=cfg_scale, remasking=remasking,
        )
        prompt_len = len(encoded)
        new_tokens = output[0, prompt_len:prompt_len + max_new_tokens].tolist()
        generated_text = tokenizer.decode(new_tokens, skip_special_tokens=True)
    else:
        input_ids = torch.tensor([encoded], dtype=torch.long, device=device)
        output_ids = model.generate(
            input_ids,
            max_new_tokens=max_new_tokens,
            temperature=temperature,
            do_sample=True if temperature > 0 else False,
            pad_token_id=tokenizer.eos_token_id
        )
        generated_ids = output_ids[0, input_ids.shape[-1]:]
        generated_text = tokenizer.decode(generated_ids, skip_special_tokens=True)
    return jsonify({"prompt": prompt, "generated_text": generated_text})


@app.route('/generate_stream', methods=['POST'])
def generate_text_stream():
    if model is None or tokenizer is None:
        return jsonify({"error": "Model workspace offline"}), 503
    data = request.get_json() or {}
    if not data or 'prompt' not in data:
        return jsonify({"error": "Missing 'prompt' operational field"}), 400
    prompt = data['prompt']
    max_new_tokens = data.get('max_new_tokens', 256)
    temperature = data.get('temperature', 0.0)
    system_prompt = data.get('system_prompt', 'You are an expert real-time translation assistant.')
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": prompt}
    ]
    encoded = tokenizer.apply_chat_template(
        messages, add_generation_prompt=True, tokenize=True,
        enable_thinking=False
    )
    if IS_DIFFUSION:
        input_ids = torch.tensor([encoded], dtype=torch.long, device=device)
        steps = data.get('steps', 256)
        block_size = data.get('block_size', 32)
        cfg_scale = data.get('cfg_scale', 0.0)
        remasking = data.get('remasking', 'low_confidence')
        capture_interval = data.get('capture_interval', 10)
        output, intermediates = generate(
            model, tokenizer, input_ids,
            steps=steps, max_new_tokens=max_new_tokens, block_size=block_size,
            temperature=temperature, cfg_scale=cfg_scale, remasking=remasking,
            capture_interval=capture_interval,
        )
        prompt_len = len(encoded)
        intermediate_states = []
        for i, intermediate in enumerate(intermediates):
            new_tokens = intermediate[0, prompt_len:prompt_len + max_new_tokens].tolist()
            text = tokenizer.decode(new_tokens, skip_special_tokens=True)
            intermediate_states.append({"step": i * capture_interval, "text": text})
        new_tokens = output[0, prompt_len:prompt_len + max_new_tokens].tolist()
        generated_text = tokenizer.decode(new_tokens, skip_special_tokens=True)
        return jsonify({"prompt": prompt, "generated_text": generated_text, "intermediate_states": intermediate_states})
    else:
        input_ids = torch.tensor([encoded], dtype=torch.long, device=device)
        output_ids = model.generate(
            input_ids, max_new_tokens=max_new_tokens, temperature=temperature,
            do_sample=True if temperature > 0 else False,
            pad_token_id=tokenizer.eos_token_id
        )
        generated_ids = output_ids[0, input_ids.shape[-1]:]
        generated_text = tokenizer.decode(generated_ids, skip_special_tokens=True)
        return jsonify({"prompt": prompt, "generated_text": generated_text, "intermediate_states": []})


@app.route('/generate_sse', methods=['POST'])
def generate_text_sse():
    if model is None or tokenizer is None:
        return jsonify({"error": "Model workspace offline"}), 503
    data = request.get_json() or {}
    if not data or 'prompt' not in data:
        return jsonify({"error": "Missing 'prompt' operational field"}), 400
    prompt = data['prompt']
    max_new_tokens = data.get('max_new_tokens', 256)
    temperature = data.get('temperature', 0.0)
    system_prompt = data.get('system_prompt', 'You are an expert real-time translation assistant.')
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": prompt}
    ]
    encoded = tokenizer.apply_chat_template(
        messages, add_generation_prompt=True, tokenize=True,
        enable_thinking=False
    )
    input_ids = torch.tensor([encoded], dtype=torch.long, device=device)

    def stream():
        if IS_DIFFUSION:
            steps = data.get('steps', 256)
            block_size = data.get('block_size', 32)
            cfg_scale = data.get('cfg_scale', 0.0)
            remasking = data.get('remasking', 'low_confidence')
            capture_interval = data.get('capture_interval', 10)
            for state in generate_stream(
                model, tokenizer, input_ids,
                steps=steps, max_new_tokens=max_new_tokens, block_size=block_size,
                temperature=temperature, cfg_scale=cfg_scale, remasking=remasking,
                capture_interval=capture_interval
            ):
                yield f"data: {json.dumps(state)}\n\n"
        else:
            streamer = TextIteratorStreamer(
                tokenizer, skip_prompt=True, skip_special_tokens=True
            )
            generation_kwargs = dict(
                input_ids=input_ids,
                streamer=streamer,
                max_new_tokens=max_new_tokens,
                temperature=temperature,
                do_sample=True if temperature > 0 else False,
                pad_token_id=tokenizer.eos_token_id,
            )

            def _generate():
                with torch.inference_mode():
                    model.generate(**generation_kwargs)

            thread = Thread(target=_generate)
            thread.start()

            accumulated = []
            for text in streamer:
                if not text:  # skip empty chunks
                    continue
                accumulated.append(text)
                current = "".join(accumulated)
                yield f"data: {json.dumps({'type': 'intermediate', 'text': current})}\n\n"

            full_text = "".join(accumulated)
            yield f"data: {json.dumps({'type': 'final', 'text': full_text, 'total_steps': 1})}\n\n"

    return Response(
        stream(), mimetype='text/event-stream',
        headers={'Cache-Control': 'no-cache', 'X-Accel-Buffering': 'no'}
    )


@app.route('/')
def index():
    return Response(HTML_UI, mimetype='text/html')

if __name__ == '__main__':
    load_model()
    app.run(host='0.0.0.0', port=int(os.getenv('PORT', 7860)))