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

Fix/adjust diffusion: round 6 race-free lifecycle + delete guards for PR #5754

Browse source 
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.
This commit is contained in:
Daniel Han-Chen 2026-05-25 01:27:57 +00:00
commit 04de106e49
10 changed files with 407 additions and 86 deletions
Download patch file Download diff file Expand all files Collapse all files

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

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

5
studio/backend/requirements/studio.txt
View file

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

62
studio/backend/routes/export.py
View file

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

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

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

25
studio/backend/routes/models.py
View file

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

11
studio/backend/routes/training.py
View file

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

131
studio/backend/tests/test_diffusion_backend.py
View file

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

38
studio/backend/tests/test_diffusion_routes.py
View file

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

17
studio/frontend/src/features/images/api.ts
View file

@ -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<DiffusionGenerateResponse> {
@ -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),
}),
);
}

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

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