Compare commits

...
Sign in to create a new pull request.

2 commits

Author SHA1 Message Date
pre-commit-ci[bot]
72200c7fce [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2026-06-11 16:03:17 +00:00
Daniel Han
fd5f0a1331 studio: serve DiffusionGemma GGUFs natively with the denoising canvas
Detect a DiffusionGemma GGUF (general.architecture=diffusion-gemma) and serve
it through the diffusion runner instead of llama-server: launch the OpenAI-compat
shim (unsloth_zoo.diffusion_studio.shim, or UNSLOTH_DG_SHIM) driving the on-device
visual decoder, resolving the visual-server binary from DG_VISUAL_BIN or next to
llama-server. Surface is_diffusion to the UI and auto-render the per-step denoising
canvas artifact for DiffusionGemma (no toggle). install_llama_prebuilt + setup.sh/ps1
place the visual-server binary in the install tree best-effort. Other models unchanged.
2026-06-11 16:00:37 +00:00
10 changed files with 356 additions and 7 deletions

View file

@ -683,6 +683,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
@ -787,6 +792,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
@ -2050,11 +2060,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",
@ -2109,6 +2124,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",
@ -2137,6 +2153,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":
@ -2203,6 +2221,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:
@ -2221,6 +2250,189 @@ 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: a shim invocation
prefix (argv), the visual-server binary, and an optional extra PYTHONPATH
dir (only for the file-based override).
Shim: UNSLOTH_DG_SHIM (a .py file) first, otherwise the installed
unsloth_zoo.diffusion_studio.shim package. Binary: DG_VISUAL_BIN first,
otherwise alongside llama-server in the install tree. The visual server
tokenizes and applies the chat template from the GGUF itself, so no
tokenizer files are needed. Returns None if the binary or a shim cannot
be found.
"""
import importlib.util
import os
import sys
# Visual-server binary: env override, else the install tree (sibling of llama-server).
visual_bin = os.environ.get("DG_VISUAL_BIN")
if not visual_bin:
base = self._find_llama_server_binary()
if base:
cand = Path(base).parent / "llama-diffusion-gemma-visual-server"
if cand.is_file():
visual_bin = str(cand)
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))
# Detect the installed unsloth_zoo.diffusion_studio.shim WITHOUT importing the
# heavy unsloth_zoo package into this backend process (find_spec on the
# top-level package does not execute 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()
# The whole [prompt | 256-canvas] must fit one ubatch; the model's native
# context (262144) is irrelevant here, so cap to a sane per-turn budget.
maxtok = n_ctx if (n_ctx and 0 < n_ctx <= 16384) else 8192
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:
env["PYTHONPATH"] = extra_pythonpath + os.pathsep + env.get("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, via its own pdeathsig, the visual
# server) dies if this backend process dies, so a Studio crash/restart
# never orphans a GPU process.
popen_kwargs = dict(_windows_hidden_subprocess_kwargs())
if os.name == "posix":
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._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
self._effective_context_length = maxtok
self._max_context_length = self._context_length or maxtok
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)
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(
@ -2908,13 +3120,10 @@ class LlamaCppBackend:
with self._lock:
self._kill_process()
# Resolve llama-server now, but defer the not-found error: a
# block-diffusion GGUF is served by the diffusion runner instead
# (the architecture is only known after the GGUF header is read).
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
@ -2969,6 +3178,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

View file

@ -160,6 +160,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)")
@ -273,6 +276,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")

View file

@ -1330,6 +1330,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),
@ -1578,6 +1579,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,
@ -2045,6 +2047,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,

View file

@ -284,9 +284,22 @@ function CodeBlockActions({
);
}
// DiffusionGemma streams its denoising visualization as a self-contained html
// canvas player; auto-render it as a sandboxed-iframe artifact (no manual toggle)
// for the diffusion model only. Matches the native-served flag (loadedIsDiffusion)
// and the external-connection model id (the "diffusiongemma" substring survives the
// external:: id encoding). Other models are unaffected.
function isDiffusionCheckpoint(checkpoint: string | null | undefined): boolean {
return !!checkpoint && checkpoint.toLowerCase().includes("diffusiongemma");
}
function StreamdownBlock(props: BlockProps) {
const shouldCollapseHtmlArtifacts = useChatRuntimeStore(
(state) => state.artifactsEnabled || state.collapseHtmlArtifacts,
(state) =>
state.artifactsEnabled ||
state.collapseHtmlArtifacts ||
state.loadedIsDiffusion ||
isDiffusionCheckpoint(state.params.checkpoint),
);
const messageHasRenderableRenderHtmlTool = useAuiState(({ message }) =>
message.parts.some(isRenderableRenderHtmlToolPart),

View file

@ -413,6 +413,7 @@ export function useChatModelRuntime() {
defaultChatTemplate: nextDefaultChatTemplate,
loadedIsMultimodal: isMultimodalResponse(statusRes),
specFallbackReason: statusRes.spec_fallback_reason ?? null,
loadedIsDiffusion: statusRes.is_diffusion ?? false,
...(prevState.loadedSpeculativeType === null && {
speculativeType: currentSpecType,
loadedSpeculativeType: currentSpecType,
@ -459,6 +460,7 @@ export function useChatModelRuntime() {
useChatRuntimeStore.setState({
modelRequiresTrustRemoteCode: false,
loadedIsMultimodal: false,
loadedIsDiffusion: false,
});
}
} catch (error) {
@ -782,6 +784,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.

View file

@ -425,6 +425,9 @@ type ChatRuntimeStore = {
specDraftNMax: number | null;
loadedSpecDraftNMax: number | null;
loadedIsMultimodal: boolean;
/** Active model is a block-diffusion model (DiffusionGemma): drives the
* denoising-canvas artifact auto-render. */
loadedIsDiffusion: boolean;
customContextLength: number | null;
defaultChatTemplate: string | null;
chatTemplateOverride: string | null;
@ -774,6 +777,7 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
specDraftNMax: null,
loadedSpecDraftNMax: null,
loadedIsMultimodal: false,
loadedIsDiffusion: false,
customContextLength: null,
defaultChatTemplate: null,
chatTemplateOverride: null,
@ -987,6 +991,7 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
specDraftNMax: null,
loadedSpecDraftNMax: null,
loadedIsMultimodal: false,
loadedIsDiffusion: false,
customContextLength: null,
defaultChatTemplate: null,
chatTemplateOverride: null,

View file

@ -109,6 +109,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;
@ -145,6 +146,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;

View file

@ -4514,6 +4514,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():
@ -6721,6 +6783,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")

View file

@ -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"

View file

@ -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