From bec81b882d6ed19e76b5b583a371cc35cb8bfb48 Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Mon, 25 May 2026 14:45:15 +0000 Subject: [PATCH] Fix/adjust diffusion: round 29 P1 + P2 batch for PR #5754 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. --- studio/backend/core/inference/diffusion.py | 61 ++++++++++++------- studio/backend/routes/inference.py | 9 ++- .../src/features/images/images-page.tsx | 12 +++- 3 files changed, 57 insertions(+), 25 deletions(-) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 3bdce54819..888e94a474 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -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 diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 636bb5cc9c..eb6ae19710 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -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, diff --git a/studio/frontend/src/features/images/images-page.tsx b/studio/frontend/src/features/images/images-page.tsx index 265786e6bd..b4468a534a 100644 --- a/studio/frontend/src/features/images/images-page.tsx +++ b/studio/frontend/src/features/images/images-page.tsx @@ -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 (