Upload app.py

app.py

@@ -1,592 +1,466 @@
-import sys
-import types
-
-# 1. Force-inject the child namespaces to prevent circular initialization locks
-try:
-    import dllm.utils
-    import dllm.pipelines
-    import dllm.data
-    import dllm.core
-except ImportError:
-    pass
-
-# 2. Re-register the primary package mock cleanly into Python's runtime cache
-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
-
-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
-
-app = Flask(__name__)
-
-model = None
-tokenizer = None
-device = None
-
-
-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)
-            if finished.all():
-                break
-
-        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
-    }
-
-
-def load_model():
-    global model, tokenizer, device
-    
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    model_name = os.getenv("MODEL_NAME", "dllm-hub/Qwen3-0.6B-diffusion-bd3lm-v0.1")
-    
-    print(f"Loading model {model_name} on {device}...")
-    model = AutoModelForMaskedLM.from_pretrained(
-        model_name, 
-        dtype=torch.bfloat16, 
-        trust_remote_code=True
-    ).to(device).eval()
-    
-    tokenizer = AutoTokenizer.from_pretrained(
-        model_name, 
-        trust_remote_code=True
-    )
-    print("Model loaded successfully!")
-
-
-@app.route('/health', methods=['GET'])
-def health():
-    return jsonify({"status": "healthy", "model_loaded": model is not None})
-
-
-@app.route('/generate', methods=['POST'])
-def generate_text():
-    if model is None or tokenizer is None:
-        return jsonify({"error": "Model not loaded"}), 503
-    
-    data = request.get_json()
-    
-    if not data or 'prompt' not in data:
-        return jsonify({"error": "Missing 'prompt' field"}), 400
-    
-    prompt = data['prompt']
-    steps = data.get('steps', 256)
-    max_new_tokens = data.get('max_new_tokens', 256)
-    block_size = data.get('block_size', 32)
-    temperature = data.get('temperature', 0.0)
-    cfg_scale = data.get('cfg_scale', 0.0)
-    remasking = data.get('remasking', 'low_confidence')
-    system_prompt = data.get('system_prompt', 'You are a helpful AI 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)
-    
-    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)
-    
-    return jsonify({
-        "prompt": prompt,
-        "generated_text": generated_text,
-        "parameters": {
-            "steps": steps,
-            "max_new_tokens": max_new_tokens,
-            "block_size": block_size,
-            "temperature": temperature,
-            "cfg_scale": cfg_scale,
-            "remasking": remasking
-        }
-    })
-
-
-@app.route('/generate_stream', methods=['POST'])
-def generate_text_stream():
-    if model is None or tokenizer is None:
-        return jsonify({"error": "Model not loaded"}), 503
-    
-    data = request.get_json()
-    
-    if not data or 'prompt' not in data:
-        return jsonify({"error": "Missing 'prompt' field"}), 400
-    
-    prompt = data['prompt']
-    steps = data.get('steps', 256)
-    max_new_tokens = data.get('max_new_tokens', 256)
-    block_size = data.get('block_size', 32)
-    temperature = data.get('temperature', 0.0)
-    cfg_scale = data.get('cfg_scale', 0.0)
-    remasking = data.get('remasking', 'low_confidence')
-    system_prompt = data.get('system_prompt', 'You are a helpful AI assistant.')
-    capture_interval = data.get('capture_interval', 10)
-    
-    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)
-    
-    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,
-        "parameters": {
-            "steps": steps,
-            "max_new_tokens": max_new_tokens,
-            "block_size": block_size,
-            "temperature": temperature,
-            "cfg_scale": cfg_scale,
-            "remasking": remasking,
-            "capture_interval": capture_interval
-        }
-    })
-
-
-@app.route('/generate_sse', methods=['POST'])
-def generate_text_sse():
-    if model is None or tokenizer is None:
-        return jsonify({"error": "Model not loaded"}), 503
-    
-    data = request.get_json()
-    
-    if not data or 'prompt' not in data:
-        return jsonify({"error": "Missing 'prompt' field"}), 400
-    
-    prompt = data['prompt']
-    steps = data.get('steps', 256)
-    max_new_tokens = data.get('max_new_tokens', 256)
-    block_size = data.get('block_size', 32)
-    temperature = data.get('temperature', 0.0)
-    cfg_scale = data.get('cfg_scale', 0.0)
-    remasking = data.get('remasking', 'low_confidence')
-    system_prompt = data.get('system_prompt', 'You are a helpful AI assistant.')
-    capture_interval = data.get('capture_interval', 10)
-    
-    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():
-        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"
-    
-    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"dLLM API Server up and running. Model: {os.getenv('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 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, (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)
+        if finished.all():
+            break
+        generated += cur_len
+    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, (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 '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 '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)))