diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index e2de5fc322..486cff1bc8 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -722,6 +722,11 @@ class LlamaCppBackend: self._spec_fallback_reason: Optional[str] = None self._hf_variant: Optional[str] = None self._is_vision: bool = False + # Block-diffusion model (e.g. DiffusionGemma): served by the diffusion + # runner, not llama-server. Set from the GGUF architecture at load. + self._architecture: Optional[str] = None + self._is_diffusion: bool = False + self._diffusion_visual_bin: Optional[str] = None self._healthy = False # Set by _classify_gpu_offload after _wait_for_health. self._gpu_offload_active: Optional[bool] = None @@ -828,6 +833,11 @@ class LlamaCppBackend: def is_vision(self) -> bool: return self._is_vision + @property + def is_diffusion(self) -> bool: + """True when the loaded GGUF is a block-diffusion model (DiffusionGemma).""" + return self._is_diffusion + @property def hf_variant(self) -> Optional[str]: return self._hf_variant @@ -2218,11 +2228,16 @@ class LlamaCppBackend: self._ssm_state_size = None self._shared_kv_layers = None self._nextn_predict_layers = None + self._architecture = None + self._is_diffusion = False try: + canvas_seen = False WANTED = { "general.architecture", "tokenizer.chat_template", + # Block-diffusion marker (DiffusionGemma); routes to the diffusion runner. + "diffusion.canvas_length", # Source-repo hints for the SWA resolver's HF fallback. "general.source.huggingface.repository", "general.source.url", @@ -2277,6 +2292,7 @@ class LlamaCppBackend: general[key] = val_s if key == "general.architecture": arch = val_s + self._architecture = val_s arch_keys = { f"{arch}.context_length": "context_length", f"{arch}.block_count": "n_layers", @@ -2305,6 +2321,8 @@ class LlamaCppBackend: if vtype == 4 else struct.unpack(" Optional[tuple[list, str, Optional[str]]]: + """Resolve how to launch the DiffusionGemma runner: (shim argv prefix, + visual-server binary, optional extra PYTHONPATH dir for the file override). + + Shim: UNSLOTH_DG_SHIM (a .py file) first, else the installed + unsloth_zoo.diffusion_studio.shim. Binary: DG_VISUAL_BIN first, else + alongside llama-server. Returns None if neither can be found. + """ + import importlib.util + import os + import sys + + # Visual-server binary: env override, else next to llama-server or in the + # install's build/bin (where the prebuilt/installer puts it). .exe on Windows. + visual_bin = os.environ.get("DG_VISUAL_BIN") + if not visual_bin: + name = "llama-diffusion-gemma-visual-server" + (".exe" if os.name == "nt" else "") + base = self._find_llama_server_binary() + if base: + base_dir = Path(base).parent + for cand in ( + base_dir / name, + base_dir / "build" / "bin" / name, + base_dir / "build" / "bin" / "Release" / name, + ): + if cand.is_file(): + visual_bin = str(cand) + break + if not (visual_bin and Path(visual_bin).is_file()): + return None + + # Shim: a file override (its dir goes on PYTHONPATH), else the zoo package via -m. + shim_file = os.environ.get("UNSLOTH_DG_SHIM") + if shim_file and Path(shim_file).is_file(): + return ([sys.executable, shim_file], visual_bin, str(Path(shim_file).parent)) + + # Find the installed shim without importing the heavy unsloth_zoo package + # (find_spec on the top-level package does not run its __init__). + try: + spec = importlib.util.find_spec("unsloth_zoo") + except Exception: + spec = None + if spec is not None and spec.submodule_search_locations: + pkg_dir = Path(list(spec.submodule_search_locations)[0]) + if (pkg_dir / "diffusion_studio" / "shim.py").is_file(): + return ( + [sys.executable, "-m", "unsloth_zoo.diffusion_studio.shim"], + visual_bin, + None, + ) + + return None + + def _start_diffusion_server( + self, + *, + model_path: str, + gguf_path: Optional[str], + hf_repo: Optional[str], + hf_variant: Optional[str], + model_identifier: str, + n_ctx: int, + extra_args: Optional[List[str]], + ) -> bool: + """Launch the OpenAI-compat diffusion shim (which drives the on-device + visual decoder) and wait for health. Presents the same /v1 + /health + interface as llama-server, so the rest of Studio is unchanged. + """ + import os + + assets = self._find_diffusion_assets() + if assets is None: + raise RuntimeError( + "DiffusionGemma runner not found. Install unsloth_zoo (which ships " + "unsloth_zoo.diffusion_studio.shim) or set UNSLOTH_DG_SHIM to a shim " + "file, and provide the visual-server binary via DG_VISUAL_BIN or next " + "to llama-server in the install tree." + ) + shim_cmd, visual_bin, extra_pythonpath = assets + self._diffusion_visual_bin = visual_bin + + self._kill_process() + self._port = self._find_free_port() + # Auto-size (0): the visual server probes the largest context that fits this GPU's VRAM + # (capped at the training context). An explicit in-range n_ctx overrides it. + maxtok = n_ctx if (n_ctx and 0 < n_ctx <= 65536) else 0 + gpu = os.environ.get("DG_GPU", "0") + + cmd = list(shim_cmd) + [ + "--gguf", + model_path, + "--host", + "127.0.0.1", + "--port", + str(self._port), + "--gpu", + gpu, + "--maxtok", + str(maxtok), + ] + + env = child_env_without_native_path_secret() + env["DG_VISUAL_BIN"] = visual_bin + env["DG_GPU"] = gpu + # The file-override shim imports its sibling visual_engine; put its dir on PYTHONPATH. + # (The zoo-package shim is an installed module and needs no PYTHONPATH change.) + if extra_pythonpath: + existing = env.get("PYTHONPATH") + env["PYTHONPATH"] = ( + (extra_pythonpath + os.pathsep + existing) if existing else extra_pythonpath + ) + + logger.info(f"Starting DiffusionGemma runner: {' '.join(cmd)}") + self._stdout_lines = [] + self._llama_log_fh = None + self._llama_log_path = None + try: + log_dir = _swa_cache_path().parent / "logs" / "diffusion-server" + log_dir.mkdir(parents = True, exist_ok = True) + self._llama_log_path = log_dir / f"diffusion-{int(time.time())}-port-{self._port}.log" + self._llama_log_fh = open(self._llama_log_path, "w", encoding = "utf-8", buffering = 1) + logger.info(f"diffusion runner stdout/stderr -> {self._llama_log_path}") + except OSError as e: + logger.debug(f"Could not open diffusion runner log file: {e}") + + # PR_SET_PDEATHSIG: the shim (and its visual server) die with this backend + # process, so a Studio crash/restart never orphans a GPU process. + popen_kwargs = dict(_windows_hidden_subprocess_kwargs()) + if sys.platform.startswith("linux"): # prctl/libc.so.6 are Linux-only + + def _pdeathsig(): + try: + import ctypes + import signal as _signal + ctypes.CDLL("libc.so.6", use_errno = True).prctl(1, _signal.SIGTERM) + except Exception: + pass + + popen_kwargs["preexec_fn"] = _pdeathsig + + self._process = subprocess.Popen( + cmd, + stdout = subprocess.PIPE, + stderr = subprocess.STDOUT, + text = True, + env = env, + **popen_kwargs, + ) + self._stdout_thread = threading.Thread( + target = self._drain_stdout, daemon = True, name = "diffusion-stdout" + ) + self._stdout_thread.start() + + # Publish state before the health wait (mirrors the llama-server path). + self._gguf_path = model_path + self._hf_repo = hf_repo + self._is_vision = False + self._is_audio = False # clear any prior TTS/audio model's routing flag + self._model_identifier = model_identifier + self._cache_type_kv = None + self._gpu_offload_active = True + if hf_variant: + self._hf_variant = hf_variant + elif gguf_path: + try: + from utils.models.model_config import _extract_quant_label + self._hf_variant = _extract_quant_label(gguf_path) + except Exception: + self._hf_variant = None + else: + self._hf_variant = None + # Provisional until the server reports the budget it resolved (auto-size picks it from VRAM). + self._effective_context_length = maxtok or self._context_length + self._max_context_length = self._context_length or maxtok or None + + healthy = self._wait_for_health(timeout = 600.0) + if healthy: + self._healthy = True + self._gpu_offload_active = True + if extra_args is not None: + self._extra_args = list(extra_args) + self._extra_args_source = (model_identifier, hf_variant) + # The visual server logs "MAXTOK=" with the context budget it actually resolved + # (auto-sized to VRAM). Read it back so the UI context bar shows the real budget. + chosen = maxtok + try: + import re as _re + for _ln in reversed(self._stdout_lines): + _m = _re.search(r"MAXTOK=(\d+)", _ln) + if _m: + chosen = int(_m.group(1)) + break + except Exception: + pass + if chosen and chosen > 0: + self._effective_context_length = chosen + self._max_context_length = chosen + self._requested_n_ctx = int(n_ctx) + else: + self._healthy = False + logger.error("DiffusionGemma runner failed to become healthy") + return healthy + # ── HF download (no lock held) ─────────────────────────────── def _download_gguf( @@ -3180,13 +3414,9 @@ class LlamaCppBackend: with self._lock: self._kill_process() + # Resolve llama-server now but defer a not-found error: a block-diffusion + # GGUF uses the diffusion runner, and its arch is only known after the header. binary = self._find_llama_server_binary() - if not binary: - raise RuntimeError( - "llama-server binary not found. " - "Run setup.sh to build it, install llama.cpp, " - "or set LLAMA_SERVER_PATH environment variable." - ) # ── Phase 2: download (NO lock held, so cancel can proceed) ── # mtp_draft_path arrives set for local Gemma loads (detected @@ -3241,6 +3471,30 @@ class LlamaCppBackend: logger.info("Load cancelled after download phase") return False + # Block-diffusion GGUFs (DiffusionGemma) cannot run on llama-server; + # serve them with the diffusion runner (same OpenAI-compat interface). + if self._is_diffusion: + with self._lock: + if self._cancel_event.is_set(): + logger.info("Load cancelled before diffusion server start") + return False + return self._start_diffusion_server( + model_path = model_path, + gguf_path = gguf_path, + hf_repo = hf_repo, + hf_variant = hf_variant, + model_identifier = model_identifier, + n_ctx = n_ctx, + extra_args = extra_args, + ) + + if not binary: + raise RuntimeError( + "llama-server binary not found. " + "Run setup.sh to build it, install llama.cpp, " + "or set LLAMA_SERVER_PATH environment variable." + ) + # Outside ``self._lock`` so /unload, /cancel, /status aren't # blocked. ``unload_model`` also records the kill, so the # frontend /unload+/load Apply path engages the wait here even @@ -5241,6 +5495,12 @@ class LlamaCppBackend: try: data = json.loads(line[6:]) + # Diffusion frame (per-step canvas) from the shim: forward untouched so + # the frontend renders it in place. No assistant text, so it never enters + # the cumulative content. + if data.get("type") == "diffusion_frame": + yield data + continue # Capture server timings/usage from final chunks. _chunk_timings = data.get("timings") if _chunk_timings: diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 5a414384d2..04c63c2247 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -169,6 +169,9 @@ class LoadResponse(BaseModel): is_vision: bool = Field(False, description = "Whether model is a vision model") is_lora: bool = Field(False, description = "Whether model is a LoRA adapter") is_gguf: bool = Field(False, description = "Whether model is a GGUF model (llama.cpp)") + is_diffusion: bool = Field( + False, description = "Whether model is a block-diffusion model (DiffusionGemma)" + ) is_audio: bool = Field(False, description = "Whether model is a TTS audio model") audio_type: Optional[str] = Field(None, description = "Audio codec type: snac, csm, bicodec, dac") has_audio_input: bool = Field(False, description = "Whether model accepts audio input (ASR)") @@ -286,6 +289,9 @@ class InferenceStatusResponse(BaseModel): ) is_vision: bool = Field(False, description = "Whether the active model is a vision model") is_gguf: bool = Field(False, description = "Whether the active model is a GGUF model (llama.cpp)") + is_diffusion: bool = Field( + False, description = "Whether the active model is a block-diffusion model (DiffusionGemma)" + ) gguf_variant: Optional[str] = Field(None, description = "GGUF quantization variant (e.g. Q4_K_M)") is_audio: bool = Field(False, description = "Whether the active model is a TTS audio model") audio_type: Optional[str] = Field(None, description = "Audio codec type: snac, csm, bicodec, dac") diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 21390d7a3e..ab8097b65e 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -1390,6 +1390,7 @@ async def load_model( is_vision = llama_backend._is_vision, is_lora = False, is_gguf = True, + is_diffusion = llama_backend.is_diffusion, is_audio = _gguf_is_audio, audio_type = _gguf_audio, has_audio_input = getattr(llama_backend, "_has_audio_input", False), @@ -1668,6 +1669,7 @@ async def load_model( is_vision = llama_backend.is_vision, is_lora = False, is_gguf = True, + is_diffusion = llama_backend.is_diffusion, is_audio = _gguf_is_audio, audio_type = _gguf_audio, has_audio_input = llama_backend._has_audio_input, @@ -2151,6 +2153,7 @@ async def get_status(current_subject: str = Depends(get_current_subject)): model_identifier = None if _native_grant_backed else _model_id, is_vision = llama_backend.is_vision, is_gguf = True, + is_diffusion = llama_backend.is_diffusion, gguf_variant = llama_backend.hf_variant, is_audio = getattr(llama_backend, "_is_audio", False), audio_type = _audio_type, @@ -3974,6 +3977,10 @@ async def openai_chat_completions( _stream_usage = cumulative.get("usage") _stream_timings = cumulative.get("timings") _stream_finish = cumulative.get("finish_reason") + elif cumulative.get("type") == "diffusion_frame": + # Diffusion frame (per-step canvas): pass through as a raw SSE line on the + # tool_status channel. No assistant text, so it never enters the cumulative diff. + yield f"data: {json.dumps(cumulative)}\n\n" else: logger.warning( "gguf_stream_chunks: unexpected dict event: %s", diff --git a/studio/frontend/src/components/assistant-ui/markdown-text.tsx b/studio/frontend/src/components/assistant-ui/markdown-text.tsx index ef9b33fced..45a0082a57 100644 --- a/studio/frontend/src/components/assistant-ui/markdown-text.tsx +++ b/studio/frontend/src/components/assistant-ui/markdown-text.tsx @@ -284,6 +284,9 @@ function CodeBlockActions({ ); } +// DiffusionGemma renders its denoising live in the bubble (see DiffusionCanvas in +// thread.tsx), so it no longer forces HTML into an iframe artifact; it follows the +// same artifact rules as every other model. function StreamdownBlock(props: BlockProps) { const shouldCollapseHtmlArtifacts = useChatRuntimeStore( (state) => state.artifactsEnabled || state.collapseHtmlArtifacts, diff --git a/studio/frontend/src/components/assistant-ui/message-timing.tsx b/studio/frontend/src/components/assistant-ui/message-timing.tsx index 447fd763e3..8d40f587ad 100644 --- a/studio/frontend/src/components/assistant-ui/message-timing.tsx +++ b/studio/frontend/src/components/assistant-ui/message-timing.tsx @@ -19,6 +19,11 @@ const formatNumber = (n: number): string => { return n.toLocaleString(); }; +const formatRate = (r: number | undefined): string => { + if (r === undefined || !Number.isFinite(r)) return "—"; + return `${Math.round(r).toLocaleString()} tok/s`; +}; + /** * Shows streaming stats as a badge with hover tooltip. * When server timings are available (GGUF), shows prompt eval, generation, @@ -51,6 +56,10 @@ export const MessageTiming: FC<{ st?.cache_n ?? custom?.contextUsage?.cachedTokens ?? 0; // Anthropic-only cache-write count. const cacheWrites = custom?.contextUsage?.cacheWriteTokens ?? 0; + // DiffusionGemma reports separately-labelled throughput (no prefill, so no "prompt + // speed"), matching the CLI: in-step parallel, effective (canvas*blocks/wall), and + // output (answer tokens/wall). + const isDiffusion = (st as { diffusion?: boolean } | undefined)?.diffusion === true; // Guard unphysical tok/s: llama.cpp emits predicted_ms=0 on no-op turns, // blowing the rate up to Infinity. Require >=1 token, a non-zero decode @@ -93,6 +102,99 @@ export const MessageTiming: FC<{ >
{st ? ( + isDiffusion ? ( + <> + {/* DiffusionGemma: honest throughput (no autoregressive prompt speed) */} + {timing.firstTokenTime !== undefined && ( +
+ First token + + {formatTimingMs(timing.firstTokenTime)} + +
+ )} + {st?.diffusion_parallel_tok_s != null && ( +
+ Speed (in-step) + + {formatRate(st.diffusion_parallel_tok_s)} + +
+ )} + {st?.diffusion_effective_tok_s != null && ( +
+ Effective + + {formatRate(st.diffusion_effective_tok_s)} + +
+ )} + {st?.diffusion_output_tok_s != null && ( +
+ Output + + {formatRate(st.diffusion_output_tok_s)} + +
+ )} + {st?.diffusion_steps != null && ( +
+ Denoising + + {formatNumber(st.diffusion_steps)} steps + {st?.diffusion_blocks != null + ? `, ${formatNumber(st.diffusion_blocks)} block${st.diffusion_blocks === 1 ? "" : "s"}` + : ""} + +
+ )} + {st?.diffusion_canvas != null && ( +
+ Canvas + + {formatNumber(st.diffusion_canvas)} tokens + +
+ )} + {(st?.diffusion_wall_ms ?? st?.predicted_ms) != null && ( +
+ Generation + + {formatTimingMs(st.diffusion_wall_ms ?? st.predicted_ms)} + +
+ )} + {timing.tokenCount !== undefined && ( +
+ Answer tokens + + {formatNumber(timing.tokenCount)} + +
+ )} + {(st?.diffusion_prompt_n ?? st?.prompt_n) != null && ( +
+ Prompt + + {formatNumber(st.diffusion_prompt_n ?? st.prompt_n)} tokens + +
+ )} +
+
+ Total + + {formatTimingMs(timing.totalStreamTime)} + +
+
+ Chunks + + {timing.totalChunks} + +
+ + ) : ( <> {/* Server-side metrics (GGUF) */} {st?.prompt_ms != null && ( @@ -135,6 +237,30 @@ export const MessageTiming: FC<{
)} + {timing.firstTokenTime !== undefined && ( +
+ First token + + {formatTimingMs(timing.firstTokenTime)} + +
+ )} + {st?.diffusion_steps != null && ( +
+ Denoising steps + + {formatNumber(st.diffusion_steps)} + +
+ )} + {st?.diffusion_blocks != null && ( +
+ Blocks + + {formatNumber(st.diffusion_blocks)} + +
+ )} {cacheHits > 0 && (
Cache hits @@ -165,6 +291,7 @@ export const MessageTiming: FC<{
+ ) ) : ( <> {/* Client-side metrics (safetensors + external provider fallback) */} diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx index ac5ae19aac..a69e493b4f 100644 --- a/studio/frontend/src/components/assistant-ui/thread.tsx +++ b/studio/frontend/src/components/assistant-ui/thread.tsx @@ -2587,6 +2587,38 @@ const ImageGenerationToolUIConfirmable = withToolConfirmation( const RenderHtmlToolUIConfirmable = withToolConfirmation(RenderHtmlToolUI); const ToolFallbackConfirmable = withToolConfirmation(ToolFallback); +// Live in-place denoising canvas for DiffusionGemma: while generating, render the +// latest per-step canvas snapshot in the bubble so the user watches the answer resolve +// out of noise. Transient (store-only, cleared on run end), so the finished message +// keeps only the committed markdown. +const DiffusionCanvas: FC = () => { + const isRunning = useAuiState( + ({ message }) => message.status?.type === "running", + ); + // A non-null canvas is set only by diffusion_frame events (diffusion models only), + // so it is a sufficient gate; loadedIsDiffusion can lag the first frame on a fresh load. + const canvas = useChatRuntimeStore((s) => s.activeDiffusionCanvas); + if (!isRunning || !canvas) { + return null; + } + const stepLabel = + canvas.total > 0 ? `step ${canvas.step + 1}/${canvas.total}` : "denoising"; + return ( +
+
+ + Denoising + + block {canvas.block + 1} - {stepLabel} + +
+ 
+        {canvas.text}
+
+
+ ); +}; + const AssistantMessage: FC = () => { return ( {
+ [0]["messages"]; @@ -2519,6 +2534,29 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { continue; } + // Diffusion frame: a transient canvas snapshot. Route it to the transient + // store (the in-bubble renderer reads it) and skip it; it has no assistant + // text, so it never enters the transcript or the counters below. + const diffusionFrame = ( + chunk as unknown as { + _diffusionFrame?: { + block?: number; + step?: number; + total?: number; + text?: string; + }; + } + )._diffusionFrame; + if (diffusionFrame !== undefined) { + runtime.setActiveDiffusionCanvas({ + block: diffusionFrame.block ?? 0, + step: diffusionFrame.step ?? 0, + total: diffusionFrame.total ?? 0, + text: diffusionFrame.text ?? "", + }); + continue; + } + // Emit tool-call content parts for assistant-ui. // tool_start: add a part (renders "running"). // tool_end: set result on the part (transitions to "complete"). @@ -3211,6 +3249,9 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { } runtime.setGeneratingStatus(null); runtime.setToolStatus(null); + // Drop the transient denoising canvas so the finished bubble shows only + // the committed markdown answer (cancellation/error included). + runtime.setActiveDiffusionCanvas(null); clearTimeout(warmupTimer); if (waitingFirstChunk) { if (firstTokenSettled) { diff --git a/studio/frontend/src/features/chat/api/chat-api.ts b/studio/frontend/src/features/chat/api/chat-api.ts index 7a2c8e7e9a..7f5cc2f974 100644 --- a/studio/frontend/src/features/chat/api/chat-api.ts +++ b/studio/frontend/src/features/chat/api/chat-api.ts @@ -753,6 +753,15 @@ export async function* streamChatCompletions( separatorIndex = buffer.search(/\r?\n\r?\n/); continue; } + // Diffusion frame: a per-step canvas snapshot. Custom SSE payload (not an OpenAI chunk) with + // no assistant text, surfaced as a transient marker for the in-place renderer, never the transcript. + if ("type" in parsed && parsed.type === "diffusion_frame") { + yield { + _diffusionFrame: parsed, + } as unknown as OpenAIChatChunk; + separatorIndex = buffer.search(/\r?\n\r?\n/); + continue; + } // Tool start/end events carry full input/output for the tool outputs panel if ( "type" in parsed && diff --git a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts index e513d351c4..55389e2bdf 100644 --- a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts +++ b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts @@ -314,6 +314,7 @@ export function useChatModelRuntime() { useChatRuntimeStore.setState({ modelRequiresTrustRemoteCode: false, loadedIsMultimodal: false, + loadedIsDiffusion: false, }); } } catch (error) { @@ -647,6 +648,7 @@ export function useChatModelRuntime() { chatTemplateOverride: effectiveChatTemplateOverride, loadedChatTemplateOverride: effectiveChatTemplateOverride, loadedIsMultimodal: isMultimodalResponse(loadResponse), + loadedIsDiffusion: loadResponse.is_diffusion ?? false, activeNativePathToken: nativePathToken ?? null, }); // Unlock attach menus for capabilities the catalog entry lacked. diff --git a/studio/frontend/src/features/chat/lib/apply-inference-status-to-store.ts b/studio/frontend/src/features/chat/lib/apply-inference-status-to-store.ts index 376329b0cc..dbe9ef4dc3 100644 --- a/studio/frontend/src/features/chat/lib/apply-inference-status-to-store.ts +++ b/studio/frontend/src/features/chat/lib/apply-inference-status-to-store.ts @@ -157,6 +157,7 @@ export function applyActiveModelStatusToStore( modelRequiresTrustRemoteCode: status.requires_trust_remote_code ?? false, defaultChatTemplate: nextDefaultChatTemplate, loadedIsMultimodal: isMultimodalResponse(status), + loadedIsDiffusion: status.is_diffusion ?? false, specFallbackReason: status.spec_fallback_reason ?? null, ...(prevState.loadedSpeculativeType === null && { speculativeType: currentSpecType, diff --git a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts index 7aecc1fcd9..e819b45ce9 100644 --- a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts +++ b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts @@ -165,6 +165,14 @@ function saveLastExternalCheckpoint(value: string | null): void { } export type ReasoningStyle = "enable_thinking" | "reasoning_effort"; +/** One live DiffusionGemma denoising snapshot: the current canvas text at a + * given step of a given block (block/step are 0-based; total = steps in block). */ +export type DiffusionCanvasFrame = { + block: number; + step: number; + total: number; + text: string; +}; export type PendingImageEditReference = { threadId: string | null; openaiImageGenerationCallId: string; @@ -524,6 +532,12 @@ type ChatRuntimeStore = { /** Backend-reported tensor-parallel state; null until first hydrated. */ loadedTensorParallel: boolean | null; loadedIsMultimodal: boolean; + /** Active model is a block-diffusion model (DiffusionGemma): drives the + * denoising-canvas artifact auto-render. */ + loadedIsDiffusion: boolean; + /** Live denoising frame for the in-progress diffusion message. Transient: set + * per step, cleared when the run ends, never persisted into the transcript. */ + activeDiffusionCanvas: DiffusionCanvasFrame | null; customContextLength: number | null; defaultChatTemplate: string | null; chatTemplateOverride: string | null; @@ -600,6 +614,7 @@ type ChatRuntimeStore = { setRagAutoInjectMinScore: (score: number) => void; setToolStatus: (status: string | null) => void; setGeneratingStatus: (status: string | null) => void; + setActiveDiffusionCanvas: (canvas: DiffusionCanvasFrame | null) => void; setAutoHealToolCalls: (enabled: boolean) => void; setMaxToolCallsPerMessage: (value: number) => void; setToolCallTimeout: (value: number) => void; @@ -875,6 +890,7 @@ export const useChatRuntimeStore = create((set, get) => ({ ), toolStatus: null, generatingStatus: null, + activeDiffusionCanvas: null, autoHealToolCalls: true, maxToolCallsPerMessage: 25, toolCallTimeout: 5, @@ -888,6 +904,7 @@ export const useChatRuntimeStore = create((set, get) => ({ tensorParallel: false, loadedTensorParallel: null, loadedIsMultimodal: false, + loadedIsDiffusion: false, customContextLength: null, defaultChatTemplate: null, chatTemplateOverride: null, @@ -1093,6 +1110,7 @@ export const useChatRuntimeStore = create((set, get) => ({ // Only the per-session enable pill resets; source/mode/top_k persist. ragEnabled: false, toolStatus: null, + activeDiffusionCanvas: null, kvCacheDtype: null, loadedKvCacheDtype: null, speculativeType: readPersistedSpeculativeType(), @@ -1103,6 +1121,7 @@ export const useChatRuntimeStore = create((set, get) => ({ tensorParallel: false, loadedTensorParallel: null, loadedIsMultimodal: false, + loadedIsDiffusion: false, customContextLength: null, defaultChatTemplate: null, chatTemplateOverride: null, @@ -1259,6 +1278,8 @@ export const useChatRuntimeStore = create((set, get) => ({ return { ragAutoInjectMinScore }; }), setToolStatus: (toolStatus) => set({ toolStatus }), + setActiveDiffusionCanvas: (activeDiffusionCanvas) => + set({ activeDiffusionCanvas }), setGeneratingStatus: (generatingStatus) => set({ generatingStatus }), setAutoHealToolCalls: (autoHealToolCalls) => set((state) => { diff --git a/studio/frontend/src/features/chat/types/api.ts b/studio/frontend/src/features/chat/types/api.ts index b3a7b16b5f..1fdd3ed721 100644 --- a/studio/frontend/src/features/chat/types/api.ts +++ b/studio/frontend/src/features/chat/types/api.ts @@ -114,6 +114,7 @@ export interface LoadModelResponse { is_vision: boolean; is_lora: boolean; is_gguf?: boolean; + is_diffusion?: boolean; is_audio?: boolean; audio_type?: string | null; has_audio_input?: boolean; @@ -152,6 +153,7 @@ export interface InferenceStatusResponse { model_identifier?: string | null; is_vision: boolean; is_gguf?: boolean; + is_diffusion?: boolean; gguf_variant?: string | null; is_audio?: boolean; audio_type?: string | null; diff --git a/studio/install_llama_prebuilt.py b/studio/install_llama_prebuilt.py index 680b621b18..92d1f84b2d 100644 --- a/studio/install_llama_prebuilt.py +++ b/studio/install_llama_prebuilt.py @@ -4548,6 +4548,68 @@ def ensure_converter_scripts(install_dir: Path, llama_tag: str) -> None: shutil.copy2(canonical, legacy) +def ensure_diffusion_visual_server( + install_dir: Path, host: HostInfo, release_tag: str | None +) -> None: + """Best-effort placement of the DiffusionGemma visual-server binary next to + llama-server in the install tree, so Studio can serve DiffusionGemma GGUFs + without any DG_* env. This is an Unsloth artifact (not a ggml-org one), so it + is optional: if it is already present we just make it executable, otherwise we + try the published release and quietly skip on absence. A source build + (setup.sh / setup.ps1) copies it from build/bin directly. Users can always + build it from llama.cpp PR #24423 and point DG_VISUAL_BIN at it. + """ + name = "llama-diffusion-gemma-visual-server" + (".exe" if host.is_windows else "") + bin_dir = install_dir / "build" / ("bin/Release" if host.is_windows else "bin") + target = bin_dir / name + + if target.exists(): + if not host.is_windows: + try: + target.chmod(0o755) + except OSError: + pass + return + + if not release_tag: + log( + "diffusion visual server not bundled (no release tag); build it from llama.cpp " + "PR #24423 and set DG_VISUAL_BIN if you want native DiffusionGemma serving" + ) + return + + try: + assets = github_release_assets(DEFAULT_PUBLISHED_REPO, release_tag) + match = None + for asset_name, url in assets.items(): + low = asset_name.lower() + if "llama-diffusion-gemma-visual-server" not in low: + continue + if host.is_windows and not low.endswith(".exe"): + continue + if (not host.is_windows) and low.endswith(".exe"): + continue + match = (asset_name, url) + break + if match is None: + log( + "diffusion visual server not found in the published release; native " + "DiffusionGemma serving needs DG_VISUAL_BIN or a source build" + ) + return + bin_dir.mkdir(parents = True, exist_ok = True) + download_file(match[1], target) + if not host.is_windows: + target.chmod(0o755) + log(f"installed diffusion visual server: {match[0]}") + except Exception as exc: + log( + "diffusion visual server fetch skipped " + f"({textwrap.shorten(str(exc), width = 160, placeholder = '...')}); " + "set DG_VISUAL_BIN or build from llama.cpp PR #24423 for native serving" + ) + + def extracted_archive_root(extract_dir: Path) -> Path: children = [path for path in extract_dir.iterdir()] if len(children) == 1 and children[0].is_dir(): @@ -6795,6 +6857,13 @@ def install_prebuilt( "converter script fetch failed after activation; install remains valid " f"({textwrap.shorten(str(exc), width = 200, placeholder = '...')})" ) + try: + ensure_diffusion_visual_server(install_dir, host, plan.release_tag) + except Exception as exc: + log( + "diffusion visual server step skipped; install remains valid " + f"({textwrap.shorten(str(exc), width = 200, placeholder = '...')})" + ) return except BusyInstallConflict as exc: log("prebuilt install path is blocked by an in-use llama.cpp install") diff --git a/studio/setup.ps1 b/studio/setup.ps1 index 94817b561a..542846ac70 100644 --- a/studio/setup.ps1 +++ b/studio/setup.ps1 @@ -3352,6 +3352,13 @@ if (-not $NeedLlamaSourceBuild) { } } + # -- Step E: Build the DiffusionGemma visual server (optional, best-effort) -- + # An example target present on llama.cpp PR #24423; lets Studio serve + # DiffusionGemma GGUFs without DG_VISUAL_BIN. No-op when not configured. + if ($BuildOk) { + $null = cmake --build $BuildDir --config Release --target llama-diffusion-gemma-visual-server -j $NumCpu 2>&1 | Out-String + } + # Swap temp build dir into final location (only if we built in a temp dir) if ($BuildOk -and $LlamaCppDir -ne $OriginalLlamaCppDir) { Assert-StudioOwnedOrAbsent -Path $OriginalLlamaCppDir -Label "llama.cpp install" diff --git a/studio/setup.sh b/studio/setup.sh index a8603bd0da..e56fab36ef 100755 --- a/studio/setup.sh +++ b/studio/setup.sh @@ -1603,6 +1603,9 @@ else if [ "$BUILD_OK" = true ]; then run_quiet_no_exit "build llama-quantize" cmake --build "$_BUILD_TMP/build" --config Release --target llama-quantize -j"$NCPU" || true + # Best-effort: the DiffusionGemma visual server (an example target, present + # on llama.cpp PR #24423). No-op when the diffusion example is not configured. + run_quiet_no_exit "build diffusion visual server" cmake --build "$_BUILD_TMP/build" --config Release --target llama-diffusion-gemma-visual-server -j"$NCPU" || true fi # Swap only after build succeeds -- preserves existing install on failure @@ -1616,6 +1619,11 @@ else if [ -f "$QUANTIZE_BIN" ]; then ln -sf build/bin/llama-quantize "$LLAMA_CPP_DIR/llama-quantize" fi + # DiffusionGemma visual server, if it was built (PR #24423): link next to + # llama-server so Studio serves DiffusionGemma GGUFs without DG_VISUAL_BIN. + if [ -f "$LLAMA_CPP_DIR/build/bin/llama-diffusion-gemma-visual-server" ]; then + ln -sf build/bin/llama-diffusion-gemma-visual-server "$LLAMA_CPP_DIR/llama-diffusion-gemma-visual-server" + fi else rm -rf "$_BUILD_TMP" fi