ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Fix/adjust diffusion: round 29 P1 + P2 batch for PR #5754

Browse source 
Five actionable findings from round 29 reviewer aggregate, plus an
origin/main merge that absorbs the chat_templates.py fix landed in
PR #5763. Skipped #4 / #5 (studio.txt + constraints.txt hub bump)
because CI evidence from round 26 contradicts that suggestion; the
real broken combo only happens via the --no-deps no-torch path
which is already bumped in no-torch-runtime.txt + pyproject.toml.

1. core/inference/diffusion.py: round 28 reordered
   _release_chat_backend_for_diffusion BEFORE
   _release_other_gpu_owners_for_diffusion to surface the helper /
   advisor busy check early, but that meant the chat unload inside
   _release_chat_backend_for_diffusion now fired before the
   training / export conflict check in the second helper. A direct
   backend caller (tests, scripts) or a route-precheck race with a
   newly-started training run would then unload the user's chat and
   then 409 with nothing loaded. Split the helper busy check into
   _raise_if_helper_advisor_busy_for_diffusion (cheap, no side
   effects), keep _release_chat_backend_for_diffusion as the
   actual chat unload with an opt-out flag, and reorder load_model
   to: (a) helper check, (b) training / export check + idle export
   shutdown, (c) chat unload. All raises now fire BEFORE any
   destructive unload.

2. Merge origin/main: absorbs af6504f9 (PR #5763
   chat_templates.py find() guards + the new
   tests/python/test_construct_chat_template_validation.py
   regression test). Removes the 101-line stale-rebase silent
   revert that round 29 reviewer 5 and 8 flagged.

3. frontend/src/features/images/images-page.tsx: supportsNegativePrompt
   now also honours customFamily when no model is loaded yet, so a
   Custom HF repo with family flux.2 / flux.2-klein correctly hides
   the negative prompt field instead of silently sending it.

4. routes/inference.py /images/generate: report the ACTUAL PNG
   width / height from PIL Image.size instead of echoing back the
   requested payload values. FLUX-family pipelines round to
   vae_scale_factor * 2, so a request for 520x520 lands as 512x512
   internally; metadata now matches the bytes on the wire.

Tests: 98 targeted (diffusion + cached_gguf + inference_validation)
and frontend npm run typecheck pass locally.
This commit is contained in:
Daniel Han-Chen 2026-05-25 14:45:15 +00:00
commit bec81b882d
3 changed files with 57 additions and 25 deletions
Download patch file Download diff file Expand all files Collapse all files

61
studio/backend/core/inference/diffusion.py
View file

