Upload app.py

app.py

@@ -1,467 +1,522 @@
-import sys
-import types
-import os
-import math
-import copy
-import json
-import torch
-import torch.nn.functional as F
-from flask import Flask, request, jsonify, Response
-from transformers import AutoTokenizer, AutoModelForMaskedLM, AutoModelForCausalLM
-
-# 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):
-    if temperature == 0:
-        return logits
-    logits = logits.to(torch.float64)
-    noise = torch.rand_like(logits, dtype=torch.float64)
-    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 diffusion_step_block(logits, x_block, mask_block, num_transfer, temperature, remasking):
-    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)
-    neg_inf = torch.full_like(conf, -float("inf"))
-    conf = torch.where(mask_block, conf, neg_inf)
-    commit = torch.zeros_like(x_block, dtype=torch.bool)
-    for i in range(B):
-        k = int(num_transfer[i].item())
-        if k > 0:
-            valid = (conf[i] > -float("inf")).sum().item()
-            k = min(k, valid)
-            _, idx = torch.topk(conf[i], k)
-            commit[i, idx] = True
-    out = x_block.clone()
-    out[commit] = x0[commit]
-    return out
-
-@torch.no_grad()
-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=copy.deepcopy(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=copy.deepcopy(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.no_grad()
-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=copy.deepcopy(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=copy.deepcopy(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,
-            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()
-        
-    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}
-    ]
-    
-    encoded = tokenizer.apply_chat_template(
-        messages,
-        add_generation_prompt=True,
-        tokenize=True,
-        enable_thinking=False if IS_DIFFUSION else None
-    )
-    
-    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:
-        # High-Speed Autoregressive Optimization Matrix
-        input_ids = torch.tensor([encoded], dtype=torch.long, device=device)
-        with torch.no_grad():
-            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 else None)
-    
-    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)
-        with torch.no_grad():
-            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 if IS_DIFFUSION else None)
-    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:
-            with torch.no_grad():
-                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]:]
-            final_text = tokenizer.decode(generated_ids, skip_special_tokens=True)
-            yield f"data: {json.dumps({'type': 'final', 'text': final_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 {
-        "status": "healthy",
-        "message": f"Multi-architecture API Router up and running. Target: {'Diffusion Framework' if IS_DIFFUSION else 'Causal Baseline Model'}", "model_loaded": MODEL_NAME}, 200
-
-if __name__ == '__main__':
-    load_model()
-    app.run(host='0.0.0.0', port=int(os.getenv('PORT', 7860)))
+import sys
+import types
+import os
+import math
+import json
+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
+
+# 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):
+    """Shallow clone of prefix KV-cache to avoid expensive copy.deepcopy."""
+    if pkv is None:
+        return None
+    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
+        )
+    return 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]
+        # Clone prefix caches once per block instead of every step
+        cond_past_clone = clone_past_key_values(cond_past)
+        uncond_past_clone = clone_past_key_values(uncond_past) if uncond_past is not None else None
+        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=cond_past_clone, 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=uncond_past_clone, 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]
+        # Clone prefix caches once per block
+        cond_past_clone = clone_past_key_values(cond_past)
+        uncond_past_clone = clone_past_key_values(uncond_past) if uncond_past is not None else None
+        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=cond_past_clone, 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=uncond_past_clone, 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 (reduce-overhead is safest for generative loops)
+    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}")
+    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 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 else None
+    )
+    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 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 else None
+    )
+    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 if IS_DIFFUSION else None
+    )
+    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()
+            for text in streamer:
+                yield f"data: {json.dumps({'type': 'token', 'text': text})}\n\n"
+            yield f"data: {json.dumps({'type': 'final', '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 {
+        "status": "healthy",
+        "message": f"Multi-architecture API Router up and running. Target: {'Diffusion Framework' if IS_DIFFUSION else 'Causal Baseline Model'}",
+        "model_loaded": MODEL_NAME
+    }, 200
+
+
+if __name__ == '__main__':
+    load_model()
+    app.run(host='0.0.0.0', port=int(os.getenv('PORT', 7860)))