From 04de106e49295dc1e2c3a4ed05fa7fa1352be009 Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Mon, 25 May 2026 01:27:57 +0000 Subject: [PATCH] Fix/adjust diffusion: round 6 race-free lifecycle + delete guards for PR #5754 Round 6 reviewers identified several races between load / unload / generate and several fail-open delete guards. This commit closes them by widening the lock scope, publishing the pending load target through status(), and switching delete guards to fail-closed. Lifecycle (P1) * core/inference/diffusion.py: load_model now also takes _generate_lock. Previous behavior released and reallocated the pipeline while a generation forward was still iterating denoising steps, corrupting scheduler state and stacking VRAM. The forward only briefly touches _lock, so taking it on the load path does not introduce a deadlock. * core/inference/diffusion.py: unload_model now also takes _generate_lock. Without it, /images/unload returned is_loaded=False while a slow forward was still running, which let chat / training / export handoffs allocate VRAM on top of the still-resident pipeline. * core/inference/diffusion.py: previous pipeline release now happens BEFORE from_single_file / from_pretrained. Switching FLUX.2 klein 4B -> 9B on a 16-24 GB GPU was failing because the new transformer allocation overlapped the old pipe's residency. * core/inference/diffusion.py: failed pipeline from_pretrained now explicitly releases the just-loaded transformer; previously its weights stayed pinned to GPU until GC and made the next load more likely to OOM. Pending-target / delete guards (P1) * core/inference/diffusion.py: load_model now publishes _pending_repo_id / _pending_base_repo / _pending_gguf_filename under _lock at the start of the call (and refreshes _pending_base_repo when the smart-base / repo defaults resolve). status() exposes those as 'repo_id' / 'base_repo' / 'gguf_filename' during is_loading=True so delete guards can see the target before _repo_id is set on success. * routes/models.py /delete-cached + /delete-finetuned: diffusion status check now fails CLOSED (HTTP 503) when status() raises. Both guards previously logged and continued, which could let a delete proceed against a repo whose status was unverifiable. * routes/models.py: is_loading is also blocked on both guards so a mid-download / mid-from_pretrained rmtree is refused. Symmetric handoffs (P1) * routes/export.py: /load-checkpoint now refuses with HTTP 409 when training is active instead of calling stop_training(). Chat and /images/load did the same after round 5; export was the remaining asymmetry that would silently kill a long training run. * routes/training.py, routes/inference.py (GGUF and standard chat), routes/export.py: diffusion handoff now treats is_loading as is_loaded. The diffusion backend's unload waits on _load_lock + _generate_lock so an in-flight load completes first. Requirements (P1) * requirements/studio.txt: pin python-multipart explicitly. The Studio routes package's eager router imports include routes/datasets.py whose FastAPI UploadFile/File validation crashes with RuntimeError without it in fresh test envs. Frontend (P2) * features/images/api.ts + images-page.tsx: seed handling now accepts the full [-2^63, 2^64 - 1] range via BigInt. The previous safe-integer cap rejected valid uint64 seeds the backend accepts. A small stringify helper emits BigInts as JSON integers without touching the rest of the payload. Tests * test_diffusion_routes.py: load routes/inference.py via importlib.spec_from_file_location to avoid triggering routes/__init__.py (which would pull in training / datasets / data_recipe imports unrelated to diffusion tests). * test_diffusion_backend.py: status() during is_loading shows pending repo + base; unload waits for in-flight generation. --- studio/backend/core/inference/diffusion.py | 139 +++++++++++++----- studio/backend/requirements/studio.txt | 5 + studio/backend/routes/export.py | 62 ++++---- studio/backend/routes/inference.py | 25 +++- studio/backend/routes/models.py | 25 +++- studio/backend/routes/training.py | 11 +- .../backend/tests/test_diffusion_backend.py | 131 +++++++++++++++++ studio/backend/tests/test_diffusion_routes.py | 38 ++++- studio/frontend/src/features/images/api.ts | 17 ++- .../src/features/images/images-page.tsx | 40 +++-- 10 files changed, 407 insertions(+), 86 deletions(-) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 29324aebfd..8ae3a13f0f 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -290,6 +290,15 @@ class DiffusionBackend: self._loaded_at: Optional[float] = None self._loading: bool = False self._last_error: Optional[str] = None + # `_pending_*` fields advertise the target of an in-flight load + # so cache- and finetuned-delete guards can refuse to rmtree a + # repo while it is being downloaded / read. They are set under + # _lock at the start of load_model and cleared on success or + # in the finally block. The route layer reads them via + # status() under _lock. + self._pending_repo_id: Optional[str] = None + self._pending_base_repo: Optional[str] = None + self._pending_gguf_filename: Optional[str] = None # ── lifecycle ───────────────────────────────────────────────── @@ -311,16 +320,24 @@ class DiffusionBackend: # POSIX layouts) to any authenticated Studio session. with self._lock: gguf_basename = Path(self._gguf_path).name if self._gguf_path else None + # During an in-flight load, expose _pending_* so cache / + # finetuned delete guards can refuse to wipe the repo + # that is mid-download. After the load completes (success + # or failure), the pending fields are cleared so status() + # reverts to publishing only the resident pipeline's id. + effective_repo = self._repo_id or self._pending_repo_id + effective_base = self._base_repo or self._pending_base_repo + effective_gguf = gguf_basename or self._pending_gguf_filename return { "is_loaded": self._pipe is not None, "is_loading": self._loading, - "repo_id": self._repo_id, + "repo_id": effective_repo, "family": self._family.name if self._family else None, "pipeline_class": ( self._family.pipeline_class if self._family else None ), - "base_repo": self._base_repo, - "gguf_filename": gguf_basename, + "base_repo": effective_base, + "gguf_filename": effective_gguf, "device": self._device, "dtype": self._dtype, "loaded_at": self._loaded_at, @@ -406,10 +423,26 @@ class DiffusionBackend: # cannot both kick off a multi-GB download + GPU upload at once. # The second caller waits behind the first and then loads on top # of the now-populated state via the normal swap path. - with self._load_lock: + # _generate_lock is also taken so we do not start swapping the + # pipeline (release old + allocate new) while a previous + # generation is still iterating denoising steps; releasing the + # pipe out from under an in-flight forward corrupts scheduler + # state. Order: _load_lock -> _generate_lock -> _lock so a + # forward (which only takes _generate_lock + briefly _lock) + # cannot block a queued load forever. + with self._load_lock, self._generate_lock: with self._lock: self._loading = True self._last_error = None + # Publish the pending target so cache / finetuned + # delete guards can see what is mid-download even + # before _repo_id / _base_repo are populated on + # success. + self._pending_repo_id = repo_id + self._pending_base_repo = base_repo + self._pending_gguf_filename = ( + Path(gguf_filename).name if gguf_filename else None + ) try: pipeline_cls = getattr(diffusers, fam.pipeline_class, None) if pipeline_cls is None: @@ -433,6 +466,10 @@ class DiffusionBackend: # 9B GGUF picks the 9B base, not the 4B fallback if base_repo: effective_base = base_repo + # Refresh pending so delete guards see the actual + # base, not just caller-supplied None. + with self._lock: + self._pending_base_repo = effective_base elif not gguf_filename: # Guard: a repo that ends in "-GGUF" (the unsloth # convention) is GGUF-only and will 500 on @@ -447,8 +484,12 @@ class DiffusionBackend: "load target." ) effective_base = repo_id + with self._lock: + self._pending_base_repo = effective_base else: effective_base = _smart_base_repo(fam, repo_id) + with self._lock: + self._pending_base_repo = effective_base logger.info( "Loading diffusion model %s (family=%s, device=%s, dtype=%s, base=%s)", repo_id, @@ -476,17 +517,41 @@ class DiffusionBackend: # pipeline / transformer class, gated download token, # transient Hub error on the GGUF download) have now # been validated. Anything past this line allocates - # GPU memory, so release competing GPU owners before - # we touch from_single_file or from_pretrained: - # * Chat backends (llama-server + safetensors) so the - # diffusion transformer does not race them for VRAM. - # * Export subprocess (also holds GB on the same GPU). + # GPU memory, so: + # 1. Release competing GPU owners (chat + export). + # 2. Release any *previous* diffusion pipeline so the + # new transformer / new from_pretrained does not + # race the old pipe for VRAM. Switching between + # FLUX.2 klein 4B and 9B on a 16-24 GB GPU OOMs + # otherwise: from_single_file allocates the new + # transformer while the old pipeline still owns + # its weights. + # 3. THEN call from_single_file / from_pretrained. # Training is *not* unloaded here: the route layer # refuses /images/load with HTTP 409 when training is # active so the user keeps their long run. _release_chat_backend_for_diffusion() _release_other_gpu_owners_for_diffusion() + old = self._pipe + if old is not None: + with self._lock: + # Clear ALL metadata together so a failed swap + # cannot leave status() reporting the previous + # repo / family / base_repo on top of an empty + # pipe. The except block below will restore + # last_error so the caller knows what happened. + self._pipe = None + self._family = None + self._repo_id = None + self._gguf_path = None + self._base_repo = None + self._device = None + self._dtype = None + self._loaded_at = None + _release(old) + old = None + if gguf_filename: quant_config = diffusers.GGUFQuantizationConfig(compute_dtype = dtype) # Diffusers-format GGUFs (FLUX.2 klein / Qwen-Image / @@ -523,26 +588,19 @@ class DiffusionBackend: if hf_token: pipe_kwargs["token"] = hf_token - old = self._pipe - if old is not None: - with self._lock: - # Clear ALL metadata together so a failed swap - # cannot leave status() reporting the previous - # repo / family / base_repo on top of an empty - # pipe. The except block below will restore - # last_error so the caller knows what happened. - self._pipe = None - self._family = None - self._repo_id = None - self._gguf_path = None - self._base_repo = None - self._device = None - self._dtype = None - self._loaded_at = None - _release(old) - old = None - - pipe = pipeline_cls.from_pretrained(effective_base, **pipe_kwargs) + try: + pipe = pipeline_cls.from_pretrained(effective_base, **pipe_kwargs) + except Exception: + # If from_pretrained fails after the transformer was + # already loaded, the transformer object holds GPU + # weights that would only be freed at GC. Drop the + # local reference and force a collect so the next + # load attempt does not stack VRAM with a phantom + # transformer. + if transformer is not None: + _release(transformer) + transformer = None + raise if enable_model_cpu_offload and device == "cuda": pipe.enable_model_cpu_offload() else: @@ -557,8 +615,6 @@ class DiffusionBackend: self._device = device self._dtype = str(dtype).replace("torch.", "") self._loaded_at = time.time() - # ``old`` was released above before the new allocation; - # nothing left to free here. return self.status() except Exception as exc: @@ -577,12 +633,25 @@ class DiffusionBackend: finally: with self._lock: self._loading = False + # Clear pending so status() falls back to publishing + # the resident pipeline (or nothing, on a failed + # swap). Keeping pending alive after the load + # finishes would falsely block deletes forever. + self._pending_repo_id = None + self._pending_base_repo = None + self._pending_gguf_filename = None def unload_model(self) -> dict[str, Any]: - # Take the load lock too so unload cannot race with an in-flight - # load_model and have the load thread overwrite the cleared state - # after we already returned {"is_loaded": false}. - with self._load_lock: + # Take the load lock and the generate lock so unload cannot: + # * race with an in-flight load_model and have the load + # thread overwrite the cleared state after we already + # returned {"is_loaded": false}. + # * return is_loaded=false while a forward pass is still + # iterating denoising steps on the soon-to-be-freed pipe. + # The generate forward only holds _generate_lock (briefly + # _lock), so acquiring _generate_lock here blocks until any + # in-flight generation completes. + with self._load_lock, self._generate_lock: with self._lock: old = self._pipe self._pipe = None diff --git a/studio/backend/requirements/studio.txt b/studio/backend/requirements/studio.txt index 96f8816b57..6628eef7f7 100644 --- a/studio/backend/requirements/studio.txt +++ b/studio/backend/requirements/studio.txt @@ -1,6 +1,11 @@ # Studio UI backend dependencies typer fastapi +# Required by FastAPI's multipart upload route validation +# (routes/datasets.py uploads files via UploadFile/File). Without +# this, importing the routes package raises RuntimeError on startup +# and CPU-only test environments fail before any test runs. +python-multipart uvicorn pydantic packaging diff --git a/studio/backend/routes/export.py b/studio/backend/routes/export.py index ff9e3d6695..d45659a386 100644 --- a/studio/backend/routes/export.py +++ b/studio/backend/routes/export.py @@ -94,42 +94,52 @@ async def load_checkpoint( except Exception as e: logger.debug("llama-server unload skipped for export: %s", e) + # Symmetric lifecycle guard: refuse to load an export + # checkpoint while training is active so we do not silently + # terminate someone's long-running training job and possibly + # fail the export load on top of that. Mirrors the + # _raise_if_training_active checks in routes/inference.py for + # chat and /images/load. Fail-closed (503) when the training + # backend can be imported but its status check raises. + try: + from core.training import get_training_backend # type: ignore + + trn = get_training_backend() + if trn.is_training_active(): + raise HTTPException( + status_code = 409, + detail = ( + "Training is currently active. Stop the training " + "run before loading an export checkpoint." + ), + ) + except HTTPException: + raise + except Exception as e: + logger.debug("training activity check skipped for export: %s", e) + # Also unload any active diffusion pipeline (Images page); it # competes for the same GPU and would survive the inference - # shutdown above. Best effort; silently skip if the module is - # absent. + # shutdown above. is_loading is treated like is_loaded so an + # in-flight load is also waited out (the diffusion unload + # acquires _load_lock + _generate_lock and blocks until the + # current load completes, then unloads). Best effort; silently + # skip if the module is absent. try: from core.inference.diffusion import get_diffusion_backend diff = get_diffusion_backend() - if diff.is_loaded: - logger.info("Unloading diffusion model to free GPU memory for export") + diff_status = diff.status() + if diff_status.get("is_loaded") or diff_status.get("is_loading"): + logger.info( + "Unloading diffusion model (loaded=%s loading=%s) for export", + diff_status.get("is_loaded"), + diff_status.get("is_loading"), + ) diff.unload_model() except Exception as e: logger.debug("diffusion unload skipped for export: %s", e) - try: - from core.training import get_training_backend - - trn = get_training_backend() - if trn.is_training_active(): - logger.info("Stopping active training to free GPU memory for export") - trn.stop_training() - # Wait for training subprocess to actually exit before proceeding, - # otherwise it may still hold GPU memory when export tries to load. - for _ in range(60): # up to 30s - if not trn.is_training_active(): - break - import time - - time.sleep(0.5) - else: - logger.warning( - "Training subprocess did not exit within 30s, proceeding anyway" - ) - except Exception as e: - logger.warning("Could not stop training: %s", e) - backend = get_export_backend() # load_checkpoint spawns and waits on a subprocess and can take # minutes. Run it in a worker thread so the event loop stays diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 1ef2a3b260..f49d28e095 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -784,13 +784,21 @@ async def load_model( # Symmetric with /images/load: drop any active diffusion # pipeline so the GGUF chat load does not race the FLUX VAE - # for VRAM. Best effort; silently continue on failure. + # for VRAM. Also handles is_loading: unload_model takes + # _load_lock + _generate_lock and will wait out an + # in-flight load before clearing state. Best effort; + # silently continue on failure. try: from core.inference.diffusion import get_diffusion_backend diff_backend = get_diffusion_backend() - if diff_backend.is_loaded: - logger.info("Unloading diffusion pipeline before GGUF load") + diff_status = diff_backend.status() + if diff_status.get("is_loaded") or diff_status.get("is_loading"): + logger.info( + "Unloading diffusion (loaded=%s loading=%s) before GGUF load", + diff_status.get("is_loaded"), + diff_status.get("is_loading"), + ) diff_backend.unload_model() except Exception as e: logger.debug("diffusion unload skipped (GGUF path): %s", e) @@ -977,14 +985,19 @@ async def load_model( llama_backend.unload_model() # Unload any active diffusion pipeline so the new chat model is - # not racing the FLUX VAE for VRAM on a 16-24 GB card. + # not racing the FLUX VAE for VRAM on a 16-24 GB card. is_loading + # is treated like is_loaded; unload waits behind _load_lock + + # _generate_lock so the in-flight load completes first. try: from core.inference.diffusion import get_diffusion_backend diff_backend = get_diffusion_backend() - if diff_backend.is_loaded: + diff_status = diff_backend.status() + if diff_status.get("is_loaded") or diff_status.get("is_loading"): logger.info( - "Unloading diffusion pipeline before loading Unsloth chat model" + "Unloading diffusion (loaded=%s loading=%s) before chat load", + diff_status.get("is_loaded"), + diff_status.get("is_loading"), ) diff_backend.unload_model() except Exception as e: diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 1029dd5bbe..a8a18c081b 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -1973,6 +1973,12 @@ async def delete_finetuned_model( # the merged repo locally, then loaded it via /images/load with a # local path as repo_id). Without this guard /delete-finetuned # could rmtree the directory the diffusion backend is reading from. + # is_loading is also blocked: status() exposes _pending_repo_id / + # _pending_base_repo during the load window so deletes during a + # mid-flight from_pretrained are refused. + # Fail-CLOSED on exception (503) like the llama.cpp / safetensors + # guards above: an unverifiable diffusion state means we cannot + # confirm the target is safe to rmtree. try: from core.inference.diffusion import get_diffusion_backend @@ -2010,6 +2016,10 @@ async def delete_finetuned_model( logger.warning( "Could not check diffusion backend loaded model before delete: %s", e ) + raise HTTPException( + status_code = 503, + detail = "Could not verify diffusion load status before deleting", + ) from e try: if export_type == "gguf" and gguf_variant: @@ -2684,6 +2694,10 @@ async def delete_cached_model( # Match exactly on repo_id (case-insensitive) instead of prefix to # avoid blocking unrelated deletes like "org/model" while # "org/model-v2" is loaded. + # Fail-CLOSED on exception (return 503) like the neighboring + # llama.cpp / safetensors guards: we cannot verify whether the + # delete is safe, so refuse rather than risk corrupting the + # pipeline's mmap. try: from core.inference.diffusion import get_diffusion_backend @@ -2700,8 +2714,15 @@ async def delete_cached_model( ) except HTTPException: raise - except Exception: - pass + except Exception as e: + logger.warning( + "Could not check diffusion backend status before cache delete: %s", + e, + ) + raise HTTPException( + status_code = 503, + detail = "Could not verify diffusion load status before deleting cache", + ) from e try: cache_scans = _all_hf_cache_scans() diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 0715a79c66..9a9f6b9761 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -314,12 +314,19 @@ async def start_training( # Also unload any loaded diffusion pipeline (Images page); it # holds the same GPU and would survive the inference shutdown. + # is_loading=True is also handled (unload_model takes + # _load_lock + _generate_lock and waits the in-flight load out). try: from core.inference.diffusion import get_diffusion_backend diff_backend = get_diffusion_backend() - if diff_backend.is_loaded: - logger.info("Unloading diffusion model to free GPU memory for training") + diff_status = diff_backend.status() + if diff_status.get("is_loaded") or diff_status.get("is_loading"): + logger.info( + "Unloading diffusion (loaded=%s loading=%s) for training", + diff_status.get("is_loaded"), + diff_status.get("is_loading"), + ) diff_backend.unload_model() except Exception as e: logger.warning("Could not unload diffusion model: %s", e) diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index 7a4312d5dd..3a73af699f 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -1018,6 +1018,137 @@ def test_generate_image_does_not_block_status(monkeypatch): t.join(timeout = 5) +def test_load_publishes_pending_target_during_loading(monkeypatch): + """status() must expose the pending repo_id / base_repo / gguf + file while is_loading=True so cache- and finetuned-delete guards + can refuse to rmtree the repo being downloaded right now.""" + import threading + import core.inference.diffusion as d + from PIL import Image + + fake = _install_fake_diffusers(monkeypatch) + + pending_seen: dict = {} + pretrained_blocked = threading.Event() + pretrained_release = threading.Event() + + class _SlowPipeline: + @classmethod + def from_pretrained(cls, base_repo, **kwargs): + pretrained_blocked.set() + # Capture status() output while the load is blocked. + backend = d.get_diffusion_backend() + pending_seen.update(backend.status()) + pretrained_release.wait(timeout = 5) + inst = cls() + inst.base_repo = base_repo + return inst + + def __call__(self, **kwargs): + class _Out: + pass + + o = _Out() + o.images = [Image.new("RGB", (kwargs["width"], kwargs["height"]))] + return o + + def enable_model_cpu_offload(self): + pass + + def to(self, device): + return self + + fake.Flux2KleinPipeline = _SlowPipeline + + backend = d.get_diffusion_backend() + backend.unload_model() + + def do_load(): + try: + backend.load_model( + "unsloth/FLUX.2-klein-4B-GGUF", + gguf_filename = "flux-2-klein-4b-Q4_K_S.gguf", + ) + except Exception: + pass + + t = threading.Thread(target = do_load) + t.start() + try: + assert pretrained_blocked.wait(timeout = 5) + # While blocked inside from_pretrained, status reads should + # already see the pending repo so deletes can be refused. + assert pending_seen.get("is_loading") is True + assert pending_seen.get("repo_id") == "unsloth/FLUX.2-klein-4B-GGUF" + assert pending_seen.get("base_repo") == "black-forest-labs/FLUX.2-klein-4B" + finally: + pretrained_release.set() + t.join(timeout = 5) + + +def test_unload_waits_for_in_flight_generation(monkeypatch): + """unload_model() must not return is_loaded=False while a + generate_image forward is still iterating; otherwise routes/... + callers see the pipe as freed while it still owns GPU memory and + can race a subsequent load.""" + import threading + import core.inference.diffusion as d + from PIL import Image + + backend = d.get_diffusion_backend() + started = threading.Event() + release = threading.Event() + generation_finished = threading.Event() + + class _SlowPipe: + def __call__(self, **kw): + started.set() + release.wait(timeout = 5) + + class _Out: + pass + + o = _Out() + o.images = [Image.new("RGB", (kw["width"], kw["height"]))] + return o + + backend._pipe = _SlowPipe() + backend._device = "cpu" + backend._family = d._FAMILIES[0] + backend._repo_id = "stub/stub" + + def do_generate(): + try: + backend.generate_image(prompt = "x", num_inference_steps = 1, + guidance_scale = 1.0, width = 64, height = 64) + finally: + generation_finished.set() + + gen_thread = threading.Thread(target = do_generate) + gen_thread.start() + try: + assert started.wait(timeout = 5) + unload_returned = threading.Event() + + def do_unload(): + backend.unload_model() + unload_returned.set() + + unload_thread = threading.Thread(target = do_unload) + unload_thread.start() + # unload should block until release sets, NOT return early. + unload_thread.join(timeout = 0.5) + assert not unload_returned.is_set(), \ + "unload_model returned while generation was still running" + release.set() + unload_thread.join(timeout = 5) + assert unload_returned.is_set() + assert generation_finished.is_set() + finally: + release.set() + gen_thread.join(timeout = 5) + + def test_bf16_falls_back_to_fp16_on_old_cuda(monkeypatch): """CUDA availability does not imply BF16 support; old GPUs report is_available()=True and is_bf16_supported()=False. The backend diff --git a/studio/backend/tests/test_diffusion_routes.py b/studio/backend/tests/test_diffusion_routes.py index 9d21f0af0b..b7bfd1e2f8 100644 --- a/studio/backend/tests/test_diffusion_routes.py +++ b/studio/backend/tests/test_diffusion_routes.py @@ -7,10 +7,17 @@ Mounts the actual ``inference_router`` on a fresh FastAPI app with the auth dependency replaced by a stub so we exercise the same FastAPI handlers Studio ships in production. The diffusion backend is replaced with an in-memory stub so we don't need diffusers / GPUs to run these. + +To stay runnable in a minimal CPU-only env, ``routes/inference.py`` +is loaded directly via ``importlib`` so we do NOT trigger +``routes/__init__.py`` -- that file eagerly imports training / +datasets / data_recipe / export and would drag in heavy deps +(matplotlib, etc.) that the diffusion tests do not need. """ from __future__ import annotations +import importlib.util import sys from pathlib import Path @@ -25,6 +32,35 @@ if str(_BACKEND_ROOT) not in sys.path: sys.path.insert(0, str(_BACKEND_ROOT)) +def _import_inference_module(): + """Load ``routes/inference.py`` without executing ``routes/__init__``. + + The package init imports training / datasets / data_recipe / export + routers, which pull in matplotlib / pandas / training stack. The + diffusion tests only need the inference module so we side-step the + package import via importlib.spec_from_file_location. + """ + # If a previous test already imported routes the normal way, reuse + # the cached module instead of re-loading. + cached = sys.modules.get("routes.inference") + if cached is not None: + return cached + target = _BACKEND_ROOT / "routes" / "inference.py" + spec = importlib.util.spec_from_file_location( + "routes.inference", + target, + # We do NOT set submodule_search_locations for routes itself + # because that would re-trigger routes/__init__.py. The module + # uses relative imports sparingly; absolute imports resolve via + # sys.path[0] = backend root. + ) + assert spec and spec.loader, "could not build spec for routes/inference.py" + module = importlib.util.module_from_spec(spec) + sys.modules["routes.inference"] = module + spec.loader.exec_module(module) + return module + + class _FakeBackend: def __init__(self) -> None: self._loaded = False @@ -71,7 +107,7 @@ class _FakeBackend: def app_with_stub(monkeypatch): """Build a FastAPI app that mounts the real inference router with auth disabled and the diffusion backend swapped for a stub.""" - from routes import inference as inf + inf = _import_inference_module() import core.inference.diffusion as d stub = _FakeBackend() diff --git a/studio/frontend/src/features/images/api.ts b/studio/frontend/src/features/images/api.ts index e576f3987e..0a09d971de 100644 --- a/studio/frontend/src/features/images/api.ts +++ b/studio/frontend/src/features/images/api.ts @@ -47,7 +47,10 @@ export interface DiffusionGenerateRequest { guidance_scale?: number; width?: number; height?: number; - seed?: number; + // bigint when the seed exceeds Number.MAX_SAFE_INTEGER, otherwise + // number. The wire format is always a JSON integer; see + // ``stringifyWithBigInt`` below. + seed?: number | bigint; } export interface DiffusionGenerateResponse { @@ -92,6 +95,16 @@ export async function unloadDiffusionModel(): Promise<{ is_loaded: boolean }> { ); } +/** JSON.stringify cannot serialise BigInt directly. We only ever + * have BigInts in the seed field, which is an integer; emit the + * literal digits so the server receives a JSON integer rather than + * a string. Pydantic v2 accepts arbitrarily large ints. */ +function stringifyWithBigInt(value: unknown): string { + return JSON.stringify(value, (_, v) => + typeof v === "bigint" ? `__bigint__:${v.toString()}` : v, + ).replace(/"__bigint__:(-?\d+)"/g, "$1"); +} + export async function generateDiffusionImage( payload: DiffusionGenerateRequest, ): Promise { @@ -99,7 +112,7 @@ export async function generateDiffusionImage( await authFetch("/api/inference/images/generate", { method: "POST", headers: { "Content-Type": "application/json" }, - body: JSON.stringify(payload), + body: stringifyWithBigInt(payload), }), ); } diff --git a/studio/frontend/src/features/images/images-page.tsx b/studio/frontend/src/features/images/images-page.tsx index 3b3891e9e6..64f06bfbb0 100644 --- a/studio/frontend/src/features/images/images-page.tsx +++ b/studio/frontend/src/features/images/images-page.tsx @@ -214,28 +214,44 @@ export function ImagesPage() { } setBusy("generating"); try { - // Reject non-integer or out-of-safe-integer-range seeds rather - // than silently rounding via Number(). The backend takes an int - // and a precision loss here would yield a different image than - // the seed the user typed. + // Reject non-integer seeds and clamp to the [-2^63, 2^64 - 1] + // range the backend's torch.Generator can actually pack. JSON + // serialises BigInts as plain integers, so we keep the wire + // format compatible and avoid the Number(seed) precision loss + // (>= 2^53 silently rounds, producing a different image than + // the seed the user typed). When the seed fits a safe integer + // it goes through unchanged; larger seeds ride along as their + // BigInt-derived string via the wire-format BigInt JSON helper + // in the api layer. const seedStr = seed.trim(); - let parsedSeed: number | undefined; + let parsedSeed: number | bigint | undefined; if (seedStr) { if (!/^-?\d+$/.test(seedStr)) { toast.error("Seed must be an integer"); return; } - const candidate = Number(seedStr); - if ( - !Number.isFinite(candidate) || - !Number.isSafeInteger(candidate) - ) { + let big: bigint; + try { + big = BigInt(seedStr); + } catch { + toast.error("Seed must be an integer"); + return; + } + const SEED_MIN = -(BigInt(2) ** BigInt(63)); + const SEED_MAX = BigInt(2) ** BigInt(64) - BigInt(1); + if (big < SEED_MIN || big > SEED_MAX) { toast.error( - "Seed must fit in a JavaScript safe integer (<= 2^53 - 1)", + "Seed must be in [-2^63, 2^64 - 1] (the torch.Generator range)", ); return; } - parsedSeed = candidate; + // Use a plain Number when it fits a safe integer so the + // existing api.ts JSON serialiser does not break on BigInt; + // otherwise pass the BigInt and let api.ts emit it as a JSON + // number via a custom replacer. + const SAFE_MAX = BigInt(Number.MAX_SAFE_INTEGER); + const SAFE_MIN = -SAFE_MAX; + parsedSeed = big >= SAFE_MIN && big <= SAFE_MAX ? Number(big) : big; } const out = await generateDiffusionImage({ prompt,