@ -1035,13 +1035,18 @@ class DiffusionBackend:
# transformer while the old pipeline still owns
# its weights.
# 4. THEN call from_single_file / from_pretrained.
# Round 28 P1 #4: helper/advisor check must fire BEFORE
# _release_other_gpu_owners_for_diffusion. Otherwise a
# blocked Images load could first tear down an idle
# export checkpoint just to then RuntimeError on the
# helper check inside _release_chat_backend_for_diffusion.
_release_chat_backend_for_diffusion()
# Round 29 P1 #1: do ALL cheap conflict checks BEFORE
# any destructive unload, so a training/export conflict
# caught inside _release_other_gpu_owners_for_diffusion
# does NOT leave the user with no chat model after we
# already unloaded it. The helper-busy check is
# split out of _release_chat_backend_for_diffusion;
# _release_other_gpu_owners_for_diffusion raises
# RuntimeError early when training/export is active
# without touching the chat backend.
_raise_if_helper_advisor_busy_for_diffusion()
_release_other_gpu_owners_for_diffusion()
_release_chat_backend_for_diffusion(check_helper_advisor = False)
old = self._pipe
if old is not None:
@ -1505,7 +1510,26 @@ def encode_png_base64(pil_image: "Any") -> str:
# ─── Helpers ──────────────────────────────────────────────────────────
def _release_chat_backend_for_diffusion() -> None:
def _raise_if_helper_advisor_busy_for_diffusion() -> None:
"""Round 29 P1 #1: split the helper-busy check out of
_release_chat_backend_for_diffusion so the diffusion load can
check ALL conflicts (helper, training, export) BEFORE doing ANY
destructive unloads. Otherwise a route-precheck race or a direct
backend call would unload the user's chat while training was
active, then 409 with the user holding no model at all.
"""
try:
from utils.datasets.llm_assist import helper_advisor_busy
except Exception:
return
if helper_advisor_busy():
raise RuntimeError(
"AI Assist (helper / advisor GGUF) is still using the GPU. "
"Wait for it to finish before loading a diffusion image model."
)
def _release_chat_backend_for_diffusion(*, check_helper_advisor: bool = True) -> None:
"""Unload any running chat backend before a diffusion load.
Diffusion pipelines on FLUX-class models can eat 12-24 GB of VRAM,
@ -1521,20 +1545,15 @@ def _release_chat_backend_for_diffusion() -> None:
diffusion ``load_model`` bails out instead of double-owning VRAM
(round 17 P1 #2).
"""
# Round 27 P1 #2: helper / advisor GGUF loads run on a PRIVATE
# LlamaCppBackend so the global llama check below cannot see them.
# Refuse the diffusion handoff while a helper / advisor still owns
# its private backend so we do not allocate FLUX VRAM on top.
try:
from utils.datasets.llm_assist import helper_advisor_busy
except Exception:
pass
else:
if helper_advisor_busy():
raise RuntimeError(
"AI Assist (helper / advisor GGUF) is still using the GPU. "
"Wait for it to finish before loading a diffusion image model."
)
# Round 27 P1 #2 / round 29 P1 #1: helper / advisor GGUF loads
# run on a PRIVATE LlamaCppBackend so the global llama check below
# cannot see them. The actual busy check now lives in
# _raise_if_helper_advisor_busy_for_diffusion so the caller can do
# ALL conflict checks BEFORE any destructive unload. Kept here as
# a default-on safety net for callers that did not run the
# standalone check.
if check_helper_advisor:
_raise_if_helper_advisor_busy_for_diffusion()
# 1. GGUF chat backend (llama-server subprocess). We unload when
# EITHER is_loaded is True (resident model) OR is_active is
# True (mid-download / startup) OR loading_model_identifier is

9
studio/backend/routes/inference.py
View file

@ -2395,11 +2395,16 @@ async def diffusion_generate(
raise HTTPException(status_code = 500, detail = str(exc))
duration_ms = int((time.time() - start) * 1000)
# Round 29 P2 #14: FLUX-family pipelines round (width, height) to
# vae_scale_factor * 2 multiples internally, so the actual PNG can
# differ from the requested dims. Report the real image size so
# the metadata caption matches the bytes on the wire.
actual_w, actual_h = (image.size if hasattr(image, "size") else (payload.width, payload.height))
return DiffusionGenerateResponse(
image_b64 = encode_png_base64(image),
image_mime = "image/png",
width = payload.width,
height = payload.height,
width = int(actual_w),
height = int(actual_h),
num_inference_steps = payload.num_inference_steps,
guidance_scale = payload.guidance_scale,
seed = payload.seed,

12
studio/frontend/src/features/images/images-page.tsx
View file

@ -308,15 +308,23 @@ export function ImagesPage() {
// FLUX.2 / FLUX.2 klein pipelines do NOT accept negative_prompt and
// would 500 if we sent one through. The backend strips the field
// defensively but hiding it client-side keeps the UI honest.
// Round 29 P2 #12: also honour the user-picked customFamily when no
// model is loaded yet, so a Custom HF repo with family flux.2 /
// flux.2-klein hides the negative-prompt field correctly.
const supportsNegativePrompt = useMemo(() => {
const family = status?.family;
if (!family) {
const candidate = useCustom ? undefined : preset.family;
let candidate: string | undefined;
if (useCustom) {
candidate = customFamily === "auto" ? undefined : customFamily;
} else {
candidate = preset.family;
}
if (!candidate) return true;
return !candidate.startsWith("flux.2");
}
return !family.startsWith("flux.2");
}, [status, useCustom, preset.family]);
}, [status, useCustom, customFamily, preset.family]);
return (
<div className="flex flex-1 flex-col gap-4 overflow-y-auto p-4 sm:p-6">