Studio: serve DiffusionGemma with live in-place denoising and honest stats (#6250)
* Studio: serve DiffusionGemma GGUFs with the on-device visual decoder * Studio: render the DiffusionGemma denoising canvas live in chat with honest stats * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: harden DiffusionGemma runner resolution (Windows .exe, build/bin lookup, clear stale audio flag, safe PYTHONPATH, Linux-only pdeathsig) * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Daniel Han
GitHub
parent
6d22829ef7
commit
90cb9499e8
15 changed files with 602 additions and 6 deletions
|
|
@ -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("<Q", f.read(8))[0]
|
||||
)
|
||||
if key == "diffusion.canvas_length":
|
||||
canvas_seen = True
|
||||
attr = arch_keys.get(key)
|
||||
if attr:
|
||||
if attr == "sliding_window_pattern":
|
||||
|
|
@ -2372,6 +2390,17 @@ class LlamaCppBackend:
|
|||
hf_repo_candidates,
|
||||
)
|
||||
|
||||
# Block-diffusion models (DiffusionGemma) report a diffusion arch
|
||||
# and/or a diffusion.canvas_length KV; they need the diffusion runner.
|
||||
self._is_diffusion = bool(
|
||||
(arch and arch.lower().startswith("diffusion")) or canvas_seen
|
||||
)
|
||||
if self._is_diffusion:
|
||||
logger.info(
|
||||
f"GGUF metadata: diffusion model detected (architecture={arch}); "
|
||||
"will serve via the diffusion runner"
|
||||
)
|
||||
|
||||
if self._context_length:
|
||||
logger.info(f"GGUF metadata: context_length={self._context_length}")
|
||||
if self._chat_template:
|
||||
|
|
@ -2390,6 +2419,211 @@ class LlamaCppBackend:
|
|||
except Exception as e:
|
||||
logger.warning(f"Failed to read GGUF metadata: {e}")
|
||||
|
||||
# ── Diffusion runner (DiffusionGemma) ──
|
||||
|
||||
def _find_diffusion_assets(self) -> 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=<N>" 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:
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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<{
|
|||
>
|
||||
<div className="grid min-w-40 gap-1.5 text-xs">
|
||||
{st ? (
|
||||
isDiffusion ? (
|
||||
<>
|
||||
{/* DiffusionGemma: honest throughput (no autoregressive prompt speed) */}
|
||||
{timing.firstTokenTime !== undefined && (
|
||||
<div className="flex items-center justify-between gap-4">
|
||||
<span className="text-muted-foreground">First token</span>
|
||||
<span className="font-mono tabular-nums">
|
||||
{formatTimingMs(timing.firstTokenTime)}
|
||||
</span>
|
||||
</div>
|
||||
)}
|
||||
{st?.diffusion_parallel_tok_s != null && (
|
||||
<div className="flex items-center justify-between gap-4">
|
||||
<span className="text-muted-foreground">Speed (in-step)</span>
|
||||
<span className="font-mono tabular-nums">
|
||||
{formatRate(st.diffusion_parallel_tok_s)}
|
||||
</span>
|
||||
</div>
|
||||
)}
|
||||
{st?.diffusion_effective_tok_s != null && (
|
||||
<div className="flex items-center justify-between gap-4">
|
||||
<span className="text-muted-foreground">Effective</span>
|
||||
<span className="font-mono tabular-nums">
|
||||
{formatRate(st.diffusion_effective_tok_s)}
|
||||
</span>
|
||||
</div>
|
||||
)}
|
||||
{st?.diffusion_output_tok_s != null && (
|
||||
<div className="flex items-center justify-between gap-4">
|
||||
<span className="text-muted-foreground">Output</span>
|
||||
<span className="font-mono tabular-nums">
|
||||
{formatRate(st.diffusion_output_tok_s)}
|
||||
</span>
|
||||
</div>
|
||||
)}
|
||||
{st?.diffusion_steps != null && (
|
||||
<div className="flex items-center justify-between gap-4">
|
||||
<span className="text-muted-foreground">Denoising</span>
|
||||
<span className="font-mono tabular-nums">
|
||||
{formatNumber(st.diffusion_steps)} steps
|
||||
{st?.diffusion_blocks != null
|
||||
? `, ${formatNumber(st.diffusion_blocks)} block${st.diffusion_blocks === 1 ? "" : "s"}`
|
||||
: ""}
|
||||
</span>
|
||||
</div>
|
||||
)}
|
||||
{st?.diffusion_canvas != null && (
|
||||
<div className="flex items-center justify-between gap-4">
|
||||
<span className="text-muted-foreground">Canvas</span>
|
||||
<span className="font-mono tabular-nums">
|
||||
{formatNumber(st.diffusion_canvas)} tokens
|
||||
</span>
|
||||
</div>
|
||||
)}
|
||||
{(st?.diffusion_wall_ms ?? st?.predicted_ms) != null && (
|
||||
<div className="flex items-center justify-between gap-4">
|
||||
<span className="text-muted-foreground">Generation</span>
|
||||
<span className="font-mono tabular-nums">
|
||||
{formatTimingMs(st.diffusion_wall_ms ?? st.predicted_ms)}
|
||||
</span>
|
||||
</div>
|
||||
)}
|
||||
{timing.tokenCount !== undefined && (
|
||||
<div className="flex items-center justify-between gap-4">
|
||||
<span className="text-muted-foreground">Answer tokens</span>
|
||||
<span className="font-mono tabular-nums">
|
||||
{formatNumber(timing.tokenCount)}
|
||||
</span>
|
||||
</div>
|
||||
)}
|
||||
{(st?.diffusion_prompt_n ?? st?.prompt_n) != null && (
|
||||
<div className="flex items-center justify-between gap-4">
|
||||
<span className="text-muted-foreground">Prompt</span>
|
||||
<span className="font-mono tabular-nums">
|
||||
{formatNumber(st.diffusion_prompt_n ?? st.prompt_n)} tokens
|
||||
</span>
|
||||
</div>
|
||||
)}
|
||||
<div className="my-0.5 border-t border-border/40" />
|
||||
<div className="flex items-center justify-between gap-4">
|
||||
<span className="text-muted-foreground">Total</span>
|
||||
<span className="font-mono tabular-nums">
|
||||
{formatTimingMs(timing.totalStreamTime)}
|
||||
</span>
|
||||
</div>
|
||||
<div className="flex items-center justify-between gap-4">
|
||||
<span className="text-muted-foreground">Chunks</span>
|
||||
<span className="font-mono tabular-nums">
|
||||
{timing.totalChunks}
|
||||
</span>
|
||||
</div>
|
||||
</>
|
||||
) : (
|
||||
<>
|
||||
{/* Server-side metrics (GGUF) */}
|
||||
{st?.prompt_ms != null && (
|
||||
|
|
@ -135,6 +237,30 @@ export const MessageTiming: FC<{
|
|||
</span>
|
||||
</div>
|
||||
)}
|
||||
{timing.firstTokenTime !== undefined && (
|
||||
<div className="flex items-center justify-between gap-4">
|
||||
<span className="text-muted-foreground">First token</span>
|
||||
<span className="font-mono tabular-nums">
|
||||
{formatTimingMs(timing.firstTokenTime)}
|
||||
</span>
|
||||
</div>
|
||||
)}
|
||||
{st?.diffusion_steps != null && (
|
||||
<div className="flex items-center justify-between gap-4">
|
||||
<span className="text-muted-foreground">Denoising steps</span>
|
||||
<span className="font-mono tabular-nums">
|
||||
{formatNumber(st.diffusion_steps)}
|
||||
</span>
|
||||
</div>
|
||||
)}
|
||||
{st?.diffusion_blocks != null && (
|
||||
<div className="flex items-center justify-between gap-4">
|
||||
<span className="text-muted-foreground">Blocks</span>
|
||||
<span className="font-mono tabular-nums">
|
||||
{formatNumber(st.diffusion_blocks)}
|
||||
</span>
|
||||
</div>
|
||||
)}
|
||||
{cacheHits > 0 && (
|
||||
<div className="flex items-center justify-between gap-4">
|
||||
<span className="text-muted-foreground">Cache hits</span>
|
||||
|
|
@ -165,6 +291,7 @@ export const MessageTiming: FC<{
|
|||
</span>
|
||||
</div>
|
||||
</>
|
||||
)
|
||||
) : (
|
||||
<>
|
||||
{/* Client-side metrics (safetensors + external provider fallback) */}
|
||||
|
|
|
|||
|
|
@ -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 (
|
||||
<div className="aui-diffusion-canvas my-1.5 overflow-hidden rounded-lg border border-primary/20 bg-primary/[0.03]">
|
||||
<div className="flex items-center gap-2 border-b border-primary/10 px-3 py-1.5 text-[11px] font-medium text-primary/80">
|
||||
<span className="inline-block size-1.5 animate-pulse rounded-full bg-primary" />
|
||||
<span>Denoising</span>
|
||||
<span className="opacity-60">
|
||||
block {canvas.block + 1} - {stepLabel}
|
||||
</span>
|
||||
</div>
|
||||
<pre className="max-h-[60vh] overflow-auto whitespace-pre-wrap px-3 py-2 font-mono text-[12.5px] leading-relaxed text-foreground/90">
|
||||
{canvas.text}
|
||||
</pre>
|
||||
</div>
|
||||
);
|
||||
};
|
||||
|
||||
const AssistantMessage: FC = () => {
|
||||
return (
|
||||
<MessagePrimitive.Root
|
||||
|
|
@ -2596,6 +2628,7 @@ const AssistantMessage: FC = () => {
|
|||
<div className="aui-assistant-message-content wrap-break-word min-w-0 text-[#0d0d0d] dark:text-foreground leading-relaxed">
|
||||
<GeneratingIndicator />
|
||||
<CancelledIndicator />
|
||||
<DiffusionCanvas />
|
||||
<MessagePrimitive.Parts
|
||||
components={{
|
||||
Text: MarkdownText,
|
||||
|
|
|
|||
|
|
@ -117,6 +117,21 @@ interface ServerTimings {
|
|||
predicted_ms: number;
|
||||
predicted_per_token_ms: number;
|
||||
predicted_per_second: number;
|
||||
// DiffusionGemma-only extras (present when serving a diffusion model; ignored otherwise).
|
||||
diffusion?: boolean;
|
||||
diffusion_blocks?: number;
|
||||
diffusion_steps?: number;
|
||||
diffusion_canvas?: number;
|
||||
diffusion_prompt_n?: number;
|
||||
diffusion_prompt_prepare_ms?: number;
|
||||
diffusion_decode_ms?: number;
|
||||
diffusion_wall_ms?: number;
|
||||
// Honest throughput, matching the standalone diffusion CLI:
|
||||
// effective = canvas*blocks/wall, parallel = canvas/per_step, output = answer tokens/wall.
|
||||
diffusion_effective_tok_s?: number;
|
||||
diffusion_parallel_tok_s?: number;
|
||||
diffusion_output_tok_s?: number;
|
||||
diffusion_steps_per_second?: number;
|
||||
}
|
||||
|
||||
type RunMessages = Parameters<ChatModelAdapter["run"]>[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) {
|
||||
|
|
|
|||
|
|
@ -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 &&
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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<ChatRuntimeStore>((set, get) => ({
|
|||
),
|
||||
toolStatus: null,
|
||||
generatingStatus: null,
|
||||
activeDiffusionCanvas: null,
|
||||
autoHealToolCalls: true,
|
||||
maxToolCallsPerMessage: 25,
|
||||
toolCallTimeout: 5,
|
||||
|
|
@ -888,6 +904,7 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
|
|||
tensorParallel: false,
|
||||
loadedTensorParallel: null,
|
||||
loadedIsMultimodal: false,
|
||||
loadedIsDiffusion: false,
|
||||
customContextLength: null,
|
||||
defaultChatTemplate: null,
|
||||
chatTemplateOverride: null,
|
||||
|
|
@ -1093,6 +1110,7 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((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<ChatRuntimeStore>((set, get) => ({
|
|||
tensorParallel: false,
|
||||
loadedTensorParallel: null,
|
||||
loadedIsMultimodal: false,
|
||||
loadedIsDiffusion: false,
|
||||
customContextLength: null,
|
||||
defaultChatTemplate: null,
|
||||
chatTemplateOverride: null,
|
||||
|
|
@ -1259,6 +1278,8 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
|
|||
return { ragAutoInjectMinScore };
|
||||
}),
|
||||
setToolStatus: (toolStatus) => set({ toolStatus }),
|
||||
setActiveDiffusionCanvas: (activeDiffusionCanvas) =>
|
||||
set({ activeDiffusionCanvas }),
|
||||
setGeneratingStatus: (generatingStatus) => set({ generatingStatus }),
|
||||
setAutoHealToolCalls: (autoHealToolCalls) =>
|
||||
set((state) => {
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue