Move Nemotron text generation to WebGPU

app.py

@@ -265,7 +265,6 @@ TTS_MODE = os.environ.get("TINY_TTS_MODE", "").strip().lower()
 VOXCPM_SPACE = os.environ.get("TINY_VOXCPM_SPACE", "").strip()
 TINY_AYA_SPACE = os.environ.get("TINY_AYA_SPACE", "").strip()
 MINICPM5_SPACE = os.environ.get("TINY_MINICPM5_SPACE", "").strip()
-NEMOTRON_SPACE = os.environ.get("TINY_NEMOTRON_SPACE", "").strip()
 _local_tts = None       # VoiceDesign model
 _local_clone = None     # Base model (voice clone) — lazy, only if a clone is requested
 _local_tts_lock = threading.Lock()
@@ -514,10 +513,6 @@ def _minicpm5_stream(system, user, max_tokens, temperature):
     yield from _space_text_stream(MINICPM5_SPACE, system, user, max_tokens, temperature)
 
 
-def _nemotron_stream(system, user, max_tokens, temperature):
-    yield from _space_text_stream(NEMOTRON_SPACE, system, user, max_tokens, temperature)
-
-
 @fastapi_app.post("/voxcpm-tts")
 async def voxcpm_tts(request: Request):
     body = await request.json()
@@ -824,13 +819,6 @@ async def text_generate_stream(request: Request):
                         if stop.is_set():
                             break
                         loop.call_soon_threadsafe(q.put_nowait, ("delta", chunk))
-                elif model == "nemotron-3-nano-4b-zerogpu":
-                    if not NEMOTRON_SPACE:
-                        raise llm.LlmUnavailable("TINY_NEMOTRON_SPACE not set")
-                    for chunk in _nemotron_stream(system, user, max_tokens, temperature):
-                        if stop.is_set():
-                            break
-                        loop.call_soon_threadsafe(q.put_nowait, ("delta", chunk))
                 else:
                     for chunk in llm.stream_chat(
                         system,

web/engineServer.js

@@ -7,7 +7,6 @@ const MODELS = [
   { id: 'server-local', label: 'Configured server model', params: 'local/remote', note: 'uses TINY_LLM_* on the Space or local app' },
   { id: 'tiny-aya-global-zerogpu', label: 'Tiny Aya Global 3.35B', params: '3.35B', note: 'ZeroGPU sidecar; multilingual' },
   { id: 'minicpm5-1b-zerogpu', label: 'MiniCPM5 1B', params: '1B', note: 'ZeroGPU sidecar; efficient MiniCPM5 text model' },
-  { id: 'nemotron-3-nano-4b-zerogpu', label: 'Nemotron 3 Nano 4B', params: '4B', note: 'ZeroGPU sidecar; GGUF llama.cpp, slower startup' },
 ]
 const get = (id) => MODELS.find((m) => m.id === id) || MODELS[0]
 

web/engineTransformers.js

@@ -6,6 +6,7 @@ const MODELS = [
   { id: 'qwen2.5-0.5b', label: 'Qwen2.5 0.5B', params: '0.5B', repo: 'onnx-community/Qwen2.5-0.5B-Instruct' },
   { id: 'smollm2-360m', label: 'SmolLM2 360M', params: '360M', repo: 'HuggingFaceTB/SmolLM2-360M-Instruct' },
   { id: 'llama3.2-1b', label: 'Llama 3.2 1B', params: '1B', repo: 'onnx-community/Llama-3.2-1B-Instruct' },
+  { id: 'nemotron-3-nano-4b', label: 'Nemotron 3 Nano 4B', params: '4B', repo: 'onnx-community/NVIDIA-Nemotron-3-Nano-4B-BF16-ONNX', webgpuOnly: true, note: 'WebGPU only; large browser download' },
 ]
 const get = (id) => MODELS.find((m) => m.id === id) || MODELS[0]
 
@@ -26,7 +27,7 @@ async function pickDevice() {
 }
 
 let _lib = null, _pipe = null, _loadedId = null, _loadingId = null, _loadPromise = null, _device = 'wasm', _chain = Promise.resolve()
-async function lib() { if (!_lib) _lib = await import('https://cdn.jsdelivr.net/npm/@huggingface/transformers@3'); return _lib }
+async function lib() { if (!_lib) _lib = await import('https://cdn.jsdelivr.net/npm/@huggingface/transformers@4.0.0-next.8'); return _lib }
 
 async function ensure(id, onProgress) {
   const m = get(id)
@@ -39,13 +40,14 @@ async function ensure(id, onProgress) {
   _loadPromise = (async () => {
     const { pipeline } = await lib()
     _device = await pickDevice()
+    if (m.webgpuOnly && _device !== 'webgpu') throw new Error(`${m.label} requires WebGPU support in this browser.`)
     const mk = (device) => pipeline('text-generation', m.repo, {
       device, dtype: 'q4',
       progress_callback: (p) => { if (onProgress && p.status === 'progress' && p.total) onProgress(p.loaded / p.total) },
     })
     let pipe
     try { pipe = await mk(_device) }
-    catch (e) { if (_device !== 'wasm') { _device = 'wasm'; pipe = await mk('wasm') } else throw e }
+    catch (e) { if (!m.webgpuOnly && _device !== 'wasm') { _device = 'wasm'; pipe = await mk('wasm') } else throw e }
     _pipe = pipe; _loadedId = m.id; return pipe
   })().catch((e) => { _loadPromise = null; _loadingId = null; throw e })
   return _loadPromise
@@ -62,8 +64,20 @@ function stream(id, system, user, { maxTokens = 200, temperature = 0.8, onToken,
       callback_function: (text) => { if (!text) return; full += text; if (onToken) onToken(text); st.tick() },
     })
     const messages = [{ role: 'system', content: system }, { role: 'user', content: user }]
-    const prompt = pipe.tokenizer.apply_chat_template(messages, { tokenize: false, add_generation_prompt: true })
-    await pipe(prompt, { max_new_tokens: maxTokens, do_sample: true, temperature, top_k: 40, top_p: 0.9, streamer })
+    if (get(id).id === 'nemotron-3-nano-4b') {
+      await pipe(messages, {
+        max_new_tokens: maxTokens,
+        do_sample: true,
+        temperature,
+        top_k: 40,
+        top_p: 0.9,
+        streamer,
+        tokenizer_encode_kwargs: { enable_thinking: false },
+      })
+    } else {
+      const prompt = pipe.tokenizer.apply_chat_template(messages, { tokenize: false, add_generation_prompt: true })
+      await pipe(prompt, { max_new_tokens: maxTokens, do_sample: true, temperature, top_k: 40, top_p: 0.9, streamer })
+    }
     return { text: full, stats: st.finish() }
   }
   const p = _chain.then(run, run); _chain = p.catch(() => {}); return p