Initial deployment: diffusion-chatbot

.gitignore

@@ -0,0 +1,7 @@
+__pycache__/
+*.pyc
+.env
+.venv
+venv/
+.git
+.ds_store

Dockerfile

@@ -0,0 +1,22 @@
+FROM pytorch/pytorch:2.6.0-cuda12.4-cudnn9-devel
+
+WORKDIR /app
+
+RUN apt-get update && apt-get install -y \
+    git \
+    && rm -rf /var/lib/apt/lists/*
+
+RUN git clone --recurse-submodules https://github.com/ZHZisZZ/dllm.git /tmp/dllm && \
+    pip install --no-cache-dir -e /tmp/dllm && \
+    pip install --no-cache-dir -e /tmp/dllm/lm-evaluation-harness
+
+RUN pip install --no-cache-dir flask
+
+COPY app.py .
+
+ENV MODEL_NAME=dllm-hub/Qwen3-0.6B-diffusion-bd3lm-v0.1
+ENV PORT=7860
+
+EXPOSE 7860
+
+CMD ["python", "app.py"]

README.md

@@ -0,0 +1,68 @@
+---
+title: Diffusion Chatbot
+emoji: 🤖
+colorFrom: blue
+colorTo: purple
+sdk: docker
+pinned: false
+license: apache-2.0
+---
+
+# Diffusion Chatbot
+
+Flask server hosting the Qwen3-0.6B-diffusion-bd3lm-v0.1 model with real-time streaming inference.
+
+## Features
+
+- **Real-time streaming**: Watch the diffusion model generate text step-by-step
+- **Three endpoints**: Simple generation, batch intermediate states, and real-time SSE streaming
+- **GPU support**: Automatically uses GPU if available, falls back to CPU
+
+## API Endpoints
+
+### Health Check
+```
+GET /health
+```
+
+### Generate Text
+```
+POST /generate
+Content-Type: application/json
+
+{
+  "prompt": "Your question here",
+  "max_new_tokens": 256
+}
+```
+
+### Generate with Real-time Streaming (SSE)
+```
+POST /generate_sse
+Content-Type: application/json
+
+{
+  "prompt": "Your question here",
+  "max_new_tokens": 100,
+  "capture_interval": 10
+}
+```
+
+## Example Usage
+
+```bash
+curl -X POST https://YOUR_USERNAME-diffusion-chatbot.hf.space/generate \
+  -H "Content-Type: application/json" \
+  -d '{"prompt": "Hello, how are you?", "max_new_tokens": 50}'
+```
+
+## Technical Details
+
+- **Model**: [dllm-hub/Qwen3-0.6B-diffusion-bd3lm-v0.1](https://huggingface.co/dllm-hub/Qwen3-0.6B-diffusion-bd3lm-v0.1)
+- **Framework**: Flask + PyTorch
+- **Diffusion Method**: Block Diffusion Language Model (BD3LM)
+
+## Environment Variables
+
+- `MODEL_NAME`: HuggingFace model name (default: dllm-hub/Qwen3-0.6B-diffusion-bd3lm-v0.1)
+- `PORT`: Server port (default: 7860 for HF Spaces)

app.py

@@ -0,0 +1,564 @@
+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',
+        }
+    )
+
+
+if __name__ == '__main__':
+    load_model()
+    app.run(host='0.0.0.0', port=int(os.getenv('PORT', 5000)))

requirements.txt

@@ -0,0 +1,6 @@
+flask
+torch
+transformers
+accelerate
+numpy<2.0
+git+https://github.com/ZHZisZZ/dllm.git