diffusion-chatbot

Sleeping

App Files Files Community

Andhs commited on Jul 3

Commit

b8d11d8

verified ·

1 Parent(s): df717ea

Upload app.py

Browse files

Files changed (1) hide show

app.py +40 -338

app.py CHANGED Viewed

@@ -10,337 +10,6 @@ from flask import Flask, request, jsonify, Response
 from transformers import AutoTokenizer, AutoModelForMaskedLM, AutoModelForCausalLM, TextIteratorStreamer
 from threading import Thread
-HTML_UI = """
-<!DOCTYPE html>
-<html lang="en">
-<head>
-    <meta charset="UTF-8">
-    <meta name="viewport" content="width=device-width, initial-scale=1.0">
-    <title>LLM API Tester</title>
-    <style>
-        * { box-sizing: border-box; }
-        body {
-            font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, sans-serif;
-            max-width: 900px;
-            margin: 0 auto;
-            padding: 20px;
-            background: #f5f5f5;
-            color: #333;
-        }
-        h1 { margin-top: 0; font-size: 1.5rem; }
-        .card {
-            background: #fff;
-            border-radius: 8px;
-            padding: 20px;
-            margin-bottom: 16px;
-            box-shadow: 0 1px 3px rgba(0,0,0,0.1);
-        }
-        label {
-            display: block;
-            font-weight: 600;
-            margin-bottom: 6px;
-            font-size: 0.9rem;
-        }
-        textarea, input, select {
-            width: 100%;
-            padding: 10px;
-            border: 1px solid #ddd;
-            border-radius: 6px;
-            font-size: 0.95rem;
-            font-family: inherit;
-        }
-        textarea { resize: vertical; min-height: 80px; }
-        .row {
-            display: grid;
-            grid-template-columns: repeat(auto-fit, minmax(140px, 1fr));
-            gap: 12px;
-            margin-bottom: 12px;
-        }
-        .field { margin-bottom: 12px; }
-        .field.inline {
-            display: flex;
-            align-items: center;
-            gap: 8px;
-        }
-        .field.inline label { margin: 0; }
-        .field.inline input, .field.inline select {
-            width: auto;
-            flex: 1;
-        }
-        button {
-            background: #2563eb;
-            color: #fff;
-            border: none;
-            padding: 10px 20px;
-            border-radius: 6px;
-            font-size: 1rem;
-            cursor: pointer;
-            font-weight: 600;
-        }
-        button:hover { background: #1d4ed8; }
-        button:disabled { background: #93c5fd; cursor: not-allowed; }
-        .output {
-            background: #1e1e1e;
-            color: #e4e4e4;
-            padding: 16px;
-            border-radius: 6px;
-            font-family: ui-monospace, SFMono-Regular, "SF Mono", Menlo, Consolas, monospace;
-            font-size: 0.9rem;
-            white-space: pre-wrap;
-            word-break: break-word;
-            min-height: 120px;
-            max-height: 500px;
-            overflow-y: auto;
-        }
-        .output:empty::before {
-            content: "Response will appear here...";
-            color: #666;
-        }
-        .status {
-            font-size: 0.85rem;
-            color: #666;
-            margin-top: 8px;
-        }
-        .error { color: #dc2626; }
-        .success { color: #16a34a; }
-        .route-badge {
-            display: inline-block;
-            background: #e5e7eb;
-            padding: 2px 8px;
-            border-radius: 4px;
-            font-size: 0.8rem;
-            font-weight: 600;
-            margin-bottom: 8px;
-        }
-    </style>
-</head>
-<body>
-    <h1>LLM API Tester</h1>
-    <div class="card">
-        <div class="route-badge" id="routeBadge">/generate</div>
-        <div class="field">
-            <label for="route">Route</label>
-            <select id="route">
-                <option value="/generate">/generate (sync JSON)</option>
-                <option value="/generate_stream">/generate_stream (sync JSON + intermediates)</option>
-                <option value="/generate_sse">/generate_sse (streaming SSE)</option>
-            </select>
-        </div>
-        <div class="field">
-            <label for="prompt">Prompt</label>
-            <textarea id="prompt" placeholder="Enter your prompt here...">Write a short poem about stars</textarea>
-        </div>
-        <div class="row">
-            <div class="field">
-                <label for="max_new_tokens">max_new_tokens</label>
-                <input type="number" id="max_new_tokens" value="150" min="1" max="2048">
-            </div>
-            <div class="field">
-                <label for="temperature">temperature</label>
-                <input type="number" id="temperature" value="0.0" min="0" max="2" step="0.1">
-            </div>
-            <div class="field">
-                <label for="steps">steps (diffusion)</label>
-                <input type="number" id="steps" value="256" min="1">
-            </div>
-            <div class="field">
-                <label for="block_size">block_size</label>
-                <input type="number" id="block_size" value="32" min="1">
-            </div>
-        </div>
-        <div class="row">
-            <div class="field">
-                <label for="cfg_scale">cfg_scale</label>
-                <input type="number" id="cfg_scale" value="0.0" min="0" step="0.1">
-            </div>
-            <div class="field">
-                <label for="remasking">remasking</label>
-                <select id="remasking">
-                    <option value="low_confidence">low_confidence</option>
-                    <option value="random">random</option>
-                </select>
-            </div>
-            <div class="field">
-                <label for="capture_interval">capture_interval</label>
-                <input type="number" id="capture_interval" value="10" min="1">
-            </div>
-        </div>
-        <button id="sendBtn">Send Request</button>
-        <div class="status" id="status"></div>
-    </div>
-    <div class="card">
-        <label>Response</label>
-        <div class="output" id="output"></div>
-    </div>
-    <script>
-        const $ = id => document.getElementById(id);
-        const routeSelect = $('route');
-        const routeBadge = $('routeBadge');
-        const sendBtn = $('sendBtn');
-        const output = $('output');
-        const status = $('status');
-        routeSelect.addEventListener('change', () => {
-            routeBadge.textContent = routeSelect.value;
-        });
-        function setStatus(msg, isError = false) {
-            status.textContent = msg;
-            status.className = 'status ' + (isError ? 'error' : 'success');
-        }
-        function appendOutput(text, clear = false) {
-            if (clear) output.textContent = '';
-            output.textContent += text;
-            output.scrollTop = output.scrollHeight;
-        }
-        function getPayload() {
-            return {
-                prompt: $('prompt').value,
-                max_new_tokens: parseInt($('max_new_tokens').value),
-                temperature: parseFloat($('temperature').value),
-                steps: parseInt($('steps').value),
-                block_size: parseInt($('block_size').value),
-                cfg_scale: parseFloat($('cfg_scale').value),
-                remasking: $('remasking').value,
-                capture_interval: parseInt($('capture_interval').value)
-            };
-        }
-        async function handleGenerate() {
-            const payload = getPayload();
-            // Remove diffusion-only fields for non-diffusion if needed, but server ignores extras
-            const t0 = performance.now();
-            const res = await fetch('/generate', {
-                method: 'POST',
-                headers: { 'Content-Type': 'application/json' },
-                body: JSON.stringify(payload)
-            });
-            const data = await res.json();
-            const ms = Math.round(performance.now() - t0);
-            if (res.ok) {
-                appendOutput(`[${ms}ms]\n${data.generated_text || JSON.stringify(data, null, 2)}\n\n`, true);
-                setStatus(`OK — ${ms}ms`);
-            } else {
-                appendOutput(`Error ${res.status}:\n${JSON.stringify(data, null, 2)}\n\n`, true);
-                setStatus(`HTTP ${res.status}`, true);
-            }
-        }
-        async function handleGenerateStream() {
-            const payload = getPayload();
-            const t0 = performance.now();
-            const res = await fetch('/generate_stream', {
-                method: 'POST',
-                headers: { 'Content-Type': 'application/json' },
-                body: JSON.stringify(payload)
-            });
-            const data = await res.json();
-            const ms = Math.round(performance.now() - t0);
-            if (res.ok) {
-                let text = `[${ms}ms]\nGenerated text:\n${data.generated_text}\n\n`;
-                if (data.intermediate_states && data.intermediate_states.length) {
-                    text += `Intermediate states (${data.intermediate_states.length}):\n`;
-                    data.intermediate_states.forEach((s, i) => {
-                        text += `  Step ${s.step}: ${s.text.substring(0, 120).replace(/\n/g, ' ')}...\n`;
-                    });
-                }
-                appendOutput(text + '\n', true);
-                setStatus(`OK — ${ms}ms, ${data.intermediate_states?.length || 0} intermediates`);
-            } else {
-                appendOutput(`Error ${res.status}:\n${JSON.stringify(data, null, 2)}\n\n`, true);
-                setStatus(`HTTP ${res.status}`, true);
-            }
-        }
-        async function handleGenerateSSE() {
-            const payload = getPayload();
-            const t0 = performance.now();
-            appendOutput('', true);
-            setStatus('Connecting SSE...');
-            const res = await fetch('/generate_sse', {
-                method: 'POST',
-                headers: { 'Content-Type': 'application/json' },
-                body: JSON.stringify(payload)
-            });
-            if (!res.ok) {
-                const data = await res.json().catch(() => ({}));
-                appendOutput(`Error ${res.status}:\n${JSON.stringify(data, null, 2)}`, true);
-                setStatus(`HTTP ${res.status}`, true);
-                return;
-            }
-            const reader = res.body.getReader();
-            const decoder = new TextDecoder();
-            let buffer = '';
-            let finalText = '';
-            let eventCount = 0;
-            while (true) {
-                const { done, value } = await reader.read();
-                if (done) break;
-                buffer += decoder.decode(value, { stream: true });
-                const lines = buffer.split('\n');
-                buffer = lines.pop(); // keep incomplete line in buffer
-                for (const line of lines) {
-                    if (!line.startsWith('data: ')) continue;
-                    const jsonStr = line.slice(6).trim();
-                    if (!jsonStr) continue;
-                    try {
-                        const event = JSON.parse(jsonStr);
-                        eventCount++;
-                        if (event.type === 'final') {
-                            finalText = event.text;
-                            const ms = Math.round(performance.now() - t0);
-                            appendOutput(`[${ms}ms | ${eventCount} events]\n${finalText}\n`, true);
-                            setStatus(`Done — ${ms}ms, ${eventCount} events, ${event.total_steps || '?'} steps`);
-                        } else if (event.type === 'intermediate' || event.type === 'token') {
-                            // Live update: overwrite with latest accumulated text
-                            appendOutput(`${event.text}`, true);
-                            setStatus(`Streaming... (${eventCount} events)`);
-                        }
-                    } catch (e) {
-                        // ignore malformed lines
-                    }
-                }
-            }
-            if (!finalText && eventCount === 0) {
-                setStatus('Stream ended with no events', true);
-            }
-        }
-        sendBtn.addEventListener('click', async () => {
-            sendBtn.disabled = true;
-            setStatus('Sending...');
-            try {
-                const route = routeSelect.value;
-                if (route === '/generate') await handleGenerate();
-                else if (route === '/generate_stream') await handleGenerateStream();
-                else if (route === '/generate_sse') await handleGenerateSSE();
-            } catch (err) {
-                appendOutput(`Network/JS Error:\n${err.message}\n\n`, true);
-                setStatus(err.message, true);
-            } finally {
-                sendBtn.disabled = false;
-            }
-        });
-    </script>
-</body>
-</html>
-"""
 # 1. Environment Parsing & Architecture Strategy Mapping
 MODEL_NAME = os.getenv("MODEL_NAME", "dllm-hub/Qwen3-0.6B-diffusion-bd3lm-v0.1")
 IS_DIFFUSION = "diffusion" in MODEL_NAME.lower()
