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 = """ LLM API Tester

LLM API Tester

/generate

Route

Prompt

max_new_tokens

temperature

steps (diffusion)

block_size

cfg_scale

remasking

capture_interval

Response

""" # 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)))