From 6ed24b061378a07408d869736b63407f370a8d15 Mon Sep 17 00:00:00 2001 From: oobabooga <112222186+oobabooga@users.noreply.github.com> Date: Fri, 3 Jul 2026 14:53:06 -0300 Subject: [PATCH 1/5] Forward declared tools to the DiffusionGemma visual server --- unsloth_zoo/diffusion_studio/shim.py | 8 ++++++-- unsloth_zoo/diffusion_studio/visual_engine.py | 20 ++++++++++++------- 2 files changed, 19 insertions(+), 9 deletions(-) diff --git a/unsloth_zoo/diffusion_studio/shim.py b/unsloth_zoo/diffusion_studio/shim.py index 6055b52bf7..5d9a0b5eb3 100644 --- a/unsloth_zoo/diffusion_studio/shim.py +++ b/unsloth_zoo/diffusion_studio/shim.py @@ -221,6 +221,10 @@ def health(): async def chat(req: Request): body = await req.json() messages = body.get("messages", []) + # Declared tools ride through to the server so the chat template can render them + # (inputs.tools); the model then emits <|tool_call> blocks the caller can heal into + # structured tool_calls. None for plain chat, so the request shape is unchanged. + tools = body.get("tools") stream = bool(body.get("stream", False)) max_blocks = _max_blocks(body) seed = int(body.get("seed", 3407)) @@ -234,7 +238,7 @@ async def chat(req: Request): def work(): with _LOCK: return V.generate_visual(srv, messages, seed=seed, max_blocks=max_blocks, - on_stats=stats_box.update) + on_stats=stats_box.update, tools=tools) try: text = await loop.run_in_executor(None, work) except V.ContextOverflow as exc: @@ -276,7 +280,7 @@ def work(): with _LOCK: full = V.generate_visual(srv, messages, seed=seed, max_blocks=max_blocks, on_frame=on_frame, on_commit=on_commit, - on_stats=stats_box.update) + on_stats=stats_box.update, tools=tools) loop.call_soon_threadsafe(q.put_nowait, ("done", full)) except V.ContextOverflow as exc: # context budget exceeded -> clean user-facing message loop.call_soon_threadsafe(q.put_nowait, ("overflow", exc)) diff --git a/unsloth_zoo/diffusion_studio/visual_engine.py b/unsloth_zoo/diffusion_studio/visual_engine.py index b7e7e13d61..d901d4aa6e 100644 --- a/unsloth_zoo/diffusion_studio/visual_engine.py +++ b/unsloth_zoo/diffusion_studio/visual_engine.py @@ -216,14 +216,20 @@ def restart(self): pass self._spawn() - def _send(self, messages, n_blocks, seed): + def _send(self, messages, n_blocks, seed, tools=None): # A previous turn may have crashed the decoder; respawn before writing so a dead child does not # strand this and every later turn with a broken pipe. if self.p is None or self.p.poll() is not None: self.restart() + req = {"seed": int(seed), "n_blocks": int(n_blocks), "messages": messages} + # Forward declared tools so the server can render them into the chat template + # (inputs.tools); the model then emits <|tool_call> blocks with schema-correct + # argument names instead of guessing. Omitted when absent so the request shape + # is unchanged for plain chat. + if tools: + req["tools"] = tools with open(self.req, "w") as f: - json.dump({"seed": int(seed), "n_blocks": int(n_blocks), "messages": messages}, f, - ensure_ascii=False) + json.dump(req, f, ensure_ascii=False) try: self.p.stdin.write(self.req + "\n") self.p.stdin.flush() @@ -256,7 +262,7 @@ def _parse_stats(line): return stats -def generate_visual(server, messages, seed=3407, max_blocks=8, on_frame=None, on_commit=None, on_stats=None): +def generate_visual(server, messages, seed=3407, max_blocks=8, on_frame=None, on_commit=None, on_stats=None, tools=None): """Stream one turn through the optimized visual server. on_frame(block, step, total, text): a denoising frame (the current argmax canvas, already decoded). @@ -282,7 +288,7 @@ def _commit(txt): for attempt in (0, 1): try: - return _generate_visual_once(server, messages, seed, max_blocks, _frame, _commit, on_stats) + return _generate_visual_once(server, messages, seed, max_blocks, _frame, _commit, on_stats, tools) except VisualServerCrashed: server.restart() # always bring a fresh server up so the next turn works regardless if attempt == 1 or progressed["emitted"]: @@ -290,8 +296,8 @@ def _commit(txt): # nothing emitted yet -> safe to transparently resend on the fresh server -def _generate_visual_once(server, messages, seed, max_blocks, on_frame, on_commit, on_stats): - server._send(messages, max_blocks, seed) +def _generate_visual_once(server, messages, seed, max_blocks, on_frame, on_commit, on_stats, tools=None): + server._send(messages, max_blocks, seed, tools) full_text = "" while True: From b8f3c31838c9d992ecf50835ad3fd279d0811a93 Mon Sep 17 00:00:00 2001 From: oobabooga <112222186+oobabooga@users.noreply.github.com> Date: Fri, 3 Jul 2026 15:11:51 -0300 Subject: [PATCH 2/5] Clarify the tools-forwarding comments --- unsloth_zoo/diffusion_studio/shim.py | 5 +---- unsloth_zoo/diffusion_studio/visual_engine.py | 7 +++---- 2 files changed, 4 insertions(+), 8 deletions(-) diff --git a/unsloth_zoo/diffusion_studio/shim.py b/unsloth_zoo/diffusion_studio/shim.py index 5d9a0b5eb3..c8b3a54fc3 100644 --- a/unsloth_zoo/diffusion_studio/shim.py +++ b/unsloth_zoo/diffusion_studio/shim.py @@ -221,10 +221,7 @@ def health(): async def chat(req: Request): body = await req.json() messages = body.get("messages", []) - # Declared tools ride through to the server so the chat template can render them - # (inputs.tools); the model then emits <|tool_call> blocks the caller can heal into - # structured tool_calls. None for plain chat, so the request shape is unchanged. - tools = body.get("tools") + tools = body.get("tools") # forwarded to the visual server (rendered once the binary reads it) stream = bool(body.get("stream", False)) max_blocks = _max_blocks(body) seed = int(body.get("seed", 3407)) diff --git a/unsloth_zoo/diffusion_studio/visual_engine.py b/unsloth_zoo/diffusion_studio/visual_engine.py index d901d4aa6e..bdd567d4b0 100644 --- a/unsloth_zoo/diffusion_studio/visual_engine.py +++ b/unsloth_zoo/diffusion_studio/visual_engine.py @@ -222,10 +222,9 @@ def _send(self, messages, n_blocks, seed, tools=None): if self.p is None or self.p.poll() is not None: self.restart() req = {"seed": int(seed), "n_blocks": int(n_blocks), "messages": messages} - # Forward declared tools so the server can render them into the chat template - # (inputs.tools); the model then emits <|tool_call> blocks with schema-correct - # argument names instead of guessing. Omitted when absent so the request shape - # is unchanged for plain chat. + # Forward tools so the server can render them into the chat template (inputs.tools), + # for schema-correct <|tool_call> args instead of guesses. No-op until the visual + # server reads the field (pending llama.cpp support); harmless to send meanwhile. if tools: req["tools"] = tools with open(self.req, "w") as f: From 12d31539d37db5c9a3faa80ac670053995082db2 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Tue, 7 Jul 2026 13:06:00 +0000 Subject: [PATCH 3/5] Honor tool_choice before forwarding tools to the DiffusionGemma visual server --- unsloth_zoo/diffusion_studio/shim.py | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/unsloth_zoo/diffusion_studio/shim.py b/unsloth_zoo/diffusion_studio/shim.py index c8b3a54fc3..04489e9133 100644 --- a/unsloth_zoo/diffusion_studio/shim.py +++ b/unsloth_zoo/diffusion_studio/shim.py @@ -217,11 +217,25 @@ def health(): return {"status": "ok", "model": MODEL_ID} +def _tools_for_choice(tools, tool_choice): + # Honor tool_choice before advertising: "none" hides tools, a forced function + # narrows to just that one, anything else forwards all. + if isinstance(tool_choice, str) and tool_choice.lower() == "none": + return None + if isinstance(tool_choice, dict): + name = (tool_choice.get("function") or {}).get("name") + if name: + return [t for t in tools or [] + if isinstance(t, dict) and (t.get("function") or {}).get("name") == name] or None + return tools + + @app.post("/v1/chat/completions") async def chat(req: Request): body = await req.json() messages = body.get("messages", []) - tools = body.get("tools") # forwarded to the visual server (rendered once the binary reads it) + # forwarded to the visual server (rendered once the binary reads it) + tools = _tools_for_choice(body.get("tools"), body.get("tool_choice")) stream = bool(body.get("stream", False)) max_blocks = _max_blocks(body) seed = int(body.get("seed", 3407)) From f6334addc24f531f6e4a73a0da2b84b8456e84de Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Wed, 8 Jul 2026 03:49:12 +0000 Subject: [PATCH 4/5] Tighten tools-forwarding comments --- unsloth_zoo/diffusion_studio/shim.py | 2 +- unsloth_zoo/diffusion_studio/visual_engine.py | 5 ++--- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/unsloth_zoo/diffusion_studio/shim.py b/unsloth_zoo/diffusion_studio/shim.py index 04489e9133..2687bfd901 100644 --- a/unsloth_zoo/diffusion_studio/shim.py +++ b/unsloth_zoo/diffusion_studio/shim.py @@ -234,7 +234,7 @@ def _tools_for_choice(tools, tool_choice): async def chat(req: Request): body = await req.json() messages = body.get("messages", []) - # forwarded to the visual server (rendered once the binary reads it) + # forwarded to the visual server, honoring tool_choice tools = _tools_for_choice(body.get("tools"), body.get("tool_choice")) stream = bool(body.get("stream", False)) max_blocks = _max_blocks(body) diff --git a/unsloth_zoo/diffusion_studio/visual_engine.py b/unsloth_zoo/diffusion_studio/visual_engine.py index bdd567d4b0..52b1ca6e7a 100644 --- a/unsloth_zoo/diffusion_studio/visual_engine.py +++ b/unsloth_zoo/diffusion_studio/visual_engine.py @@ -222,9 +222,8 @@ def _send(self, messages, n_blocks, seed, tools=None): if self.p is None or self.p.poll() is not None: self.restart() req = {"seed": int(seed), "n_blocks": int(n_blocks), "messages": messages} - # Forward tools so the server can render them into the chat template (inputs.tools), - # for schema-correct <|tool_call> args instead of guesses. No-op until the visual - # server reads the field (pending llama.cpp support); harmless to send meanwhile. + # Forward tools so the server renders them into the chat template (inputs.tools) + # for schema-correct <|tool_call> args. if tools: req["tools"] = tools with open(self.req, "w") as f: From 741b83f186da6e791dee7ea119ce8b398fd5ab11 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Wed, 8 Jul 2026 04:01:44 +0000 Subject: [PATCH 5/5] Write the visual-server request file as UTF-8 --- unsloth_zoo/diffusion_studio/visual_engine.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/unsloth_zoo/diffusion_studio/visual_engine.py b/unsloth_zoo/diffusion_studio/visual_engine.py index 52b1ca6e7a..5fbeabdbb9 100644 --- a/unsloth_zoo/diffusion_studio/visual_engine.py +++ b/unsloth_zoo/diffusion_studio/visual_engine.py @@ -226,7 +226,7 @@ def _send(self, messages, n_blocks, seed, tools=None): # for schema-correct <|tool_call> args. if tools: req["tools"] = tools - with open(self.req, "w") as f: + with open(self.req, "w", encoding="utf-8") as f: json.dump(req, f, ensure_ascii=False) try: self.p.stdin.write(self.req + "\n")