@@ -410,13 +79,11 @@ 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__()
@@ -428,7 +95,6 @@ def clone_past_key_values(pkv):
             return new_cache
         except Exception:
             pass
-    # Fallback
     return copy.deepcopy(pkv)
@@ -666,7 +332,6 @@ def load_model():
             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)
@@ -679,7 +344,7 @@ def load_model():
         MODEL_NAME,
         trust_remote_code=IS_DIFFUSION
     )
-    print("Model compilation completed and loaded into memory workspace.")
 @app.route('/health', methods=['GET'])
@@ -702,7 +367,6 @@ def generate_text():
         {"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,
@@ -847,7 +511,7 @@ def generate_text_sse():
             accumulated = []
             for text in streamer:
-                if not text:  # skip empty chunks
                     continue
                 accumulated.append(text)
                 current = "".join(accumulated)
@@ -862,10 +526,48 @@ def generate_text_sse():
     )
 @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)))

 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()
     """Clone KV-cache. Fast path for tuples and Cache objects; falls back to 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
         )
     if hasattr(pkv, 'key_cache') and hasattr(pkv, 'value_cache'):
         try:
             new_cache = pkv.__class__()
             return new_cache
         except Exception:
             pass
     return copy.deepcopy(pkv)
             torch_dtype=torch.bfloat16,
             trust_remote_code=False
         ).to(device).eval()
     if not IS_DIFFUSION:
         try:
             model = torch.compile(model, mode="reduce-overhead", fullgraph=False)
         MODEL_NAME,
         trust_remote_code=IS_DIFFUSION
     )
+    print("Model loaded into memory workspace.")
 @app.route('/health', methods=['GET'])
         {"role": "system", "content": system_prompt},
         {"role": "user", "content": prompt}
     ]
     encoded = tokenizer.apply_chat_template(
         messages,
         add_generation_prompt=True,
             accumulated = []
             for text in streamer:
+                if not text:
                     continue
                 accumulated.append(text)
                 current = "".join(accumulated)
     )
+HTML_UI = """<!DOCTYPE html>
+<html lang="en">
+<head>
+<meta charset="UTF-8">
+<meta name="viewport" content="width=device-width, initial-scale=1.0">
+<title>LLM API Tester</title>
+<style>
+*{box-sizing:border-box}body{font-family:-apple-system,BlinkMacSystemFont,"Segoe UI",Roboto,sans-serif;max-width:900px;margin:0 auto;padding:20px;background:#f5f5f5;color:#333}h1{margin-top:0;font-size:1.5rem}.card{background:#fff;border-radius:8px;padding:20px;margin-bottom:16px;box-shadow:0 1px 3px rgba(0,0,0,.1)}label{display:block;font-weight:600;margin-bottom:6px;font-size:.9rem}textarea,input,select{width:100%;padding:10px;border:1px solid #ddd;border-radius:6px;font-size:.95rem;font-family:inherit}textarea{resize:vertical;min-height:80px}.row{display:grid;grid-template-columns:repeat(auto-fit,minmax(140px,1fr));gap:12px;margin-bottom:12px}.field{margin-bottom:12px}button{background:#2563eb;color:#fff;border:none;padding:10px 20px;border-radius:6px;font-size:1rem;cursor:pointer;font-weight:600}button:hover{background:#1d4ed8}button:disabled{background:#93c5fd;cursor:not-allowed}.output{background:#1e1e1e;color:#e4e4e4;padding:16px;border-radius:6px;font-family:ui-monospace,SFMono-Regular,"SF Mono",Menlo,Consolas,monospace;font-size:.9rem;white-space:pre-wrap;word-break:break-word;min-height:120px;max-height:500px;overflow-y:auto}.output:empty::before{content:"Response will appear here...";color:#666}.status{font-size:.85rem;color:#666;margin-top:8px}.error{color:#dc2626}.success{color:#16a34a}.route-badge{display:inline-block;background:#e5e7eb;padding:2px 8px;border-radius:4px;font-size:.8rem;font-weight:600;margin-bottom:8px}
+</style>
+</head>
+<body>
+<h1>LLM API Tester</h1>
+<div class="card">
+<div class="route-badge" id="routeBadge">/generate</div>
+<div class="field"><label for="route">Route</label><select id="route"><option value="/generate">/generate (sync JSON)</option><option value="/generate_stream">/generate_stream (sync JSON + intermediates)</option><option value="/generate_sse">/generate_sse (streaming SSE)</option></select></div>
+<div class="field"><label for="prompt">Prompt</label><textarea id="prompt" placeholder="Enter your prompt here...">Write a short poem about stars</textarea></div>
+<div class="row"><div class="field"><label for="max_new_tokens">max_new_tokens</label><input type="number" id="max_new_tokens" value="150" min="1" max="2048"></div><div class="field"><label for="temperature">temperature</label><input type="number" id="temperature" value="0.0" min="0" max="2" step="0.1"></div><div class="field"><label for="steps">steps (diffusion)</label><input type="number" id="steps" value="256" min="1"></div><div class="field"><label for="block_size">block_size</label><input type="number" id="block_size" value="32" min="1"></div></div>
+<div class="row"><div class="field"><label for="cfg_scale">cfg_scale</label><input type="number" id="cfg_scale" value="0.0" min="0" step="0.1"></div><div class="field"><label for="remasking">remasking</label><select id="remasking"><option value="low_confidence">low_confidence</option><option value="random">random</option></select></div><div class="field"><label for="capture_interval">capture_interval</label><input type="number" id="capture_interval" value="10" min="1"></div></div>
+<button id="sendBtn">Send Request</button>
+<div class="status" id="status"></div>
+</div>
+<div class="card"><label>Response</label><div class="output" id="output"></div></div>
+<script>
+const $=id=>document.getElementById(id);
+const routeSelect=$('route');const routeBadge=$('routeBadge');const sendBtn=$('sendBtn');const output=$('output');const status=$('status');
+routeSelect.addEventListener('change',()=>{routeBadge.textContent=routeSelect.value;});
+function setStatus(msg,isError=false){status.textContent=msg;status.className='status '+(isError?'error':'success');}
+function appendOutput(text,clear=false){if(clear)output.textContent='';output.textContent+=text;output.scrollTop=output.scrollHeight;}
+function getPayload(){return{prompt:$('prompt').value,max_new_tokens:parseInt($('max_new_tokens').value),temperature:parseFloat($('temperature').value),steps:parseInt($('steps').value),block_size:parseInt($('block_size').value),cfg_scale:parseFloat($('cfg_scale').value),remasking:$('remasking').value,capture_interval:parseInt($('capture_interval').value)};}
+async function handleGenerate(){const payload=getPayload();const t0=performance.now();const res=await fetch('/generate',{method:'POST',headers:{'Content-Type':'application/json'},body:JSON.stringify(payload)});const data=await res.json();const ms=Math.round(performance.now()-t0);if(res.ok){appendOutput(`[${ms}ms]\n${data.generated_text||JSON.stringify(data,null,2)}\n\n`,true);setStatus(`OK — ${ms}ms`);}else{appendOutput(`Error ${res.status}:\n${JSON.stringify(data,null,2)}\n\n`,true);setStatus(`HTTP ${res.status}`,true);}}
+async function handleGenerateStream(){const payload=getPayload();const t0=performance.now();const res=await fetch('/generate_stream',{method:'POST',headers:{'Content-Type':'application/json'},body:JSON.stringify(payload)});const data=await res.json();const ms=Math.round(performance.now()-t0);if(res.ok){let text=`[${ms}ms]\nGenerated text:\n${data.generated_text}\n\n`;if(data.intermediate_states&&data.intermediate_states.length){text+=`Intermediate states (${data.intermediate_states.length}):\n`;data.intermediate_states.forEach((s,i)=>{text+=`  Step ${s.step}: ${s.text.substring(0,120).replace(/\n/g,' ')}...\n`;});}appendOutput(text+'\n',true);setStatus(`OK — ${ms}ms, ${data.intermediate_states?.length||0} intermediates`);}else{appendOutput(`Error ${res.status}:\n${JSON.stringify(data,null,2)}\n\n`,true);setStatus(`HTTP ${res.status}`,true);}}
+async function handleGenerateSSE(){const payload=getPayload();const t0=performance.now();appendOutput('',true);setStatus('Connecting SSE...');const res=await fetch('/generate_sse',{method:'POST',headers:{'Content-Type':'application/json'},body:JSON.stringify(payload)});if(!res.ok){const data=await res.json().catch(()=>({}));appendOutput(`Error ${res.status}:\n${JSON.stringify(data,null,2)}`,true);setStatus(`HTTP ${res.status}`,true);return;}const reader=res.body.getReader();const decoder=new TextDecoder();let buffer='';let finalText='';let eventCount=0;while(true){const {done,value}=await reader.read();if(done)break;buffer+=decoder.decode(value,{stream:true});const lines=buffer.split('\n');buffer=lines.pop();for(const line of lines){if(!line.startsWith('data: '))continue;const jsonStr=line.slice(6).trim();if(!jsonStr)continue;try{const event=JSON.parse(jsonStr);eventCount++;if(event.type==='final'){finalText=event.text;const ms=Math.round(performance.now()-t0);appendOutput(`[${ms}ms | ${eventCount} events]\n${finalText}\n`,true);setStatus(`Done — ${ms}ms, ${eventCount} events, ${event.total_steps||'?'} steps`);}else if(event.type==='intermediate'||event.type==='token'){appendOutput(`${event.text}`,true);setStatus(`Streaming... (${eventCount} events)`);}}catch(e){}}}if(!finalText&&eventCount===0){setStatus('Stream ended with no events',true);}}
+sendBtn.addEventListener('click',async()=>{sendBtn.disabled=true;setStatus('Sending...');try{const route=routeSelect.value;if(route==='/generate')await handleGenerate();else if(route==='/generate_stream')await handleGenerateStream();else if(route==='/generate_sse')await handleGenerateSSE();}catch(err){appendOutput(`Network/JS Error:\n${err.message}\n\n`,true);setStatus(err.message,true);}finally{sendBtn.disabled=false;}});
+</script>
+</body>
+</html>"""
 @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)))