studio: opt-in model auto-switch for the image and video apis (#8766)

* add opt-in model auto-switch for the image and video apis

POST /v1/images/generations answered 503 unless a model had already been picked on the
Images page, and its `model` field was documented as informational, so an agent could not
choose what it generated on. Chat has had this since openai_api_auto_switch_model; media had
no equivalent, and media_auto_unload_idle_seconds could unload a pipeline but never load one.

core/inference/media_auto_switch.py resolves the requested name against the downloaded image
and video models, waits for in-flight work to drain, then runs the same load /images/load and
/video/load run. Names come from the local-model scan the Images picker already uses, so a
repo id, the scanner id, a display label, and `<id>:<QUANT>` for a GGUF all resolve. A cached
GGUF loads by repo id rather than its snapshot directory, whose entries are symlinks into
blobs/ that the loader's containment check rejects; a cached pipeline loads from
snapshots/<sha>, which is where its model_index.json lives.

POST /api/inference/video/generate gains an optional `model` for the same purpose. The Video
page never sends it, so that route is unchanged for existing callers.

Gated on media_api_auto_switch_model, off by default, with its own toggle under Settings >
API. With it off nothing here runs and `model` keeps its informational meaning. Nothing
starts a download, so a name that resolves to no downloaded model is refused with the ids
that would work rather than answered by whichever model happens to be resident.

Both waits are bounded because secure mode's tunnel caps an origin response near 100 seconds:
a drain that does not clear in 30s returns 409, and a load still running after 90s returns 503
with Retry-After, leaving the load to finish for the retry.

"Only unload models loaded by the API" changes with it. That setting used to zero the media
TTL outright, on the grounds that nothing but the user could load an image or video model.
Auto-switch breaks that premise, so it is now a per-model rule: the load routes record who
asked, and media_keepwarm spares a user-loaded pipeline while still collecting an API-loaded
one, matching what chat already does with _loaded_by_user_action.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* close six gaps in the media auto-switch path

Refuse a pick whose companions are not on disk. The resolver only indexes downloaded
checkpoints, but a GGUF loads its text encoders and VAE from a base repo the loader prefetches,
so a request could pull tens of gigabytes despite the setting promising it never downloads.
The same planner /images/download-plan serves now answers first, and a nonzero remainder is a
409 naming the missing gigabytes. A plan that cannot be built is treated as nothing missing:
that is almost always an unreachable Hub, where no download can happen either.

Discount requests queued on the switch lock from the drain. Two concurrent requests for the
same absent model are both counted by the media middleware, so the lock holder saw the queued
one as work to drain, waited the full 30 seconds and returned 409, and so did the other. The
chat switch already excludes its own waiters for this reason.

Give the drain and the load one shared budget. Separate 30s and 90s allowances added up to
roughly 120s, past the ~100s origin window the bounds exist to stay inside, so a slow switch
lost the socket instead of returning the retryable 503. The drain is now capped inside a 90s
end-to-end budget and the load inherits what is left.

Match a GGUF on quant in the serving check. Loose .gguf files in one scan folder share that
folder as their model_path, so the path alone reported a sibling as already serving and the
API generated with the wrong weights. The quant is the per-file identity the backend
publishes, and the same one media_keepwarm counts as part of the build.

Key load provenance to the target it was recorded for. A load is recorded when it is accepted
and can still fail during prefetch with the previous model resident, so an API load that failed
marked a user-loaded pipeline as API-loaded and "only unload models loaded by the API" then
evicted it. An origin now only answers for the model it was written against.

Hold the keep-warm gate from the final drain check through load registration. Nothing stopped
a generation entering that gap, and the load path cancels active work as it tears the pipeline
down, so a swap that had just waited for the queue to clear could cut one short.

* tighten the media auto-switch guards

Plan the download against the engine that will load the pick, not the resident one. With
native sd.cpp active and a diffusers pipeline requested, the sd.cpp planner refuses the pick
and that refusal read as nothing missing, so the guard let the load download after all.
predict_engine picks the planner, as /images/download-plan already does.

Hold the admission gate across the final drain observation instead of after it. A request
admitted between a passing drain and the gate is tracked but has not marked the backend
active, so it read as idle and the load's teardown could still cancel it. The drain condition
is now evaluated once more under the gate.

Recheck the target after draining. A retry can acquire the switch lock while the earlier
attempt's load is still running; draining waits that out, and without this the retry tore
down the model that had just landed and reloaded it.

Confirm the requested model actually served. _await_loaded returned success as soon as
anything was resident, so a user load accepted between two polls superseded ours and the
request generated on the replacement while naming the requested model. The final status is
matched against the pick, GGUF quant included.

* sharpen media auto-switch identity and locality checks

Compare a GGUF on the variant lister's full label rather than the backend's quant token. The
token collapses IQ4_XS-3.53bpw and IQ4_XS-3.97bpw to IQ4_XS and reduces an unlabelled file to
nothing, so either build reported the other as already serving. A label that cannot match the
published token now reloads instead of risking the wrong weights.

Key load provenance by the build, not the path. A user-loaded Q4 and an API load of Q8 from
one repo shared a target string, so a failed API load marked the resident Q4 as API-loaded and
"only unload models loaded by the API" freed it.

Refuse when locality cannot be established. video.download_plan returns zero bytes on any
metadata error because its own caller falls back to an inline pull, so the guard read that as
a complete cache and allowed the download it exists to prevent. The planner now flags the
failure and the switch treats an unverifiable plan as a refusal.

Re-resolve the image engine under the admission gate. A concurrent load can activate the other
engine while this request drains, leaving the captured backend idle and both checks passing
against something nothing is using.

Start the budget before resolution and bound the scan and the plan with it. It began after the
cold scan and never covered the planner's Hub calls, so either could consume the response
window before a deadline was consulted.

Serve an exact resident match without consulting the index. A scan that failed caches an empty
result, which turned a request naming the loaded model into a 404 for that window.

* compare media builds on the identity the backend publishes

The previous round compared a GGUF against the variant lister's full label, which no backend
publishes: status derives gguf_variant with extract_quant_token, so IQ4_XS-3.97bpw and every
unlabelled checkpoint never matched and the switch returned 503 for a model that had loaded
correctly. The comparison is back on the published token, and the case that motivated the
label is handled where it belongs: a pick whose token another build under the same path also
publishes is marked ambiguous at index time and never short-circuits, so the switch reloads
instead of assuming. The load a request started is checked without that rule, since it is its
own load and nothing is being assumed.

Drop an alias two models answer to. Cached repos advertise their final component, so
org-a/model and org-b/model both offer "model", and binding whichever the scan reached first
would load arbitrary weights for a name the resolver documents as usable. Full ids are
unaffected.

Keep the exact-resident shortcut away from a GGUF. A bare repo id means the preferred quant,
which a repo id comparison cannot see, so a resident non-preferred quant answered the request.

Normalize a bare single-file directory before planning. Both load routes reinterpret it as a
single_file load and resolve that family's companions, while the plan described a local
pipeline with nothing to fetch.

Treat plan entries as missing, not just their byte total. Both planners keep an entry whose
size could not be read and coerce that size to zero, so a pending multi-GB fetch reported as
zero bytes passed the guard.

Flag a partial image plan. _estimate_download_bytes swallows a companion metadata failure and
returns whatever it had accumulated, so a plan missing its companions looked complete; it now
records the failure and download_plan reports plan_failed, which the switch already refuses.

Plan on the card the load will rank for itself, since automatic precision can select a
different hosted pre-quantized artifact per card.

Re-plan under the admission gate. The drain can last 30 seconds, and a cache deletion during
it sees a target that is neither loaded nor loading, so files verified before the wait could
be gone by the time the load starts.

* stop the download guard from refusing local models

Flagging a partial image plan in the previous round turned every local full pipeline into a
refusal: the planner asks HfApi about the absolute path, that failure is now recorded, and the
switch reads it as unverifiable. A directory on disk is what from_pretrained loads, so it is
complete by definition and is no longer planned at all.

Verify the diffusers fallback when the native engine is predicted. predict_engine treats an
absent sd.cpp binary as available whenever its installation is allowed, while activation falls
back to diffusers if that install produces nothing runnable, and the two engines read different
companion sets. Both plans are checked when they can differ.

Compare filesystem paths without folding case. /models/Foo and /models/foo are different models
where the filesystem says so, and folding them reported one as already serving the other.
Aliases and repo ids still fold.

Bound the post-drain re-plan by the switch budget. It runs while the admission gate is held, so
a stalled Hub call held every new generation and load off that backend as well as overrunning
the response window.

* narrow the media auto-switch discovery and engine checks

Verify the diffusers fallback only when the prediction assumed an install. select_and_activate_engine
falls back to diffusers when no runnable sd.cpp binary exists, so on a host that already has one the
load stays native and demanding the diffusers shards refused a model sd.cpp can serve, including
after an idle unload. predict_engine's own binary probe is extracted as native_binary_installed so
both callers share one implementation.

Skip partial catalog rows. A cancelled or incomplete pull still lists, and registering it advertised
an id whose load fails predictably and reported it in the "not found" listing.

Detect the HF cache layout rather than the scanner's source label. A cache tree inside a user-added
scan folder is relabelled custom, while its snapshot entries stay symlinks into blobs/ that the
loader's containment check refuses, so those rows have to load by repo id like any other cached repo.

Refuse an edit-only family before switching. The local catalog tags Kontext and Qwen-Image-Edit
text-to-image, so the switch loaded a multi-GB pipeline that this endpoint then rejected for lacking
txt2img, with the previously useful model already evicted.

* translate the media auto-switch settings and tighten its refusals

Locale parity runs strict in CI, so the three new keys are translated into every overlay rather
than left to the English fallback. mediaIdlePaused is restored with its new wording, since the
setting it named no longer vetoes the media TTL outright.

Refuse an incompatible plan before switching. A FLUX.2 GGUF paired with a different-size base is
fully cached and still unloadable, and the route's cheap validation misses it, so the mismatch
surfaced only from the background loader once the resident pipeline was already torn down.

Bound load setup by the switch budget. Companion preflight and a first-run sd.cpp install both
run before begin_load registers, while the admission gate is held, so a stall blocked every new
generation and load on that backend as well as overrunning the response window.

Preserve case in the resident shortcut. It compared a filesystem path case-insensitively, which
the _satisfied_by comparison had already stopped doing.

Do not advertise a directory the load route rejects. Several checkpoints and no model_index.json
is ambiguous, and both routes refuse rather than choose, so indexing it only cost a failed switch.

* close the remaining media auto-switch identity and timing gaps

Shield load setup from the deadline. Cancelling the await does not stop the worker, so a timeout
landing inside select_and_activate_engine left the engine switched, the resident model unloaded
and no begin_load coming, with the gate already released. The setup now always reaches
registration; the timeout only frees the caller.

Acquire the switch lock within the budget. A request that spent most of it resolving could queue
behind another full switch and pass the response window before any inner wait noticed.

Flag the native MiniMax-H3 plan failure. That planner has its own metadata-error return, which
the switch read as a verified cache and would have downloaded the missing components.

Compare the H3 partition. A switch sends no h3_task and therefore loads the default keyframe
denoiser, so a resident ref2va does not answer a plain request: serving it accepted a generation
that then failed for missing references.

Accept modular pipeline indexes in discovery. A fully downloaded dense MiniMax-H3 carries
modular_model_index.json, and rejecting it 404'd every named request for it.

Rank root variants for a bare id. A repo holding both root and subdirectory builds could hand a
bare id a qualified build, where a plain local load resolves non-recursively and the picker and
chat resolver both take the root.

Keep case in provenance keys, matching the other filesystem comparisons.

* hand the gate and switch lock to the shielded load setup

Shielding kept the setup running but not the contexts around it: a timeout unwound the admission
gate and the switch lock while the task was still short of begin_load, so a newly admitted
generation could start on an engine that setup was about to activate or tear down. The gated
section is now its own task that acquires the gate and releases the lock itself, so giving up on
the wait frees only the caller and nothing else is admitted until registration completes.

Check the MiniMax-H3 partition in the resident shortcut too. It returns before _resident_is_pick
runs, so a resident ref2va still answered a plain request for that repo id and the generation
then failed for missing references.

* correct the switch budget, waiter accounting and identity keys

Cancel a shielded future rather than closing it. An already-expired budget reached the shield
branch and raised AttributeError, since a Future has no close(), instead of the retryable 503
the branch exists to return.

Bound the final plan inside the gated task. It holds the admission gate and the switch lock, so
a stalled planner blocked every generation and load on that backend indefinitely while the
caller had long since returned. The step has no side effects, so timing out can give both back.

Count a request as a waiter only while it waits for the switch lock. It stayed marked through
load polling, so a second switch discounted a request that was about to generate and could
replace its model out from under it.

Ignore requests parked behind the gate the task owns. A newcomer is counted pending and then
blocks on that same gate, so counting it aborted an otherwise idle switch with model_busy.

Normalize paths in the ambiguity scan and record the MiniMax-H3 partition in provenance, so a
case-sensitive filesystem stops merging distinct directories and a failed API load of the
default partition no longer unpins a resident ref2va the user loaded.

* accept the partition in the provenance API and derive it from the checkpoint

note_load_origin never grew the partition parameter both load routes had started passing, so
every /video/load raised TypeError after its background load was already accepted. It takes and
stores the partition now, and a route-level test drives the real load route so a signature that
drifts from its callers fails here rather than in production.

Derive the expected MiniMax-H3 partition from the selected checkpoint. The native backend
publishes ref2va for a minimax_h3_ref2va denoiser, so assuming the keyframe default rejected the
checkpoint that had just loaded and reloaded a resident one that already matched.

Discount every recorded switch waiter. The holder leaves the marker when it acquires the lock,
so subtracting one fewer left a phantom request and the holder drained the full window before
returning busy against an idle backend.

Plan local video pipelines rather than trusting the directory. A local MiniMax-H3 modular
pipeline still substitutes a hosted quantized conditioner during assembly, tens of gigabytes the
shortcut would have waved through.

Record a GGUF variant only for a gguf load. A single-file checkpoint with a precision token in
its name stored a variant that status never publishes for that kind, so the idle policy read the
mismatch as unknown and never freed an API-loaded pipeline.

* recognise a native GGUF, name the H3 partition and carry the caller's token

The native sd.cpp status reports a GGUF through dtype and a quant rather than model_kind, so the
resident shortcut treated every native checkpoint as a plain pipeline and served whichever quant
happened to be up instead of the preferred one a bare id names.

Send the resolved MiniMax-H3 partition with the load. The route records what the request carried
while the backend publishes what it resolved, so leaving it unset gave every API-loaded H3
pipeline an origin key that never matched and the idle policy never freed it.

Thread the caller's HF token into planning and loading. A cached gated repo resolves by repo id,
so the planner still makes authenticated metadata reads; without the token they failed and the
switch refused a model whose files were all local. Only the caller's header is used, never the
server owner's ambient credential.

* account for the other media backend and for pipelines with external parts

Wait for the other media backend before switching. The load route takes the GPU through the
arbiter, whose cross-owner handoff unloads whoever holds it, so an image switch could cancel a
running video generation and the reverse. The drain now counts that backend's tracked requests
and its in-flight work, and refuses rather than triggering the handoff.

Plan a HiDream pipeline even when it is local. Its pipeline loads a separate encoder repo of
about 16 GB, so a directory on disk is not evidence that nothing will be fetched. Detection runs
on the model id as well as the path, and an unrecognised family keeps the shortcut, since
planning a local path always fails.

Take the MiniMax-H3 partition from the checkpoint basename and the family. A qualified variant
lives under ref2va/, which the prefix test missed, and a modular pipeline resolves to a local
directory that is neither bundle repo id, so its provenance recorded no partition at all and the
idle policy never freed it.

* protect every owner the GPU handoff can evict

Count chat in the drain. The arbiter evicts whoever owns the GPU, so a media switch terminated a
streaming completion that had nothing to do with the request.

Hold both media gates through registration. Draining the other backend was not atomic with
admission to it, so a request arriving in the gap started generating and the handoff cancelled
it. Both gates are taken in a fixed order, image first, so two switches cannot deadlock.

Bound the busy probes by the switch budget rather than the drain loop's deadline.
loading_repo_ids takes the backend lock, which the loader holds across pipeline assembly, so an
unbounded probe outlived the response window; the in-gate check evaluates once with no time to
wait, which must not be read as every backend being busy.

Verify HiDream's encoder against the cache instead of planning its local path. The planner
cannot be handed an absolute pipeline directory, so refusing on a failed plan rejected a model
whose pipeline and encoder were both present.

Keep the user's provenance when an API load targets an indistinguishable build. Sibling GGUFs
can share a quant token, so a load that is accepted and then fails would otherwise reclassify
the user's surviving model as API-loaded.

* make the handoff checks atomic and scope them to loads that take the GPU

Hold the chat lifecycle gate alongside the two media gates. Counting chat was not atomic with
registration, so a completion starting in that gap was evicted by the handoff anyway. With the
gate held, the in-gate drain no longer probes chat at all: that counter includes media requests,
and one arriving while the gates are held is already blocked in the middleware, so counting it
aborted an otherwise idle switch.

Skip the cross-owner waits when the load will not take the GPU. A CPU diffusion device releases
ownership rather than acquiring it, so such a switch interrupts nothing and had no reason to
refuse while chat or the other backend was busy.

Check every shard of HiDream's encoder. A cache hit was any single weight file, while the
pipeline opens the whole repository, so an interrupted sharded pull read as local and the load
fetched the rest.

Derive the partition for a user H3 GGUF load that leaves h3_task unset, since the backend
publishes what it derived; recording nothing let a later API switch overwrite the user's mark.

* bound gate acquisition and close two more download paths

Acquire the media and chat gates under the switch budget. A gate held by an unrelated stall
would otherwise keep the setup task, and with it the per-owner switch lock, alive indefinitely,
so every later switch on that backend timed out too. Acquisition happens through an exit stack
before the non-cancellable phase, where cancelling releases whatever was already taken.

Skip the other owner's tracked requests for a CPU load as well. The previous round gated only
the probes, and the test covered only those, so a CPU switch still waited on a request the load
route can never evict.

Refuse a renamed LTX-2.3 checkpoint whose extras are missing. The planner judges 2.3 by name
while the loader reads the checkpoint header, so a generically named file plans as 2.0, reports
nothing missing, and pulls the 2.3 VAE, audio and connector artifacts during assembly.

* stop two switches refusing each other, and name the LTX-2.3 companions exactly

Serialize GPU-taking media switches and discount the switchers. An image switch and a video
switch each counted the other as cross-owner work, and chat's counter saw both, so neither
started and both returned busy. A request performing a switch is now marked for its whole
duration and subtracted from the cross-owner and chat counts, and one cross-backend lock lets
them proceed in turn rather than deadlocking on each other. Every acquired lock transfers to the
setup task, as the per-owner one already did.

Check the exact LTX-2.3 companion files. The extras repo also holds checkpoints, so "any weight
file is cached" proved nothing about the variant-specific connector, video VAE and audio VAE the
assembly reads; the check now asks for that file list by name.

* count parked requests once and check cached checkpoints on disk

The chat probe discounted a queued switcher twice, once as a switcher and once as a waiter, so
an active chat stream could read as idle and the load took the GPU out from under it. A waiter
is marked inside its own switch, so the switcher count alone is the right one.

Header-check a cached LTX-2.3 pick. A GGUF discovered in the Hub cache is named by its repo id,
so the path does not exist and the check was skipped, while the checkpoint is on disk and the
loader reads 2.3 straight out of it and then fetches the companions.

Treat a missing shard index as incomplete for the encoder that is always sharded. One shard
without model.safetensors.index.json is an interrupted pull, not a single-file layout, and
from_pretrained would fetch the index and the remaining shards.

* share one family probe and stop guessing on a spent budget

Family detection was done three different ways: the edit-only guard passed the un-normalized
pick and only its path, so a sole-checkpoint directory named nothing like its model resolved to
no family and an edit-only pipeline was loaded and then refused for lacking txt2img. The planner
had the same one-needle gap and fell back to the resident engine, which is exactly the engine it
must not plan against. Both now go through one helper that tries the id and the path with the
normalized pick, so the checkpoint filename participates.

A probe with no budget left knows nothing about the backend, and reporting "busy" sent callers
after a generation that does not exist. It raises the slow-switch 503 instead, and the drain
probes answer to the switch budget rather than their own window so a genuinely busy backend is
still a 409.

Re-read the backend status after resolution rather than reusing the pre-resolution snapshot: the
index build can run for the whole budget, and an idle unload landing inside it left the switch
reporting a model that was no longer there. The 404 listing shares the budget for the same
reason, since it can trigger a second cold scan.

Resolve the GPU-handoff question once per drain instead of on every 0.2s poll, and drop
_held_within, MediaModelPick.quant and _variant_label: nothing read them, since every identity
decision derives the token from the filename.

* split the media auto-switch into focused modules

media_auto_switch.py had grown to 1447 lines covering name resolution, locality
verification, backend probing, lock accounting, refusal shapes and the switch itself.
It is now the orchestration only, with each of those beside it: media_model_index,
media_locality, media_switch_backends, media_switch_locks and media_switch_errors.

Behaviour is unchanged. The locality refusal was written out twice and the busy
refusal three times, so both collapse into one helper; the pre-index resident
shortcut and the lock acquisition become named helpers; one dead import goes.

Two defects fixed on the way. A lone unlabelled GGUF was marked ambiguous because its
published quant token is empty, so every request naming it reloaded the pipeline that
was already resident; only a token two builds share is ambiguous. The handed-over setup
task's exception was never retrieved once the budget expired, so an ordinary refusal
from it was reported by the loop as an unhandled task exception.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* close three more no-download and no-eviction gaps

The HiDream encoder check counted shards only, so a cache holding every shard but not
config.json or the tokenizer files read as complete and the accepted switch then fetched
them: from_pretrained is called on the whole repository, not on the weights.

The LTX-2.3 extras check resolved the family by name, so a checkpoint whose repo and
filename both carry no family token was exempt from the very check its general.architecture
would have triggered. It now resolves the family the way the loader does, header fallback
included.

A video request whose shape the target family cannot render was judged only by
begin_generate, after the switch had evicted the resident model and spent minutes loading
the target. maybe_auto_switch_media_model takes a before_switch hook that runs on the
resolved pick while the resident model is still up, and the generate route uses it to raise
the same 422 up front.

* judge video conditioning against the target before the switch

A request whose keyframes or references the target checkpoint cannot take was rejected only
by begin_generate, after the switch had evicted the resident pipeline and loaded a multi-GB
target for a request that was already unservable.

The five rules _resolve_keyframes and _resolve_references applied inline move into
validate_video_keyframe_conditioning and validate_video_reference_conditioning, which are
pure in the family and the MiniMax-H3 partition. The backend calls them where it raised
before, so its under-lock answer is unchanged, and the generate route's before_switch hook
calls them on the resolved pick's family and expected partition.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* consolidate the media auto-switch tests

83 tests had grown one per review finding, with the same four-line FLUX setup written out 26
times and the same FastAPI wiring five times. Shared fixtures and a client helper replace the
boilerplate, and the pairs that differed only in an input are parametrized: resident by repo id
or by path, the H3 partition through the index and through the pre-index shortcut, the other
media backend and chat, the three plan refusals, the two before_switch route refusals.

Eight tests are dropped where a broader test already covers the same code: two units of
errors.probe and errors.bounded, one that belongs to the keep-warm suite, two that stubbed the
very function under test, and the LTX and partition cases that duplicate a sibling's branch.

67 tests over 1570 lines, from 83 over 2002.

* plan a pick the way the loader reads it

detect_family_for_pick reads a local model_index.json ahead of any guess made from a name, and
the load route is only ever handed the pick's path. Asking about the catalog id first therefore
answered FLUX for a HiDream pipeline sitting in a directory called flux.1, so the locality check
found no outside dependency while the loader identified HiDream and fetched its 16 GB encoder.
The path is now the first needle, with the id as the fallback.

A non-active HF cache is scanned with the snapshot directory as the entry path, so it is never
unwrapped to a repo id, and its entries are symlinks into blobs/ that both load validators refuse
for escaping the directory. The repo id is no answer either, since it would send the loader to
the active cache and download the model again, so the index no longer advertises a GGUF build the
loader cannot open.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* hold only the gates a switch can evict behind

A CPU diffusion load releases GPU ownership instead of acquiring it, so it can evict neither
chat nor the other media backend, yet the gated section took all three admission gates
unconditionally. An unrelated chat teardown could therefore time a CPU switch out, and the
switch held chat and video admission closed across the re-plan and the load registration. It
now takes those two only when the load actually takes the GPU, from the same decision that
picks up the cross-backend switch lock.

The pre-switch hook also judged only shape and conditioning, so a flow_shift the target family
does not expose still cost an eviction and a full load before _resolve_flow_shifts refused it.
Its family half moves to validate_video_flow_controls, called from the backend where it raised
and from the hook; the engine half stays behind, since a target's engine is not chosen yet.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* check a local pipeline's parts and the target's engine before switching

A directory carrying model_index.json was treated as complete by definition, which holds only
if its components are there. A hand-copied or interrupted pipeline passed, and the loader tore
the resident pipeline down before from_pretrained found the gap, leaving the API with no image
model. Each component the index names must now have a directory with something in it, and a
sharded one must hold every shard its own weight index lists.

MiniMax-H3 GGUFs always load through sd.cpp, which derives the audio schedule against a fixed
shift, so that rule is knowable from the pick. validate_video_flow_controls takes the engine and
now carries the whole of _resolve_flow_shifts's validation, with the route passing sd_cpp for a
native H3 target so an audio_flow_shift it cannot apply is refused before the eviction.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* close four more no-download and no-eviction gaps

The external-encoder check ran only for local image pipelines, so a standalone HiDream
checkpoint, which normalized_pick gives a filename, skipped it. Its single-file assembly still
calls hidream_te4_kwargs unconditionally and the planner never covers that repo, so a cached
base and an absent encoder meant an API switch could start a 16 GB download. It now runs for
every image pick.

The component check ran only for image pipelines, while the video planner also omits base files
whenever the local path exists, so an incomplete video directory reported nothing missing. It
now runs for either owner, and reads modular_model_index.json as well.

A component holding its config and none of its weights passed the nonempty-directory test.
Weight-bearing components declare config.json, where schedulers, tokenizers and processors
declare their own *_config.json, so one that declares the former must now hold a weight file.

The reference-sizing rules join the flow controls in video_families, so a native MiniMax-H3
target refuses reference_image_size=max before the switch rather than after the load.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* read modular entries, and bound the last unbounded wait

A modular_model_index.json entry is [library, class, spec], so the two-element test skipped
every component of a MiniMax-H3 pipeline. Three-element entries now count, except where the spec
names another repository: that component is never expected in this directory and the planner is
what covers it.

A directory holding exactly one checkpoint is reinterpreted as a single_file load, which resolves
the name through the same containment check a GGUF gets, so a cache snapshot's symlink into
blobs/ is refused. Such a directory is no longer advertised, matching the GGUF rule.

load_progress walks cache directories to count bytes, so on a stalled filesystem a single poll
outlived the budget the check at the bottom of the loop enforces. Both probes now run through
the same bounded helper as every other wait in the switch.

Also drops the H3_REF_SIZE_MAX import left unused by the previous commit, which the source lint
job flagged as a blocker.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* verify hosted modular components and tokenizer vocabularies

load_components pulls every repository a modular index names, and the video planner omits its
base manifest whenever the selected path already exists, so a local modular directory whose
hosted VAE or processor was absent verified clean and downloaded it after the eviction. Each
hosted entry is now resolved against the cache: a spec naming a subfolder is checked in that
repo's cached snapshot with the same component rules a local one gets, and one without a
subfolder falls back to _upstream_is_cached.

A tokenizer directory holding only tokenizer_config.json also passed, since it declares no
config.json and so was read as weightless. Which vocabulary file a tokenizer class needs varies,
so any one of the known spellings answers for all of them, and a directory with none is
refused.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* read the active revision, the whole component, and the whole split set

_cached_snapshot_root took whichever snapshot sorted first, so a superseded revision holding a
complete component could vouch for the partial one from_pretrained actually reads. It now
resolves refs/main the way the cache does, falling back to any revision only for a
commit-pinned download that has no ref.

A hosted component named as a whole repo went through _upstream_is_cached, whose no-manifest
branch is satisfied by one weight file. Both shapes now go through the same component rules, so
an interrupted sharded pull no longer reads as complete.

A split GGUF opens its siblings implicitly and the planners read a local checkpoint as already
present, so half a set would evict the resident model and fail at startup. colocated_split_shards
decides that before the variant is indexed.

Load setup held the gates for as long as it took to reach begin_load, which for a first-run
native install is minutes of blocked chat and video admission. It now runs as its own task with
a grace period, after which the gates go back and the load carries on without them.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* honour what the spec pins and what the index declares

A modular spec carries revision and variant, and ComponentSpec.load is handed both, so a
complete default snapshot could approve a component whose pinned commit or named weight set was
never downloaded. The revision now selects the snapshot, with no fallback to main once one is
pinned, and a variant requires files carrying its name.

A spec pointing straight at a local directory was accepted for existing at all, so an empty or
half-copied source passed; it now goes through the same component rules as a cached one.

A shard index declaring an empty weight_map named nothing missing and was still taken as proof
that a component was complete. An index has to declare at least one shard to count.

A cached repo-id GGUF skipped the split-set check, since the containment rule does not apply
there. The child is resolved out of the cache and its set checked the same way.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* tell two h3 partitions apart, and recheck chat under the gate

Two MiniMax-H3 denoisers in one directory share a quant token, so both were marked
indistinguishable and every request naming one reloaded a multi-GB checkpoint that status
already identifies through h3_task. The collision key now carries the expected partition, which
is exactly what partition_matches reads.

A chat request admitted between the outer drain's last probe and the switch taking the lifecycle
gate is already running, and the in-gate drain skipped chat entirely, so the handoff could
terminate it. That drain now asks again with include_pending=False: a request blocked behind the
held gate has not started inference and must not abort the switch, one already running must.

A metadata-only component is its config, so a scheduler or processor directory holding a stray
file and nothing else no longer passes as complete.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* make the no-download promise the loader's own rule

Auto-switch has been predicting what a load would fetch: which config, tokenizer, shard,
variant and revision each component needs, enumerated by hand against what from_pretrained and
load_components happen to open. That enumeration can never be complete, and every gap in it is a
download the feature promised would not happen.

An API-initiated load now passes local_files_only through to the calls that reach the Hub during
assembly: the pipeline and transformer builds, the HiDream Llama encoder, the modular
ModularPipeline load and load_components, which forwards it into every ComponentSpec.load. A
load nobody asked for can no longer fetch anything, whatever the locality check failed to
predict. The picker's own load is untouched, since that is what downloads a model in the first
place.

The prediction stays as the first line: it is what turns a would-be failure into a 409 that
names the model and leaves the resident pipeline up.

* split the ambiguity group by h3 partition only

Adding the expected partition to the collision key let a non-H3 build escape its group: it has
no partition to be told apart by, and partition_matches reads a resident fl2va as answering for
it, so a Wan request could be served by a resident MiniMax-H3 checkpoint. Grouping is back on
what status publishes, the path and the quant token, and the partition only splits a group whose
members are all H3 builds with distinct partitions.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio: accept and honour local_files_only on the native image engine

The image load route passes local_files_only to whichever engine was activated, through one
call site, unconditionally. SdCppDiffusionBackend.begin_load did not declare it and has no
**kwargs catch-all, so every load on a host that selects the native engine raised

    TypeError: begin_load() got an unexpected keyword argument 'local_files_only'

That is CPU-only hosts, opted-in MPS and UNSLOTH_DIFFUSION_ENGINE=sd_cpp, and it is not
limited to the new auto-switch: the ordinary Images-page load passes the flag too, as False.
The raise lands after select_and_activate_engine has already switched engines and, on the GPU
path, after chat was evicted, and the route catches only ValueError, FileNotFoundError and
RuntimeError, so it surfaced as an unhandled 500.

The parameter is now declared in the same position and with the same default as the diffusers
twin, and honoured rather than swallowed: the FLUX.2 header probe, the companion-repo
preflight, the size probe and the asset fetch each resolve from the cache or stand down, and a
miss is reported as a RuntimeError naming the repo and file rather than a raw hub error. It
covers model assets only; the sd-cli binary tree keeps its own install policy, which the new
tests pin as deliberate.

hf_hub_download_with_xet_fallback grows the same parameter. When set it goes straight to
huggingface_hub rather than through the shared unsloth_zoo ladder, because that ladder is an
optional runtime dependency whose degraded stub drops keywords it does not recognise, and a
dropped local_files_only downloads.

The reason fifteen review rounds did not catch this is that every fake backend in the test tree
is written as begin_load(self, *a, **kwargs), which accepts anything. The new contract test
binds the keywords the route actually passes against both real signatures, and drives the
native path through create_autospec so a signature drift fails the suite.

* studio: bind local_files_only in the dense-quant image build

_load_dense_quant_pipeline forwarded local_files_only to both of its _assemble_pipe calls but
never declared it and was never passed it, so the name resolved as a module global that does
not exist. Every load taking the transformer-quant fast path raised

    NameError: name 'local_files_only' is not defined

which the caller turns into "transformer_quant='fp8' could not be used: the quantised
transformer build failed", so an explicit fp8 or nvfp4 pin 409s and Auto silently drops to the
GGUF build. Nothing about it is gated behind the auto-switch setting, and pyflakes reports it
as an undefined name on the branch and on nothing before it.

The parameter is declared and threaded from load_pipeline, and while it was being threaded it
was carried the rest of the way rather than stopping at the pipeline assembly: the dense bf16
transformer download, the pre-quantised checkpoint fetch and its config read are the largest
fetches on this path, so a load nobody asked for has to be refused at each of them rather than
allowed to pull several GB. A cache miss falls back exactly as an absent checkpoint already
does.

Six tests in test_diffusion_backend.py and one in the saved-metadata contract were failing on
this and now pass. The prequant doubles gain the parameter, since they are hand-written fakes
with exact signatures and the production signature moved.

* studio: keep the no-download promise across the video prefetch

An API-initiated video load reached load_pipeline with local_files_only set and everything
before it without. _run_load fetched the checkpoint, the pre-cast text encoders, the H3
conditioner and the scoped base snapshot, and sized all of them off model_info, none of which
consulted the flag; load_pipeline re-entered _predownload_base without it; and
_run_load_h3_native swallowed it into **_ and then downloaded its runtime. Those calls are
where the multi-GB pulls actually happen, so the promise the setting makes was carried only by
the locality check the flag exists to stop relying on.

The flag now reaches every one of them. The verification probes stand down rather than refuse:
the denoiser-prequant question becomes a cache probe over both roots, the size estimate returns
its existing "metadata unavailable" answer, and the scoped base predownload returns None
because from_pretrained resolves the cached snapshot itself. A fully cached model still loads
offline, which is the case the flag exists for, and each miss lands in the handler a failed
fetch already had. Offline the native H3 path no longer installs its runtime, since an install
is a download.

The new suite installs sentinels that raise on any model_info and on any download not asking
for the cache, runs a full API-initiated load against a cached fake, and asserts the
user-initiated mirror still does all of it.

* studio: prove media locality from files, not directory entries

_component_present judged three of its checks on entry names from the directory listing with no
test that the entry is a file. A Hugging Face snapshot holds symlinks into blobs/, and a blob
removed by a cache sweep or an aborted pull leaves the link behind, so such a component read as
complete: missing_download_bytes answered zero, the switch accepted, the resident pipeline was
evicted, and from_pretrained then failed with nothing loaded. That is the outcome the module
exists to prevent.

It also split by host. Windows without developer mode copies the blob into the snapshot instead
of linking, and a copy layout cannot express a dangling entry at all, so the same repository
verified as downloaded on one machine and not on another. A directory that merely ends in
.safetensors passed the same way.

The shard-index detection deliberately stays on the full listing: a dangling index must still
route into _shards_declared, which refuses on an index it cannot read, rather than fall through
to the weight test and pass on whichever sibling shard survived.

Found by a cross-OS simulation that builds both cache materialisations under tmp_path and
requires them to agree; the new tests are that comparison, reduced to the cases that moved.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio: pin the gpu-taking path in the media switch tests that assume it

Four tests in test_media_auto_switch.py assert the cross-owner wait: the switch refuses while
chat or the other media backend is busy, and a gate held elsewhere times the switch out rather
than pinning the lock. All four depend on load_takes_the_gpu, and none of them said so, so they
read the running host: green on a CUDA box, red on every CPU-only runner, where the switch
correctly waits on nobody because a CPU load evicts nobody.

Reproduced by running the file with CUDA_VISIBLE_DEVICES emptied, which is what the CI runners
look like, and confirmed on the parent commit so it is not a regression from the recent fixes.
The neighbouring CPU test already pins the same function to False in both bindings, so the new
fixture is that pattern named and given the other value.

* studio: take chat's lifecycle gate before the media admission gates

A media generation route is counted on chat's in-flight counter as well as its own, and the
middleware takes chat's gate, notes the request and releases it BEFORE it parks on the media
gate. With the media gates taken first, a request arriving in that window passed the still-open
chat gate, incremented chat's _inflight, and only then blocked on the held media gate. The
in-gate drain discounts it correctly on the media side, but chat_busy(count_pending=False) read
the same request as running chat work, so an otherwise idle switch answered 409 and loaded
nothing.

Reproduced with the real middleware and a real gate holder, which is what media_keepwarm._tick
does on every idle poll:

    chat _inflight 2, _pending 0 | media inflight/pending 1/1
    chat_busy(count_pending=False) -> True -> 409 model_busy

Taking chat's gate first parks such a request in _note_pending instead, where both counters
ignore it. The order is safe because this function is the only place in the tree that holds a
media admission gate and chat's lifecycle gate at once: the middleware holds at most one at a
time, the idle loop calls idle_unload_step outside its own gate rather than nested, and the four
lifecycle-gate sites in routes/inference.py take no media gate. AsyncExitStack unwinds in
reverse, so chat's gate now also releases last, instead of waking every parked chat request onto
media gates this switch still holds.

* studio: refuse the hosted h3 denoiser on a load nobody asked for

_load_h3_modular_pipeline threads local_files_only into ModularPipeline.from_pretrained and
load_components but skipped the multi-GB load_prequantized_transformer between them, even though
that callee accepts the flag and the image twin already passes it for exactly this reason.

No cache race is needed to reach it, and the code says so itself: the download plan settles auto
against the card's capacity, while this branch re-decides against live free memory once the
previous pipeline is gone. So a pick the plan sized as "the released bf16 denoiser fits", with
nothing hosted staged and the locality gate reporting zero missing bytes, can be re-decided here
as "take the hosted int8 checkpoint" and pull it on a load that promised to fetch nothing.

The failure path was already right: a cache miss returns None and the build continues dense,
which load_components refuses under the same flag a few lines below.

* studio: keep the no-download promise across the image prefetch

The image twin of the video fix. begin_load takes local_files_only and hands it to the worker,
but _run_load never read it, so the whole staging phase ran unrestricted: _prefetch_files pulled
the checkpoint and every base companion through the xet wrapper without the flag, and
_estimate_download_bytes, _te_prequant_plan_files and _assert_base_repo_accessible each asked the
Hub. The flag only took effect at load_pipeline, by which point the bytes had moved.

So an API-initiated load could fetch multi-GB assets on the user's connection whenever a file
vanished after the switch's locality check, or that check read the cache as more complete than it
was, which is precisely the prediction the flag exists to stop relying on.

The metadata probes now return their existing "unavailable" answers rather than refusing: the
estimate gives (0, []), the pre-cast plan gives {}, and the base preflight keeps only its
other-root cache read, which is what still lets a base staged under huggingface_hub's import-time
root load off disk. An empty file list means the prefetch stages nothing and from_pretrained
resolves the cached snapshot itself.

Two probes are deliberately left alone: the base_model card-tag read, because dropping it offline
would resolve a different base than the one whose weights are cached, and the FLUX.2 pick
compatibility check, which already prefers the on-disk header and whose refusal is real.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio: anchor the offline-load ast read on the package, not the cwd

The keyword assertion opened core/inference/video.py relative to the process working directory,
which only exists when pytest runs from studio/backend. CI runs it from the repo root with the
backend merely on PYTHONPATH, so it raised FileNotFoundError there. Resolved off the module's own
__file__ instead, and read with an explicit encoding like every other checked-in file read.

* studio: build the offline preflight's expected path with os.path, not a posix literal

The assertion compared against a hardcoded /snap/... spelling while the function strips the
file's own relative path with os.path, so on Windows it failed for the separator alone.

* studio: carry local-only mode into the family assemblers

Four more places an API-initiated load reached the Hub, all of them the per-family assembly the
generic path hands off to.

The MiniMax-H3 conditioner loader took neither the flag nor the other-root reuse its own stager
uses, so a ~27 GB artifact the staging phase had already accepted from the alternate cache root
was missed in the live root and pulled again, after the resident pipeline was evicted. It now
resolves exactly the way the stager that cleared the load did.

_resolve_gguf_path re-resolved the revision on every call, and huggingface_hub only skips that
HEAD when local_files_only is set, so this was not merely a deletion race: a checkpoint
republished upstream since the cache was filled re-downloaded on an ordinary cache hit, under
the generation lock, after the old pipeline was gone and with progress already reading 100%.

The Krea and LTX 2.3 assemblers bypass the guarded pipe_kwargs entirely and open the base repo
through their own from_pretrained and load_config calls. Krea has three call sites, not the one
the review named, including the one in _assemble_pipe that passes a repo id whenever nothing was
staged. LTX 2.3 is handed no staged snapshot at all by design, which makes the missing flag worse
rather than better: every component resolves the hub id.

Scope: local_files_only only, not cache_dir. That these two assemblers ignore Studio's configured
cache root is a real defect, but it predates this and breaks user-initiated loads the same way, so
it is not part of an opt-in switch's no-download promise. The one widening is deliberate, and is
what makes the H3 fix correct rather than a refusal of an artifact the load was already cleared on.

Three hand-written test doubles carrying exact production signatures move with it, the same class
of breakage that hid the native-engine TypeError for fifteen rounds.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio: importorskip the libraries the family assembler tests actually reach

Three of these drive the real transformers, diffusers and safetensors entry points rather than a
stub, so on a runner that ships none of them they failed rather than skipped. The rest of the file
is unaffected and still runs there.

* Studio: pin the family assemblers to the live HF cache root

An API-initiated media switch passes local_files_only=True, and the locality
gate that cleared the switch reads Studio's LIVE cache root (media_locality
passes cache_dir = hub_cache_dir()). The Krea 2, LTX 2.3 and HiDream assemblers
build every component themselves from a repo id, bypassing the guarded
pipe_kwargs, and left cache_dir unset -- which resolves through huggingface_hub's
import-time constant.

Studio's cache folder is a setting, and set_hf_cache_home only writes the DB, so
after a mid-session change the two roots differ. On main that was survivable: a
miss in the stale root simply downloaded again. Under local_files_only it is not,
because the load raises after the resident pipeline was already evicted, for a
model that is fully present.

Pins every repo-id component load to active_hf_hub_cache(), imported inside the
function to avoid the diffusion.py import cycle. diffusion_ideogram4 and
diffusion_te_prequant are deliberately untouched: neither takes local_files_only,
so both behave exactly as before.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: keep the single-file transformer config lookup local-only

from_single_file(config = <repo id>, subfolder = "transformer") is a Hub
resolution, not a local read: diffusers 0.39 pops local_files_only and forwards
it into the load_config() that resolves that id, and an unset flag is None, which
permits the network. The pipeline assembly after it was already guarded, so this
was the last unguarded Hub read on the GGUF/safetensors image and video paths,
and it runs after the resident pipeline has been evicted. (cache_dir does not
help: diffusers forwards it to the checkpoint read only, never to that
load_config.)

The flag alone is not enough. transformer/config.json was excluded from the
staged base file set on both paths, correctly for the shards, which the
checkpoint supplies, but that leaves the config systematically absent rather
than absent in a race: the locality gate clears the pick and local_files_only
would then fail essentially every API-initiated GGUF load.

So the config, and only the config, is now admitted to the base file set. A
complete install loads offline as promised; an install missing it is refused
with a 409 before eviction instead of losing the resident model to a late
LocalEntryNotFound.

The two in-tree tests that pinned the blanket exclusion are updated to assert
the intended shape (config in, shards out).

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
This commit is contained in:
Maheswar Kumar 2026-08-16 17:04:16 +12:00 committed by GitHub
parent 321c3ee019
commit cfee13795e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
55 changed files with 7903 additions and 301 deletions

View file

@ -617,6 +617,8 @@ def _assert_base_repo_accessible(
base_repo: str,
hf_token: Optional[str],
probe_file: str = "model_index.json",
*,
local_files_only: bool = False,
) -> Optional[str]:
"""Fail up front, with the licence URL, when a companion base cannot be read.
@ -629,7 +631,14 @@ def _assert_base_repo_accessible(
Returns the base's snapshot dir when an ACCESS verdict was excused by a copy living only under
huggingface_hub's import-time root, so the caller can load off disk: ``from_pretrained`` is
pinned to ``hub_cache_dir()`` and cannot see it, and the failure that earned the escape also
empties the size estimate. None otherwise."""
empties the size estimate. None otherwise.
Under ``local_files_only`` the Hub half stands down and only the cache probe runs. This whole
function exists to name the repo whose BYTES a download is about to fail on; a load that is
forbidden to download has no such fetch to pre-empt, so the verdict it would compute is about a
request that will never be made. The one thing worth keeping is the other-root escape, which is
a pure cache read, and it is exactly what lets a base staged under huggingface_hub's
import-time root still load off disk."""
repo = (base_repo or "").strip()
# Only a remote 'org/name' can be gated; a local base is already on disk.
if not repo or repo.count("/") != 1:
@ -685,6 +694,15 @@ def _assert_base_repo_accessible(
pass
return False
if local_files_only:
# Cache read only: no model_info, no metadata HEAD. ``_already_downloaded`` is called for
# its side effect (it sets ``other_root_snapshot``), and its bool is deliberately dropped
# -- a miss here is not a refusal. The loader is the one that raises for a file that is not
# on disk, with the file's own name in it, which is a better answer than this probe's
# "add a token" text for a load that was never going to fetch anything.
_already_downloaded()
return other_root_snapshot
def _is_auth_error(exc: Any) -> bool:
"""A 401/403 that hf_raise_for_status did not classify: an expired token 401s "Invalid
credentials in Authorization header", which _http.py excludes from its RepoNotFound branch
@ -1326,7 +1344,21 @@ class DiffusionBackend:
pin_cuda_ordinal(state.placed_ordinal)
return target
def _resolve_gguf_path(self, repo_id: str, gguf_filename: str, hf_token: Optional[str]) -> str:
def _resolve_gguf_path(
self,
repo_id: str,
gguf_filename: str,
hf_token: Optional[str],
local_files_only: bool = False,
) -> str:
"""The local path of this pick's checkpoint, downloading it when it is not on disk.
``local_files_only`` makes both resolutions below cache lookups. The prefetch already
staged this file under the same flag, so the promise looks kept -- but this call re-resolves
the revision against the Hub, and a checkpoint republished upstream since the cache was
filled (or removed between the staging and here) is a multi-GB pull taken under the
generation lock, after the resident pipeline was evicted, where unload cannot preempt it and
progress already reads 100%. It is the last unrestricted byte-mover on the image path."""
local_root = Path(repo_id).expanduser()
if local_root.exists():
return str(resolve_local_gguf_child(local_root, gguf_filename))
@ -1349,13 +1381,23 @@ class DiffusionBackend:
if isinstance(elsewhere, str) and Path(elsewhere).is_file():
try:
return hf_hub_download(
repo_id, gguf_filename, token = hf_token, cache_dir = None
repo_id,
gguf_filename,
token = hf_token,
cache_dir = None,
local_files_only = local_files_only,
)
except Exception: # noqa: BLE001 — revalidation is a bonus, never a new failure
return elsewhere
except Exception: # noqa: BLE001 — an unreadable cache is not a verdict, just download
pass
return hf_hub_download(repo_id, gguf_filename, token = hf_token, cache_dir = cache_dir)
return hf_hub_download(
repo_id,
gguf_filename,
token = hf_token,
cache_dir = cache_dir,
local_files_only = local_files_only,
)
def _dense_quant_prefetch_needed(
self,
@ -1545,6 +1587,7 @@ class DiffusionBackend:
hf_token: Optional[str],
cancel_event: Optional[threading.Event] = None,
fetch_base: Optional[str] = None,
local_files_only: bool = False,
) -> Optional[str]:
"""Pre-download the GGUF + the given ``base_files`` into the HF cache,
WITHOUT the lock and honoring ``cancel_event`` (this load's own event, so a
@ -1556,7 +1599,14 @@ class DiffusionBackend:
the pipeline manifest, so from_pretrained can load from disk instead of
re-sweeping the hub (its own sweep also pulls files the scoped list skips,
e.g. the 24 GB packaged root singles in each FLUX.1 repo); None otherwise
(estimate failure, config-only base, local repo) -> hub id as before."""
(estimate failure, config-only base, local repo) -> hub id as before.
``local_files_only`` makes every resolution below a CACHE LOOKUP: the file is returned when
it is on disk and ``LocalEntryNotFoundError`` is raised when it is not. This is the one call
in the staging phase that moves multi-GB bytes, so a dropped flag here is the whole
no-download promise -- a checkpoint that vanished after the switch's locality check, or a
check that read the cache as more complete than it is, becomes a pull on the user's
connection that nobody asked for."""
from utils.hf_xet_fallback import hf_hub_download_with_xet_fallback
# Only the BYTES move; the caller keeps the upstream id. ``fetch_base`` is the load-wide
@ -1576,6 +1626,7 @@ class DiffusionBackend:
hf_token,
cancel_event = cancel,
reuse_other_cache_root = True,
local_files_only = local_files_only,
)
# Base repo (VAE / text-encoder / scheduler); list comes from the estimate.
snapshot_root: Optional[str] = None
@ -1587,7 +1638,12 @@ class DiffusionBackend:
if cancel.is_set():
raise RuntimeError("Cancelled")
local = hf_hub_download_with_xet_fallback(
base, rfilename, hf_token, cancel_event = cancel, reuse_other_cache_root = True
base,
rfilename,
hf_token,
cancel_event = cancel,
reuse_other_cache_root = True,
local_files_only = local_files_only,
)
# The resolved path minus the file's own relative path, so a subfolder entry yields the
# same root as a top-level one. Not resolve()d: that follows the symlink into blobs/.
@ -1752,6 +1808,7 @@ class DiffusionBackend:
self,
repo_id: str,
*,
local_files_only: bool = False,
gguf_filename: Optional[str] = None,
base_repo: Optional[str] = None,
family_override: Optional[str] = None,
@ -1825,6 +1882,7 @@ class DiffusionBackend:
target = self._run_load,
kwargs = dict(
repo_id = repo_id,
local_files_only = local_files_only,
gguf_filename = gguf_filename,
base_repo = base_repo,
family_override = family_override,
@ -1853,6 +1911,13 @@ class DiffusionBackend:
token = kwargs.get("_load_token")
# This load's own event: a later load replaces self._cancel_event rather than clearing it.
cancel_event = kwargs.pop("_cancel_event", None) or self._cancel_event
# An API-initiated load promises to download NOTHING: it may only open what is already
# cached. Read once here and threaded into the staging phase below, the same way the video
# path does it -- the flag reached load_pipeline and nothing before it, so the prefetch that
# actually moves the bytes ran unrestricted and the promise was carried only by the locality
# check the flag exists to stop relying on.
# READ, not popped: load_pipeline takes it too (it is in this thread's kwargs by contract).
local_files_only = bool(kwargs.get("local_files_only"))
try:
# Resolve the base repo and estimate sizes here (both network) so begin_load returns instantly.
fam = detect_family_for_pick(
@ -1873,6 +1938,7 @@ class DiffusionBackend:
kwargs.get("text_encoder_quant"),
kwargs.get("hf_token"),
kwargs.get("gpu_ordinal"),
local_files_only = local_files_only,
)
expected, base_files = self._estimate_download_bytes(
kwargs["repo_id"],
@ -1898,6 +1964,7 @@ class DiffusionBackend:
else False
),
skip_te_components = tuple(te_prequant_files),
local_files_only = local_files_only,
)
# Only shards this prefetch staged may be materialised by the dense fallback, so read it
# off the staged list: a failed size estimate drops every base file too. A LOCAL base
@ -1915,7 +1982,9 @@ class DiffusionBackend:
# refusing the upstream id would reject the gated picks the ungated mirror rescues.
# Everything above is metadata, so no byte has moved yet. Returns a snapshot dir when it
# excused the base off a copy only the import-time root holds.
base_snapshot = _assert_base_repo_accessible(fetch_base, kwargs.get("hf_token"))
base_snapshot = _assert_base_repo_accessible(
fetch_base, kwargs.get("hf_token"), local_files_only = local_files_only
)
# And the same size-pairing preflight the plan and the route run, here because a direct
# begin_load (a saved config, the deploy path) reaches neither. Still metadata only, so
# it lands before _prefetch_files stages the base and before load_pipeline unloads the
@ -1942,6 +2011,7 @@ class DiffusionBackend:
kwargs.get("hf_token"),
cancel_event = cancel_event,
fetch_base = fetch_base,
local_files_only = local_files_only,
)
or base_snapshot
)
@ -2025,12 +2095,20 @@ class DiffusionBackend:
text_encoder_quant: Optional[str],
hf_token: Optional[str],
gpu_ordinal: Optional[int] = None,
local_files_only: bool = False,
) -> dict[str, tuple[str, list[tuple[str, int]]]]:
"""``{component: (repo_id, [(rfilename, size)])}`` for the text encoders this pick will
take PRE-CAST from a hosted checkpoint instead of the base repo's dense weights.
Empty unless the request asked for a scheme with a hosted artifact AND that artifact
really resolves, so a plan can never drop a dense encoder the load still wants."""
really resolves, so a plan can never drop a dense encoder the load still wants.
Empty under ``local_files_only``: ``te_prequant_hub_files`` is a ``model_info`` per source,
and the only consumer is the size estimate, which has already stood down. Empty is the
answer an unresolvable pre-cast gives today, and it is the safe one -- it keeps the dense
shards in the list rather than dropping weights the load may still open."""
if local_files_only:
return {}
try:
from huggingface_hub import HfApi
@ -2174,6 +2252,8 @@ class DiffusionBackend:
file_sizes_out: Optional[dict[str, dict[str, int]]] = None,
revisions_out: Optional[dict[str, str]] = None,
skip_te_components: tuple[str, ...] = (),
failures_out: Optional[list] = None,
local_files_only: bool = False,
) -> tuple[int, list[str]]:
"""Total download size for the progress bar, plus the base-repo files to
fetch (the prefetch reuses this list, so the base is listed only once).
@ -2201,7 +2281,15 @@ class DiffusionBackend:
encoder for a pre-cast load wastes tens of GB (FLUX.2-dev's Mistral-24B is ~48 GB,
Qwen-Image's Qwen2.5-VL ~16.6 GB) and nothing ever opens them. Everything else in the
component folder (config, shard index, tokenizer) is kept -- the pre-cast loader
meta-inits the encoder from the base repo's config."""
meta-inits the encoder from the base repo's config.
``local_files_only`` returns the "metadata unavailable" answer -- ``(0, [])`` -- instead of
asking. Every byte this counts belongs to a download that is not permitted, and the empty
file list is what makes the prefetch below stage nothing: ``from_pretrained`` resolves the
cached snapshot itself. Callers already handle it, because it is the same pair a failed
lookup produces."""
if local_files_only:
return 0, []
from huggingface_hub import HfApi
# No swap on purpose: the Hub gates only the BYTE endpoint, so model_info answers
@ -2308,6 +2396,11 @@ class DiffusionBackend:
_record_revision(revisions_out, base_repo, base_info)
except Exception as exc: # noqa: BLE001 — estimate is best-effort
logger.warning("diffusion.size_estimate_failed: %s", exc)
# Recorded so a caller that must not guess (media auto-switch refuses rather than
# download) can tell a partial file list from a complete one; the estimate itself
# stays best-effort for the UI, whose fallback is the inline pull.
if failures_out is not None:
failures_out.append(exc)
return total, base_files
def download_plan(
@ -2359,6 +2452,7 @@ class DiffusionBackend:
sizes: dict[str, int] = {}
file_sizes: dict[str, dict[str, int]] = {}
revisions: dict[str, str] = {}
plan_failures: list = []
required_total, base_files = self._estimate_download_bytes(
repo_id,
gguf_filename,
@ -2387,6 +2481,7 @@ class DiffusionBackend:
file_sizes_out = file_sizes,
revisions_out = revisions,
skip_te_components = tuple(te_files),
failures_out = plan_failures,
)
# Decided once, from the staged file list, and both probed and reported: a gated base
# answers model_info anonymously, so the plan would otherwise be confident and the 401 land
@ -2557,6 +2652,10 @@ class DiffusionBackend:
"required_bytes": int(required_total),
"checkpoint_bytes": checkpoint_bytes,
"incompatible_reason": incompatible,
# Companion discovery failed, so the file list above is partial. The UI stages what
# it can and the loader pulls the rest inline; a caller that must not download at
# all has to treat this plan as unverified rather than as nothing missing.
"plan_failed": bool(plan_failures),
}
@staticmethod
@ -3000,6 +3099,7 @@ class DiffusionBackend:
self,
repo_id: str,
*,
local_files_only: bool = False,
gguf_filename: Optional[str] = None,
base_repo: Optional[str] = None,
family_override: Optional[str] = None,
@ -3106,7 +3206,9 @@ class DiffusionBackend:
# Single-file kinds resolve a checkpoint path; the pipeline kind has none.
single_file_path = (
self._resolve_gguf_path(repo_id, gguf_filename, hf_token)
self._resolve_gguf_path(
repo_id, gguf_filename, hf_token, local_files_only = local_files_only
)
if kind in ("gguf", "single_file")
else None
)
@ -3727,6 +3829,7 @@ class DiffusionBackend:
lora_specs = loras,
text_encoder_quant = text_encoder_quant,
fetch_base = fetch_base,
local_files_only = local_files_only,
)
except Exception as exc: # noqa: BLE001 — fall back to the GGUF build
logger.warning(
@ -3797,6 +3900,9 @@ class DiffusionBackend:
fetch_base,
dtype,
hf_token = hf_token,
# The branch never sees pipe_kwargs, so the one keyword that keeps
# the no-download promise has to be handed over with the rest.
local_files_only = local_files_only,
text_encoder = te_prequant_pipe_kwargs(
fam,
fetch_base,
@ -3812,6 +3918,7 @@ class DiffusionBackend:
pipe = load_ideogram4_pipeline(fetch_base, dtype, hf_token = hf_token)
else:
pipe_kwargs: dict[str, Any] = {
"local_files_only": local_files_only,
"torch_dtype": dtype,
"cache_dir": hub_cache_dir(),
}
@ -3826,6 +3933,7 @@ class DiffusionBackend:
fam = fam,
te_quant_mode = text_encoder_quant,
target = target,
local_files_only = local_files_only,
)
)
# A hosted pre-cast fp8 text encoder skips the dense TE download; the cast re-applies idempotently.
@ -3847,6 +3955,7 @@ class DiffusionBackend:
elif kind == "single_file" and fam.single_file_is_pipeline:
# A single-file SDXL-style checkpoint is the WHOLE pipeline: load it through the pipeline class with ``config`` on the base repo.
sf_pipe_kwargs: dict[str, Any] = {
"local_files_only": local_files_only,
"torch_dtype": dtype,
# ``config`` is a REPO FETCH ahead of the mirrored load, so a gated id
# would 401 here first.
@ -3865,6 +3974,11 @@ class DiffusionBackend:
"subfolder": "transformer",
"token": hf_token,
"cache_dir": hub_cache_dir(),
# config is a REPO ID, and diffusers forwards this into the
# load_config() that resolves it (single_file_model.py), so without the
# flag this branch reaches the Hub on a load nobody asked for. The
# pipeline assembly below was already guarded; this call was not.
"local_files_only": local_files_only,
}
if kind == "gguf":
# Dequantise the GGUF transformer on-device at the compute dtype.
@ -3884,6 +3998,10 @@ class DiffusionBackend:
dtype,
hf_token = hf_token,
transformer = transformer,
# Same reason as the full-pipeline branch: the single file supplies
# only the denoiser, so the encoder, VAE and tokenizer below are
# still GB this load promised not to fetch.
local_files_only = local_files_only,
# Same pre-cast TE hand-in as the full-pipeline branch.
text_encoder = te_prequant_pipe_kwargs(
fam,
@ -3897,6 +4015,7 @@ class DiffusionBackend:
)
else:
pipe_kwargs = {
"local_files_only": local_files_only,
"torch_dtype": dtype,
"transformer": transformer,
"cache_dir": hub_cache_dir(),
@ -3912,6 +4031,7 @@ class DiffusionBackend:
fam = fam,
te_quant_mode = text_encoder_quant,
target = target,
local_files_only = local_files_only,
)
)
# Same pre-cast TE injection as above: the GGUF supplies the transformer, so the TE is the big download.
@ -4304,6 +4424,7 @@ class DiffusionBackend:
lora_specs: Optional[list[tuple[str, float]]] = None,
text_encoder_quant: Optional[str] = None,
fetch_base: Optional[str] = None,
local_files_only: bool = False,
) -> tuple[Any, str]:
"""Build the opt-in fast pipeline and return ``(pipe, engaged_scheme)``.
@ -4356,6 +4477,11 @@ class DiffusionBackend:
scheme = scheme,
# Reject a checkpoint with a different Linear filter so prequant matches runtime-quant.
min_features = DEFAULT_MIN_LINEAR_FEATURES,
# A load nobody asked for may not fetch this checkpoint either: it is a
# multi-GB download like any other, and the switch verified only what the
# plan listed. A cache miss falls out to the dense path, which is refused
# for the same reason a line below.
local_files_only = local_files_only,
# Only enforced when the caller forces fp8 fast-accum; a checkpoint that baked the other choice falls to the dense path.
fast_accum = fast_accum,
# The root _uncached_prequant_repo cleared this load against, so its hit is
@ -4376,6 +4502,7 @@ class DiffusionBackend:
te_quant_mode = text_encoder_quant,
target = target,
fetch_base = fetch_base,
local_files_only = local_files_only,
)
return pipe, scheme
@ -4396,6 +4523,9 @@ class DiffusionBackend:
torch_dtype = dtype,
token = hf_token,
cache_dir = hub_cache_dir(),
# The dense bf16 transformer is the largest single fetch on this path, so an
# API-initiated load has to be refused here rather than allowed to pull it.
local_files_only = local_files_only,
)
pipe = self._assemble_pipe(
pipeline_cls,
@ -4409,6 +4539,7 @@ class DiffusionBackend:
te_quant_mode = text_encoder_quant,
target = target,
fetch_base = fetch_base,
local_files_only = local_files_only,
)
if _has_active_lora(lora_specs):
# Bake the adapters BEFORE quantize_: peft wraps the dense Linears (post-quant torchao dispatch would TypeError),
@ -4456,6 +4587,7 @@ class DiffusionBackend:
te_quant_mode: Optional[str] = None,
target: Any = None,
fetch_base: Optional[str] = None,
local_files_only: bool = False,
) -> Any:
"""Assemble the diffusers pipeline around ``transformer`` and place it on ``device``
(a no-op for an already-placed pre-quantized transformer; it moves the companions).
@ -4483,10 +4615,14 @@ class DiffusionBackend:
hf_token = hf_token,
transformer = transformer,
text_encoder = krea_te,
# ``base_local_dir`` is None whenever nothing was staged, and then this is a repo
# id: the same guard the pipe_kwargs below carry for every other family.
local_files_only = local_files_only,
)
pipe.to(device)
return pipe
pipe_kwargs: dict[str, Any] = {
"local_files_only": local_files_only,
"torch_dtype": dtype,
"transformer": transformer,
"cache_dir": hub_cache_dir(),
@ -4502,6 +4638,7 @@ class DiffusionBackend:
fam = fam,
te_quant_mode = te_quant_mode,
target = target,
local_files_only = local_files_only,
)
)
# Same pre-cast TE injection as the other branches: the dense path supplies only the transformer.
@ -6037,6 +6174,15 @@ def _base_file_downloaded(rfilename: str, *, include_transformer: bool = False)
``include_transformer`` admits the ``transformer/`` shards for loads where the
dense transformer-quant path will fetch them anyway (see
``_dense_quant_prefetch_needed``)."""
if rfilename == "transformer/config.json":
# The one transformer/ file that is ALWAYS needed, shards or not: from_single_file(config =
# <repo id>, subfolder = "transformer") resolves it through the Hub, so a load that promised
# to download nothing cannot keep that promise unless this ~1 KB file is already on disk.
# Excluding it made the locality gate pass and the load fetch it afterwards, AFTER the
# resident pipeline was evicted. Counting it here is what lets the gate refuse up front
# (cleanly, before eviction) or clear a pick that really can load offline. Video keeps the
# same exception at video.py's snapshot filter.
return True
if rfilename.startswith("transformer/"):
return include_transformer
if "/" not in rfilename: # top-level: only the pipeline manifest is fetched

View file

@ -226,6 +226,20 @@ def select_and_activate_engine(
return _activate(ENGINE_DIFFUSERS, reason)
def native_binary_installed() -> bool:
"""Whether a RUNNABLE sd.cpp binary is already on disk, installing nothing to find out.
Separated from the prediction because the two answers differ where it matters: prediction
counts an absent binary as available whenever installing one is allowed, and a caller that
must know whether selection could still fall back to diffusers needs the unassumed answer.
"""
server_binary = ensure_sd_server_binary(allow_install = False)
if server_binary and _server_binary_runnable(server_binary):
return True
binary = ensure_sd_cpp_binary(allow_install = False)
return bool(binary and SdCppEngine(binary = binary).version() is not None)
def predict_engine(fam: DiffusionFamily, *, model_kind: Optional[str] = None) -> str:
"""The engine a load of ``fam`` would select on this host, WITHOUT any side effect.
@ -256,13 +270,7 @@ def predict_engine(fam: DiffusionFamily, *, model_kind: Optional[str] = None) ->
if not (policy_eligible and family_sd_cpp_supported(fam)):
return ENGINE_DIFFUSERS
server_binary = ensure_sd_server_binary(allow_install = False)
if server_binary and not _server_binary_runnable(server_binary):
server_binary = None
binary = ensure_sd_cpp_binary(allow_install = False)
if binary and SdCppEngine(binary = binary).version() is None:
binary = None
native_available = bool(binary or server_binary) or _install_allowed()
native_available = native_binary_installed() or _install_allowed()
return select_diffusion_engine(
backend, native_available = native_available, prefer_native = prefer_native
)

View file

@ -34,6 +34,7 @@ def hidream_te4_kwargs(
fam: Any = None,
te_quant_mode: Optional[str] = None,
target: Any = None,
local_files_only: bool = False,
) -> dict[str, Any]:
"""``{text_encoder_4, tokenizer_4}`` kwargs for a HiDream pipeline ``from_pretrained``.
@ -44,11 +45,27 @@ def hidream_te4_kwargs(
TE4 -- HiDream's HEAVIEST encoder -- is handled here: when the requested TE quant is
layerwise fp8 (and the device/family qualify, same gates as the runtime cast), TE4 is
fp8-cast too, preferring the hosted pre-cast checkpoint (~half the download) and
falling back to dense-load-then-cast. Any other mode keeps today's dense bf16 TE4."""
falling back to dense-load-then-cast. Any other mode keeps today's dense bf16 TE4.
``local_files_only`` is set by a load no user asked for, where fetching this repo is the
thing the caller promised would not happen: it raises here instead of downloading 16 GB."""
import torch # noqa: F401 -- dtype values are torch dtypes; import keeps parity with callers
from transformers import AutoTokenizer, LlamaForCausalLM
tokenizer_4 = AutoTokenizer.from_pretrained(HIDREAM_LLAMA_REPO, token = hf_token)
# Pinned to the LIVE hub root: ``encoder_repo_complete`` verifies these assets there, so an
# unpinned lookup after a mid-session cache-folder change searches huggingface_hub's
# import-time root instead and fails under local_files_only for a 16 GB encoder that is
# present, after the resident image pipeline was evicted.
from utils.hf_cache_settings import active_hf_hub_cache
cache_dir = active_hf_hub_cache()
tokenizer_4 = AutoTokenizer.from_pretrained(
HIDREAM_LLAMA_REPO,
token = hf_token,
local_files_only = local_files_only,
cache_dir = cache_dir,
)
fp8_engages = False
if target is not None:
@ -102,6 +119,8 @@ def hidream_te4_kwargs(
output_attentions = True,
torch_dtype = dtype,
token = hf_token,
local_files_only = local_files_only,
cache_dir = cache_dir,
)
if fp8_engages:
try:
@ -122,6 +141,8 @@ def hidream_te4_kwargs(
output_hidden_states = True,
output_attentions = True,
torch_dtype = dtype,
local_files_only = local_files_only,
token = hf_token,
cache_dir = cache_dir,
)
return {"text_encoder_4": text_encoder_4, "tokenizer_4": tokenizer_4}

View file

@ -36,11 +36,34 @@ logger = get_logger(__name__)
KREA2_FAMILY_NAME = "krea-2"
def load_krea2_tokenizer(repo_id: str, hf_token: Optional[str] = None):
def _live_cache_dir() -> str:
"""Studio's LIVE hub cache root, which every component load here must be pinned to.
An unset ``cache_dir`` resolves through huggingface_hub's import-time constant, and Studio's
cache folder is a setting: after a mid-session change the two roots differ. This assembler is
reached with a repo id, and the locality gate that cleared the switch reads the live root
(``media_locality`` passes ``cache_dir = hub_cache_dir()``), so an unpinned load looks in the
OTHER root -- which under ``local_files_only`` raises after the resident pipeline was already
evicted, for a model that is fully downloaded. Read from utils rather than
``diffusion.hub_cache_dir`` to avoid a circular import, the same way diffusion_auto_policy does.
"""
from utils.hf_cache_settings import active_hf_hub_cache
return active_hf_hub_cache()
def load_krea2_tokenizer(
repo_id: str,
hf_token: Optional[str] = None,
local_files_only: bool = False,
):
"""The Krea 2 tokenizer, tolerating the repo's transformers-5.x tokenizer config."""
from transformers import AutoTokenizer
kwargs: dict[str, Any] = {"subfolder": "tokenizer"}
kwargs: dict[str, Any] = {
"subfolder": "tokenizer",
"local_files_only": local_files_only,
"cache_dir": _live_cache_dir(),
}
if hf_token:
kwargs["token"] = hf_token
try:
@ -64,11 +87,16 @@ def load_krea2_text_encoder(
repo_id: str,
dtype,
hf_token: Optional[str] = None,
local_files_only: bool = False,
):
"""The Qwen3-VL text encoder, remapping 5.x ``rope_parameters`` for a 4.x runtime."""
from transformers import AutoConfig, Qwen3VLModel
kwargs: dict[str, Any] = {"subfolder": "text_encoder"}
kwargs: dict[str, Any] = {
"subfolder": "text_encoder",
"local_files_only": local_files_only,
"cache_dir": _live_cache_dir(),
}
if hf_token:
kwargs["token"] = hf_token
config = AutoConfig.from_pretrained(repo_id, **kwargs)
@ -76,7 +104,11 @@ def load_krea2_text_encoder(
return Qwen3VLModel.from_pretrained(repo_id, config = config, dtype = dtype, **kwargs)
def _load_model_index(repo_id: str, hf_token: Optional[str] = None) -> dict[str, Any]:
def _load_model_index(
repo_id: str,
hf_token: Optional[str] = None,
local_files_only: bool = False,
) -> dict[str, Any]:
"""model_index.json as a dict, from a local path or the Hub cache."""
is_local_dir = False
try:
@ -92,7 +124,13 @@ def _load_model_index(repo_id: str, hf_token: Optional[str] = None) -> dict[str,
raise FileNotFoundError(f"model_index.json not found in local model dir {repo_id}")
from huggingface_hub import hf_hub_download
path = hf_hub_download(repo_id, "model_index.json", token = hf_token or None)
path = hf_hub_download(
repo_id,
"model_index.json",
token = hf_token or None,
local_files_only = local_files_only,
cache_dir = _live_cache_dir(),
)
return json.loads(Path(path).read_text(encoding = "utf-8"))
@ -103,6 +141,7 @@ def load_krea2_pipeline(
transformer = None,
with_transformer: bool = True,
text_encoder = None,
local_files_only: bool = False,
):
"""A ready ``Krea2Pipeline`` for ``repo_id`` (still on CPU; caller places it).
@ -112,6 +151,13 @@ def load_krea2_pipeline(
pre-cast TE path (diffusion_te_prequant) hand in an already-built encoder, skipping
the dense Qwen3-VL download. The remaining components (VAE, tokenizer, scheduler)
come from the repo.
``local_files_only`` is a load nobody asked for. This assembler is reached with a REPO ID
rather than a staged snapshot dir and builds every component itself, so without the flag a
switch that verified locality from the outside can still pull the 26 GB transformer, the
8.88 GB Qwen3-VL encoder and the VAE here, after the resident pipeline was evicted. Every
component load below therefore resolves from the cache or raises, which is what the
caller's ``pipe_kwargs`` already does for every non-Krea family.
"""
import diffusers
@ -124,20 +170,37 @@ def load_krea2_pipeline(
)
token = hf_token or None
tokenizer = load_krea2_tokenizer(repo_id, hf_token = token)
cache_dir = _live_cache_dir()
tokenizer = load_krea2_tokenizer(repo_id, hf_token = token, local_files_only = local_files_only)
if text_encoder is None:
text_encoder = load_krea2_text_encoder(repo_id, dtype, hf_token = token)
text_encoder = load_krea2_text_encoder(
repo_id, dtype, hf_token = token, local_files_only = local_files_only
)
scheduler = diffusers.FlowMatchEulerDiscreteScheduler.from_pretrained(
repo_id, subfolder = "scheduler", token = token
repo_id,
subfolder = "scheduler",
token = token,
local_files_only = local_files_only,
cache_dir = cache_dir,
)
vae = diffusers.AutoencoderKLQwenImage.from_pretrained(
repo_id, subfolder = "vae", torch_dtype = dtype, token = token
repo_id,
subfolder = "vae",
torch_dtype = dtype,
token = token,
local_files_only = local_files_only,
cache_dir = cache_dir,
)
if transformer is None and with_transformer:
transformer = diffusers.Krea2Transformer2DModel.from_pretrained(
repo_id, subfolder = "transformer", torch_dtype = dtype, token = token
repo_id,
subfolder = "transformer",
torch_dtype = dtype,
token = token,
local_files_only = local_files_only,
cache_dir = cache_dir,
)
model_index = _load_model_index(repo_id, hf_token = token)
model_index = _load_model_index(repo_id, hf_token = token, local_files_only = local_files_only)
return diffusers.Krea2Pipeline(
scheduler = scheduler,
vae = vae,

View file

@ -601,6 +601,7 @@ def load_prequantized_transformer(
cache_dir: Optional[str] = None,
prepare_model: Optional[Any] = None,
config_subfolder: str = "transformer",
local_files_only: bool = False,
logger: Any = None,
) -> Optional[Any]:
"""Load the pre-quantized transformer described by ``source`` onto ``device``.
@ -649,7 +650,9 @@ def load_prequantized_transformer(
)
return None
path = _resolve_checkpoint_path(source, hf_token, cache_dir)
path = _resolve_checkpoint_path(
source, hf_token, cache_dir, local_files_only = local_files_only
)
if path is None:
return None
@ -670,7 +673,13 @@ def load_prequantized_transformer(
# the pinned root may be gone or read-only, and load_config's raise is swallowed below into
# a None return, silently dropping a prequant whose checkpoint is cached and already loaded.
config = _load_transformer_config(
transformer_cls, base, hf_token, cache_dir, path, config_subfolder
transformer_cls,
base,
hf_token,
cache_dir,
path,
config_subfolder,
local_files_only = local_files_only,
)
from accelerate import init_empty_weights
@ -765,6 +774,7 @@ def _download_checkpoint_name(
cache_dir: Optional[str],
*,
propagate_missing: bool,
local_files_only: bool = False,
) -> str:
"""Download ONE checkpoint filename, reusing a copy that sits under the other cache root.
@ -788,6 +798,7 @@ def _download_checkpoint_name(
filename = name,
token = hf_token,
cache_dir = None,
local_files_only = local_files_only,
)
except LocalEntryNotFoundError: # offline with the copy right there: use it
return elsewhere
@ -802,6 +813,7 @@ def _download_checkpoint_name(
filename = name,
token = hf_token,
cache_dir = cache_dir,
local_files_only = local_files_only,
)
@ -809,8 +821,13 @@ def _resolve_checkpoint_path(
source: PrequantSource,
hf_token: Optional[str],
cache_dir: Optional[str] = None,
*,
local_files_only: bool = False,
) -> Optional[str]:
"""The local file path for ``source``, downloading from the Hub if needed; None if absent."""
"""The local file path for ``source``, downloading from the Hub if needed; None if absent.
``local_files_only`` is the caller's promise that this load may not fetch anything, so a cache
miss answers None and the build falls back rather than pulling several GB nobody asked for."""
if source.kind == "path":
import os
@ -829,6 +846,7 @@ def _resolve_checkpoint_path(
hf_token,
cache_dir,
propagate_missing = has_fallback,
local_files_only = local_files_only,
)
except EntryNotFoundError:
if not has_fallback:
@ -840,6 +858,7 @@ def _resolve_checkpoint_path(
hf_token,
cache_dir,
propagate_missing = False,
local_files_only = local_files_only,
)
return None
@ -873,13 +892,22 @@ def _load_transformer_config(
cache_dir: Optional[str],
checkpoint_path: str,
subfolder: str = "transformer",
*,
local_files_only: bool = False,
) -> Any:
"""``transformer_cls.load_config`` against the checkpoint's cache root, then the other one."""
"""``transformer_cls.load_config`` against the checkpoint's cache root, then the other one.
The config is a few KB, but it is still a Hub fetch, and a load that promised to reach nothing
has to keep that promise for the small files too."""
last: Optional[BaseException] = None
for root in _config_cache_roots(checkpoint_path, cache_dir):
try:
return transformer_cls.load_config(
base, subfolder = subfolder, token = hf_token, cache_dir = root
base,
subfolder = subfolder,
token = hf_token,
cache_dir = root,
local_files_only = local_files_only,
)
except Exception as exc: # noqa: BLE001 — try the other root before giving up
last = exc

View file

@ -0,0 +1,522 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Opt-in model auto-switch for the image and video generation APIs.
The chat twin lives in ``local_model_resolver`` + ``routes.inference``: a ``/v1`` request
naming a downloaded GGUF loads it before serving. Media had no equivalent, so
``POST /v1/images/generations`` answered 503 unless someone had already picked a model on the
Images page, and ``model`` was documented as informational. This resolves that name against
the downloaded image/video models, drains what the backend is doing, and runs the load the
picker would run.
Off by default (``media_api_auto_switch_model``), so existing clients see no change.
Only downloaded models resolve, and an unknown name is refused rather than answered by
whatever is resident. Nothing here starts a download: the media equivalent of the chat
auto-download setting would let one API key spend tens of GB, which is its own decision.
Both waits are bounded, because Studio's secure-mode tunnel caps an origin response near 100
seconds. Exceeding a bound leaves the work running and asks the caller to retry, the contract
``begin_load`` already gives the UI.
This module is the orchestration. The pieces it drives live next to it: ``media_model_index``
resolves a name and recognises the resident model, ``media_locality`` proves a pick is already
downloaded, ``media_switch_backends`` waits out work a switch would interrupt,
``media_switch_locks`` serializes switches, and ``media_switch_errors`` holds the refusals.
"""
from __future__ import annotations
import asyncio
import contextlib
import functools
import time
from typing import Any, Callable, Optional
from core.inference.gpu_arbiter import DIFFUSION, VIDEO
from core.inference.media_locality import is_edit_only, missing_download_bytes
from core.inference.media_model_index import (
IMAGE_TASK,
VIDEO_TASK,
MediaModelPick,
available_media_model_ids,
expected_partition,
invalidate_index,
partition_matches,
resident_is_gguf,
resident_is_pick,
resolve_local_media_model,
same_identity,
satisfied_by,
)
from core.inference.media_switch_backends import (
POLL_S,
backend_for,
drain,
load_takes_the_gpu,
)
from core.inference.media_switch_errors import (
EDIT_ONLY_MSG,
LOADING_MSG,
RETRY_AFTER_S,
UNVERIFIED_MSG,
bounded,
busy,
format_available,
incomplete_message,
refuse,
)
from core.inference.media_switch_locks import (
gpu_switch_lock,
note_switcher,
note_waiter,
switch_lock,
)
from loggers import get_logger
logger = get_logger(__name__)
# one end-to-end budget for the whole switch, under the ~100s tunnel window
_SWITCH_BUDGET_S = 90.0
# a generation the caller cannot see is yielded to rather than cut short, capped inside the budget
_DRAIN_WAIT_S = 30.0
# how long the gates are kept for a load that has not reached begin_load yet
_SETUP_GRACE_S = 120.0
def _resident_answers_exactly(resident: dict[str, Any], name: str) -> bool:
"""Whether the resident model is this exact name, needing no discovery at all.
A scan that failed or skipped an entry would otherwise 404 the very model that is loaded,
for as long as the empty index stays cached. Never true for a resident GGUF: a bare repo id
means the preferred quant, which this comparison cannot see, so it would serve whichever
quant happens to be up.
"""
return (
bool(resident.get("loaded"))
and not resident_is_gguf(resident)
and partition_matches(resident)
and same_identity(name, str(resident.get("repo_id") or ""))
)
async def _require_local(
owner: str,
pick: MediaModelPick,
deadline: float,
*,
kind: str,
openai_errors: bool,
hf_token: Optional[str],
) -> None:
"""Refuse unless *pick* is provably downloaded in full.
Bounded inside the switch budget, and free of side effects, so a planner that stalls can
safely give back whatever locks and gates the caller is holding while it runs.
"""
missing = await bounded(
asyncio.to_thread(missing_download_bytes, owner, pick, hf_token),
deadline,
kind = kind,
openai_errors = openai_errors,
)
if missing is None:
raise refuse(
UNVERIFIED_MSG.format(model = pick.model_id, kind = kind),
status_code = 409,
openai_errors = openai_errors,
code = "model_not_downloaded",
)
if missing:
raise refuse(
incomplete_message(pick.model_id, missing, kind),
status_code = 409,
openai_errors = openai_errors,
code = "model_not_downloaded",
)
async def _acquire_all(locks: list, deadline: float, *, kind: str, openai_errors: bool) -> None:
"""Take every lock within the budget, releasing what was taken if one cannot be had.
A request that spent most of its budget resolving would otherwise queue behind another full
switch and blow past the response window before any of the inner waits could notice.
"""
acquired: list = []
try:
for held in locks:
await bounded(held.acquire(), deadline, kind = kind, openai_errors = openai_errors)
acquired.append(held)
except BaseException:
for held in reversed(acquired):
held.release()
raise
def _consume_detached_error(task: "asyncio.Task") -> None:
"""Retrieve a handed-over task's exception, since the caller may have stopped awaiting it.
``_gated_start_load`` refuses on ordinary paths (a backend still busy at the in-gate drain,
a cache deletion during it), and once the budget expires nothing awaits the task again. An
unretrieved exception is reported by the loop at collection time, so a routine slow switch
would log a traceback for a refusal that was handled correctly.
"""
if task.cancelled():
return
exc = task.exception()
if exc is not None:
logger.debug("Media auto-switch: setup finished after the caller stopped waiting: %s", exc)
async def _await_loaded(
backend: Any,
name: str,
pick: MediaModelPick,
deadline: float,
*,
kind: str,
openai_errors: bool,
) -> bool:
"""Poll the background load until the REQUESTED model is resident; False if still going.
Checked against the pick, not merely "something is loaded": a user load accepted between
two polls supersedes this one, and returning success there would generate on the
replacement while reporting the requested model.
The probes are bounded like every other wait here: ``load_progress`` walks cache directories
to count bytes, so on a slow or stalled filesystem a single poll can outlive the budget that
the check at the bottom of the loop is meant to enforce.
"""
probe = functools.partial(bounded, deadline = deadline, kind = kind, openai_errors = openai_errors)
while True:
progress = await probe(asyncio.to_thread(backend.load_progress)) or {}
phase = progress.get("phase")
if phase == "error":
raise RuntimeError(progress.get("error") or "The model failed to load.")
if phase in (None, "ready"):
status = await probe(asyncio.to_thread(backend.status))
# the landed check, not the skip check: this load is ours, so ambiguity is settled
if resident_is_pick(status, name, pick):
return True
# loaded, but not this pick: a load that landed after ours replaced it
raise RuntimeError(f"'{pick.model_id}' was replaced by another load before it served.")
if time.monotonic() >= deadline:
return False
await asyncio.sleep(POLL_S)
async def _start_load(
owner: str,
pick: MediaModelPick,
current_subject: str,
hf_token: Optional[str] = None,
) -> None:
"""Run the load its own route would run, as an API load rather than a user one."""
partition = expected_partition(pick)
if owner == DIFFUSION:
from models.inference import DiffusionLoadRequest
from routes.inference import load_diffusion_model_gated
await load_diffusion_model_gated(
DiffusionLoadRequest(
model_path = pick.model_path,
gguf_filename = pick.gguf_filename,
model_kind = pick.model_kind,
hf_token = hf_token,
),
current_subject,
user_initiated = False,
)
else:
from models.inference import VideoLoadRequest
from routes.video import load_video_model_gated
await load_video_model_gated(
VideoLoadRequest(
model_path = pick.model_path,
gguf_filename = pick.gguf_filename,
model_kind = pick.model_kind,
h3_task = partition,
hf_token = hf_token,
),
current_subject,
user_initiated = False,
)
logger.info("Media auto-switch: loading %s on the %s backend", pick.model_id, owner)
async def _gated_start_load(
owner: str,
name: str,
pick: MediaModelPick,
current_subject: str,
locks: list,
deadline: float,
*,
kind: str,
openai_errors: bool,
hf_token: Optional[str],
takes_the_gpu: bool,
) -> bool:
"""Run the final checks and start the load, owning the gates and *locks* throughout.
Returns True when the resident model already answers the request, so the caller can stop.
Ownership is the point. The caller shields this and may stop waiting on it, and the work
from the last drain observation through ``begin_load`` must not be interruptible: engine
activation unloads the resident pipeline on its way, so anything admitted before
registration would be cut short by a load that no longer has a request behind it.
Ownership is bounded all the same: a load that has not registered within ``_SETUP_GRACE_S``
gives the gates back and carries on without them, since an installer running for minutes
behind them costs more than the race they close.
The gates held are the ones this load could evict behind, entered in a fixed order so two
switches cannot deadlock, and one at a time under the budget: a stalled holder elsewhere
would otherwise pin this task, and with it the switch lock, indefinitely. Cancelling during
that acquisition is free, and the stack releases whatever was already entered; nothing past
it may be interrupted.
Chat's lifecycle gate is the FIRST of them, not the last. Every media generation route is
counted on chat's in-flight counter as well as its own, and the middleware takes chat's gate
and releases it before it parks on the media one. With the media gates taken first, a request
arriving in between passed the still-open chat gate, incremented chat's ``_inflight``, and
only then blocked on the held media gate: the in-gate drain discounts it on the media side
(``count_pending=False``) but ``chat_busy(count_pending=False)`` still read it as running chat
work, and an otherwise idle switch answered 409 without loading anything. Taking chat's gate
first parks such a request in ``_note_pending`` instead, where both counters ignore it, and
the middleware never holds a media gate while it waits for chat's, so the order is safe.
A load that does not take the GPU holds its own backend's gate only. It cannot evict chat or
the other media backend, so waiting on their gates would let an unrelated chat teardown time
the switch out, and holding them would block new chat and video requests for as long as the
re-plan and the load registration take.
"""
from core.inference.media_keepwarm import admission_gate
from core.inference.llama_keepwarm import inference_lifecycle_gate
needed = (
(inference_lifecycle_gate(), admission_gate(DIFFUSION), admission_gate(VIDEO))
if takes_the_gpu
else (admission_gate(owner),)
)
try:
async with contextlib.AsyncExitStack() as gates:
for gate in needed:
await bounded(
gates.enter_async_context(gate),
deadline,
kind = kind,
openai_errors = openai_errors,
)
# re-resolved under the gate: a concurrent load can activate the other image engine
backend = backend_for(owner)
# what the drain waited out may have been the very load this request wanted
if satisfied_by(await asyncio.to_thread(backend.status), name, pick):
return True
if not await drain(
owner,
backend,
time.monotonic(),
count_pending = False,
probe_deadline = deadline,
kind = kind,
openai_errors = openai_errors,
):
raise busy(kind, openai_errors)
# re-planned because a cache deletion during the drain sees no load to guard against
await _require_local(
owner,
pick,
deadline,
kind = kind,
openai_errors = openai_errors,
hf_token = hf_token,
)
# given its own task and waited on with a cap: a first-run native install runs for
# minutes before begin_load, and holding both media gates and chat's that long
# blocks every unrelated request. On expiry the load keeps going without them.
setup = asyncio.ensure_future(_start_load(owner, pick, current_subject, hf_token))
setup.add_done_callback(_consume_detached_error)
with contextlib.suppress(asyncio.TimeoutError):
await asyncio.wait_for(asyncio.shield(setup), _SETUP_GRACE_S)
return False
finally:
for held in reversed(locks):
held.release()
async def maybe_auto_switch_media_model(
requested_model: Optional[str],
*,
owner: str,
current_subject: str,
openai_errors: bool,
hf_token: Optional[str] = None,
before_switch: Optional[Callable[[MediaModelPick], None]] = None,
) -> None:
"""Load the image or video model a generation request names, if it is not resident.
No-op when the setting is off or nothing was named, so ``model`` keeps its old
informational meaning for every existing client. With the setting on, a name that resolves
to no downloaded model is refused: answering it would return one model's output under
another's name.
``before_switch`` is the caller's last say on the resolved pick, run only when a switch is
actually going to happen. It exists so a request the target model cannot serve is refused
while the resident one is still loaded, rather than after a multi-minute load; a request the
resident model already answers skips it, since the generate route judges that one anyway.
"""
from utils.openai_auto_switch_settings import get_media_auto_switch_enabled
if not isinstance(requested_model, str) or not requested_model.strip():
return
if not get_media_auto_switch_enabled():
return
# started before resolution: the cold scan is part of the wait the caller experiences
deadline = time.monotonic() + _SWITCH_BUDGET_S
name = requested_model.strip()
task = IMAGE_TASK if owner == DIFFUSION else VIDEO_TASK
kind = "image" if owner == DIFFUSION else "video"
if _resident_answers_exactly(await asyncio.to_thread(backend_for(owner).status), name):
return
# off the loop: a cold index walks the model roots and reads gguf headers
pick = await bounded(
asyncio.to_thread(resolve_local_media_model, name, task = task),
deadline,
kind = kind,
openai_errors = openai_errors,
)
if pick is None:
available = format_available(
await bounded(
asyncio.to_thread(available_media_model_ids, task),
deadline,
kind = kind,
openai_errors = openai_errors,
)
)
raise refuse(
f"No downloaded {kind} model matches '{name}'."
+ (f" Downloaded {kind} models: {available}." if available else ""),
status_code = 404,
openai_errors = openai_errors,
code = "model_not_found",
)
# before anything is evicted: the load would otherwise finish and be refused for lacking txt2img
if owner == DIFFUSION and await asyncio.to_thread(is_edit_only, pick):
raise refuse(
EDIT_ONLY_MSG.format(model = pick.model_id),
status_code = 400,
openai_errors = openai_errors,
code = "invalid_value",
)
# re-read: the index build can run for the whole budget, and an idle unload can land in it
if satisfied_by(await asyncio.to_thread(backend_for(owner).status), name, pick):
return
if before_switch is not None:
await bounded(
asyncio.to_thread(before_switch, pick),
deadline,
kind = kind,
openai_errors = openai_errors,
)
lock = switch_lock(owner)
# held only when the load takes the gpu, since a cpu-only switch takes it from nobody
takes_the_gpu = await asyncio.to_thread(load_takes_the_gpu)
gpu_lock = gpu_switch_lock() if takes_the_gpu else None
locks = [held for held in (gpu_lock, lock) if held is not None]
with note_switcher(owner):
# the marker covers only the wait: once this request holds the lock it is real work
with note_waiter(owner):
await _acquire_all(locks, deadline, kind = kind, openai_errors = openai_errors)
handed_over = False
try:
backend = backend_for(owner)
# re-read under the lock: a concurrent request may have just loaded this model
if satisfied_by(await asyncio.to_thread(backend.status), name, pick):
return
await _require_local(
owner,
pick,
deadline,
kind = kind,
openai_errors = openai_errors,
hf_token = hf_token,
)
if not await drain(
owner,
backend,
min(deadline, time.monotonic() + _DRAIN_WAIT_S),
# probes answer to the switch budget: only a spent budget is the slow-switch 503
probe_deadline = deadline,
kind = kind,
openai_errors = openai_errors,
):
raise busy(kind, openai_errors)
# its own task, so a timeout below frees the caller without unwinding gate or lock
setup = asyncio.ensure_future(
_gated_start_load(
owner,
name,
pick,
current_subject,
locks,
deadline,
kind = kind,
openai_errors = openai_errors,
hf_token = hf_token,
takes_the_gpu = takes_the_gpu,
)
)
setup.add_done_callback(_consume_detached_error)
handed_over = True
if await bounded(
asyncio.shield(setup), deadline, kind = kind, openai_errors = openai_errors
):
return
finally:
if not handed_over:
for held in reversed(locks):
held.release()
try:
# re-resolved: an engine switch (diffusers <-> sd.cpp) replaces the object
ready = await _await_loaded(
backend_for(owner), name, pick, deadline, kind = kind, openai_errors = openai_errors
)
except RuntimeError as exc:
# the loader already redacts this text; a bare raise would 500 with it
raise refuse(
f"'{pick.model_id}' could not be loaded: {exc}",
status_code = 503,
openai_errors = openai_errors,
code = "model_load_failed",
)
if not ready:
raise refuse(
LOADING_MSG.format(model = pick.model_id),
status_code = 503,
openai_errors = openai_errors,
code = "model_loading",
retry_after = RETRY_AFTER_S,
)
__all__ = [
"IMAGE_TASK",
"VIDEO_TASK",
"MediaModelPick",
"available_media_model_ids",
"invalidate_index",
"maybe_auto_switch_media_model",
"resolve_local_media_model",
]

View file

@ -22,6 +22,7 @@ from __future__ import annotations
import asyncio
import contextlib
import os
import sys
import threading
import time
@ -77,6 +78,10 @@ class _Tracker:
with self._lock:
self._last_active = time.monotonic()
def outstanding(self, *, count_pending: bool = True) -> int:
with self._lock:
return self._inflight + (self._pending if count_pending else 0)
def is_idle(self, ttl_seconds: float) -> bool:
with self._lock:
return (
@ -88,6 +93,95 @@ class _Tracker:
_TRACKERS = {DIFFUSION: _Tracker(DIFFUSION), VIDEO: _Tracker(VIDEO)}
# Per backend, not per engine object: a diffusers <-> sd.cpp switch replaces that object.
# Keyed by the target a load was started for, since a load route records this the moment the
# background load is accepted and that load can still fail with the previous model resident.
_LOAD_ORIGINS: dict[str, tuple[tuple[str, str, str], bool]] = {}
_LOAD_ORIGINS_GUARD = threading.Lock()
def _origin_key(
target: Optional[str],
variant: Optional[str],
partition: Optional[str] = None,
) -> tuple[str, str, str]:
"""The build a provenance record answers for: the repo/path AND its GGUF variant.
The path alone is not the build: a user-loaded Q4 and an API load of Q8 from the same repo
share it, so a failed API load would mark the resident Q4 as API-loaded and free it.
The variant is the token ``status()`` publishes, not a fuller filename label. Two builds the
token cannot separate therefore share a key, which errs toward reporting a model as
user-loaded and sparing it. Keying on more would never match what the resident model
publishes, and would spare everything.
"""
text = str(target or "").strip()
# A repo id folds case; a path does not, or /models/Foo and /models/foo share an origin.
key = os.path.normcase(text) if os.path.isabs(text) else text.lower()
return (key, str(variant or "").strip().lower(), str(partition or "").strip().lower())
def note_load_origin(
owner: str,
target: Optional[str],
variant: Optional[str] = None,
partition: Optional[str] = None,
*,
user_action: bool,
) -> None:
"""Record who asked for the model a load route is bringing up.
An API load over a user-loaded build with the same key keeps the user's mark: the two are
indistinguishable to ``status()`` (sibling GGUFs can share a quant token), and a load that
is accepted and then fails leaves the user's model resident, which this must not reclassify.
"""
key = _origin_key(target, variant, partition)
with _LOAD_ORIGINS_GUARD:
previous = _LOAD_ORIGINS.get(owner)
if not user_action and previous is not None and previous[0] == key and previous[1]:
return
_LOAD_ORIGINS[owner] = (key, user_action)
def loaded_by_user_action(
owner: str,
resident: Optional[str] = None,
variant: Optional[str] = None,
partition: Optional[str] = None,
) -> bool:
"""Whether the RESIDENT model was loaded from Studio rather than by an API request.
The record only answers for the build it was written against: a load that was accepted and
then failed leaves the previous model resident, and reading its origin off the failed
load would let the idle unload free a model the user had pinned. Anything unrecognised
reads as user-loaded, which is the direction that spares a model.
"""
with _LOAD_ORIGINS_GUARD:
entry = _LOAD_ORIGINS.get(owner)
if entry is None:
return True
key, user_action = entry
if resident is not None and key[0] and key != _origin_key(resident, variant, partition):
return True
return user_action
def other_request_count(
owner: str,
*,
current_request_counted: bool = False,
count_pending: bool = True,
) -> int:
"""Tracked media requests in flight on *owner*, excluding this one when it is counted.
The auto-switch drain reads this from inside a tracked request, so its own entry must
not make the backend look permanently busy. ``count_pending`` False drops requests that
have registered but are still blocked on the gate, which a gate holder must not wait for.
"""
total = _TRACKERS[owner].outstanding(count_pending = count_pending)
return max(0, total - 1) if current_request_counted else total
# The concrete mounted paths, not any path that ends in one of them: FastAPI answers
# /v1/anything/images/generations with a 404 without ever running an endpoint, and stamping
# that as activity would let unauthenticated 404s keep a pipeline resident past every TTL.
@ -118,6 +212,19 @@ def owner_for_path(path: str) -> Optional[str]:
return _TRACKED_PATHS.get(path)
@contextlib.asynccontextmanager
async def admission_gate(owner: str):
"""Hold new tracked media requests off *owner* for the duration of the block.
The media auto-switch keeps this closed from its final drain check through registering
the load: the load path cancels active work as it tears the pipeline down, so a
generation admitted in that gap would be cut short by a swap that just waited for the
queue to clear. Requests arriving meanwhile park in ``begin_request`` until it reopens.
"""
async with _gate(_TRACKERS[owner]):
yield
@contextlib.asynccontextmanager
async def _gate(tracker: _Tracker):
# Polled non-blocking acquire, exactly like llama_keepwarm's: it keeps the wait off this
@ -265,10 +372,19 @@ async def _tick(tracker: _Tracker, ttl: float) -> None:
return
# Re-read the effective setting immediately before the teardown. One step covers both
# backends and an unload frees several GB, so a residency veto applied while it runs
# (Model Memory, API-only, or the TTL itself moved) would otherwise be ignored by
# every teardown left in the step, freeing a model the settings page now calls pinned.
# (Model Memory, or the TTL itself moved) would otherwise be ignored by every teardown
# left in the step, freeing a model the settings page now calls pinned.
ttl = _effective_ttl()
if ttl <= 0 or not tracker.is_idle(ttl):
if (
ttl <= 0
or not tracker.is_idle(ttl)
or _user_pinned(
tracker.owner,
status.get("repo_id"),
status.get("gguf_variant"),
status.get("h3_task"),
)
):
return
await asyncio.to_thread(backend.unload)
# Drop ownership only if nothing came back meanwhile, and check it under the arbiter
@ -283,11 +399,23 @@ async def _tick(tracker: _Tracker, ttl: float) -> None:
def _effective_ttl() -> float:
"""The media TTL with the residency vetoes applied: 0 means nothing is unloaded."""
"""The media TTL with the residency veto applied: 0 means nothing is unloaded."""
from utils.openai_auto_switch_settings import get_media_auto_unload_idle_seconds
return float(get_media_auto_unload_idle_seconds())
def _user_pinned(
owner: str, resident: Optional[str], variant: Optional[str], partition: Optional[str]
) -> bool:
"""Whether "only unload models loaded by the API" spares this backend's model.
Read immediately before the teardown, like the TTL: the setting can be turned on
while a step is running, and a model it now pins must not be freed by the rest of it.
"""
from utils.openai_auto_switch_settings import get_auto_unload_api_only
return get_auto_unload_api_only() and loaded_by_user_action(owner, resident, variant, partition)
async def idle_unload_step() -> None:
"""The media half of one idle_unload_loop tick. Inert when the TTL is off."""
ttl = _effective_ttl()

View file

@ -0,0 +1,539 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Proving a media model is already downloaded, before the switch evicts anything for it.
Auto-switch never downloads. That promise is easy to state and hard to keep, because the index
only sees CHECKPOINTS: a GGUF or single-file pick loads its text encoders and VAE from a
companion base repo, HiDream-I1 fetches a 16 GB Llama encoder no amount of pipeline on disk
accounts for, and an LTX-2.3 checkpoint pulls VAE, audio and connector artifacts the planner
only recognises by name. Any of those would let one API request spend tens of gigabytes.
So locality is verified through the same download planner ``/images/download-plan`` serves, and
the answer is tri-state: complete, incomplete by some number of bytes, or unverifiable. Zero
bytes from a planner that failed is not evidence of a complete cache, and the switch refuses on
anything short of proof.
"""
from __future__ import annotations
from dataclasses import replace
from pathlib import Path
from typing import Any, Optional
from core.inference.gpu_arbiter import DIFFUSION, VIDEO
from core.inference.media_model_index import MediaModelPick
from core.inference.media_switch_backends import backend_for
from core.inference.media_switch_errors import UNSIZED_MISSING
from loggers import get_logger
logger = get_logger(__name__)
# the image family whose pipeline loads a separate encoder repo its own directory cannot hold
_EXTERNAL_ENCODER_FAMILIES = frozenset({"hidream-i1"})
# encoder repos that always ship sharded, where a missing index means an interrupted download
_SHARDED_ENCODER_REPOS = frozenset({"unsloth/Meta-Llama-3.1-8B-Instruct"})
# what from_pretrained reads besides the weights, tiny next to them but still a download
_ENCODER_METADATA_FILES = ("config.json", "tokenizer.json", "tokenizer_config.json")
# suffixes a weight-bearing pipeline component can satisfy from_pretrained with
_WEIGHT_SUFFIXES = (".safetensors", ".bin", ".pt", ".pth", ".ckpt", ".msgpack", ".onnx")
# any one of these is the vocabulary a tokenizer class builds itself from
_TOKENIZER_ASSETS = (
"tokenizer.json",
"vocab.json",
"vocab.txt",
"merges.txt",
"spiece.model",
"tokenizer.model",
"sentencepiece.bpe.model",
)
def detected_image_family(pick: MediaModelPick) -> Any:
"""The diffusion family for *pick*, tried against its path and then its id.
The path comes first because it is the only needle the load route is ever handed, and the
only one that can carry a local ``model_index.json``: ``detect_family_for_pick`` reads that
index ahead of any guess made from a name, so asking about the id first answers FLUX for a
HiDream pipeline in a directory called ``flux.1`` while the loader answers HiDream and
fetches its 16 GB encoder. The id is kept as a fallback for a pick whose path says nothing.
"""
from core.inference.diffusion_families import detect_family_for_pick
for needle in (pick.model_path, pick.model_id):
if not needle:
continue
try:
fam = detect_family_for_pick(needle, pick.gguf_filename, None)
except Exception: # noqa: BLE001 -- a probe failure must not refuse a loadable pick
continue
if fam is not None:
return fam
return None
def normalized_pick(pick: MediaModelPick) -> MediaModelPick:
"""The pick as the LOAD route will read it, with a bare single-file directory reinterpreted.
Both load routes turn a kindless directory holding exactly one checkpoint into a
``single_file`` load and then resolve that family's companions. Planning the un-normalized
pick describes a local pipeline with nothing to fetch, and misses those companions.
"""
from core.inference.diffusion import resolve_local_single_file
if pick.model_kind or pick.gguf_filename:
return pick
sole = resolve_local_single_file(pick.model_path)
if sole is None:
return pick
return replace(pick, gguf_filename = sole, model_kind = "single_file")
def is_edit_only(pick: MediaModelPick) -> bool:
"""Whether *pick* is an instruction-editing family, which has no text-to-image mode.
The local catalog tags these text-to-image, so without this the switch would evict a working
model for a multi-GB pipeline that /v1/images/generations then refuses for lacking txt2img.
"""
from core.inference.diffusion import _family_workflows
fam = detected_image_family(normalized_pick(pick))
if fam is None:
return False
return "txt2img" not in _family_workflows(fam)
def _needs_external_encoder(pick: MediaModelPick) -> bool:
"""Whether this pick's pipeline fetches an encoder that its own directory cannot hold."""
# an unrecognised family keeps the shortcut, since refusing every on-device model is worse
fam = detected_image_family(pick)
return fam is not None and getattr(fam, "name", "") in _EXTERNAL_ENCODER_FAMILIES
def _cached_snapshot_file(repo_id: str, filename: str) -> Optional[str]:
"""The cached path of ``filename`` in ``repo_id``, or None when it is not downloaded."""
from huggingface_hub import try_to_load_from_cache
from core.inference.diffusion import hub_cache_dir
hit = try_to_load_from_cache(repo_id, filename, cache_dir = hub_cache_dir())
return hit if isinstance(hit, str) else None
def encoder_repo_complete(repo_id: str) -> bool:
"""Whether every shard of a cached encoder repo is present, not merely one of them.
``_upstream_is_cached`` counts any single weight file, while the pipeline calls
from_pretrained on the whole repository, so an interrupted sharded pull would otherwise
read as local and the load would fetch the rest.
The config and tokenizer files count as much as the shards. They are kilobytes rather than
gigabytes, but the encoder is built with ``AutoTokenizer.from_pretrained`` and
``LlamaForCausalLM.from_pretrained`` on the whole repository, so a cache holding every shard
and none of those still reaches the Hub during an accepted switch.
"""
import json
from core.inference.diffusion_families import _upstream_is_cached, cache_holds_files
if not _upstream_is_cached(repo_id):
return False
if not cache_holds_files(repo_id, list(_ENCODER_METADATA_FILES)):
return False
index = _cached_snapshot_file(repo_id, "model.safetensors.index.json")
if index is None:
# a repo known to be sharded has no unsharded reading, so a missing index means a partial
return repo_id not in _SHARDED_ENCODER_REPOS
with open(index, encoding = "utf-8") as handle:
shards = sorted(set((json.load(handle).get("weight_map") or {}).values()))
return bool(shards) and cache_holds_files(repo_id, shards)
def _missing_external_encoder(pick: MediaModelPick) -> Optional[int]:
"""0 when this local pipeline needs nothing more, else what its outside dependency costs.
HiDream-I1 loads unsloth/Meta-Llama-3.1-8B-Instruct unconditionally, around 16 GB, which no
amount of the pipeline being on disk accounts for. Checked against the cache directly rather
than through the planner, which cannot be handed an absolute pipeline path.
"""
if not _needs_external_encoder(pick):
return 0
from core.inference.diffusion_hidream import HIDREAM_LLAMA_REPO
try:
if encoder_repo_complete(HIDREAM_LLAMA_REPO):
return 0
except Exception as exc: # noqa: BLE001 -- an unreadable cache is not proof of locality
logger.debug("media auto-switch: hidream encoder probe failed: %s", exc)
return None
return UNSIZED_MISSING
def hidden_ltx23_extras(owner: str, pick: MediaModelPick) -> bool:
"""Whether this local video pick is an LTX-2.3 checkpoint the plan did not treat as one.
The planner judges 2.3 by name, while the loader reads the checkpoint header and then pulls
the 2.3 VAE, audio and connector artifacts. A renamed checkpoint therefore plans as 2.0,
reports nothing missing, and downloads those extras during assembly.
The family is resolved the way the loader resolves it, which falls back to the checkpoint's
``general.architecture`` where neither the repo nor the filename carries a family token.
Deciding by name alone left a generically named LTX checkpoint exempt from the very check
its header would have triggered.
"""
if owner != VIDEO or not pick.gguf_filename:
return False
try:
from core.inference.diffusion_families import resolve_local_gguf_child
from core.inference.video import _detect_load_family
from core.inference.video_ltx2 import LTX23_EXTRAS_REPO, is_ltx23_checkpoint
except Exception: # noqa: BLE001 -- no ltx support here means nothing to hide
return False
fam = _detect_load_family(pick.model_path, pick.gguf_filename, None) or (
_detect_load_family(pick.model_id, pick.gguf_filename, None) if pick.model_id else None
)
if fam is None or getattr(fam, "name", None) != "ltx-2":
return False
root = Path(pick.model_path).expanduser()
try:
if root.exists():
checkpoint = resolve_local_gguf_child(root, pick.gguf_filename)
else:
# a cached repo id: the checkpoint is on disk all the same, and its header decides
cached = _cached_snapshot_file(pick.model_path, pick.gguf_filename)
if cached is None:
return False
checkpoint = Path(cached)
except Exception: # noqa: BLE001 -- an unreadable pick is refused by the load itself
return False
if not is_ltx23_checkpoint(checkpoint):
return False
from core.inference.diffusion_families import cache_holds_files
from core.inference.video_ltx2 import ltx23_extras_files
extras = ltx23_extras_files(checkpoint)
# the exact three artifacts, since the repo also holds checkpoints that prove nothing here
return bool(extras) and not cache_holds_files(LTX23_EXTRAS_REPO, list(extras))
def planners_for(owner: str, pick: MediaModelPick) -> list:
"""Every engine whose plan this pick could end up loading through.
Usually one. ``predict_engine`` treats an absent sd.cpp binary as available whenever its
installation is allowed, while ``select_and_activate_engine`` falls back to diffusers when
that install produces nothing runnable, and the two engines read different companion sets.
Both are verified only in that case: with a runnable binary already on disk the load stays
native, and demanding the diffusers shards too would refuse a model sd.cpp can serve.
"""
if owner != DIFFUSION:
return [backend_for(owner)]
from core.inference.diffusion import resolve_model_kind
from core.inference.diffusion_engine_router import (
engine_for,
native_binary_installed,
predict_engine,
)
from core.inference.sd_cpp_engine import ENGINE_DIFFUSERS, ENGINE_SD_CPP
fam = detected_image_family(pick)
if fam is None:
return [backend_for(owner)]
kind = resolve_model_kind(pick.gguf_filename, pick.model_kind)
predicted = predict_engine(fam, model_kind = kind)
names = [predicted]
if predicted == ENGINE_SD_CPP and not native_binary_installed():
names.append(ENGINE_DIFFUSERS)
return [engine_for(name) for name in names]
def plan_gpu_ordinal() -> Optional[int]:
"""The card the load route will rank for itself, so the plan sizes the same file set.
Automatic precision is chosen per card, and a different card can select a different hosted
pre-quantized artifact, which a plan plotted against the default device would omit.
"""
from core.inference.diffusion_device import (
resolve_diffusion_device_target,
resolve_selected_cuda_ordinal,
)
if resolve_diffusion_device_target().device != "cuda":
return None
return resolve_selected_cuda_ordinal(None)
def _pipeline_components_present(root: Path) -> bool:
"""Whether every component a local pipeline's own index names is on disk.
A directory carrying a pipeline index is treated as complete by definition, because
from_pretrained reads it off disk and the planner cannot be asked about an absolute path.
That holds only if the components are actually there: a hand-copied or interrupted pipeline
passes the index check, and the loader then tears the resident pipeline down before
from_pretrained discovers the gap, leaving the API with no model at all.
Judged on what can be seen without reading weights. A directory carrying neither index is
not this function's business. A modular entry whose spec names another repository is checked
against the cache instead, since the load pulls that repository itself.
"""
import json
for name in ("model_index.json", "modular_model_index.json"):
index_file = root / name
if not index_file.is_file():
continue
try:
with open(index_file, encoding = "utf-8-sig") as handle:
index = json.load(handle)
except Exception as exc: # noqa: BLE001 -- an index the loader cannot read is not complete
logger.debug("media auto-switch: unreadable pipeline index under %s: %s", root, exc)
return False
if not isinstance(index, dict):
return False
for component, entry in index.items():
if component.startswith("_") or not isinstance(entry, (list, tuple)):
continue
# [null, null] marks a component this pipeline deliberately ships without.
if len(entry) not in (2, 3) or not entry[1]:
continue
# a modular entry is [library, class, spec], and its spec can name another repo,
# which this directory is never expected to hold but the load still pulls
hosted = _hosted_source(entry[2]) if len(entry) == 3 else None
if hosted is not None:
if not _hosted_component_cached(*hosted):
return False
continue
if not _component_present(root / component):
return False
return True
def _hosted_source(spec: Any) -> Optional[tuple[str, str, str, str]]:
"""What a modular index entry asks the loader for: repo, subfolder, revision and variant.
``ComponentSpec.load`` is handed the whole spec, so a component can pin a commit or a named
weight variant. Checking the default snapshot for those would approve a switch and then
download the pinned files after the resident pipeline is gone.
"""
if not isinstance(spec, dict):
return None
source = spec.get("pretrained_model_name_or_path") or spec.get("repo")
if not isinstance(source, str) or not source.strip():
return None
def _text(key: str) -> str:
value = spec.get(key)
return value.strip() if isinstance(value, str) else ""
return source.strip(), _text("subfolder"), _text("revision"), _text("variant")
def _cached_snapshot_root(repo_id: str, revision: str = "") -> Optional[Path]:
"""The cached snapshot a load of *repo_id* would read, or None when it is not downloaded.
The revision the loader asks for: the pinned one where a spec names it, else the one
``refs/main`` resolves to, rather than whichever snapshot sorts first. A superseded revision
can hold a complete component while the active one is partial, and approving the old copy is
how the load ends up fetching the new one.
"""
from core.inference.diffusion import hub_cache_dir
repo_dir = Path(hub_cache_dir()) / f"models--{repo_id.replace('/', '--')}"
snapshots = repo_dir / "snapshots"
# a pinned revision is a commit sha, or a branch or tag the cache records under refs/, and
# it is the only candidate: falling back to main is how the default snapshot approves a pin
for candidate in [revision] if revision else ["main"]:
pinned = snapshots / candidate
if pinned.is_dir():
return pinned
try:
ref = (repo_dir / "refs" / candidate).read_text(encoding = "utf-8").strip()
except OSError:
continue
resolved = snapshots / ref if ref else None
if resolved is not None and resolved.is_dir():
return resolved
if revision:
# pinned and not cached under that name: the default snapshot is not what will load
return None
try:
# no ref file means a commit-pinned download, where any cached revision is the one
return next((child for child in sorted(snapshots.iterdir()) if child.is_dir()), None)
except OSError:
return None
def _hosted_component_cached(source: str, subfolder: str, revision: str, variant: str) -> bool:
"""Whether a modular component the index sources elsewhere is already on disk.
``load_components`` pulls each repository the index names, and the video planner omits its
base manifest whenever the selected path exists, so a local modular directory with a missing
hosted component would otherwise verify clean and download it after the eviction.
"""
local = Path(source).expanduser()
try:
if local.is_dir():
return _component_present(local / subfolder if subfolder else local, variant)
except OSError:
return False
snapshot = _cached_snapshot_root(source, revision)
if snapshot is None:
return False
# the same component rules either way: _upstream_is_cached's no-manifest branch is satisfied
# by a single weight file, which an interrupted sharded pull leaves behind
return _component_present(snapshot / subfolder if subfolder else snapshot, variant)
def _component_present(component: Path, variant: str = "") -> bool:
"""Whether one named pipeline component holds what from_pretrained will ask it for.
``variant`` is the named weight set a modular spec can pin (``fp16`` and the like), which
from_pretrained requires by name rather than falling back to the default files.
Judged on entries that are real FILES, not merely names in the directory listing. An HF cache
snapshot holds symlinks into ``blobs/``, and a deleted blob leaves the link behind: matching
on the name alone reads such a component as complete, evicts the resident pipeline, and then
fails in from_pretrained with nothing loaded. The same listing on Windows without developer
mode holds copies rather than links and cannot express that state at all, so the two hosts
disagreed about the very same repository. A directory that merely ends in ``.safetensors``
is excluded by the same test.
"""
try:
if not component.is_dir():
return False
entries = list(component.iterdir())
files = [entry for entry in entries if entry.is_file()]
except OSError:
return False
if not entries:
return False
if not _shards_present(component):
return False
if variant and not any(f".{variant}." in entry.name for entry in files):
return False
# kept on the full listing: a shard index whose blob is gone must still route here, where
# _shards_declared reads the unreadable index and refuses, rather than fall through to the
# weight test below and pass on whichever sibling shard did survive
if any(entry.name.endswith(".index.json") for entry in entries):
# an index is proof only once it declares something; an empty weight_map declares nothing
return _shards_declared(component)
# a weight-bearing component declares config.json; schedulers, tokenizers and processors
# carry their own *_config.json instead and ship no weights at all
if (component / "config.json").is_file():
return any(entry.suffix.lower() in _WEIGHT_SUFFIXES for entry in files)
# a tokenizer ships no weights but is still useless without its vocabulary, and which file
# that is varies by class, so any one of the known spellings answers for all of them
if (component / "tokenizer_config.json").is_file():
return any((component / name).is_file() for name in _TOKENIZER_ASSETS)
# a metadata-only component is its config: a scheduler or processor directory holding
# anything else at all (a stray README) builds nothing and is fetched at load time
return any(entry.name.endswith("config.json") for entry in files)
def _shards_declared(component: Path) -> bool:
"""Whether any shard index in *component* names at least one weight file."""
import json
for index_file in component.glob("*.index.json"):
try:
with open(index_file, encoding = "utf-8-sig") as handle:
if (json.load(handle) or {}).get("weight_map"):
return True
except Exception: # noqa: BLE001 -- an unreadable index declares nothing
return False
return False
def _shards_present(component: Path) -> bool:
"""Whether a sharded component holds every file its own weight index names."""
import json
for index_file in component.glob("*.index.json"):
try:
with open(index_file, encoding = "utf-8-sig") as handle:
weight_map = (json.load(handle) or {}).get("weight_map") or {}
except Exception: # noqa: BLE001 -- an unreadable shard index is not evidence of presence
return False
if any(not (component / shard).is_file() for shard in set(weight_map.values())):
return False
return True
def missing_download_bytes(
owner: str,
pick: MediaModelPick,
hf_token: Optional[str] = None,
) -> Optional[int]:
"""Bytes this pick would still have to fetch, or 0 when nothing is missing.
Planned against the engine that will LOAD this pick, the way /images/download-plan does:
the resident engine can be native sd.cpp while the target loads through diffusers, and its
planner refuses the pick, which the catch below would read as nothing missing.
A local full IMAGE pipeline is complete by definition, since from_pretrained reads it off
disk and the planner would ask the Hub about an absolute path and fail, which reads as
unverifiable and would refuse every on-device model. Video is excluded: a local MiniMax-H3
modular pipeline still substitutes a hosted quantized conditioner, tens of GB the loader
fetches during assembly, so it has to be planned like any other pick.
Returns None when locality could not be established: the image planner raises, and the
video one returns zero bytes with ``plan_failed`` because its own caller falls back to an
inline pull. Either way zero is not evidence of a complete cache, and treating it as such
would allow exactly the download this exists to prevent, so the switch refuses instead.
"""
target = normalized_pick(pick)
local_pipeline = not target.gguf_filename and Path(target.model_path).is_dir()
if local_pipeline and not _pipeline_components_present(Path(target.model_path)):
return UNSIZED_MISSING
if owner == DIFFUSION:
# asked of every image pick, not only local pipelines: a single-file HiDream checkpoint
# plans clean and its assembly still loads the encoder repo unconditionally
external = _missing_external_encoder(target)
if external is None or external:
return external
if local_pipeline:
# the pipeline is present, so only a dependency outside it could still be fetched
return 0
try:
ordinal = plan_gpu_ordinal()
plans = [
planner.download_plan(
target.model_path,
gguf_filename = target.gguf_filename,
model_kind = target.model_kind,
gpu_ordinal = ordinal,
hf_token = hf_token,
)
or {}
for planner in planners_for(owner, target)
]
except Exception as exc: # noqa: BLE001 -- see the docstring
logger.debug("media auto-switch: download plan for %s failed: %s", pick.model_id, exc)
return None
if any(plan.get("plan_failed") for plan in plans):
return None
# cached in full and still unloadable (a flux.2 gguf on a different-size base) shows up here
if any(plan.get("incompatible_reason") for plan in plans):
return None
if hidden_ltx23_extras(owner, target):
return UNSIZED_MISSING
missing = max((max(0, int(plan.get("total_bytes") or 0)) for plan in plans), default = 0)
# both planners coerce an unknown size to zero, so entries decide and bytes only describe
if not missing and any(plan.get("entries") for plan in plans):
return UNSIZED_MISSING
return missing
__all__ = [
"detected_image_family",
"encoder_repo_complete",
"hidden_ltx23_extras",
"is_edit_only",
"missing_download_bytes",
"normalized_pick",
"plan_gpu_ordinal",
"planners_for",
]

View file

@ -0,0 +1,526 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Which downloaded media model a requested name means, and whether it is the resident one.
Two halves of one question. The index walks the model roots once per few seconds and maps every
name a downloaded image or video model answers to onto the load spec its route takes. The
matching half then decides whether the backend already holds that exact build, which is what
lets a switch be skipped rather than reloading the model that is already serving.
Only downloaded models are indexed. A name that resolves to nothing is refused by the caller
rather than answered by whichever model happens to be resident.
"""
from __future__ import annotations
import os
import threading
import time
from dataclasses import dataclass, replace
from pathlib import Path
from typing import Any, Optional
from loggers import get_logger
logger = get_logger(__name__)
IMAGE_TASK = "text-to-image"
VIDEO_TASK = "text-to-video"
# the scan walks several roots and reads gguf headers, and this runs per request
_INDEX_TTL_S = 5.0
_index_lock = threading.Lock()
_index: dict[str, tuple[float, dict[str, "MediaModelPick"]]] = {}
# the video family whose partitions are a load-time choice rather than a property of the files
_H3_FAMILY = "minimax-h3"
@dataclass(frozen = True)
class MediaModelPick:
"""A downloaded media model, in the shape its load route takes."""
model_id: str
model_path: str
gguf_filename: Optional[str] = None
model_kind: Optional[str] = None
# true when a sibling build publishes the same quant token, so identity cannot be proven
ambiguous: bool = False
# sentinel for a name two different models answer to; resolution treats it as no match
_AMBIGUOUS = MediaModelPick("", "")
# ── resolving a name to a downloaded model ──────────────────────────
def _resolve_load_dir(p: Path) -> Path:
"""The directory holding the weights, unwrapping an HF cache repo to its snapshot.
The chat resolver's helper, reused so both surfaces resolve a cached repo to the same
local directory rather than to the download-capable repo id.
"""
from core.inference.local_model_resolver import _resolve_load_dir as _chat_resolve
return Path(_chat_resolve(p))
def _register(index: dict[str, MediaModelPick], keys, pick: MediaModelPick) -> None:
"""Bind every name *pick* answers to, dropping any that two different models share.
Display labels collide readily: a cached repo advertises its final component, so
``org-a/model`` and ``org-b/model`` both offer ``model``. Taking whichever the scan
reached first would load arbitrary weights for a name the docs say is usable, and the
full ids stay available either way.
"""
for key in keys:
if not isinstance(key, str) or not key.strip():
continue
normalized = key.strip().lower()
existing = index.get(normalized)
if existing is None:
index[normalized] = pick
elif existing is not _AMBIGUOUS and (existing.model_path, existing.gguf_filename) != (
pick.model_path,
pick.gguf_filename,
):
index[normalized] = _AMBIGUOUS
def _name_keys(info) -> tuple[str, ...]:
"""Names a request may use for *info*: its repo id, scanner id and label.
An absolute path is excluded: the ./models and LM Studio scanners report one as ``id``,
and a host path is not something an API caller should have to send.
"""
from core.inference.local_model_resolver import _is_abs_path_id
return tuple(
value
for value in (
getattr(info, "model_id", None),
getattr(info, "id", None),
getattr(info, "display_name", None),
)
if isinstance(value, str) and value and not _is_abs_path_id(value)
)
def _gguf_load_path(info, on_disk: Path, load_dir: Path) -> str:
"""What ``/images/load`` takes as ``model_path`` for a GGUF under *info*.
An HF cache repo is named by its repo id, as the picker names a Hub pick. Its snapshot
entries are symlinks into ``blobs/``, and the loader's local branch resolves a symlink
before its containment check, so a snapshot directory refuses its own file. Anything else
is a real directory and loads by path.
Keyed on the layout rather than the scanner's ``source``, which is rewritten to ``custom``
for a cache tree sitting inside a user-added scan folder while the symlinks stay exactly
as fragile.
"""
repo_id = getattr(info, "model_id", None)
if load_dir != on_disk and isinstance(repo_id, str) and repo_id:
return repo_id
return str(load_dir)
def _loader_can_open(load_path: str, filename: str) -> bool:
"""Whether the load routes will resolve *filename* under *load_path*, by their own rule.
A repo id is opened from the cache by id and has nothing to check here. A directory does:
an HF cache snapshot's entries are symlinks into ``blobs/``, and both validators resolve a
symlink before their containment check, so such a directory refuses its own file. The
scanner hands one over already unwrapped for a non-active cache, where naming the repo id
instead would only send the loader to the active cache and download the model again.
A split checkpoint counts as openable only when its whole set is beside it: the loader opens
the siblings implicitly, and the planners read a local checkpoint as already present, so a
half-copied set would evict the resident model and then fail.
Advertising a name the loader then refuses costs a 400 on every request for a model the
lister shows as downloaded, so an unopenable build is left out of the index.
"""
from utils.models.model_config import colocated_split_shards
root = Path(load_path)
if not root.is_dir():
# a repo id loads from the cache, where the containment rule does not apply but the
# split set still has to be whole; an uncached child is the download guard's business
cached = _cached_repo_file(load_path, filename)
return True if cached is None else bool(colocated_split_shards(cached)[1])
from core.inference.diffusion_families import resolve_local_gguf_child
try:
child = resolve_local_gguf_child(root, filename)
except Exception: # noqa: BLE001 -- whatever the loader refuses, the index does not advertise
return False
# a split checkpoint opens its siblings implicitly, so an incomplete set fails at load time
return bool(colocated_split_shards(child)[1])
def _cached_repo_file(repo_id: str, filename: str) -> Optional[Path]:
"""The cached path of *filename* in *repo_id*, or None when it is not downloaded."""
from huggingface_hub import try_to_load_from_cache
from core.inference.diffusion import hub_cache_dir
try:
hit = try_to_load_from_cache(repo_id, filename, cache_dir = hub_cache_dir())
except Exception: # noqa: BLE001 -- an unreadable cache is not an answer about the shards
return None
return Path(hit) if isinstance(hit, str) else None
def _add_gguf_picks(
index: dict[str, MediaModelPick], info, keys: tuple[str, ...], on_disk: Path, load_dir: Path
) -> bool:
"""Index every GGUF quant under *info*, bare and as ``<id>:<QUANT>``; False if it holds none.
A bare id means the quant a plain load takes, ranked by the ``preferred_quant`` the chat
resolver and /v1/models already share, so one id cannot mean different weights per surface.
Root checkpoints are ranked alone when there are any: a plain local load resolves
non-recursively and always takes the root, so ranking a qualified ``distilled/...`` build
alongside them would let one id mean different weights here than in the picker.
"""
from core.inference.openai_auto_download import preferred_quant
from utils.models.model_config import list_local_gguf_variants
if load_dir.is_file():
if load_dir.suffix.lower() != ".gguf":
return False
if _loader_can_open(str(load_dir.parent), load_dir.name):
_register(
index,
keys,
MediaModelPick(
keys[0],
str(load_dir.parent),
load_dir.name,
"gguf",
),
)
return True
# filenames come back relative to this directory, which is what the loader joins them onto
variants, _ = list_local_gguf_variants(str(load_dir))
by_quant = {v.quant: v for v in variants if v.quant}
if not by_quant:
return False
load_path = _gguf_load_path(info, on_disk, load_dir)
openable = {
quant: variant
for quant, variant in by_quant.items()
if _loader_can_open(load_path, variant.filename)
}
if not openable:
return True
for quant, variant in openable.items():
# model_id stays the bare id so a "not found" error lists models, not one row per quant
_register(
index,
[f"{key}:{quant}" for key in keys],
MediaModelPick(keys[0], load_path, variant.filename, "gguf"),
)
unqualified = [quant for quant in openable if "/" not in quant]
best = preferred_quant(unqualified or list(openable)) or next(iter(unqualified or openable))
_register(
index,
keys,
MediaModelPick(keys[0], load_path, openable[best].filename, "gguf"),
)
return True
def _loadable_directory(load_dir: Path) -> bool:
"""Whether a non-GGUF directory is something the load routes can actually open.
Either a full diffusers pipeline, or a directory holding exactly one checkpoint, which both
routes reinterpret as a single_file load. Several checkpoints and no index is ambiguous, and
the routes reject it rather than choose, so advertising one would only cost a failed switch.
Both index layouts count: a Modular Diffusers pipeline (a dense MiniMax-H3) carries
``modular_model_index.json`` instead, and the video loader opens either.
"""
from core.inference.diffusion import resolve_local_single_file
try:
if any(
(load_dir / name).is_file() for name in ("model_index.json", "modular_model_index.json")
):
return True
except OSError:
return False
# a sole checkpoint is reinterpreted as a single_file load, which resolves the name through
# the same containment check a gguf goes through, so a cache snapshot's symlink is refused
sole = resolve_local_single_file(str(load_dir))
return sole is not None and _loader_can_open(str(load_dir), sole)
def _build_index(task: str) -> dict[str, MediaModelPick]:
"""Map every name a downloaded *task* model answers to onto its load spec."""
from routes.models import _local_model_task, collect_local_models
index: dict[str, MediaModelPick] = {}
try:
candidates = collect_local_models(Path("./models").resolve())
except Exception as exc: # noqa: BLE001 -- a failed scan must not 500 the generation
logger.debug("media auto-switch: local model scan failed: %s", exc)
return index
for info in candidates:
try:
# a cancelled or incomplete pull still lists, and loading it fails predictably
if getattr(info, "partial", False):
continue
if _local_model_task(info) != task:
continue
keys = _name_keys(info)
if not keys:
continue
# an hf cache repo keeps its weights, and its model_index.json, under snapshots/<sha>
on_disk = Path(info.path).expanduser()
load_dir = _resolve_load_dir(on_disk)
if _add_gguf_picks(index, info, keys, on_disk, load_dir):
continue
if not _loadable_directory(load_dir):
continue
_register(index, keys, MediaModelPick(keys[0], str(load_dir)))
except Exception as exc: # noqa: BLE001 -- one unreadable model must not hide the rest
logger.debug("media auto-switch: skipped %s: %s", getattr(info, "id", "?"), exc)
return index
def _partition_of(pick: MediaModelPick) -> Optional[str]:
"""The MiniMax-H3 partition this pick brings up, or None when it is not an H3 build."""
return expected_partition(pick)
def _mark_ambiguous_builds(index: dict[str, MediaModelPick]) -> dict[str, MediaModelPick]:
"""Flag every GGUF pick another build under its path cannot be told apart from.
Grouped on what the backend publishes about a resident model, which is the path and the
quant token. The H3 partition then splits a group, because status publishes ``h3_task`` and
``resident_is_pick`` compares it: two H3 denoisers sharing a quant are distinguishable, and
marking them ambiguous reloads a multi-GB checkpoint on every request.
It splits nothing else. A non-H3 sibling has no partition to be told apart by, and
``partition_matches`` reads a resident ``fl2va`` as answering for it, so it has to stay in
the group with everything else that shares its token.
"""
groups: dict[tuple[str, str], list[MediaModelPick]] = {}
for pick in index.values():
if pick is _AMBIGUOUS or pick.model_kind != "gguf":
continue
groups.setdefault((identity_key(pick.model_path), published_token(pick)), []).append(pick)
collides = set()
for key, picks in groups.items():
files = {pick.gguf_filename for pick in picks}
if len(files) < 2:
continue
partitions = [_partition_of(pick) for pick in picks]
# every member an H3 build with a partition of its own: status tells them apart
if all(partitions) and len(set(partitions)) == len(files):
continue
collides.add(key)
if not collides:
return index
return {
name: (
pick
if pick is _AMBIGUOUS
or pick.model_kind != "gguf"
or (identity_key(pick.model_path), published_token(pick)) not in collides
else replace(pick, ambiguous = True)
)
for name, pick in index.items()
}
def _cached_index(task: str) -> dict[str, MediaModelPick]:
now = time.monotonic()
with _index_lock:
hit = _index.get(task)
if hit is not None and now - hit[0] < _INDEX_TTL_S:
return hit[1]
built = _mark_ambiguous_builds(_build_index(task))
with _index_lock:
# stamped after the scan, so one slower than the ttl is not already expired
_index[task] = (time.monotonic(), built)
return built
def invalidate_index() -> None:
"""Drop the cached scan. For tests and anything that changes what is downloaded."""
with _index_lock:
_index.clear()
def resolve_local_media_model(name: str, *, task: str) -> Optional[MediaModelPick]:
"""The downloaded *task* model *name* refers to, or None."""
if not isinstance(name, str) or not name.strip():
return None
pick = _cached_index(task).get(name.strip().lower())
return None if pick is _AMBIGUOUS else pick
def available_media_model_ids(task: str) -> list[str]:
"""Sorted ids a request may name for *task*, for a "not found" error to list."""
return sorted(
{pick.model_id for pick in _cached_index(task).values() if pick is not _AMBIGUOUS}
)
# ── recognising the resident model ──────────────────────────────────
def published_token(pick: MediaModelPick) -> str:
"""The ``gguf_variant`` the backend will publish once *pick* is loaded, lowercased."""
from hub.utils.gguf import extract_quant_token
if not pick.gguf_filename:
return ""
token = extract_quant_token(pick.gguf_filename)
return (token or "").strip().lower()
def identity_key(value: str) -> str:
"""A model identity normalized for comparison: a repo id folds case, a path does not."""
text = str(value or "").strip()
return os.path.normcase(text) if os.path.isabs(text) else text.lower()
def same_identity(requested: str, resident: str) -> bool:
"""Whether two model identities name the same thing.
A repo id folds case; a filesystem path does not, since /models/Foo and /models/foo are
different models where the filesystem says so.
"""
requested, resident = requested.strip(), resident.strip()
if not requested or not resident:
return False
return identity_key(requested) == identity_key(resident)
def resident_is_gguf(status: dict[str, Any]) -> bool:
"""Whether the resident build is a GGUF, however its engine says so.
The native sd.cpp status publishes ``dtype="gguf"`` and a quant but no ``model_kind``, so a
model_kind test alone reads every native checkpoint as a plain pipeline.
"""
return (
status.get("model_kind") == "gguf"
or str(status.get("dtype") or "").strip().lower() == "gguf"
or bool(status.get("gguf_variant"))
)
def resident_is_pick(status: dict[str, Any], name: str, pick: MediaModelPick) -> bool:
"""Whether the resident build is the one *pick* names, on the identity status publishes.
A modular MiniMax-H3 build is its partition too: an auto-load of this name selects the
default keyframe denoiser, so a resident ``ref2va`` does not answer for it.
"""
if not status.get("loaded"):
return False
resident = str(status.get("repo_id") or "").strip().lower()
if not resident:
return False
aliases = {name.strip().lower(), pick.model_id.strip().lower()}
# not case-folded: /models/Foo and /models/foo are different models where the filesystem says so
same_path = os.path.normcase(str(status.get("repo_id") or "").strip()) == os.path.normcase(
pick.model_path.strip()
)
if resident not in aliases and not same_path:
return False
if not partition_matches(status, pick):
return False
if pick.model_kind != "gguf" and not resident_is_gguf(status):
return True
loaded_quant = str(status.get("gguf_variant") or "").strip().lower()
return loaded_quant == published_token(pick)
def satisfied_by(status: dict[str, Any], name: str, pick: MediaModelPick) -> bool:
"""Whether the resident model already answers this request.
Matched on the requested name AND the pick's on-disk path: a model loaded from the Images
page reports its repo id while one loaded here reports the local path it was given, and
either has to count as already serving or every request reswaps. Never on ``base_repo``,
which is a companion encoder/VAE repo and would answer a request for that full pipeline
with whichever GGUF happens to borrow it.
A GGUF also has to match on quant. Loose ``.gguf`` files in one scan folder share that
folder as their ``model_path``, so the path alone would report a sibling as already
serving and generate on the wrong weights.
The comparison uses the token the backend actually publishes. Where that token cannot tell
two indexed builds apart (``IQ4_XS-3.53bpw`` and ``-3.97bpw`` both publish ``IQ4_XS``), the
pick is marked ambiguous at index time and this answers False: reloading costs a load,
serving the sibling returns the wrong image.
"""
if not resident_is_pick(status, name, pick):
return False
# ambiguity only blocks the skip, never the "did my load land" check: the reload settles it
return not pick.ambiguous
def expected_partition(pick: MediaModelPick) -> Optional[str]:
"""The MiniMax-H3 partition this pick will come up on, or None when it is not an H3 model.
Sent with the load so the recorded provenance matches what status publishes: a GGUF takes
the partition its filename names, and a modular pipeline takes the keyframe default.
"""
try:
from core.inference.video_families import detect_video_family
from core.inference.video_minimax_h3 import H3_TASK_KEYFRAMES, h3_transformer_task
except Exception: # noqa: BLE001 -- no h3 support here means no partition to name
return None
# the basename, since a qualified variant lives at ref2va/minimax_h3_ref2va-*.gguf
name = Path(pick.gguf_filename or "").name.lower()
if name.startswith("minimax_h3_"):
return h3_transformer_task(name)
try:
# keyed on the family: a modular pipeline resolves to a directory, not a bundle repo id
for needle in (pick.model_id, pick.model_path):
fam = detect_video_family(needle) if needle else None
if fam is not None and getattr(fam, "name", "") == _H3_FAMILY:
return H3_TASK_KEYFRAMES
except Exception: # noqa: BLE001 -- a probe failure must not name a partition
return None
return None
def partition_matches(status: dict[str, Any], pick: Optional[MediaModelPick] = None) -> bool:
"""Whether the resident MiniMax-H3 partition is the one this pick would bring up.
Derived from the checkpoint, not assumed: the native backend publishes ``ref2va`` for a
``minimax_h3_ref2va`` denoiser, so hardcoding the keyframe default rejected the very
checkpoint that had just loaded. Absent a filename the switch sends no ``h3_task`` and the
load takes the family default.
"""
resident = str(status.get("h3_task") or "").strip().lower()
if not resident:
return True
try:
from core.inference.video_minimax_h3 import H3_TASK_KEYFRAMES, h3_transformer_task
except Exception: # noqa: BLE001 -- no h3 support here means nothing to compare
return True
filename = (pick.gguf_filename if pick else None) or ""
expected = h3_transformer_task(filename) if filename else H3_TASK_KEYFRAMES
return resident == str(expected or "").strip().lower()
__all__ = [
"IMAGE_TASK",
"VIDEO_TASK",
"MediaModelPick",
"available_media_model_ids",
"expected_partition",
"identity_key",
"invalidate_index",
"partition_matches",
"published_token",
"resident_is_gguf",
"resident_is_pick",
"resolve_local_media_model",
"same_identity",
"satisfied_by",
]

View file

@ -0,0 +1,189 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Reaching the media backends, and waiting until nothing a switch would interrupt is running.
Loading a media model is not a private act. The load route takes the GPU through the arbiter,
whose cross-owner handoff unloads whichever owner holds it, so a switch can cancel a video
generation or terminate a streaming chat completion that has nothing to do with the request that
asked for it. The drain here waits all of that out first, and refuses rather than interrupting.
"""
from __future__ import annotations
import asyncio
import functools
import time
from typing import Any, Optional
from core.inference.gpu_arbiter import DIFFUSION, VIDEO
from core.inference.media_switch_errors import probe
from core.inference.media_switch_locks import switcher_count, waiter_count
POLL_S = 0.2
def backend_for(owner: str) -> Any:
"""The live backend object for *owner*, resolved on each call rather than cached."""
if owner == DIFFUSION:
from core.inference.diffusion_engine_router import get_active_diffusion_engine
return get_active_diffusion_engine()
from core.inference.video import get_video_backend
return get_video_backend()
def other_owner(owner: str) -> str:
"""The media backend this one would take the GPU from."""
return VIDEO if owner == DIFFUSION else DIFFUSION
def load_takes_the_gpu() -> bool:
"""Whether this load will go through the arbiter and evict the current owner.
A CPU-only diffusion device releases ownership instead of acquiring it, so such a switch
interrupts nothing and must not wait on chat or the other media backend.
"""
try:
from core.inference.diffusion_device import resolve_diffusion_device_target
return resolve_diffusion_device_target().device != "cpu"
except Exception: # noqa: BLE001 -- assume the handoff, which is the careful direction
return True
def chat_busy(count_pending: bool = True) -> bool:
"""Whether a chat request or load is in flight, so the GPU handoff would interrupt it.
The arbiter evicts chat unconditionally for the current owner, terminating a streaming
completion that has nothing to do with this switch.
``count_pending`` is False once the lifecycle gate is held: a request blocked in the
middleware behind that gate has not started inference and cannot be interrupted, while one
admitted just before the gate was taken is already running and still can be.
"""
try:
from core.inference.llama_keepwarm import other_inference_request_count
except Exception: # noqa: BLE001 -- no chat stack means no chat work
return False
try:
# chat's counter covers media requests too, and none of those is using chat
parked = switcher_count()
counted = other_inference_request_count(
current_request_counted = True, include_pending = count_pending
)
# counted once, since a request parked on a switch lock is a waiter inside its own switch
return max(0, counted - max(0, parked - 1)) > 0
except Exception: # noqa: BLE001
return False
def backend_busy(backend: Any) -> bool:
"""One off-loop read of whether a load or generation is running. Mirrors media_keepwarm."""
if backend.loading_repo_ids():
return True
return bool((backend.generate_progress() or {}).get("active"))
def other_backend_busy(owner: str) -> bool:
"""Whether the other media backend is loading or generating, off the loop.
Guarded and lazy: a Studio that never opened the other page has no backend to ask, and
importing one just to find that out would drag torch in for nothing.
"""
import sys
other = other_owner(owner)
wanted = (
{"core.inference.video"}
if other == VIDEO
else {"core.inference.diffusion", "core.inference.sd_cpp_backend"}
)
if not wanted & set(sys.modules):
return False
try:
return backend_busy(backend_for(other))
except Exception: # noqa: BLE001 -- an unavailable backend is not busy work
return False
async def drain(
owner: str,
backend: Any,
deadline: float,
*,
count_pending: bool = True,
probe_deadline: Optional[float] = None,
check_chat: bool = True,
kind: str = "image",
openai_errors: bool = True,
) -> bool:
"""Wait out other tracked requests and any in-flight load or generation.
A request queued on this backend's switch lock is counted by the middleware but is not
doing any work, so it is discounted here: two concurrent requests for the same absent
model would otherwise each wait the other out and both return 409. Mirrors the chat
switch, which excludes its own waiters from ``_wait_for_model_switch_idle``.
The other media backend counts too, because the arbiter's cross-owner handoff unloads
whatever holds the GPU. So does chat, whether or not it is streaming. Both are skipped
entirely when this load does not take the GPU at all.
``count_pending`` is False for the check made while holding the admission gate. A request
arriving then is counted pending and immediately blocks on that gate, so counting it would
abort a switch over a newcomer that cannot be touching the backend.
``probe_deadline`` bounds the busy probes themselves, and is the switch budget rather than
this loop's deadline: the in-gate check evaluates the condition once with no time to wait,
and reusing that as the probe bound would report every backend busy. The probes need a bound
at all because ``loading_repo_ids`` takes the backend lock, which the loader holds across
pipeline assembly, so an unbounded probe outlives the response window.
``check_chat`` stays on for the in-gate check, where ``count_pending`` is what makes it safe:
a chat request blocked behind the held lifecycle gate has not started inference and must not
abort the switch, but one admitted between the outer drain's last probe and this gate being
taken is already running and would be terminated by the handoff.
"""
from core.inference.media_keepwarm import other_request_count
# device configuration is resolved once, not on every poll: that cost ~150 round-trips a switch
cross_owner = await asyncio.to_thread(load_takes_the_gpu)
while True:
# this request is itself tracked and itself a waiter, so it counts as neither
others = other_request_count(
owner, current_request_counted = True, count_pending = count_pending
)
others -= waiter_count(owner)
if cross_owner:
other = other_owner(owner)
others += max(
0,
other_request_count(other, count_pending = count_pending) - switcher_count(other),
)
probe_by = deadline if probe_deadline is None else probe_deadline
bounded_probe = functools.partial(probe, kind = kind, openai_errors = openai_errors)
if (
others <= 0
and not await bounded_probe(backend_busy, backend, probe_by)
and not (cross_owner and await bounded_probe(other_backend_busy, owner, probe_by))
and not (
cross_owner
and check_chat
and await bounded_probe(functools.partial(chat_busy, count_pending), None, probe_by)
)
):
return True
if time.monotonic() >= deadline:
return False
await asyncio.sleep(POLL_S)
__all__ = [
"POLL_S",
"backend_busy",
"backend_for",
"chat_busy",
"drain",
"load_takes_the_gpu",
"other_backend_busy",
"other_owner",
]

View file

@ -0,0 +1,171 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Refusals the media model auto-switch answers a generation request with.
Every one of these is an ordinary outcome rather than a fault. The switch never downloads,
never cuts a running generation short, and never outlives the ~100 second window Studio's
secure-mode tunnel gives an origin response, so where it cannot serve the request it says which
of those it hit and asks the caller to retry.
Each refusal is rendered in the error shape of the route that raised it: the OpenAI-compatible
body for ``/v1``, a plain string for the native media routes.
"""
from __future__ import annotations
import asyncio
import time
from typing import Any, Optional
RETRY_AFTER_S = 15
MAX_LISTED_MODELS = 8
# stands for "entries are missing but their size is unknown", so the refusal reports no figure
UNSIZED_MISSING = -1
BUSY_MSG = (
"The {kind} model is busy with another request, so it could not be switched in time. "
"Retry once the current generation finishes."
)
LOADING_MSG = (
"Loading '{model}'. It was not resident when this request arrived and is still coming up; "
"retry shortly."
)
SLOW_MSG = (
"Selecting the {kind} model took too long to answer inside this request. It is still "
"being prepared; retry shortly."
)
UNVERIFIED_MSG = (
"Could not verify that '{model}' is fully downloaded, so it was not switched in. "
"Auto-switch never downloads; load it once from the {kind} page and retry."
)
EDIT_ONLY_MSG = (
"'{model}' is an edit-only model: it requires an input image, which this endpoint cannot "
"supply. Name a text-to-image model instead."
)
UNSIZED_MSG = (
"'{model}' is missing some of its weights. Auto-switch never downloads, so load it once "
"from the {kind} page and retry."
)
INCOMPLETE_MSG = (
"'{model}' is not fully downloaded: about {gb:.1f} GB of its companion weights are missing. "
"Auto-switch never downloads, so load it once from the {kind} page and retry."
)
def refuse(
message: str,
*,
status_code: int,
openai_errors: bool,
code: str,
retry_after: int = 0,
):
"""The HTTPException to raise, in the error shape the calling route publishes."""
from fastapi import HTTPException
from utils.api_errors import openai_error_body
detail: Any = message
if openai_errors:
detail = openai_error_body(message, status = status_code, code = code, param = "model")
return HTTPException(
status_code = status_code,
detail = detail,
headers = {"Retry-After": str(retry_after)} if retry_after else None,
)
def slow_switch(kind: str, openai_errors: bool):
"""The refusal for a switch that ran out of budget before it could answer."""
return refuse(
SLOW_MSG.format(kind = kind),
status_code = 503,
openai_errors = openai_errors,
code = "model_loading",
retry_after = RETRY_AFTER_S,
)
def busy(kind: str, openai_errors: bool):
"""The refusal for a backend that stayed busy for the whole drain."""
return refuse(
BUSY_MSG.format(kind = kind),
status_code = 409,
openai_errors = openai_errors,
code = "model_busy",
retry_after = RETRY_AFTER_S,
)
def incomplete_message(model_id: str, missing: int, kind: str) -> str:
"""The refusal text, which only quotes a size when the plan could size what it is missing."""
if missing == UNSIZED_MISSING:
return UNSIZED_MSG.format(model = model_id, kind = kind)
return INCOMPLETE_MSG.format(model = model_id, gb = missing / 1e9, kind = kind)
def format_available(ids: list[str]) -> str:
if not ids:
return ""
shown = ", ".join(ids[:MAX_LISTED_MODELS])
extra = len(ids) - MAX_LISTED_MODELS
return f"{shown} and {extra} more" if extra > 0 else shown
async def bounded(coro, deadline: float, *, kind: str, openai_errors: bool):
"""Await *coro* within the switch budget, refusing rather than outliving the response window.
The worker thread behind a ``to_thread`` keeps running after this returns; what matters is
that the request stops waiting on it, since the caller's connection is the thing on a clock.
"""
remaining = deadline - time.monotonic()
if remaining <= 0:
# shield() yields a Future, which has no close(); a bare coroutine has no cancel().
if hasattr(coro, "cancel"):
coro.cancel()
else:
coro.close()
raise slow_switch(kind, openai_errors)
try:
return await asyncio.wait_for(coro, timeout = remaining)
except asyncio.TimeoutError:
raise slow_switch(kind, openai_errors)
async def probe(fn, arg: Optional[Any], deadline: float, *, kind: str, openai_errors: bool) -> bool:
"""Run a blocking busy probe off the loop, refusing rather than guessing on an overrun.
A spent budget is not a busy backend: reporting one sends the caller after a generation
that does not exist, where the slow-switch 503 says what actually happened.
"""
remaining = deadline - time.monotonic()
call = asyncio.to_thread(fn, arg) if arg is not None else asyncio.to_thread(fn)
if remaining <= 0:
call.close()
raise slow_switch(kind, openai_errors)
try:
return bool(await asyncio.wait_for(call, timeout = remaining))
except asyncio.TimeoutError:
raise slow_switch(kind, openai_errors)
__all__ = [
"BUSY_MSG",
"EDIT_ONLY_MSG",
"INCOMPLETE_MSG",
"LOADING_MSG",
"MAX_LISTED_MODELS",
"RETRY_AFTER_S",
"SLOW_MSG",
"UNSIZED_MISSING",
"UNSIZED_MSG",
"UNVERIFIED_MSG",
"bounded",
"busy",
"format_available",
"incomplete_message",
"probe",
"refuse",
"slow_switch",
]

View file

@ -0,0 +1,113 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Serialization and bookkeeping for media model switches.
One switch at a time per backend, so two requests cannot race the single pipeline slot, plus a
single cross-backend lock every GPU-taking switch queues on. Without the latter two switchers
each see the other as work they would interrupt and refuse each other; queueing makes the second
a waiter instead.
The counters exist because the request that is switching is itself tracked by the middleware.
A request parked on a switch lock holds no work, and a request performing a switch is not using
the backend it is counted against, so both are discounted when a switch asks whether anything
else is running.
The locks are per running loop, like ``_auto_switch_lock`` in ``routes.inference``: a
module-level ``asyncio.Lock`` binds to the loop that first awaited it and hangs a second one.
"""
from __future__ import annotations
import asyncio
import contextlib
import threading
import weakref
from typing import Optional
# not an owner: the key the cross-backend gpu switch lock is stored under
_GPU_SWITCH_KEY = "gpu-switch"
_switch_locks: "weakref.WeakKeyDictionary" = weakref.WeakKeyDictionary()
_switch_locks_guard = threading.Lock()
_waiters: dict[str, int] = {}
_waiters_guard = threading.Lock()
_switching: dict[str, int] = {}
_switching_guard = threading.Lock()
@contextlib.contextmanager
def note_switcher(owner: str):
"""Mark this request as performing a switch on *owner*, for its whole duration."""
with _switching_guard:
_switching[owner] = _switching.get(owner, 0) + 1
try:
yield
finally:
with _switching_guard:
remaining = _switching.get(owner, 0) - 1
if remaining > 0:
_switching[owner] = remaining
else:
_switching.pop(owner, None)
def switcher_count(owner: Optional[str] = None) -> int:
"""Requests currently switching *owner*, or across every backend when it is None."""
with _switching_guard:
if owner is None:
return sum(_switching.values())
return _switching.get(owner, 0)
@contextlib.contextmanager
def note_waiter(owner: str):
"""Mark this request as parked on *owner*'s switch lock, doing no work of its own."""
with _waiters_guard:
_waiters[owner] = _waiters.get(owner, 0) + 1
try:
yield
finally:
with _waiters_guard:
remaining = _waiters.get(owner, 0) - 1
if remaining > 0:
_waiters[owner] = remaining
else:
_waiters.pop(owner, None)
def waiter_count(owner: str) -> int:
"""Requests parked on *owner*'s switch lock."""
with _waiters_guard:
return _waiters.get(owner, 0)
def gpu_switch_lock() -> asyncio.Lock:
"""The single lock every GPU-taking media switch queues on, per running loop."""
return switch_lock(_GPU_SWITCH_KEY)
def switch_lock(owner: str) -> asyncio.Lock:
"""The switch lock for *owner* on the running loop, created on first use."""
loop = asyncio.get_running_loop()
# weakkeydictionary mutation is not thread-safe, so guard the get-or-create
with _switch_locks_guard:
per_owner = _switch_locks.get(loop)
if per_owner is None:
per_owner = _switch_locks[loop] = {}
lock = per_owner.get(owner)
if lock is None:
lock = per_owner[owner] = asyncio.Lock()
return lock
__all__ = [
"gpu_switch_lock",
"note_switcher",
"note_waiter",
"switch_lock",
"switcher_count",
"waiter_count",
]

View file

@ -977,6 +977,24 @@ def _fetch_repo_map(assets: list[tuple[str, str, str]], hf_token: Optional[str])
}
class _NeverRaised(Exception):
"""Placeholder ``except`` target for a hub layout with no LocalEntryNotFoundError."""
def _local_entry_not_found_error() -> type[BaseException]:
"""huggingface_hub's "not cached and downloads are disabled" error, or an unraisable stand-in.
Resolved lazily and defensively for the same reason the rest of this module imports
``huggingface_hub`` inside functions: an unexpected hub layout must degrade to today's error,
never break the import or swallow an unrelated exception. The stand-in matches nothing, so a
missing class simply leaves the raw hub error on load-progress."""
try:
from huggingface_hub.errors import LocalEntryNotFoundError
return LocalEntryNotFoundError
except Exception: # noqa: BLE001 -- an unexpected hub layout keeps the raw error
return _NeverRaised
def _with_mirrors(repo_ids) -> tuple[str, ...]:
"""``repo_ids`` plus the ungated mirror and the community repack of each, de-duplicated, order
preserved.
@ -1148,6 +1166,17 @@ class SdCppDiffusionBackend:
self,
repo_id: str,
*,
# Same name, position and default as DiffusionBackend.begin_load: the route calls whichever
# engine was activated through ONE call site and passes this unconditionally, so an engine
# that does not declare it TypeErrors every load on the hosts that select it (CPU-only,
# opted-in MPS, UNSLOTH_DIFFUSION_ENGINE=sd_cpp) -- including the ordinary user-initiated
# ones, which pass False.
#
# Covers the MODEL ASSETS only: the GGUF, the VAE and the text encoders this pick fetches
# from the Hub. It deliberately says nothing about the sd-cli/sd-server BINARY, which is a
# separate managed tree with its own install policy (_install_allowed / ensure_sd_*_binary);
# a background load may still install one, exactly as it does today.
local_files_only: bool = False,
gguf_filename: Optional[str] = None,
base_repo: Optional[str] = None,
family_override: Optional[str] = None,
@ -1256,6 +1285,7 @@ class SdCppDiffusionBackend:
target = self._run_load,
kwargs = dict(
repo_id = repo_id,
local_files_only = local_files_only,
gguf_filename = gguf_filename,
base = base,
fam = fam,
@ -1279,6 +1309,9 @@ class SdCppDiffusionBackend:
base: str,
fam: DiffusionFamily,
hf_token: Optional[str],
# Cache-only when set: every Hub call below is either skipped or told to resolve from disk,
# so a load nobody asked for cannot pull bytes. See begin_load for what it does not cover.
local_files_only: bool = False,
cpu_offload: bool = False,
memory_mode: Optional[str] = None,
speed_mode: Optional[str] = None,
@ -1337,7 +1370,13 @@ class SdCppDiffusionBackend:
# Swap ONCE so the size probe and the download agree: sizes come from paths-info, which
# -- unlike model_info -- 401s anonymously on a gated repo, so probing the upstream
# drops the VAE from the progress total the mirror then pulls.
inner_dim = self._flux2_inner_dim(repo_id, gguf_filename, fam, hf_token)
# The probe is a RANGE READ off the Hub when the checkpoint is not on disk, so an
# offline load asks it the way begin_load does: memo or local header or nothing. A
# None here only falls back to the filename heuristic for the encoder pick, and a
# cache-only load can fetch nothing the heuristic did not already have.
inner_dim = self._flux2_inner_dim(
repo_id, gguf_filename, fam, hf_token, allow_network = not local_files_only
)
specs = self._asset_specs(repo_id, gguf_filename, fam, inner_dim)
fetch_repo = _fetch_repo_map(specs, hf_token)
assets = [(fetch_repo[repo], fn, kind) for repo, fn, kind in specs]
@ -1370,10 +1409,24 @@ class SdCppDiffusionBackend:
# 15 GiB into the prefetch, without refusing one an ungated mirror stands in for. The
# plan alone is not enough: the images page falls back to this load when it fails.
self._preflight_companion_repos(
self._assets_by_repo(assets), fetch_repo.get(repo_id, repo_id), hf_token
self._assets_by_repo(assets),
fetch_repo.get(repo_id, repo_id),
hf_token,
local_files_only = local_files_only,
)
# Skipped outright offline: the size probe is get_paths_info, a Hub round trip, and its
# only product is the progress bar's denominator. A cache-only load resolves every
# asset from disk in milliseconds, so 0 (the value this method already reports for any
# size the Hub will not answer) costs nothing and asking would be the one network call
# left on the path.
if not local_files_only:
self._set_expected_bytes(assets, hf_token)
paths = self._fetch_assets(
assets,
hf_token,
cancel_event = cancel_event,
local_files_only = local_files_only,
)
self._set_expected_bytes(assets, hf_token)
paths = self._fetch_assets(assets, hf_token, cancel_event = cancel_event)
files = SdCppModelFiles(
diffusion_model = paths["diffusion_model"],
@ -1788,7 +1841,11 @@ class SdCppDiffusionBackend:
@staticmethod
def _preflight_companion_repos(
by_repo: dict[str, list[str]], repo_id: str, hf_token: Optional[str]
by_repo: dict[str, list[str]],
repo_id: str,
hf_token: Optional[str],
*,
local_files_only: bool = False,
) -> None:
"""Refuse a companion repo this pick cannot read, before any byte is fetched.
@ -1796,8 +1853,19 @@ class SdCppDiffusionBackend:
black-forest-labs/FLUX.1-schnell), and neither ``_plan_file_sizes`` nor the size probe
surfaces the 401: the entry is planned at 0 bytes and the fetch dies on the bare Hub token
error this replaces. Run from BOTH the plan and ``_run_load``, as the diffusers backend
does, because the UI falls back to /images/load when the plan call fails."""
does, because the UI falls back to /images/load when the plan call fails.
``local_files_only`` skips it entirely. The probe is a ``model_info`` call plus, for a
gated repo, a metadata HEAD -- pure network, whose whole purpose is to turn a 401 that
would otherwise arrive mid-download into a licence URL up front. A cache-only load never
starts that download: it either resolves the companion from disk (in which case the probe
would only have excused it anyway, via ``_already_downloaded``) or fails on the local
miss, which is the clearer error of the two. Skipping is therefore strictly what the
offline contract asks for and never hides a refusal a network load would have made."""
if local_files_only:
return
from core.inference.diffusion import _assert_base_repo_accessible
for repo, names in by_repo.items():
# Companions only: the picker only lists repos it could already read.
if repo != repo_id and names:
@ -1925,9 +1993,15 @@ class SdCppDiffusionBackend:
assets: list[tuple[str, str, str]],
hf_token: Optional[str],
cancel_event: Optional[threading.Event] = None,
local_files_only: bool = False,
) -> dict[str, str]:
"""Download every asset (cancellable via this load's own ``cancel_event``, so
a replacement load cannot un-cancel this pull), returning kind -> local path."""
a replacement load cannot un-cancel this pull), returning kind -> local path.
``local_files_only`` resolves each asset from the HF cache and never from the network; an
asset that is not there fails HERE, with the repo and filename named, rather than being
quietly pulled. This is the last and only network call left on an offline load's path, so
it is the one that has to honour the flag rather than merely accept it."""
from utils.hf_xet_fallback import hf_hub_download_with_xet_fallback
# Callers without a per-load event (tests, direct use) fall back to the current one.
@ -1948,9 +2022,28 @@ class SdCppDiffusionBackend:
# Resolve an asset cached only under huggingface_hub's import-time root through
# that root, as the preflight does. Pinned to the live root, a cache-folder change
# re-downloads every moved asset and 401s on an already-downloaded gated base.
path = hf_hub_download_with_xet_fallback(
repo, fn, hf_token, cancel_event = cancel, reuse_other_cache_root = True
)
try:
path = hf_hub_download_with_xet_fallback(
repo,
fn,
hf_token,
cancel_event = cancel,
reuse_other_cache_root = True,
local_files_only = local_files_only,
)
except _local_entry_not_found_error() as exc:
# Raised by huggingface_hub for exactly "not cached and outgoing traffic is
# disabled", so it can only fire under local_files_only. Its own text names
# neither the repo nor the file, and this string is what /images/load-progress
# toasts, so restate it with both. Re-raised untouched in the (unreachable)
# online case rather than relabelled, so nothing changes when the flag is off.
if not local_files_only:
raise
raise RuntimeError(
f"'{fn}' is not in the local cache for '{repo}', and this load may not "
f"download (it was not user-initiated). Open the model from the Images "
f"page to fetch it."
) from exc
paths[kind] = path
with self._lock:
if self._loading is not None:

View file

@ -127,6 +127,9 @@ from .video_families import (
snap_num_frames,
snap_video_size,
supported_video_family_names,
validate_video_flow_controls,
validate_video_keyframe_conditioning,
validate_video_reference_conditioning,
validate_video_request_shape,
video_family_prequant_available,
video_family_prequant_repo,
@ -138,7 +141,6 @@ from .video_minimax_h3 import (
H3_CANVAS_MAX_PIXELS,
H3_CANVAS_SHORT_EDGE,
H3_REF_SIZE_MATCH,
H3_REF_SIZE_MAX,
H3_TASK_REFERENCES,
fit_h3_keyframe,
fit_h3_reference_image,
@ -1271,6 +1273,7 @@ class VideoBackend:
self,
repo_id: str,
*,
local_files_only: bool = False,
gguf_filename: Optional[str] = None,
base_repo: Optional[str] = None,
family_override: Optional[str] = None,
@ -1353,6 +1356,7 @@ class VideoBackend:
target = self._run_load,
kwargs = dict(
repo_id = repo_id,
local_files_only = local_files_only,
gguf_filename = gguf_filename,
base_repo = base_repo,
family_override = family_override,
@ -1378,6 +1382,11 @@ class VideoBackend:
token = kwargs.get("_load_token")
# This load's own event: a later load replaces self._cancel_event rather than clearing it.
cancel_event = kwargs.pop("_cancel_event", None) or self._cancel_event
# An API-initiated load promises to download NOTHING: it may only open what is already
# cached. Read once here and threaded into every helper below -- the metadata probes as
# much as the fetches, since a model_info call reaches the Hub just as a weight pull does.
# READ, not popped: load_pipeline takes it too (it is in this thread's kwargs by contract).
local_files_only = bool(kwargs.get("local_files_only"))
try:
fam = _detect_load_family(
kwargs["repo_id"], kwargs.get("gguf_filename"), kwargs.get("family_override")
@ -1386,6 +1395,8 @@ class VideoBackend:
from .video_minimax_h3 import is_h3_native
if is_h3_native(fam, kind):
# local_files_only rides in kwargs and is now a NAMED parameter over there, so the
# native path binds it explicitly instead of swallowing it.
self._run_load_h3_native(
fam = fam,
token = token,
@ -1432,6 +1443,7 @@ class VideoBackend:
base,
kwargs.get("h3_task"),
kwargs.get("hf_token"),
local_files_only = local_files_only,
)
# Handed to the loader only when the dense shards really are gone from the pull. Then
# the choice is already committed and the loader must not re-take it against a reading
@ -1445,7 +1457,11 @@ class VideoBackend:
# Verified, not just name-matched: this scheme decides whether the base pull DROPS the
# dense encoder, and a derivative that the seed will decline must keep its own.
h3_te_scheme = self._h3_te_quant_scheme_verified(
fam, kwargs.get("text_encoder_quant"), base, kwargs.get("hf_token")
fam,
kwargs.get("text_encoder_quant"),
base,
kwargs.get("hf_token"),
local_files_only = local_files_only,
)
expected = self._estimate_download_bytes(
kwargs["repo_id"],
@ -1457,6 +1473,7 @@ class VideoBackend:
skip_transformer_weights = skip_transformer_weights,
h3_task = kwargs.get("h3_task"),
h3_te_scheme = h3_te_scheme,
local_files_only = local_files_only,
)
with self._lock:
if self._load_token == token and self._loading is not None:
@ -1482,6 +1499,9 @@ class VideoBackend:
# The plan counts a file cached under EITHER root, so the load has to
# resolve through both or it re-pulls what the planner skipped.
reuse_other_cache_root = True,
# An API-initiated load takes the cached checkpoint or fails; it never
# pulls the multi-GB file itself.
local_files_only = local_files_only,
)
)
# An LTX-2.3 checkpoint supplies the VAEs/vocoder/connectors, so the base pull shrinks to scheduler + TE + tokenizer; recompute the estimate.
@ -1501,6 +1521,7 @@ class VideoBackend:
kwargs["repo_id"],
kwargs.get("gguf_filename"),
kwargs.get("hf_token"),
local_files_only = local_files_only,
)
except Exception: # noqa: BLE001 -- surfaced by load_pipeline
probe = None
@ -1517,19 +1538,26 @@ class VideoBackend:
skip_transformer_weights = skip_transformer_weights,
h3_task = kwargs.get("h3_task"),
h3_te_scheme = h3_te_scheme,
local_files_only = local_files_only,
)
with self._lock:
if self._load_token == token and self._loading is not None:
self._loading.expected_bytes = expected
# Only a pre-cast checkpoint actually on disk earns the dense skip below.
te_skipped = self._fetch_te_prequant(
te_sources, kwargs.get("hf_token"), cancel_event = cancel_event
te_sources,
kwargs.get("hf_token"),
cancel_event = cancel_event,
local_files_only = local_files_only,
)
kwargs["_te_prequant_skipped"] = te_skipped
# Same rule for the H3 conditioner: the artifact has to be on disk before the base pull
# is allowed to leave the dense encoder behind.
h3_te_skipped = self._fetch_h3_te_quant(
h3_te_scheme, kwargs.get("hf_token"), cancel_event = cancel_event
h3_te_scheme,
kwargs.get("hf_token"),
cancel_event = cancel_event,
local_files_only = local_files_only,
)
base_local = self._predownload_base(
base,
@ -1542,6 +1570,7 @@ class VideoBackend:
# separate 66.28 GB transformer_ref/ that load_components(workflow="ref2va") opens.
h3_task = kwargs.get("h3_task"),
cancel_event = cancel_event,
local_files_only = local_files_only,
)
# The 2.3 assembly pulls per component from the hub id (its snapshot lacks the base VAEs), so it only gets the warmed cache.
kwargs["_base_local_dir"] = None if ltx23 else base_local
@ -1577,9 +1606,16 @@ class VideoBackend:
hf_token: Optional[str] = None,
memory_mode: Optional[str] = None,
gpu_ordinal: Optional[int] = None,
# NAMED, not left to the ``**_`` swallow below: an API-initiated load hands this in
# through _run_load's kwargs, and swallowed it meant the four-file bundle, the sizing
# metadata and the sd-cli install were all fetched by a load that promised no downloads.
local_files_only: bool = False,
**_: Any,
) -> None:
"""Download and commit the four-file stable-diffusion.cpp H3 runtime."""
"""Download and commit the four-file stable-diffusion.cpp H3 runtime.
``local_files_only`` restricts this to what is already on disk: no Hub sizing metadata, no
prebuilt install, and every file resolved from the cache."""
from huggingface_hub import HfApi
from .sd_cpp_args import (
@ -1634,7 +1670,11 @@ class VideoBackend:
if cancel_event.is_set():
raise RuntimeError(VIDEO_CANCELLED_MSG)
target = self._device_target(gpu_ordinal)
allow_install = _install_allowed()
# An install DOWNLOADS and extracts the sd-cli prebuilt, so an offline load may not start
# one. A build already on disk (managed or user-supplied) is still discovered and used;
# when there is none, the ensure returns None and the refusal below names it, which is the
# honest answer for a load that was told not to fetch anything.
allow_install = _install_allowed() and not local_files_only
binary = ensure_h3_sd_cpp_binary(
allow_install = allow_install,
accelerator = _install_accelerator_for(target.backend),
@ -1700,8 +1740,14 @@ class VideoBackend:
)
total = 0
try:
api = HfApi(token = hf_token or None)
# Skipped wholesale offline: model_info is a Hub call, and the number it produces is
# the size of a download that is not going to happen -- every file below either
# resolves from the cache or fails the load. total stays 0, which load_progress already
# reads as "no estimate" and reports as a bare phase.
api = HfApi(token = hf_token or None) if not local_files_only else None
for repo, wanted in requests:
if api is None:
break
if Path(repo).expanduser().exists():
continue
info = api.model_info(repo, files_metadata = True)
@ -1733,6 +1779,9 @@ class VideoBackend:
hf_token,
cancel_event = cancel_event,
reuse_other_cache_root = True,
# Offline this resolves from the cache and raises when the file is not
# there; h3_download_error below still names the missing repo/file.
local_files_only = local_files_only,
)
)
except Exception as exc: # noqa: BLE001 -- re-raised below, narrowed by name
@ -2024,7 +2073,11 @@ class VideoBackend:
return None
@staticmethod
def _h3_te_base_index_source(base: Optional[str], hf_token: Optional[str]) -> Optional[str]:
def _h3_te_base_index_source(
base: Optional[str],
hf_token: Optional[str],
local_files_only: bool = False,
) -> Optional[str]:
"""The text-encoder source ``base``'s own ``modular_model_index.json`` names, or None.
The staging twin of ``_h3_te_index_source``, which can only run once the pipeline exists.
@ -2033,7 +2086,12 @@ class VideoBackend:
62 GB inline. Reading the index first makes the skip as exact as the substitution.
A few KB, from the repo the load is about to pull anyway, pinned to the live cache root.
None on anything unanswerable, which keeps the dense shards."""
None on anything unanswerable, which keeps the dense shards.
Under ``local_files_only`` the index is read from the cache instead of fetched. A load that
gets this far offline has the base staged, so the index is there and the answer is the same
one the online read would give; an unstaged base raises inside the try and keeps the dense
shards, which is the existing unanswerable path."""
if not base:
return None
try:
@ -2051,6 +2109,7 @@ class VideoBackend:
hf_token,
cache_dir = hub_cache_dir(),
reuse_other_cache_root = True,
local_files_only = local_files_only,
)
)
with open(path, "r", encoding = "utf-8") as handle:
@ -2069,6 +2128,7 @@ class VideoBackend:
base: Optional[str],
h3_task: Optional[str],
hf_token: Optional[str],
local_files_only: bool = False,
) -> bool:
"""``_denoiser_prequant_covered`` plus the Hub check, for the decisions that COMMIT.
@ -2081,7 +2141,16 @@ class VideoBackend:
disk budget).
Same rule as ``_h3_te_quant_scheme_verified`` next door, and the same fail-closed
direction: unanswerable keeps the dense shards."""
direction: unanswerable keeps the dense shards.
``local_files_only`` swaps the Hub probe for a cache probe rather than refusing the load.
Refusing would break the one case the flag exists for -- a fully cached model coming up
offline -- and the question the probe answers is not actually about the Hub: it is
"is there a replacement denoiser to open instead of the dense shards?". Offline nothing is
downloaded either way, so the honest reading is whether the artifact is already on disk,
and that is exactly as strict as the online one in the direction that matters (an absent
checkpoint keeps the dense shards, so the loader's bf16 fallback still has something to
open)."""
if not self._denoiser_prequant_covered(fam, transformer_quant, base, h3_task):
return False
from .diffusion_prequant import restricted_prequant_load_supported
@ -2096,13 +2165,16 @@ class VideoBackend:
transformer_quant,
)
return False
try:
from huggingface_hub import HfApi
repo, _files = self._denoiser_prequant_hub_files(
fam, transformer_quant, base, HfApi(token = hf_token), h3_task
)
except Exception: # noqa: BLE001 -- an unanswerable probe keeps the dense shards
return False
if local_files_only:
repo = self._denoiser_prequant_cached_repo(fam, transformer_quant, base, h3_task)
else:
try:
from huggingface_hub import HfApi
repo, _files = self._denoiser_prequant_hub_files(
fam, transformer_quant, base, HfApi(token = hf_token), h3_task
)
except Exception: # noqa: BLE001 -- an unanswerable probe keeps the dense shards
return False
if repo is None:
logger.info(
"video.denoiser_prequant: no hosted %s checkpoint resolved for %s %s; keeping its "
@ -2176,17 +2248,23 @@ class VideoBackend:
text_encoder_quant: Optional[str],
base: Optional[str],
hf_token: Optional[str],
local_files_only: bool = False,
) -> Optional[str]:
"""``_h3_te_quant_scheme`` plus the exact index check, for the decisions that COMMIT.
The pure resolver compares repo NAMES, which is right for a plan (it must not raise and
must not touch the network) but not for dropping a derivative's dense encoder from the
pull. This one is allowed the few-KB index read, so the staged snapshot and the seed agree
on the same base."""
on the same base.
Offline the read comes from the cache (see ``_h3_te_base_index_source``) rather than being
skipped: this check is what stops a derivative being conditioned on someone else's encoder,
and an offline load is no safer to get that wrong. An unreadable index keeps the dense
encoder, exactly as it does online."""
scheme = self._h3_te_quant_scheme(fam, text_encoder_quant, base)
if scheme is None:
return None
source = self._h3_te_base_index_source(base, hf_token)
source = self._h3_te_base_index_source(base, hf_token, local_files_only = local_files_only)
if source is None or _h3_te_canonical(source) != _h3_te_canonical(
getattr(fam, "base_repo", None)
):
@ -2252,6 +2330,7 @@ class VideoBackend:
hf_token: Optional[str],
*,
cancel_event: Optional[threading.Event] = None,
local_files_only: bool = False,
) -> tuple[str, ...]:
"""Pre-fetch the hosted quantized conditioner; return ``("text_encoder",)`` when it landed.
@ -2263,7 +2342,11 @@ class VideoBackend:
A file that lands and then fails to LOAD (corrupt, or a re-upload this loader does not
implement) is slow rather than broken: the base snapshot has no dense encoder, so
``load_components`` pulls it from the modular index's own repo id inline. Degrading to the
download this skip avoided is the right failure -- the alternative is refusing the load."""
download this skip avoided is the right failure -- the alternative is refusing the load.
``local_files_only`` turns the pull into a cache lookup. An artifact already staged still
earns the skip; one that is not raises inside the same handler a failed fetch lands in, so
the load keeps the dense encoder rather than being refused."""
from .video_minimax_h3_te import h3_te_quant_filename
filename = h3_te_quant_filename(scheme)
@ -2280,6 +2363,7 @@ class VideoBackend:
cancel_event = cancel,
cache_dir = hub_cache_dir(),
reuse_other_cache_root = True,
local_files_only = local_files_only,
)
except Exception as exc: # noqa: BLE001 -- no artifact just means the dense encoder
if cancel.is_set():
@ -2378,6 +2462,51 @@ class VideoBackend:
return source.location, [(name, by_name[name])]
return None, []
@staticmethod
def _denoiser_prequant_cached_repo(
fam: Any,
transformer_quant: Optional[str],
base: Optional[str],
h3_task: Optional[str] = None,
) -> Optional[str]:
"""The hosted pre-quantized denoiser repo when its checkpoint is ALREADY cached, else None.
The offline twin of ``_denoiser_prequant_hub_files``, for a load that may not reach the
Hub. ``resolve_prequant_source`` is pure registry work and ``try_to_load_from_cache``
(behind ``_hub_file_is_cached``) is network-free, so this answers the same question -- is
there a replacement denoiser to open? -- from disk alone.
Both candidate names are tried, in the order the load tries them, and both cache roots are
searched, because that is what the fetch would have resolved through. No size to
corroborate against without the Hub, so a cached name is taken at face value: the loader
opening a damaged checkpoint falls back to dense, which is the same failure an online
size mismatch would have produced one step later."""
try:
from .diffusion_prequant import resolve_prequant_source
source = resolve_prequant_source(
fam,
(transformer_quant or "").strip().lower(),
base_repo = base,
task = h3_task,
)
except Exception as exc: # noqa: BLE001 -- a bad registry entry keeps the dense shards
logger.warning("video.denoiser_prequant_unresolved: %s", exc)
return None
# A local override is on disk by definition; only a hosted checkpoint has a cache to probe.
if source is None or getattr(source, "kind", None) != "repo":
return None
from core.inference.diffusion import DiffusionBackend
for name in (
getattr(source, "filename", None),
getattr(source, "fallback_filename", None),
):
if name and DiffusionBackend._hub_file_is_cached(source.location, name):
return source.location
# No log here: the caller reports the same "keeping its dense denoiser shards" outcome for
# a miss, and logging it twice would read as two separate decisions.
return None
@staticmethod
def _te_prequant_hub_files(
sources: dict[str, Any], api: Any
@ -2484,7 +2613,15 @@ class VideoBackend:
continue
if "/" not in name and name.endswith(".safetensors"):
continue
if kind != "pipeline" and name.startswith("transformer/"):
if (
kind != "pipeline"
and name.startswith("transformer/")
# transformer/config.json is the exception: from_single_file(config = <repo id>,
# subfolder = "transformer") reads it off the Hub, so a load that promised to
# download nothing needs it staged. Without it the locality gate cleared the pick
# and the fetch happened after eviction.
and name != "transformer/config.json"
):
continue
# is_prequant_covered_weight is component-agnostic (it matches "<component>/" against a
# weight suffix), so the encoder helper serves the denoiser verbatim.
@ -2511,6 +2648,7 @@ class VideoBackend:
skip_transformer_weights: bool = False,
h3_task: Optional[str] = None,
h3_te_scheme: Optional[str] = None,
local_files_only: bool = False,
) -> Optional[int]:
"""Total bytes this load will pull (checkpoint + companions), or None.
@ -2519,7 +2657,17 @@ class VideoBackend:
pre-quantized checkpoint replaces (``skip_transformer_weights``) and H3's dense
``text_encoder/`` under ``h3_te_scheme``. None of those hosted checkpoints' own bytes are
added: the progress bar counts cached bytes for the checkpoint and base repos only, so a
third repo in the total would leave the bar permanently short of 100%."""
third repo in the total would leave the bar permanently short of 100%.
Under ``local_files_only`` this is None without asking anything. The function is nothing
but ``model_info`` calls, and the number they produce is the size of a download that is not
going to happen: every file below either resolves from the cache or fails the load. None is
the value the estimate ALREADY takes whenever the metadata is unavailable, and
``load_progress`` reads it as "no estimate" and reports the bare phase -- so the offline
path costs nothing beyond a bar that does not fill, which is right for a load that fetches
nothing."""
if local_files_only:
return None
try:
from huggingface_hub import HfApi
@ -2764,7 +2912,16 @@ class VideoBackend:
)
except Exception as exc: # noqa: BLE001 -- an unavailable plan falls back to the inline pull
logger.warning("video.download_plan_failed: %s", exc)
return {"entries": [], "total_bytes": 0, "required_bytes": 0, "checkpoint_bytes": 0}
# Flagged, because zero here means "unknown", not "nothing to fetch": a caller that
# reads it as verified locality (media auto-switch) would allow the very download
# the inline fallback then performs.
return {
"entries": [],
"total_bytes": 0,
"required_bytes": 0,
"checkpoint_bytes": 0,
"plan_failed": True,
}
for repo, where in located.items():
# Files STRADDLING the two roots cannot be handed to from_pretrained as one snapshot:
# _predownload_base returns no directory then and the assembly is pinned to
@ -2888,11 +3045,14 @@ class VideoBackend:
entry["gguf_filename"] = filename
except Exception as exc: # noqa: BLE001 -- inline loading remains the fallback
logger.warning("video.h3_native_download_plan_failed: %s", exc)
# Flagged like the outer planner's failure: zero here means "unknown", and a caller
# that must not download (media auto-switch) has to tell the two apart.
return {
"entries": [],
"total_bytes": 0,
"required_bytes": 0,
"checkpoint_bytes": 0,
"plan_failed": True,
}
entries = []
for repo, entry in grouped.items():
@ -2932,6 +3092,7 @@ class VideoBackend:
hf_token: Optional[str],
*,
cancel_event: Optional[threading.Event] = None,
local_files_only: bool = False,
) -> tuple[str, ...]:
"""Pre-fetch the hosted pre-cast encoder checkpoints; return the components whose
dense weights the base pull can therefore skip.
@ -2940,7 +3101,11 @@ class VideoBackend:
is what makes that skip safe: only a checkpoint already on disk earns the right to
drop the dense shards, and the pull becomes cancellable and resumable like every
other load download instead of an untracked stall. A component whose fetch fails
keeps its dense weights, so the load still has an encoder to fall back to."""
keeps its dense weights, so the load still has an encoder to fall back to.
``local_files_only`` makes each fetch a cache lookup: a checkpoint already staged still
earns the skip, and one that is not lands in the same handler a failed fetch does, so the
component keeps its dense weights instead of the load being refused."""
cancel = cancel_event if cancel_event is not None else self._cancel_event
from utils.hf_xet_fallback import hf_hub_download_with_xet_fallback
@ -2957,6 +3122,7 @@ class VideoBackend:
cancel_event = cancel,
# Pre-quant encoders are planned with the same both-roots probe.
reuse_other_cache_root = True,
local_files_only = local_files_only,
)
except Exception as exc: # noqa: BLE001 -- no pre-cast file just means the dense encoder
if cancel.is_set():
@ -2982,6 +3148,7 @@ class VideoBackend:
skip_transformer_weights: bool = False,
h3_task: Optional[str] = None,
cancel_event: Optional[threading.Event] = None,
local_files_only: bool = False,
) -> Optional[str]:
"""Pull exactly the base-repo files the load needs; return the local snapshot dir.
@ -2991,8 +3158,18 @@ class VideoBackend:
also cancellable per file, and handing the local dir to from_pretrained skips
diffusers' own expected-files sweep. None -> caller keeps the hub id (local
path, non-diffusers layout, or any metadata failure: from_pretrained then
resolves the repo exactly as before)."""
resolves the repo exactly as before).
``local_files_only`` skips the whole prefetch and returns None. The scoped file list comes
from ``model_info``, which is a Hub call, and there is nothing to scope: the load below
passes ``local_files_only=True`` to ``from_pretrained``, which resolves the cached snapshot
itself and fetches nothing. Enumerating the cache here instead would only reconstruct what
diffusers is about to do offline anyway, and getting it wrong would hand it an incomplete
directory. None is the value this already returns whenever the metadata is unavailable, so
the caller path is the tested one."""
cancel = cancel_event if cancel_event is not None else self._cancel_event
if local_files_only:
return None
try:
if not base or Path(base).expanduser().exists():
return None
@ -3142,6 +3319,7 @@ class VideoBackend:
self,
repo_id: str,
*,
local_files_only: bool = False,
gguf_filename: Optional[str] = None,
base_repo: Optional[str] = None,
family_override: Optional[str] = None,
@ -3184,6 +3362,9 @@ class VideoBackend:
gguf_filename = gguf_filename,
hf_token = hf_token,
memory_mode = memory_mode,
# Carried, not defaulted: load_pipeline is also reached directly (no _run_load),
# and dropping it here would let an offline load fetch the four-file bundle.
local_files_only = local_files_only,
)
return self.status()
@ -3235,6 +3416,7 @@ class VideoBackend:
return self._load_h3_modular_pipeline(
diffusers = diffusers,
torch = torch,
local_files_only = local_files_only,
fam = fam,
target = target,
repo_id = repo_id,
@ -3295,7 +3477,9 @@ class VideoBackend:
te_scale = te_prequant_budget_scale(fam, te_quant_mode = text_encoder_quant, target = target)
transformer_mib: Optional[int] = None
if kind != "pipeline":
checkpoint_path = self._resolve_checkpoint_path(repo_id, gguf_filename, hf_token)
checkpoint_path = self._resolve_checkpoint_path(
repo_id, gguf_filename, hf_token, local_files_only = local_files_only
)
size_mib = file_size_mib(str(checkpoint_path))
if kind == "gguf":
transformer_mib = estimate_gguf_resident_mib(size_mib)
@ -3412,7 +3596,11 @@ class VideoBackend:
# ── build the pipeline.
pipeline_cls = getattr(diffusers, fam.pipeline_class)
# cache_dir pins every loader call to the live cache root, so a mid-session change cannot split one model across roots.
pipe_kwargs: dict[str, Any] = {"torch_dtype": dtype, "cache_dir": hub_cache_dir()}
pipe_kwargs: dict[str, Any] = {
"torch_dtype": dtype,
"cache_dir": hub_cache_dir(),
"local_files_only": local_files_only,
}
if getattr(fam, "vae_force_fp32", False):
# Wan VAE must decode in float32: a scalar torch_dtype truncates its fp32 weights to bf16 and a later widen only
# restores lossy values (banding / black frames). "default" MUST be set or unlisted components fall back to fp32.
@ -3453,7 +3641,17 @@ class VideoBackend:
", ".join(missing_dense),
)
_base_local_dir = (
self._predownload_base(base, hf_token, kind, ltx23 = False) or _base_local_dir
self._predownload_base(
base,
hf_token,
kind,
ltx23 = False,
# An offline load never had a scoped snapshot to top up (the prefetch returns
# None), so this is a no-op there -- but the flag is carried rather than
# defaulted so the top-up can never become the one call that reaches the Hub.
local_files_only = local_files_only,
)
or _base_local_dir
)
if kind == "pipeline":
# The pre-downloaded snapshot dir keeps from_pretrained off the hub; hub id when pre-download was skipped.
@ -3467,6 +3665,10 @@ class VideoBackend:
"subfolder": "transformer",
"token": hf_token,
"cache_dir": hub_cache_dir(),
# base is a REPO ID that diffusers resolves through load_config(), so the flag has
# to ride along or the non-LTX single-file path fetches the config after eviction.
# The LTX 2.3 branch below takes local_files_only through its own assembler.
"local_files_only": local_files_only,
}
if kind == "gguf":
sf_kwargs["quantization_config"] = diffusers.GGUFQuantizationConfig(
@ -3484,6 +3686,11 @@ class VideoBackend:
hf_token = hf_token,
# The 2.3 assembly builds every component itself and never sees pipe_kwargs, so hand the pre-cast encoder over explicitly.
text_encoder = pipe_kwargs.get("text_encoder"),
# And the no-download promise with it, for the same reason: it is in
# pipe_kwargs, and this branch is the one path that does not read them. There is
# no staged snapshot to fall back on either -- _base_local_dir is None for 2.3
# by design -- so every component below resolves the hub id.
local_files_only = local_files_only,
)
else:
transformer = transformer_cls.from_single_file(str(checkpoint_path), **sf_kwargs)
@ -3953,6 +4160,7 @@ class VideoBackend:
_load_token: Optional[int] = None,
_base_local_dir: Optional[str] = None,
_h3_auto_denoiser_planned: Optional[str] = None,
local_files_only: bool = False,
) -> dict[str, Any]:
"""Load MiniMax-H3 through its official Modular Diffusers workflow.
@ -3983,6 +4191,7 @@ class VideoBackend:
load_kwargs: dict[str, Any] = {
"components_manager": manager,
"cache_dir": hub_cache_dir(),
"local_files_only": local_files_only,
}
if hf_token:
load_kwargs["token"] = hf_token
@ -4095,6 +4304,13 @@ class VideoBackend:
# Reject a checkpoint baked under a different Linear filter, exactly as the
# image path does, so prequant and runtime-quant stay the same model.
min_features = DEFAULT_MIN_LINEAR_FEATURES,
# A load nobody asked for may not fetch this checkpoint either, exactly as the
# image twin refuses it. The switch verified only what the PLAN listed, and
# `auto` is re-decided here against live free memory rather than the card's
# capacity, so this branch can ask for a hosted checkpoint the plan never
# cleared and never staged. A cache miss returns None and the load continues
# dense, which load_components refuses under the same flag a few lines below.
local_files_only = local_files_only,
# Pin the live cache root, as every other loader call does: unset, the fetch
# lands under huggingface_hub's import-time constant and a later cache change
# re-downloads multiple GB into a root Studio no longer reads.
@ -4232,6 +4448,11 @@ class VideoBackend:
# already holds text_encoder/config.json: read it from there rather than sending
# the config resolution back to the hub.
local_base = _base_local_dir,
# And the same refusal the denoiser and load_components take: this artifact is
# ~27 GB, the flag protected only the calls either side of it, and the staging
# fetch that cleared the skip accepted a copy under EITHER cache root while this
# lookup searched only the live one.
local_files_only = local_files_only,
logger = logger,
)
if text_encoder is None:
@ -4306,6 +4527,7 @@ class VideoBackend:
workflow = workflow,
dtype = dtype,
cache_dir = hub_cache_dir(),
local_files_only = local_files_only,
**({"token": hf_token} if hf_token else {}),
)
# The video VAE loads at float32 and the decode runs under float16 autocast, so both
@ -4612,9 +4834,16 @@ class VideoBackend:
@staticmethod
def _resolve_checkpoint_path(
repo_id: str, gguf_filename: Optional[str], hf_token: Optional[str]
repo_id: str,
gguf_filename: Optional[str],
hf_token: Optional[str],
local_files_only: bool = False,
) -> Path:
"""The local checkpoint file for a gguf/single_file load (downloads if hub)."""
"""The local checkpoint file for a gguf/single_file load (downloads if hub).
``local_files_only`` resolves it from the cache instead: an API-initiated load has already
been told the checkpoint is staged, and this call runs a second time inside load_pipeline
(after _run_load's own fetch), where a cache miss must fail rather than pull the file."""
from .diffusion_families import resolve_local_gguf_child
root = Path(repo_id).expanduser()
@ -4631,6 +4860,7 @@ class VideoBackend:
hf_token,
# Matches the planner's both-roots cache probe, as the diffusion fetches already do.
reuse_other_cache_root = True,
local_files_only = local_files_only,
)
)
@ -5323,17 +5553,9 @@ class VideoBackend:
Unsupported keyframes are rejected. Omitting both dimensions matches the source aspect.
"""
if first_frame or last_frame:
if not fam.supports_keyframes:
raise ValueError(
f"{fam.name} generates from the prompt alone; it takes no first or last frame."
)
if h3_task == H3_TASK_REFERENCES:
raise ValueError(
"The loaded MiniMax-H3 checkpoint is the Ref2VA partition, which conditions "
"on references rather than keyframes. Load a minimax_h3_fl2va checkpoint to "
"generate from a first or last frame."
)
validate_video_keyframe_conditioning(
fam, h3_task, has_keyframes = bool(first_frame or last_frame)
)
first = last = None
if first_frame or last_frame:
from .diffusion import decode_b64_image
@ -5374,24 +5596,11 @@ class VideoBackend:
sd.cpp rejects non-default audio shifts because it cannot apply them.
"""
if flow_shift is not None and fam.default_flow_shift is None:
raise ValueError(f"{fam.name} does not expose a video flow_shift control.")
if audio_flow_shift is not None and fam.default_audio_flow_shift is None:
raise ValueError(f"{fam.name} does not expose an audio_flow_shift control.")
validate_video_flow_controls(fam, flow_shift, audio_flow_shift, engine = engine)
shift = flow_shift if flow_shift is not None else fam.default_flow_shift
audio_shift = (
audio_flow_shift if audio_flow_shift is not None else fam.default_audio_flow_shift
)
if (
audio_flow_shift is not None
and engine == "sd_cpp"
and audio_flow_shift != fam.default_audio_flow_shift
):
raise ValueError(
"stable-diffusion.cpp derives the audio schedule against a fixed "
f"{fam.default_audio_flow_shift:g} shift, so audio_flow_shift needs the "
"Diffusers engine."
)
return shift, audio_shift
@staticmethod
@ -5423,34 +5632,16 @@ class VideoBackend:
images = list(reference_images or [])
videos = list(reference_videos or [])
audios = list(reference_audios or [])
validate_video_reference_conditioning(
fam,
h3_task,
has_references = bool(images or videos or audios),
reference_image_size = reference_image_size,
engine = engine,
)
if not (images or videos or audios):
if h3_task == H3_TASK_REFERENCES:
# Ref2VA has no text-only denoiser, so every request needs an image or video.
raise ValueError(
"The loaded MiniMax-H3 checkpoint is the Ref2VA partition, which generates "
"from references. Add at least one reference image or video, or load a "
"minimax_h3_fl2va checkpoint for text-to-video."
)
return MiniMaxH3References()
if not fam.supports_references:
raise ValueError(f"{fam.name} takes no reference images, videos or audio.")
if h3_task != H3_TASK_REFERENCES:
raise ValueError(
"The loaded MiniMax-H3 checkpoint is the FL2VA partition, which conditions on "
"keyframes rather than references. Load a minimax_h3_ref2va checkpoint to "
"generate from references."
)
policy = (reference_image_size or H3_REF_SIZE_MATCH).strip().lower()
if policy not in (H3_REF_SIZE_MATCH, H3_REF_SIZE_MAX):
raise ValueError(
f"reference_image_size must be '{H3_REF_SIZE_MATCH}' or '{H3_REF_SIZE_MAX}'."
)
if policy == H3_REF_SIZE_MAX and engine == "sd_cpp":
raise ValueError(
"stable-diffusion.cpp scales every reference to the generation's pixel area, so "
f"'{H3_REF_SIZE_MAX}' reference sizing needs the Diffusers engine. Use "
f"'{H3_REF_SIZE_MATCH}' with this checkpoint."
)
from .diffusion import decode_b64_image

View file

@ -621,6 +621,110 @@ def validate_video_request_shape(
)
def validate_video_keyframe_conditioning(
fam: VideoFamily, h3_task: Optional[str], *, has_keyframes: bool
) -> None:
"""Raise ``ValueError`` when a checkpoint cannot take the keyframes a request supplies.
Pure in the family and the MiniMax-H3 partition, which is what lets the generate route judge
the checkpoint it is about to SWITCH TO by the same rules the backend applies to the loaded
one. Without that, an auto-switch evicts a working pipeline and spends minutes loading a
target for a request that was already known to be unservable.
"""
if not has_keyframes:
return
from .video_minimax_h3 import H3_TASK_REFERENCES
if not fam.supports_keyframes:
raise ValueError(
f"{fam.name} generates from the prompt alone; it takes no first or last frame."
)
if h3_task == H3_TASK_REFERENCES:
raise ValueError(
"The MiniMax-H3 checkpoint is the Ref2VA partition, which conditions on references "
"rather than keyframes. Load a minimax_h3_fl2va checkpoint to generate from a first "
"or last frame."
)
def validate_video_flow_controls(
fam: VideoFamily,
flow_shift: Optional[float],
audio_flow_shift: Optional[float],
*,
engine: Optional[str] = None,
) -> None:
"""Raise ``ValueError`` when a request sets a shift the checkpoint cannot honour.
The backend's flow-shift rules, kept here so the generate route can judge the checkpoint it
is about to switch TO by the same ones. ``engine`` is optional because a target's engine is
normally not chosen until the load runs; where it IS determined by the pick, as MiniMax-H3
GGUFs are, passing it refuses an unservable request before anything is evicted.
"""
if flow_shift is not None and fam.default_flow_shift is None:
raise ValueError(f"{fam.name} does not expose a video flow_shift control.")
if audio_flow_shift is not None and fam.default_audio_flow_shift is None:
raise ValueError(f"{fam.name} does not expose an audio_flow_shift control.")
if (
audio_flow_shift is not None
and engine == "sd_cpp"
and audio_flow_shift != fam.default_audio_flow_shift
):
raise ValueError(
"stable-diffusion.cpp derives the audio schedule against a fixed "
f"{fam.default_audio_flow_shift:g} shift, so audio_flow_shift needs the "
"Diffusers engine."
)
def validate_video_reference_conditioning(
fam: VideoFamily,
h3_task: Optional[str],
*,
has_references: bool,
reference_image_size: Optional[str] = None,
engine: Optional[str] = None,
) -> None:
"""Raise ``ValueError`` when a checkpoint cannot be conditioned on the request's references.
The absence of references is a rule too: the Ref2VA partition has no text-only denoiser. See
``validate_video_keyframe_conditioning`` for why these live here rather than inline.
``engine`` is optional for the same reason it is on the flow controls: a target's engine is
normally unknown before the load, but where the pick decides it, passing it refuses an
unservable sizing policy before anything is evicted.
"""
from .video_minimax_h3 import H3_REF_SIZE_MATCH, H3_REF_SIZE_MAX, H3_TASK_REFERENCES
if not has_references:
if h3_task == H3_TASK_REFERENCES:
raise ValueError(
"The MiniMax-H3 checkpoint is the Ref2VA partition, which generates from "
"references. Add at least one reference image or video, or load a "
"minimax_h3_fl2va checkpoint for text-to-video."
)
return
if not fam.supports_references:
raise ValueError(f"{fam.name} takes no reference images, videos or audio.")
if h3_task != H3_TASK_REFERENCES:
raise ValueError(
"The MiniMax-H3 checkpoint is the FL2VA partition, which conditions on keyframes "
"rather than references. Load a minimax_h3_ref2va checkpoint to generate from "
"references."
)
policy = (reference_image_size or H3_REF_SIZE_MATCH).strip().lower()
if policy not in (H3_REF_SIZE_MATCH, H3_REF_SIZE_MAX):
raise ValueError(
f"reference_image_size must be '{H3_REF_SIZE_MATCH}' or '{H3_REF_SIZE_MAX}'."
)
if policy == H3_REF_SIZE_MAX and engine == "sd_cpp":
raise ValueError(
"stable-diffusion.cpp scales every reference to the generation's pixel area, so "
f"'{H3_REF_SIZE_MAX}' reference sizing needs the Diffusers engine. Use "
f"'{H3_REF_SIZE_MATCH}' with this checkpoint."
)
# Default (steps, guidance) per checkpoint variant, matched by substring (picked id then base repo), most specific first.
_VIDEO_GENERATION_DEFAULTS: tuple[tuple[str, int, float], ...] = (
("distilled", 8, 1.0),

View file

@ -32,6 +32,15 @@ logger = get_logger(__name__)
# Companion files (text projections, VAEs incl. vocoder) beside the quants in unsloth's GGUF repo: the official Lightricks weights split out of the combined checkpoint. Keyed by variant.
LTX23_EXTRAS_REPO = "unsloth/LTX-2.3-GGUF"
def _live_cache_dir() -> str:
"""Studio's LIVE hub cache root. Read from utils rather than ``diffusion.hub_cache_dir`` to
avoid a circular import, the same way diffusion_auto_policy does."""
from utils.hf_cache_settings import active_hf_hub_cache
return active_hf_hub_cache()
_EXTRAS_TEXT_PROJ = "text_encoders/ltx-2.3-22b-{variant}_embeddings_connectors.safetensors"
_EXTRAS_VIDEO_VAE = "vae/ltx-2.3-22b-{variant}_video_vae.safetensors"
_EXTRAS_AUDIO_VAE = "vae/ltx-2.3-22b-{variant}_audio_vae.safetensors"
@ -333,7 +342,11 @@ def _split_checkpoint(state: dict[str, Any]) -> dict[str, dict[str, Any]]:
return groups
def _load_extras_file(filename: str, hf_token: Optional[str]) -> dict[str, Any]:
def _load_extras_file(
filename: str,
hf_token: Optional[str],
local_files_only: bool = False,
) -> dict[str, Any]:
from safetensors.torch import load_file
from utils.hf_xet_fallback import hf_hub_download_with_xet_fallback
@ -345,6 +358,9 @@ def _load_extras_file(filename: str, hf_token: Optional[str]) -> dict[str, Any]:
# The plan counts an extras file cached under EITHER root and stages neither, so this has to
# resolve both or it re-pulls what the planner skipped, inline and outside the manager.
reuse_other_cache_root = True,
# And a load nobody asked for takes the cached copy or fails: the switch's locality gate
# cleared these three artifacts by name, so a miss here is a promise it cannot keep.
local_files_only = local_files_only,
)
return load_file(path)
@ -446,6 +462,7 @@ def load_ltx23_transformer(
torch_dtype: Any,
is_gguf: bool,
hf_token: Optional[str],
local_files_only: bool = False,
) -> Any:
import diffusers
from diffusers import LTX2VideoTransformer3DModel
@ -459,6 +476,8 @@ def load_ltx23_transformer(
"subfolder": "transformer",
"torch_dtype": torch_dtype,
"token": hf_token,
# ``config`` is the BASE REPO, so the 2.0 transformer config is a hub read here.
"local_files_only": local_files_only,
**LTX_2_3_TRANSFORMER_CONFIG_OVERRIDES,
}
if is_gguf:
@ -467,7 +486,12 @@ def load_ltx23_transformer(
def load_ltx23_connectors(
connector_state: dict[str, Any], *, variant: str, torch_dtype: Any, hf_token: Optional[str]
connector_state: dict[str, Any],
*,
variant: str,
torch_dtype: Any,
hf_token: Optional[str],
local_files_only: bool = False,
) -> Any:
from diffusers.pipelines.ltx2.connectors import LTX2TextConnectors
@ -475,7 +499,7 @@ def load_ltx23_connectors(
if not any(k.startswith("text_embedding_projection") for k in connector_state):
connector_state = dict(connector_state)
connector_state.update(
_load_extras_file(_EXTRAS_TEXT_PROJ.format(variant = variant), hf_token)
_load_extras_file(_EXTRAS_TEXT_PROJ.format(variant = variant), hf_token, local_files_only)
)
return _build_from_config(
LTX2TextConnectors,
@ -487,11 +511,18 @@ def load_ltx23_connectors(
def load_ltx23_vae(
vae_state: dict[str, Any], *, variant: str, torch_dtype: Any, hf_token: Optional[str]
vae_state: dict[str, Any],
*,
variant: str,
torch_dtype: Any,
hf_token: Optional[str],
local_files_only: bool = False,
) -> Any:
from diffusers import AutoencoderKLLTX2Video
if not vae_state:
vae_state = _load_extras_file(_EXTRAS_VIDEO_VAE.format(variant = variant), hf_token)
vae_state = _load_extras_file(
_EXTRAS_VIDEO_VAE.format(variant = variant), hf_token, local_files_only
)
return _build_from_config(
AutoencoderKLLTX2Video,
_VIDEO_VAE_CONFIG,
@ -509,12 +540,15 @@ def load_ltx23_audio_vae_and_vocoder(
variant: str,
torch_dtype: Any,
hf_token: Optional[str],
local_files_only: bool = False,
) -> tuple[Any, Any]:
from diffusers import AutoencoderKLLTX2Audio
from diffusers.pipelines.ltx2.vocoder import LTX2VocoderWithBWE
if not audio_vae_state or not vocoder_state:
combined = _load_extras_file(_EXTRAS_AUDIO_VAE.format(variant = variant), hf_token)
combined = _load_extras_file(
_EXTRAS_AUDIO_VAE.format(variant = variant), hf_token, local_files_only
)
audio_vae_state = {
k[len("audio_vae.") :]: v for k, v in combined.items() if k.startswith("audio_vae.")
}
@ -551,6 +585,7 @@ def load_ltx23_pipeline(
is_gguf: bool,
hf_token: Optional[str] = None,
text_encoder: Optional[Any] = None,
local_files_only: bool = False,
) -> Any:
"""Full LTX-2.3 pipeline from a single-file/GGUF checkpoint. Assembled per-component
(constructor, not from_pretrained) because the base model_index pins LTX2Vocoder while 2.3
@ -558,7 +593,14 @@ def load_ltx23_pipeline(
``text_encoder`` supplies an already-built encoder (the caller's pre-cast fp8 Gemma3);
None builds it dense from the base repo. Because the assembly bypasses
``from_pretrained``, this is the only way an fp8 request reaches the 2.3 path."""
``from_pretrained``, this is the only way an fp8 request reaches the 2.3 path.
``local_files_only`` is a load nobody asked for. Because the assembly bypasses
``from_pretrained`` it also bypasses the caller's guarded ``pipe_kwargs``, and it is handed the
base REPO ID rather than a staged snapshot (the 2.3 snapshot lacks the base VAEs, so
``_base_local_dir`` is deliberately None here), so without the flag the base config, the
scheduler, the tokenizer, the dense Gemma3 encoder and the companion VAE/vocoder artifacts are
all fetched by a load that promised to fetch nothing."""
import transformers
from diffusers import LTX2Pipeline
from diffusers.loaders.single_file_utils import load_single_file_checkpoint
@ -589,30 +631,53 @@ def load_ltx23_pipeline(
torch_dtype = torch_dtype,
is_gguf = is_gguf,
hf_token = hf_token,
local_files_only = local_files_only,
)
connectors = load_ltx23_connectors(
groups["connectors"],
variant = variant,
torch_dtype = torch_dtype,
hf_token = hf_token,
local_files_only = local_files_only,
)
vae = load_ltx23_vae(
groups["vae"],
variant = variant,
torch_dtype = torch_dtype,
hf_token = hf_token,
local_files_only = local_files_only,
)
vae = load_ltx23_vae(groups["vae"], variant = variant, torch_dtype = torch_dtype, hf_token = hf_token)
audio_vae, vocoder = load_ltx23_audio_vae_and_vocoder(
groups["audio_vae"],
groups["vocoder"],
variant = variant,
torch_dtype = torch_dtype,
hf_token = hf_token,
local_files_only = local_files_only,
)
# Shared 2.0/2.3 components from the base repo via model_index, so upstream class renames break loudly here rather than drift.
index = LTX2Pipeline.load_config(base_repo, token = hf_token)
# Pinned to the LIVE hub root, not huggingface_hub's import-time constant: Studio's cache
# folder is a setting, and the locality gate that cleared this switch reads the live root. An
# unpinned lookup after a mid-session change searches the OTHER root, so under
# local_files_only it raises for a base that is fully downloaded, after eviction.
cache_dir = _live_cache_dir()
index = LTX2Pipeline.load_config(
base_repo, token = hf_token, local_files_only = local_files_only, cache_dir = cache_dir
)
def _sub(name: str, **extra: Any) -> Any:
library, class_name = index[name]
module = transformers if library == "transformers" else __import__("diffusers")
return getattr(module, class_name).from_pretrained(
base_repo, subfolder = name, token = hf_token, **extra
base_repo,
subfolder = name,
token = hf_token,
# The dense Gemma3 encoder below is the largest of these by far, and every one of them
# resolves the hub id: the flag is what keeps each a cache read.
local_files_only = local_files_only,
cache_dir = cache_dir,
**extra,
)
scheduler = _sub("scheduler")

View file

@ -292,6 +292,7 @@ def load_h3_quantized_text_encoder(
hf_token: Optional[str] = None,
cache_dir: Optional[str] = None,
local_base: Optional[str] = None,
local_files_only: bool = False,
logger: Any = None,
) -> Optional[Any]:
"""The hosted quantized Qwen3-VL conditioner for ``scheme``, on CPU, ready to seed into the
@ -307,6 +308,14 @@ def load_h3_quantized_text_encoder(
resolves through huggingface_hub's import-time constant instead and can re-download into a root
Studio no longer reads (or fail outright on an offline host that has already staged it).
``local_files_only`` is a load nobody asked for, which may not fetch anything. The artifact is
~27 GB, and the caller's staging phase (``_fetch_h3_te_quant``) has already accepted it -- so
without the flag this is where that promise is broken, after the resident pipeline was evicted.
It rides with the same other-root reuse the stager uses: the stager accepts a copy living only
under huggingface_hub's import-time root, so a lookup pinned to ``cache_dir`` alone would refuse
an artifact the load was cleared on and drop to the dense encoder the base pull already left
behind. A genuine miss still returns None through the handler below.
CPU on purpose: ``enable_auto_cpu_offload`` owns placement for every component, and a
pre-placed encoder would only be moved again."""
try:
@ -322,7 +331,14 @@ def load_h3_quantized_text_encoder(
from utils.hf_xet_fallback import hf_hub_download_with_xet_fallback
path = hf_hub_download_with_xet_fallback(
H3_TE_QUANT_REPO, filename, hf_token, cache_dir = cache_dir
H3_TE_QUANT_REPO,
filename,
hf_token,
cache_dir = cache_dir,
# Resolve the artifact through whichever root holds it, exactly as the stager that
# cleared this load did; pinned to cache_dir alone a moved cache folder re-pulls 27 GB.
reuse_other_cache_root = True,
local_files_only = local_files_only,
)
config = transformers.AutoConfig.from_pretrained(
@ -330,6 +346,9 @@ def load_h3_quantized_text_encoder(
subfolder = "text_encoder",
token = hf_token,
cache_dir = cache_dir,
# ``local_base`` is None on an offline load (the scoped base predownload stands down),
# so this reads the hub id and would go out for the config without the flag.
local_files_only = local_files_only,
)
text_config = getattr(config, "text_config", config)
released_layers = int(getattr(text_config, "num_hidden_layers", 0))

View file

@ -3660,6 +3660,12 @@ class VideoGenerateRequest(BaseModel):
negative_prompt: Optional[str] = Field(
None, description = "What to avoid (if the model supports it)"
)
model: Optional[str] = Field(
None,
description = "Video model to generate on. Only read when media auto-switch is on, "
"where a downloaded model that is not the resident one is loaded first; omit to use "
"whatever is loaded. The Video page never sends it.",
)
# Width/height/num_frames/fps default per loaded family, so they are optional here. These bounds stay a COARSE outer
# guard only -- they are family-agnostic, and a request that clears them can still be one no checkpoint can render. The
# enforced rule is the LOADED family's own (its resolution presets and k * frame_step + frame_offset lattice), which the

View file

@ -23344,6 +23344,20 @@ def _assert_native_precision_unset(
async def load_diffusion_model(
request: DiffusionLoadRequest, current_subject: str = Depends(get_current_subject)
):
return await load_diffusion_model_gated(request, current_subject, user_initiated = True)
async def load_diffusion_model_gated(
request: DiffusionLoadRequest,
current_subject: str,
*,
user_initiated: bool = False,
):
"""Everything ``POST /images/load`` does, plus who asked for it.
Media auto-switch awaits this rather than the route so the idle unload can tell an
API-loaded pipeline from one the user picked on the Images page.
"""
from core.inference.diffusion import (
get_diffusion_backend,
resolve_local_single_file,
@ -23362,6 +23376,8 @@ async def load_diffusion_model(
select_and_activate_engine,
)
from core.inference.gpu_arbiter import acquire_for, release, DIFFUSION
from core.inference.media_keepwarm import note_load_origin as note_media_load_origin
from hub.utils.gguf import extract_quant_token
from core.inference.sd_cpp_engine import ENGINE_DIFFUSERS, ENGINE_SD_CPP
from utils.native_path_leases import redact_native_paths
@ -23488,6 +23504,9 @@ async def load_diffusion_model(
# Kicks the slow load onto a background thread and returns at once (the client polls images/load-progress).
return engine.begin_load(
request.model_path,
# a load nobody asked for may not reach the hub: the switch verified locality
# from the outside, and this is what makes that promise the loader's own rule
local_files_only = not user_initiated,
gguf_filename = request.gguf_filename,
base_repo = request.base_repo,
family_override = request.family_override,
@ -23526,6 +23545,14 @@ async def load_diffusion_model(
# A CPU-only native load never touches the GPU, but switching FROM a previous GPU load leaves DIFFUSION marked as owner, so release (owner-guarded).
await asyncio.to_thread(release, DIFFUSION)
status_dict = await asyncio.to_thread(_begin_load)
# Keyed to the target: this load can still fail with the previous model resident, and
# its origin must not be read off that model.
note_media_load_origin(
DIFFUSION,
request.model_path,
extract_quant_token(request.gguf_filename) if kind == "gguf" else None,
user_action = user_initiated,
)
return DiffusionStatusResponse(**annotate_status(status_dict))
except (ValueError, FileNotFoundError) as exc:
raise HTTPException(status_code = 400, detail = redact_native_paths(str(exc)))
@ -24078,15 +24105,21 @@ async def openai_image_generations(
body: ImageGenerationRequest,
request: Request,
current_subject: str = Depends(get_current_subject),
hf_token: Optional[str] = Depends(get_hf_token),
):
"""OpenAI-compatible text-to-image (POST /v1/images/generations).
Generates ``n`` images from ``prompt`` on the loaded diffusion model and
returns them as URLs (default) or base64 PNGs per ``response_format``. Steps
and guidance have no OpenAI knob, so they default per loaded model."""
and guidance have no OpenAI knob, so they default per loaded model.
With media auto-switch on, ``model`` names the image model to serve on and is loaded
when it is not the resident one; with it off ``model`` stays informational."""
from core.inference import image_gallery
from core.inference.diffusion_engine_router import get_active_diffusion_engine
from core.inference.diffusion_families import default_generation_params
from core.inference.gpu_arbiter import DIFFUSION
from core.inference.media_auto_switch import maybe_auto_switch_media_model
if body.stream:
raise HTTPException(
@ -24102,6 +24135,15 @@ async def openai_image_generations(
status_code = 400, detail = openai_error_body(str(exc), status = 400, param = "size")
)
# Before the loaded check: the requested model may be the one this brings up.
await maybe_auto_switch_media_model(
body.model,
owner = DIFFUSION,
current_subject = current_subject,
openai_errors = True,
hf_token = hf_token,
)
# Use the active engine (diffusers OR native sd.cpp), the same accessor /images/generate uses.
backend = get_active_diffusion_engine()
status = backend.status()

View file

@ -66,6 +66,7 @@ from utils.openai_auto_switch_settings import (
BATCH_SIZE_MIN,
DEFAULT_AUTO_UNLOAD_API_ONLY,
DEFAULT_AUTO_UNLOAD_KEEP_KV,
DEFAULT_MEDIA_AUTO_SWITCH_ENABLED,
DEFAULT_MEDIA_AUTO_UNLOAD_IDLE_SECONDS,
DEFAULT_OPENAI_AUTO_DOWNLOAD_ENABLED,
DEFAULT_OPENAI_AUTO_SWITCH_ENABLED,
@ -76,6 +77,7 @@ from utils.openai_auto_switch_settings import (
get_auto_unload_api_only,
get_auto_unload_idle_seconds,
get_auto_unload_keep_kv,
get_media_auto_switch_enabled,
get_media_auto_unload_idle_seconds,
get_model_overrides,
get_openai_auto_switch_enabled,
@ -625,6 +627,8 @@ class OpenAIAutoSwitchPayload(BaseModel):
auto_unload_api_only: Optional[bool] = None
# The image/video TTL is its own setting, not a share of the chat one.
media_auto_unload_idle_seconds: Optional[int] = Field(default = None, ge = 0)
# And so is image/video auto-switch, for the same reason.
media_auto_switch_model: Optional[bool] = None
class OpenAIAutoSwitchResponse(BaseModel):
@ -644,6 +648,8 @@ class OpenAIAutoSwitchResponse(BaseModel):
# (residency, or API-loaded only) is holding the image/video unload off.
media_auto_unload_idle_seconds: int = DEFAULT_MEDIA_AUTO_UNLOAD_IDLE_SECONDS
media_idle_unload_active: bool = False
# When true, a media request may load the image or video model it names.
media_auto_switch_model: bool = DEFAULT_MEDIA_AUTO_SWITCH_ENABLED
# A quant suffix, as modelOverrideKey builds it. Matched against the loader's quant pattern,
@ -1027,6 +1033,7 @@ def get_openai_auto_switch(
auto_unload_api_only = get_auto_unload_api_only(),
media_auto_unload_idle_seconds = get_stored_media_auto_unload_idle_seconds(),
media_idle_unload_active = get_media_auto_unload_idle_seconds() > 0,
media_auto_switch_model = get_media_auto_switch_enabled(),
)
@ -1042,6 +1049,7 @@ def update_openai_auto_switch(
auto_download,
api_only,
media_idle_seconds,
media_auto_switch,
) = set_openai_auto_switch(
payload.enabled,
payload.auto_unload_idle_seconds,
@ -1049,6 +1057,7 @@ def update_openai_auto_switch(
payload.auto_download_model,
payload.auto_unload_api_only,
payload.media_auto_unload_idle_seconds,
payload.media_auto_switch_model,
)
except ValueError as exc:
raise log_and_http_error(
@ -1073,6 +1082,7 @@ def update_openai_auto_switch(
auto_unload_api_only = api_only,
media_auto_unload_idle_seconds = media_idle_seconds,
media_idle_unload_active = get_media_auto_unload_idle_seconds() > 0,
media_auto_switch_model = media_auto_switch,
)

View file

@ -28,6 +28,7 @@ from fastapi import APIRouter, Depends, HTTPException, Query, Response
from pydantic import ValidationError
from auth.authentication import get_current_subject
from hub.dependencies import get_hf_token
from loggers import get_logger
from models.inference import (
DiffusionDownloadPlanResponse,
@ -59,6 +60,20 @@ def _training_is_active() -> bool:
return _images_training_is_active()
def _derived_h3_task(gguf_filename: Optional[str], kind: str) -> Optional[str]:
"""The MiniMax-H3 partition a GGUF load resolves to from its filename, else None."""
if kind != "gguf" or not gguf_filename:
return None
try:
from core.inference.video_minimax_h3 import h3_transformer_task
from pathlib import Path as _Path
name = _Path(gguf_filename).name.lower()
return h3_transformer_task(name) if name.startswith("minimax_h3_") else None
except Exception: # noqa: BLE001 -- a probe failure must not fail the load
return None
def _guard_video_load_against_training() -> None:
"""Refuse loading a video model while a training run is active. Unlike chat,
a video pipeline's VRAM can't be cheaply estimated before the load, so the
@ -191,12 +206,28 @@ async def video_download_plan(
async def load_video_model(
request: VideoLoadRequest, current_subject: str = Depends(get_current_subject)
):
return await load_video_model_gated(request, current_subject, user_initiated = True)
async def load_video_model_gated(
request: VideoLoadRequest,
current_subject: str,
*,
user_initiated: bool = False,
):
"""Everything ``POST /video/load`` does, plus who asked for it.
Media auto-switch awaits this rather than the route so the idle unload can tell an
API-loaded pipeline from one the user picked on the Video page.
"""
from core.inference.diffusion import resolve_local_single_file
from core.inference.diffusion_device import (
resolve_diffusion_device_target,
resolve_selected_cuda_ordinal,
)
from core.inference.gpu_arbiter import VIDEO, acquire_for, release
from core.inference.media_keepwarm import note_load_origin
from hub.utils.gguf import extract_quant_token
from core.inference.video import (
assert_video_precision_available,
get_video_backend,
@ -258,6 +289,9 @@ async def load_video_model(
# Kicks the (slow) load onto a background thread and returns at once; begin_load itself validates network-free.
return backend.begin_load(
request.model_path,
# a load nobody asked for may not reach the hub: the switch verified locality
# from the outside, and this makes that promise the loader's own rule
local_files_only = not user_initiated,
gguf_filename = request.gguf_filename,
base_repo = request.base_repo,
family_override = request.family_override,
@ -289,6 +323,16 @@ async def load_video_model(
else:
await asyncio.to_thread(release, VIDEO)
status_dict = await asyncio.to_thread(_begin_load)
# Keyed to the target: this load can still fail with the previous model resident, and
# its origin must not be read off that model.
note_load_origin(
VIDEO,
request.model_path,
extract_quant_token(request.gguf_filename) if kind == "gguf" else None,
# Derived when the caller left it unset, since that is what the backend publishes.
request.h3_task or _derived_h3_task(request.gguf_filename, kind),
user_action = user_initiated,
)
return VideoStatusResponse(**status_dict)
except (ValueError, FileNotFoundError) as exc:
raise HTTPException(status_code = 400, detail = redact_native_paths(str(exc)))
@ -305,13 +349,20 @@ async def video_load_progress(current_subject: str = Depends(get_current_subject
@router.post("/video/generate", response_model = VideoGenerateResponse)
async def generate_video(
request: VideoGenerateRequest, current_subject: str = Depends(get_current_subject)
request: VideoGenerateRequest,
current_subject: str = Depends(get_current_subject),
hf_token: Optional[str] = Depends(get_hf_token),
):
"""Start a generation job and return at once (the begin_load pattern): a clip
takes minutes, and secure mode's tunnel caps the origin response window near
100 seconds, so the response must not span the generation. The worker runs the
generate + gallery-persist pipeline; the terminal outcome (completed with the
saved record / failed with a client-safe error) arrives via generate-progress."""
saved record / failed with a client-safe error) arrives via generate-progress.
With media auto-switch on, ``model`` names the video model to generate on and is loaded
when it is not the resident one."""
from core.inference.gpu_arbiter import VIDEO
from core.inference.media_auto_switch import maybe_auto_switch_media_model
from core.inference.video import get_video_backend
from core.inference.video_families import (
VIDEO_GENERATION_BUSY_MSG,
@ -319,6 +370,64 @@ async def generate_video(
VideoShapeError,
)
def _refuse_unservable_request(pick) -> None:
"""Judge the request against the family being switched TO, before it evicts anything.
begin_generate judges it against the loaded family under the lock, which is what makes
the answer race-proof, but by then a request no model could have served has already cost
the resident pipeline and a multi-minute load. The same rules, applied to the target's
family and MiniMax-H3 partition, both of which the pick already determines.
"""
from core.inference.media_model_index import expected_partition
from core.inference.video import _detect_load_family, resolve_video_model_kind
from core.inference.video_minimax_h3 import is_h3_native
from core.inference.video_families import (
validate_video_flow_controls,
validate_video_keyframe_conditioning,
validate_video_reference_conditioning,
validate_video_request_shape,
)
fam = _detect_load_family(pick.model_path, pick.gguf_filename, None)
if fam is None:
return
validate_video_request_shape(fam, request.width, request.height, request.num_frames)
h3_task = expected_partition(pick)
validate_video_keyframe_conditioning(
fam, h3_task, has_keyframes = bool(request.first_frame or request.last_frame)
)
# the engine is only knowable up front where the pick decides it, as an h3 gguf does
kind = resolve_video_model_kind(pick.gguf_filename, pick.model_kind)
engine = "sd_cpp" if is_h3_native(fam, kind) else None
validate_video_reference_conditioning(
fam,
h3_task,
has_references = bool(
request.reference_images or request.reference_videos or request.reference_audios
),
reference_image_size = request.reference_image_size,
engine = engine,
)
validate_video_flow_controls(
fam, request.flow_shift, request.audio_flow_shift, engine = engine
)
# Before the backend is resolved: the requested model may be the one this brings up.
try:
await maybe_auto_switch_media_model(
request.model,
owner = VIDEO,
current_subject = current_subject,
openai_errors = False,
hf_token = hf_token,
before_switch = _refuse_unservable_request,
)
except VideoShapeError as exc:
raise HTTPException(status_code = 422, detail = str(exc))
except ValueError as exc:
# the conditioning rules, which begin_generate reports the same way below
raise HTTPException(status_code = 400, detail = str(exc))
backend = get_video_backend()
# The request bounds on VideoGenerateRequest are a coarse outer guard; the real rule is the LOADED
# family's (its presets and frame lattice), and begin_generate applies it under the same lock that

View file

@ -4132,9 +4132,14 @@ def test_assemble_pipe_routes_krea2_per_component(monkeypatch):
hf_token = None,
transformer = None,
text_encoder = None,
# Spelled out rather than swallowed by a **kwargs: this double exists to pin the exact
# production signature, and this branch never sees the guarded pipe_kwargs, so the keyword
# that keeps the no-download promise has to be one _assemble_pipe really passes.
local_files_only = False,
):
calls["base"] = base
calls["transformer"] = transformer
calls["local_files_only"] = local_files_only
return Pipe()
monkeypatch.setattr(dmod, "load_krea2_pipeline", fake_loader)
@ -4158,7 +4163,14 @@ def test_assemble_pipe_routes_krea2_per_component(monkeypatch):
)
assert isinstance(pipe, Pipe)
# _assemble_pipe reads base only to FETCH, so the loader gets the ungated mirror.
assert calls == {"base": "unsloth/Krea-2-Turbo", "transformer": marker, "device": "cuda:0"}
assert calls == {
"base": "unsloth/Krea-2-Turbo",
"transformer": marker,
"device": "cuda:0",
# Default here (a direct call), but PASSED rather than left to the loader's own default:
# the parameter this test binds is what an API-initiated load flips.
"local_files_only": False,
}
def test_dense_quant_unusable_prequant_path_runs_dense_refit(

View file

@ -0,0 +1,360 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""The load contract the two image engines share, and the native engine's offline half.
``POST /images/load`` calls ``begin_load`` through ONE call site for whichever engine was
activated, so every keyword it passes has to be accepted by both. That is not a style rule: the
native engine is what a CPU-only host, an opted-in MPS host and ``UNSLOTH_DIFFUSION_ENGINE=sd_cpp``
select, so a keyword only the diffusers engine declares TypeErrors every single load on those
hosts -- including the ordinary user-initiated ones from the Images page, which pass the flag's
default. ``local_files_only`` shipped exactly that way.
The engine doubles here are ``create_autospec`` mocks on purpose. A hand-written fake with
``**kwargs`` accepts anything, which is why the existing route tests passed against an engine that
could not be called at all; autospec binds against the real signature and raises the TypeError the
user would have seen.
"""
from __future__ import annotations
import ast
import asyncio
import inspect
import textwrap
import threading
import types
from pathlib import Path
from unittest.mock import create_autospec
import pytest
from core.inference.diffusion import DiffusionBackend
from core.inference.sd_cpp_backend import SdCppDiffusionBackend
# ── What the route actually passes ─────────────────────────────────────────
def _route_begin_load_keywords() -> list[str]:
"""The keyword names ``_start_engine_load`` hands ``engine.begin_load``, read off the route.
Parsed rather than duplicated so this test cannot drift: the next keyword added to that call
is covered the moment it is added, which is the whole failure mode here.
"""
import routes.inference as route_module
source = textwrap.dedent(inspect.getsource(route_module.load_diffusion_model_gated))
for node in ast.walk(ast.parse(source)):
if not (isinstance(node, ast.FunctionDef) and node.name == "_start_engine_load"):
continue
for call in ast.walk(node):
if (
isinstance(call, ast.Call)
and isinstance(call.func, ast.Attribute)
and call.func.attr == "begin_load"
):
# ``**kwargs`` would arrive as a None-named keyword; the route spells every one out.
return [kw.arg for kw in call.keywords if kw.arg is not None]
raise AssertionError("_start_engine_load no longer calls engine.begin_load")
def test_the_route_still_passes_local_files_only():
# Guards the guard: if the route stopped passing it, every assertion below would still hold
# while the offline promise had quietly gone.
assert "local_files_only" in _route_begin_load_keywords()
@pytest.mark.parametrize("engine", [DiffusionBackend, SdCppDiffusionBackend])
def test_both_engines_accept_every_keyword_the_route_passes(engine):
"""``inspect.signature().bind`` is the exact check the interpreter makes at call time."""
keywords = _route_begin_load_keywords()
signature = inspect.signature(engine.begin_load)
# Bound against the UNBOUND function, so ``self`` is just the first positional and no engine
# has to be constructed. bind checks names and arity, never values.
signature.bind(
None,
"unsloth/FLUX.1-dev-GGUF",
**{name: None for name in keywords},
)
def test_the_two_begin_load_signatures_declare_local_files_only_alike():
"""Same name, same keyword-only kind, same default on both engines.
A native ``**kwargs`` catch-all would satisfy the bind test above while silently DROPPING the
flag, so the shape is asserted, not just the acceptance.
"""
params = {
engine: inspect.signature(engine.begin_load).parameters
for engine in (DiffusionBackend, SdCppDiffusionBackend)
}
for engine, parameters in params.items():
assert "local_files_only" in parameters, engine
declared = parameters["local_files_only"]
assert declared.kind is inspect.Parameter.KEYWORD_ONLY, engine
assert declared.default is False, engine
assert not any(
p.kind is inspect.Parameter.VAR_KEYWORD for p in params[SdCppDiffusionBackend].values()
), "a **kwargs catch-all would accept the flag and ignore it"
# ── The route, driven onto the native engine ───────────────────────────────
def _drive_the_images_load(monkeypatch, *, user_initiated: bool):
"""Run ``POST /images/load``'s body with the NATIVE engine selected; return the mock engine.
Autospec'd off the real class, so the call the route makes is bound against the real
``begin_load`` signature: this is what turns the shipped TypeError into a test failure.
"""
import core.inference.diffusion_device as device_module
import core.inference.diffusion_engine_router as router_module
from core.inference.sd_cpp_engine import ENGINE_SD_CPP
from models.inference import DiffusionLoadRequest
from routes.inference import load_diffusion_model_gated
engine = create_autospec(SdCppDiffusionBackend, instance = True)
engine.begin_load.return_value = {"loaded": False, "repo_id": None}
engine.preflight_base_access.return_value = None
monkeypatch.setattr(router_module, "predict_engine", lambda *a, **k: ENGINE_SD_CPP)
monkeypatch.setattr(router_module, "active_engine_name", lambda: ENGINE_SD_CPP)
monkeypatch.setattr(router_module, "engine_for", lambda name: engine)
monkeypatch.setattr(router_module, "select_and_activate_engine", lambda *a, **k: engine)
monkeypatch.setattr(router_module, "begin_load_on", lambda _engine, start: start())
monkeypatch.setattr(router_module, "annotate_status", lambda status: status)
# A CPU-only host is where the native engine is selected in the first place.
monkeypatch.setattr(
device_module,
"resolve_diffusion_device_target",
lambda: types.SimpleNamespace(device = "cpu"),
)
monkeypatch.setattr("routes.inference._guard_diffusion_load_against_training", lambda: None)
async def _no_ordinal(_gpu_ids):
return None
monkeypatch.setattr("routes.inference._selected_gpu_ordinal", _no_ordinal)
asyncio.run(
load_diffusion_model_gated(
DiffusionLoadRequest(
model_path = "unsloth/FLUX.1-dev-GGUF",
gguf_filename = "flux1-dev-Q4_K_M.gguf",
),
"test-user",
user_initiated = user_initiated,
)
)
return engine
@pytest.mark.parametrize("user_initiated", [True, False])
def test_the_images_page_can_load_on_the_native_engine(monkeypatch, user_initiated):
# The regression: this raised TypeError for BOTH values, so the Images page could not load a
# model at all on any host that selects sd.cpp. The parametrisation keeps the user-initiated
# case explicit, because that is the one nobody expects an offline flag to break.
engine = _drive_the_images_load(monkeypatch, user_initiated = user_initiated)
engine.begin_load.assert_called_once()
assert engine.begin_load.call_args.kwargs["local_files_only"] is (not user_initiated)
# ── The native loader honours it ───────────────────────────────────────────
def _no_hub(monkeypatch):
"""Make every huggingface_hub API call this load could reach an outright failure."""
import huggingface_hub
def _forbidden(*_a, **_k):
raise AssertionError("a cache-only load reached the Hub")
monkeypatch.setattr(huggingface_hub.HfApi, "model_info", _forbidden)
monkeypatch.setattr(huggingface_hub.HfApi, "get_paths_info", _forbidden)
monkeypatch.setattr(huggingface_hub, "get_hf_file_metadata", _forbidden)
def test_a_cache_only_native_load_makes_no_hub_call(monkeypatch):
"""The size probe and the companion preflight are both pure network; neither may run.
Their failure mode is quiet -- ``_set_expected_bytes`` swallows everything and the preflight
fails open -- so an unguarded call would not fail the load, it would just download.
"""
from core.inference.diffusion_families import detect_family
from core.inference.sd_cpp_backend import SdCppDiffusionBackend as Native
_no_hub(monkeypatch)
backend = Native(engine = None)
monkeypatch.setattr(
Native,
"_resolve_backend",
lambda self: ("oneshot", None, types.SimpleNamespace(version = lambda: "master")),
)
fetched: list = []
def _fetch(
self,
assets,
token,
cancel_event = None,
local_files_only = False,
):
fetched.append(local_files_only)
raise RuntimeError("stop here; the Hub calls under test all precede the fetch")
monkeypatch.setattr(Native, "_fetch_assets", _fetch)
repo = "unsloth/FLUX.1-dev-GGUF"
Native._run_load(
backend,
repo_id = repo,
gguf_filename = "flux1-dev-Q4_K_M.gguf",
base = "black-forest-labs/FLUX.1-dev",
fam = detect_family(repo),
hf_token = None,
local_files_only = True,
_load_token = 1,
)
# Reached the fetch (so the probe and preflight were skipped, not merely tolerated) and the
# flag arrived there, which is the only call that can still pull bytes.
assert fetched == [True]
def test_the_native_fetch_resolves_from_cache_only(monkeypatch, tmp_path):
"""``local_files_only`` reaches huggingface_hub, where it is the only thing that stops a pull."""
import utils.hf_xet_fallback as xet
from core.inference.sd_cpp_backend import SdCppDiffusionBackend as Native
seen: list = []
cached = tmp_path / "flux1-dev-Q4_K_M.gguf"
cached.write_bytes(b"")
def _download(repo_id, filename, token, **kwargs):
seen.append((repo_id, filename, kwargs.get("local_files_only")))
return str(cached)
monkeypatch.setattr(xet, "hf_hub_download_with_xet_fallback", _download)
Native(engine = None)._fetch_assets(
[("unsloth/FLUX.1-dev-GGUF", "flux1-dev-Q4_K_M.gguf", "diffusion_model")],
None,
local_files_only = True,
)
assert seen == [("unsloth/FLUX.1-dev-GGUF", "flux1-dev-Q4_K_M.gguf", True)]
def test_an_uncached_asset_fails_with_a_local_error_naming_it(monkeypatch):
"""The miss must READ as a miss. huggingface_hub's own text names neither repo nor file, and
this string is what /images/load-progress puts in front of the user."""
from huggingface_hub.errors import LocalEntryNotFoundError
import utils.hf_xet_fallback as xet
from core.inference.sd_cpp_backend import SdCppDiffusionBackend as Native
def _download(*_a, **_k):
raise LocalEntryNotFoundError("Cannot find the requested files in the disk cache")
monkeypatch.setattr(xet, "hf_hub_download_with_xet_fallback", _download)
with pytest.raises(RuntimeError) as caught:
Native(engine = None)._fetch_assets(
[("black-forest-labs/FLUX.1-dev", "ae.safetensors", "vae")],
None,
local_files_only = True,
)
message = str(caught.value)
assert "ae.safetensors" in message
# The FETCH repo, which is where the bytes were looked for: the gated vendor base is swapped
# to its ungated mirror before the lookup, so naming the upstream id would misdirect.
assert "unsloth/FLUX.1-dev" in message
def test_the_default_still_takes_the_xet_fallback_ladder(monkeypatch, tmp_path):
"""Nothing changes with the flag off: the shared Xet -> HTTP path is still the one used, and
``local_files_only`` is not forwarded to a shared layer that may predate it."""
import utils.hf_xet_fallback as xet
seen: list = []
def _shared(repo_id, filename, token, **kwargs):
seen.append(kwargs)
return str(tmp_path / filename)
monkeypatch.setattr(xet, "_shared_hf_hub_download_with_xet_fallback", _shared)
xet.hf_hub_download_with_xet_fallback(
"unsloth/FLUX.1-dev-GGUF", "flux1-dev-Q4_K_M.gguf", None, cache_dir = str(tmp_path)
)
assert len(seen) == 1
assert "local_files_only" not in seen[0]
def test_the_offline_download_never_reaches_the_shared_ladder(monkeypatch, tmp_path):
"""And with the flag on it goes straight to huggingface_hub.
Deliberately NOT forwarded to unsloth_zoo: ``start_watchdog`` already showed that an older
installed zoo silently drops kwargs it does not declare, and a dropped ``local_files_only``
downloads -- the one outcome the flag exists to prevent.
"""
import huggingface_hub
import utils.hf_xet_fallback as xet
def _forbidden(*_a, **_k):
raise AssertionError("the shared Xet ladder must not run for a cache-only download")
monkeypatch.setattr(xet, "_shared_hf_hub_download_with_xet_fallback", _forbidden)
seen: list = []
def _hub(**kwargs):
seen.append(kwargs)
return str(tmp_path / "flux1-dev-Q4_K_M.gguf")
monkeypatch.setattr(huggingface_hub, "hf_hub_download", _hub)
xet.hf_hub_download_with_xet_fallback(
"unsloth/FLUX.1-dev-GGUF",
"flux1-dev-Q4_K_M.gguf",
None,
cache_dir = str(tmp_path),
local_files_only = True,
)
assert seen and seen[0]["local_files_only"] is True
def test_a_cancelled_offline_download_still_stops(monkeypatch, tmp_path):
"""The cancellation contract is the ladder's, not huggingface_hub's, so the bypass keeps it."""
import utils.hf_xet_fallback as xet
cancel = threading.Event()
cancel.set()
with pytest.raises(RuntimeError):
xet.hf_hub_download_with_xet_fallback(
"unsloth/FLUX.1-dev-GGUF",
"flux1-dev-Q4_K_M.gguf",
None,
cache_dir = str(tmp_path),
local_files_only = True,
cancel_event = cancel,
)
def test_the_binary_install_is_not_covered_by_the_flag():
"""Stated as a test so the boundary is not re-litigated by accident.
``local_files_only`` is about MODEL ASSETS. The sd-cli / sd-server binary lives in a separate
managed tree with its own install policy, and ``_run_load`` resolves it before any asset is
fetched; a background load may still install one, exactly as before. If that ever needs to
change it is a deliberate decision, not a side effect of this flag.
"""
source = inspect.getsource(SdCppDiffusionBackend._run_load)
resolve = source.index("self._resolve_backend()")
fetch = source.index("self._fetch_assets(")
assert resolve < fetch, "the binary is resolved before the assets; the comment above assumes it"
assert Path(inspect.getsourcefile(SdCppDiffusionBackend)).name == "sd_cpp_backend.py"

View file

@ -671,6 +671,7 @@ def _native_backend_ready(monkeypatch):
assets,
token,
cancel_event = None,
local_files_only = False,
):
fetched.append(assets)
raise AssertionError("the gated companion must be caught before any byte is fetched")
@ -974,6 +975,9 @@ def test_a_base_excused_by_the_other_root_is_loaded_from_that_snapshot(monkeypat
hf_token,
cancel_event = None,
fetch_base = None,
# Tracks the real signature: the staging phase now threads the no-download flag into the
# prefetch, and a double that refuses it turns the load into a TypeError.
local_files_only = False,
):
staged.append(base_files)
return None

View file

@ -101,11 +101,15 @@ def test_load_krea2_pipeline_threads_init_config(monkeypatch, tmp_path):
monkeypatch.setitem(sys.modules, "diffusers", fake_diffusers)
monkeypatch.setattr(
"core.inference.diffusion_krea2.load_krea2_tokenizer",
lambda repo_id, hf_token = None: SimpleNamespace(tag = "tokenizer"),
# Hand-written fakes with EXACT signatures, so they have to follow the production one:
# load_krea2_pipeline now passes local_files_only down to every component load.
lambda repo_id, hf_token = None, local_files_only = False: SimpleNamespace(tag = "tokenizer"),
)
monkeypatch.setattr(
"core.inference.diffusion_krea2.load_krea2_text_encoder",
lambda repo_id, dtype, hf_token = None: SimpleNamespace(tag = "text_encoder"),
lambda repo_id, dtype, hf_token = None, local_files_only = False: SimpleNamespace(
tag = "text_encoder"
),
)
pipe = load_krea2_pipeline(str(tmp_path), "bf16")

View file

@ -0,0 +1,216 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""An API-initiated IMAGE load downloads NOTHING.
The video twin of this suite lives in ``test_video_offline_load.py``; this is the same promise on
the diffusers image path, where ``local_files_only`` reached ``load_pipeline`` and nothing in the
``_run_load`` staging phase that runs before it -- so the byte estimate and the base preflight still
asked the Hub, and ``_prefetch_files`` (the one call on that path that moves multi-GB weights)
fetched without the flag. Every network-capable helper the staging phase reaches is replaced with a
sentinel that RAISES when it is asked to fetch, so a load that regains the network is a failing test
rather than a multi-GB surprise on a user's connection. The mirror test proves the user-initiated
(UI) path still calls exactly those helpers, which is the pre-PR behaviour nothing here changes.
"""
from __future__ import annotations
import os
import types
import pytest
import utils.hf_xet_fallback as xet
from core.inference import diffusion as diffusion_mod
from core.inference.diffusion import DiffusionBackend
from core.inference.diffusion_families import detect_family_for_pick
# A plain FLUX.1 GGUF pick: it walks the shared staging path every image pick walks, and its family
# name keeps the FLUX.2 pairing preflight out of the way (that guard is a header read, not a fetch).
FLUX_GGUF = "unsloth/FLUX.1-dev-GGUF"
FLUX_BASE = "black-forest-labs/FLUX.1-dev"
FLUX_FILE = "flux1-dev-Q4_K_M.gguf"
class _Calls:
"""Every Hub call the load made, and how it made it."""
def __init__(self):
self.model_info: list[str] = []
self.downloads: list[tuple[str, str, bool]] = []
def _install_sentinels(monkeypatch, calls, tmp_path, *, offline):
"""Replace every network helper the staging phase can reach.
``offline`` is the assertion: a metadata probe is refused outright (there is no offline form of
``model_info``), and a download is refused unless it carries ``local_files_only=True``, which is
what makes it a cache lookup rather than a fetch. Online they only record, so the same fake
serves both directions and the two tests differ by one flag.
"""
import huggingface_hub
def _model_info(self, repo_id, **_kwargs):
calls.model_info.append(repo_id)
if offline:
raise AssertionError(f"model_info({repo_id!r}) reached the Hub on an offline load")
return types.SimpleNamespace(siblings = [], sha = "deadbeef", gated = False, cardData = {})
def _download(
repo_id,
filename,
token = None,
**kwargs,
):
local_files_only = bool(kwargs.get("local_files_only"))
calls.downloads.append((repo_id, filename, local_files_only))
if offline and not local_files_only:
raise AssertionError(
f"{repo_id}/{filename} was fetched without local_files_only on an offline load"
)
path = tmp_path / filename
path.parent.mkdir(parents = True, exist_ok = True)
path.write_bytes(b"")
return str(path)
monkeypatch.setattr(huggingface_hub.HfApi, "model_info", _model_info, raising = False)
monkeypatch.setattr(xet, "hf_hub_download_with_xet_fallback", _download)
# The wrapper's own offline branch calls this directly; a sentinel here catches a bypass.
monkeypatch.setattr(
huggingface_hub,
"hf_hub_download",
lambda **kwargs: _download(
kwargs.get("repo_id"), kwargs.get("filename"), kwargs.get("token"), **kwargs
),
raising = False,
)
# Deterministic fetch target: the mirror swap is a pure local-cache test, and which side it
# picks depends on the developer's own HF cache. Pinning it keeps both directions readable.
monkeypatch.setenv("UNSLOTH_DIFFUSION_NO_MIRROR", "1")
def _backend(monkeypatch, calls_seen):
"""A backend whose family detection is pinned and whose pipeline build is a capture."""
backend = DiffusionBackend()
backend._load_token = 1
backend._loading = diffusion_mod._LoadingState(repo_id = FLUX_GGUF, base_repo = FLUX_BASE)
fam = detect_family_for_pick(FLUX_GGUF, FLUX_FILE, None)
assert fam is not None
monkeypatch.setattr(diffusion_mod, "detect_family_for_pick", lambda *_a, **_k: fam)
monkeypatch.setattr(backend, "load_pipeline", lambda **kwargs: calls_seen.update(kwargs))
return backend
def test_an_api_initiated_image_load_opens_the_cache_and_downloads_nothing(monkeypatch, tmp_path):
"""The whole promise, end to end: every helper the image staging phase reaches either stays off
the Hub or asks it for a cached file only."""
calls = _Calls()
_install_sentinels(monkeypatch, calls, tmp_path, offline = True)
seen: dict = {}
backend = _backend(monkeypatch, seen)
backend._run_load(
repo_id = FLUX_GGUF,
gguf_filename = FLUX_FILE,
# Carried by the request the way a saved image config carries it, so the card-tag lookup in
# _resolve_base_repo is out of the picture: that read is metadata that FAILS OPEN, and
# dropping it offline would resolve a DIFFERENT base than the load that cached the weights.
base_repo = FLUX_BASE,
local_files_only = True,
_load_token = 1,
)
# _run_load swallows failures onto load_progress rather than raising, so the state IS the
# result: cleared means the load ran through, an error string means a sentinel fired.
assert backend._loading is None, getattr(backend._loading, "error", None)
assert seen.get("local_files_only") is True
# Not one metadata probe: the byte estimate, the pre-cast plan and the base preflight all stand
# down offline.
assert calls.model_info == []
# The checkpoint is still resolved -- as a cache lookup. THIS is the multi-GB call.
assert calls.downloads == [(FLUX_GGUF, FLUX_FILE, True)]
# And nothing was staged for from_pretrained, which resolves the cached snapshot itself.
assert seen.get("_base_local_dir") is None
def test_a_user_initiated_image_load_still_calls_every_one_of_them(monkeypatch, tmp_path):
"""The pre-PR path, unchanged: the UI load asks the Hub for sizes and PULLS the checkpoint."""
calls = _Calls()
_install_sentinels(monkeypatch, calls, tmp_path, offline = False)
seen: dict = {}
backend = _backend(monkeypatch, seen)
backend._run_load(
repo_id = FLUX_GGUF,
gguf_filename = FLUX_FILE,
base_repo = FLUX_BASE,
_load_token = 1,
)
assert backend._loading is None, getattr(backend._loading, "error", None)
assert seen.get("local_files_only") in (False, None)
# The byte estimate probes the checkpoint repo and the base; the preflight probes the base too.
assert FLUX_GGUF in calls.model_info and FLUX_BASE in calls.model_info
# And the checkpoint is FETCHED, not looked up.
assert calls.downloads == [(FLUX_GGUF, FLUX_FILE, False)]
def test_the_estimate_and_the_pre_cast_plan_stand_down_offline(monkeypatch):
"""Both are pure Hub metadata, and both already have a "could not tell" answer their callers
handle, so offline they take it rather than inventing a probe."""
class _Boom:
def __init__(self, *_a, **_k):
pass
def model_info(self, *_a, **_k):
raise AssertionError("the Hub was asked about an offline load")
import huggingface_hub
monkeypatch.setattr(huggingface_hub, "HfApi", _Boom)
backend = DiffusionBackend()
assert backend._estimate_download_bytes(
FLUX_GGUF, FLUX_FILE, FLUX_BASE, None, local_files_only = True
) == (0, [])
assert backend._te_prequant_plan_files(None, "fp8", None, None, local_files_only = True) == {}
def test_the_base_preflight_reads_the_cache_and_never_the_hub_offline(monkeypatch):
"""The preflight exists to name the repo a DOWNLOAD is about to 401 on. Offline there is no
such download, so the Hub half stands down -- but the other-root escape it computes is a pure
cache read and still runs, since that is what lets a base staged under huggingface_hub's
import-time root load off disk."""
import huggingface_hub
def _boom(*_a, **_k):
raise AssertionError("the Hub was asked about an offline load")
monkeypatch.setattr(huggingface_hub.HfApi, "model_info", _boom, raising = False)
monkeypatch.setattr(huggingface_hub, "get_hf_file_metadata", _boom, raising = False)
# Cached only under the import-time root: the live root misses, the fallback hits.
# Built with os.path.join rather than a "/" literal: the function strips the file's own
# relative path with os.path, so a POSIX spelling here would compare against a
# backslash-separated answer on Windows and fail for the separator alone.
snapshot = os.path.join(os.sep + "snap", *FLUX_BASE.split("/"))
monkeypatch.setattr(
huggingface_hub,
"try_to_load_from_cache",
lambda repo, name, cache_dir = None: (
None if cache_dir is not None else os.path.join(snapshot, *name.split("/"))
),
raising = False,
)
assert (
diffusion_mod._assert_base_repo_accessible(FLUX_BASE, None, local_files_only = True)
== snapshot
)
def test_the_prefetch_signature_declares_the_flag():
"""A default-True or missing parameter here is the bug itself: this is the call that moves the
weights, so the flag has to be a named, default-False parameter of it."""
import inspect
param = inspect.signature(DiffusionBackend._prefetch_files).parameters["local_files_only"]
assert param.default is False

View file

@ -1002,6 +1002,7 @@ def test_load_repo_source_allowed_without_optin(monkeypatch, tmp_path):
filename,
token = None,
cache_dir = None,
local_files_only = False,
):
roots.append(cache_dir)
return str(downloaded)
@ -1048,6 +1049,7 @@ def test_load_repo_source_falls_back_to_legacy_filename(monkeypatch, tmp_path):
filename,
token = None,
cache_dir = None,
local_files_only = False,
):
requested.append(filename)
if filename != "transformer_fp8.pt":
@ -1304,6 +1306,7 @@ def test_a_live_root_hit_still_goes_through_the_hub_so_it_revalidates(monkeypatc
filename,
token = None,
cache_dir = None,
local_files_only = False,
):
asked.append((filename, cache_dir))
return str(ckpt) # the same blob, revalidated
@ -1352,6 +1355,7 @@ def test_a_hit_only_in_the_other_root_is_revalidated_through_that_root(monkeypat
filename,
token = None,
cache_dir = None,
local_files_only = False,
):
asked.append((filename, cache_dir))
return str(ckpt) # unchanged upstream: the same blob, revalidated
@ -1436,6 +1440,9 @@ def test_an_uncached_checkpoint_downloads_into_the_live_root(monkeypatch):
"filename": "Z-Image-Turbo-FP8.pt",
"token": "tok",
"cache_dir": "/live-hub",
# An ordinary user-initiated load still downloads; only an API-initiated one is pinned
# to the cache, and that is the caller's flag rather than this helper's default.
"local_files_only": False,
}
]
@ -1482,6 +1489,7 @@ def test_a_cached_legacy_file_does_not_pre_empt_the_canonical_one(monkeypatch, t
filename,
token = None,
cache_dir = None,
local_files_only = False,
):
asked.append(filename)
return str(tmp_path / "downloaded-canonical.pt")
@ -1513,6 +1521,7 @@ def test_the_legacy_name_is_still_used_once_the_canonical_one_is_absent(monkeypa
filename,
token = None,
cache_dir = None,
local_files_only = False,
):
asked.append(filename)
if filename == "Z-Image-Turbo-FP8.pt":
@ -1559,6 +1568,7 @@ def test_a_legacy_copy_in_the_other_root_is_reused_after_the_primary_404s(monkey
filename,
token = None,
cache_dir = None,
local_files_only = False,
):
asked.append((filename, cache_dir))
if filename == "Z-Image-Turbo-FP8.pt":
@ -1600,6 +1610,7 @@ def test_a_primary_404_during_revalidation_still_reaches_the_legacy_fallback(mon
filename,
token = None,
cache_dir = None,
local_files_only = False,
):
asked.append((filename, cache_dir))
if filename == "Z-Image-Turbo-FP8.pt":
@ -1641,6 +1652,7 @@ def test_offline_revalidation_still_returns_the_other_root_copy(monkeypatch, tmp
filename,
token = None,
cache_dir = None,
local_files_only = False,
):
asked.append((filename, cache_dir))
raise LocalEntryNotFoundError("offline and not in this root")
@ -1693,7 +1705,7 @@ def test_load_config_reads_the_same_cache_root_as_the_checkpoint(monkeypatch, tm
ckpt = live_root / "ck.pt" # the checkpoint came from the LIVE root
ckpt.write_bytes(b"x")
monkeypatch.setattr(P, "_resolve_checkpoint_path", lambda s, t, c = None: str(ckpt))
monkeypatch.setattr(P, "_resolve_checkpoint_path", lambda s, t, c = None, **_: str(ckpt))
monkeypatch.setattr(P, "_validate_checkpoint", lambda *a, **k: True)
monkeypatch.setattr(P, "_pin_kernel_preference", lambda *a, **k: 0)
monkeypatch.setattr(torch, "load", lambda *a, **k: {"state_dict": {}, "scheme": "int8"})
@ -1705,6 +1717,7 @@ def test_load_config_reads_the_same_cache_root_as_the_checkpoint(monkeypatch, tm
subfolder = None,
token = None,
cache_dir = None,
local_files_only = False,
):
seen.append(cache_dir)
raise RuntimeError("stop right after the config fetch")
@ -1741,7 +1754,7 @@ def test_the_config_follows_the_checkpoint_into_the_other_cache_root(monkeypatch
ckpt = other_root / "ck.pt"
ckpt.write_bytes(b"x")
monkeypatch.setattr(P, "_resolve_checkpoint_path", lambda s, t, c = None: str(ckpt))
monkeypatch.setattr(P, "_resolve_checkpoint_path", lambda s, t, c = None, **_: str(ckpt))
monkeypatch.setattr(P, "_validate_checkpoint", lambda *a, **k: True)
monkeypatch.setattr(P, "_pin_kernel_preference", lambda *a, **k: 0)
monkeypatch.setattr(torch, "load", lambda *a, **k: {"state_dict": {}, "scheme": "int8"})
@ -1767,6 +1780,7 @@ def test_the_config_follows_the_checkpoint_into_the_other_cache_root(monkeypatch
subfolder = None,
token = None,
cache_dir = None,
local_files_only = False,
):
seen.append(cache_dir)
if cache_dir is not None:

View file

@ -202,6 +202,10 @@ def test_base_config_filter_skips_weights():
assert not keep("unet/diffusion_pytorch_model.safetensors")
assert not keep("vae/diffusion_pytorch_model.bin")
assert not keep("text_encoder/model.onnx")
# transformer/ and assets/ stay excluded (inherited from _base_file_downloaded).
assert not keep("transformer/config.json")
# transformer/ and assets/ stay excluded (inherited from _base_file_downloaded), except for
# transformer/config.json: from_single_file(config = <repo id>, subfolder = "transformer")
# resolves that one off the Hub, so an offline load needs it staged and the locality gate has
# to count it. The shards stay excluded -- the single file supplies those.
assert keep("transformer/config.json")
assert not keep("transformer/diffusion_pytorch_model.safetensors")
assert not keep("assets/x.png")

File diff suppressed because it is too large Load diff

View file

@ -0,0 +1,595 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""The FAMILY-SPECIFIC assemblers keep the no-download promise too.
``test_diffusion_offline_load.py`` and ``test_video_offline_load.py`` pin the SHARED path: the
staging phase, the byte estimate, the prefetch, the guarded ``pipe_kwargs`` every ordinary family
is built from. Four assemblers do not go through that dict, and each of them was still reaching
the Hub on a load nobody asked for:
- the MiniMax-H3 hosted conditioner (~27 GB), fetched by ``load_h3_quantized_text_encoder``;
- the image checkpoint, REOPENED by ``_resolve_gguf_path`` under the generation lock;
- Krea 2, assembled per-component because the repo ships transformers-5.x configs;
- LTX 2.3, assembled per-component because its vocoder class differs from the base pin.
Krea and LTX are reachable with no race at all: both are handed a REPO ID rather than a staged
snapshot (``_base_local_dir`` is None for 2.3 by design), so an assembler that resolves it without
the flag downloads whatever the caller's cache root does not hold. Every test below therefore
records what each component load was actually asked for, and the mirror tests keep the
user-initiated path fetching exactly as it did before.
"""
from __future__ import annotations
import ast
import json
import pathlib
import sys
from types import SimpleNamespace
import pytest
from core.inference import diffusion as diffusion_mod
def _call_keyword_sets(module_path: str, function: str, callee: str) -> list[set[str]]:
"""The keywords EVERY call to *callee* inside *function* spells out, one set per call site.
Read from the source rather than driven: these branches need a real Krea single-file build, a
real 2.3 checkpoint header or a Modular Diffusers H3 pipeline to reach, which no unit test can
stage -- but the keyword either is written there or is not. ``callee`` matches a bare name
(``load_krea2_pipeline``) or the last attribute segment (``self._resolve_gguf_path``,
``LTX2Pipeline.load_config``), so a call site that moves onto or off a receiver still counts.
Anchored on the package, not on the process CWD: CI runs pytest from the repo root with the
backend merely on PYTHONPATH, where a relative open raises FileNotFoundError.
"""
backend_root = pathlib.Path(diffusion_mod.__file__).resolve().parents[2]
tree = ast.parse((backend_root / module_path).read_text(encoding = "utf-8"))
found: list[set[str]] = []
for node in ast.walk(tree):
if not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) or node.name != function:
continue
for call in ast.walk(node):
if not isinstance(call, ast.Call):
continue
name = getattr(call.func, "id", None) or getattr(call.func, "attr", None)
if name == callee:
found.append({kw.arg for kw in call.keywords if kw.arg})
if not found:
raise AssertionError(f"no call to {callee}() inside {function}()")
return found
# ── [A] the MiniMax-H3 hosted conditioner ────────────────────────────────────
def _h3_te_module():
from core.inference import video_minimax_h3_te as te_mod
return te_mod
def _drive_h3_conditioner(monkeypatch, *, local_files_only):
"""Call the conditioner loader far enough to record its two Hub reads.
The loader is best-effort by contract and swallows everything into a None return, so the
RECORD is the result: the artifact fetch and the config read are the only two calls that can
leave the process, and a stub that raises after recording stops the 62 GB meta-init below.
"""
# The bare CI runners ship neither, and this driver reaches the real library rather than
# a stub, so the honest answer there is a skip.
transformers = pytest.importorskip("transformers")
import utils.hf_xet_fallback as xet
seen: dict = {}
# accelerate is imported at the top of the loader body, ahead of both Hub reads, and it is not
# a hard dependency of this backend; without the stub the whole function degrades to its
# best-effort None return before it asks for anything and the test would pass vacuously.
monkeypatch.setitem(
sys.modules, "accelerate", SimpleNamespace(init_empty_weights = lambda **_k: None)
)
def _download(
repo_id,
filename,
token = None,
**kwargs,
):
seen["download"] = (repo_id, filename, kwargs)
return "/nowhere/artifact.safetensors"
def _config(*args, **kwargs):
seen["config"] = kwargs
raise RuntimeError("stop after the two Hub reads")
monkeypatch.setattr(xet, "hf_hub_download_with_xet_fallback", _download)
monkeypatch.setattr(transformers.AutoConfig, "from_pretrained", _config, raising = False)
assert (
_h3_te_module().load_h3_quantized_text_encoder(
"MiniMaxAI/MiniMax-H3",
"int8",
dtype = None,
cache_dir = "/live/root",
local_files_only = local_files_only,
)
is None
)
return seen
def test_the_h3_conditioner_is_opened_from_the_cache_on_a_load_nobody_asked_for(monkeypatch):
"""The artifact is ~27 GB and the staging fetch already accepted it, so the loader may only
look it up -- through the SAME both-roots rule the stager used, or a moved cache folder makes
it re-pull what the load was cleared on."""
seen = _drive_h3_conditioner(monkeypatch, local_files_only = True)
repo, filename, kwargs = seen["download"]
te_mod = _h3_te_module()
assert (repo, filename) == (te_mod.H3_TE_QUANT_REPO, te_mod.H3_TE_QUANT_FILES["int8"])
assert kwargs["local_files_only"] is True
assert kwargs["reuse_other_cache_root"] is True
# The component config is a hub read too: _base_local_dir is None on an offline load, because
# the scoped base predownload stands down, so `local_base or base` resolves the repo id.
assert seen["config"]["local_files_only"] is True
def test_a_user_initiated_h3_load_still_fetches_the_conditioner(monkeypatch):
"""The pre-PR behaviour, unchanged: a load the user asked for pulls the artifact."""
seen = _drive_h3_conditioner(monkeypatch, local_files_only = False)
assert seen["download"][2]["local_files_only"] is False
assert seen["config"]["local_files_only"] is False
def test_the_h3_modular_build_hands_the_flag_to_the_conditioner_loader():
# The flag protected load_components() on one side and load_prequantized_transformer() on the
# other; the conditioner load between them was the remaining multi-GB fetch on that path.
for keywords in _call_keyword_sets(
"core/inference/video.py",
"_load_h3_modular_pipeline",
"load_h3_quantized_text_encoder",
):
assert "local_files_only" in keywords
# ── [B] reopening the image checkpoint under the generation lock ─────────────
def _drive_resolve_gguf(monkeypatch, *, cached_here, local_files_only):
"""``_resolve_gguf_path`` with the cache answering ``cached_here`` for the live root."""
import huggingface_hub
seen: list[dict] = []
monkeypatch.setattr(
huggingface_hub,
"try_to_load_from_cache",
lambda repo, name, cache_dir = None: (
"/live/checkpoint.gguf" if cache_dir is not None and cached_here else "/other/ck.gguf"
),
raising = False,
)
def _download(repo_id, filename, **kwargs):
seen.append(kwargs)
return "/resolved/checkpoint.gguf"
monkeypatch.setattr(huggingface_hub, "hf_hub_download", _download, raising = False)
diffusion_mod.DiffusionBackend()._resolve_gguf_path(
"unsloth/FLUX.1-dev-GGUF",
"flux1-dev-Q4_K_M.gguf",
None,
local_files_only = local_files_only,
)
return seen
@pytest.mark.parametrize("cached_here", [True, False])
def test_reopening_the_image_checkpoint_is_a_cache_lookup_offline(monkeypatch, cached_here):
"""Both resolutions, the live root and the other-root reuse. Neither is a no-op even on a hit:
``hf_hub_download`` re-resolves the revision against the Hub, so a checkpoint republished since
the cache was filled is a multi-GB pull taken AFTER the resident pipeline was evicted, inside
the generation lock, with progress already reading 100%."""
for kwargs in _drive_resolve_gguf(monkeypatch, cached_here = cached_here, local_files_only = True):
assert kwargs["local_files_only"] is True
@pytest.mark.parametrize("cached_here", [True, False])
def test_a_user_initiated_image_load_still_revalidates_the_checkpoint(monkeypatch, cached_here):
"""Unchanged for the UI load: it still goes to the Hub, which is how a republished GGUF is
picked up."""
for kwargs in _drive_resolve_gguf(monkeypatch, cached_here = cached_here, local_files_only = False):
assert kwargs["local_files_only"] is False
def test_the_image_assembly_hands_the_flag_to_the_checkpoint_resolver():
for keywords in _call_keyword_sets(
"core/inference/diffusion.py", "load_pipeline", "_resolve_gguf_path"
):
assert "local_files_only" in keywords
# ── [C] the Krea 2 per-component assembler ───────────────────────────────────
def _drive_krea(
monkeypatch,
tmp_path,
*,
local_files_only,
with_transformer = True,
):
"""Assemble a Krea pipeline against fakes, recording what each component was asked for."""
from core.inference.diffusion_krea2 import load_krea2_pipeline
import huggingface_hub
index = tmp_path / "model_index.json"
index.write_text(json.dumps({"patch_size": 2}), encoding = "utf-8")
seen: dict = {}
monkeypatch.setattr(
huggingface_hub,
"hf_hub_download",
lambda repo_id, filename, **kwargs: seen.update(model_index = kwargs) or str(index),
raising = False,
)
class _Component:
def __init__(self, tag):
self.tag = tag
def from_pretrained(self, repo_id, **kwargs):
seen[self.tag] = kwargs
return SimpleNamespace(tag = self.tag)
monkeypatch.setitem(
sys.modules,
"diffusers",
SimpleNamespace(
FlowMatchEulerDiscreteScheduler = _Component("scheduler"),
AutoencoderKLQwenImage = _Component("vae"),
Krea2Transformer2DModel = _Component("transformer"),
Krea2Pipeline = lambda **kwargs: SimpleNamespace(**kwargs),
),
)
for name in ("tokenizer", "text_encoder"):
monkeypatch.setattr(
f"core.inference.diffusion_krea2.load_krea2_{name}",
(
lambda *_a, tag = name, **kwargs: (
seen.__setitem__(tag, kwargs),
SimpleNamespace(tag = tag),
)[1]
),
)
load_krea2_pipeline(
# A hub id, not the local dir the other Krea tests use: the branch that reaches this passes
# ``fetch_base`` (or ``base_local_dir or base``, which is the id whenever nothing staged),
# and a local dir would resolve every component off disk and prove nothing.
"krea/Krea-2-Turbo",
"bf16",
with_transformer = with_transformer,
local_files_only = local_files_only,
)
return seen
def test_the_krea_assembler_opens_every_component_from_the_cache_offline(monkeypatch, tmp_path):
"""The 26 GB transformer, the 8.88 GB Qwen3-VL encoder, the VAE, the tokenizer and the
scheduler: this branch never sees the guarded pipe_kwargs, so each one has to carry the flag
itself or a load that promised nothing pulls it."""
seen = _drive_krea(monkeypatch, tmp_path, local_files_only = True)
assert set(seen) == {
"scheduler",
"vae",
"transformer",
"tokenizer",
"text_encoder",
"model_index",
}
for tag, kwargs in seen.items():
assert kwargs.get("local_files_only") is True, tag
def test_a_user_initiated_krea_load_still_fetches_every_component(monkeypatch, tmp_path):
seen = _drive_krea(monkeypatch, tmp_path, local_files_only = False)
for tag, kwargs in seen.items():
assert kwargs.get("local_files_only") is False, tag
def test_the_krea_model_index_read_is_a_cache_lookup_offline(monkeypatch):
"""A few KB, but still a fetch: the assembly reads the init config (``is_distilled`` carries
Turbo's mu shift) straight off the hub id when the repo is not a local directory."""
import huggingface_hub
from core.inference.diffusion_krea2 import _load_model_index
seen: dict = {}
def _download(repo_id, filename, **kwargs):
seen.update(kwargs)
raise RuntimeError("stop before the read")
monkeypatch.setattr(huggingface_hub, "hf_hub_download", _download, raising = False)
with pytest.raises(RuntimeError):
_load_model_index("krea/Krea-2-Turbo", None, local_files_only = True)
assert seen["local_files_only"] is True
def test_every_krea_call_site_hands_over_the_flag():
# Three: the full-pipeline branch, the transformer-only single-file branch, and _assemble_pipe.
sites = _call_keyword_sets(
"core/inference/diffusion.py", "load_pipeline", "load_krea2_pipeline"
) + _call_keyword_sets("core/inference/diffusion.py", "_assemble_pipe", "load_krea2_pipeline")
assert len(sites) == 3
for keywords in sites:
assert "local_files_only" in keywords
# ── [D] the LTX 2.3 per-component assembler ──────────────────────────────────
def test_the_ltx23_extras_fetch_is_a_cache_lookup_offline(monkeypatch):
"""The text projections, the video VAE and the audio VAE/vocoder: the switch's locality gate
clears these three by name, so a miss here is a promise it cannot keep."""
# monkeypatched by dotted path below, which imports the module to patch it.
pytest.importorskip("safetensors")
import utils.hf_xet_fallback as xet
from core.inference import video_ltx2
seen: dict = {}
monkeypatch.setattr(
xet,
"hf_hub_download_with_xet_fallback",
lambda repo_id, filename, token = None, **kwargs: seen.update(kwargs) or "/nowhere.st",
)
monkeypatch.setattr("safetensors.torch.load_file", lambda _path: {}, raising = False)
video_ltx2._load_extras_file("vae/x.safetensors", None, True)
assert seen["local_files_only"] is True
assert seen["reuse_other_cache_root"] is True
seen.clear()
video_ltx2._load_extras_file("vae/x.safetensors", None)
assert seen["local_files_only"] is False
def _drive_ltx23(monkeypatch, *, local_files_only):
"""Assemble a 2.3 pipeline against fakes, recording what the base repo reads were asked for."""
diffusers = pytest.importorskip("diffusers")
from core.inference import video_ltx2
seen: dict = {}
class _Sub:
@classmethod
def from_pretrained(cls, repo_id, **kwargs):
seen[kwargs["subfolder"]] = kwargs
return SimpleNamespace(tag = kwargs["subfolder"])
class _FakePipeline:
def __init__(self, **kwargs):
self.kwargs = kwargs
@classmethod
def load_config(cls, repo_id, **kwargs):
seen["load_config"] = kwargs
return {
name: ["diffusers", "_LtxOfflineSub"]
for name in ("scheduler", "tokenizer", "text_encoder")
}
monkeypatch.setattr(diffusers, "LTX2Pipeline", _FakePipeline, raising = False)
monkeypatch.setattr(diffusers, "_LtxOfflineSub", _Sub, raising = False)
monkeypatch.setattr(
"diffusers.loaders.single_file_utils.load_single_file_checkpoint",
lambda _path: {},
raising = False,
)
monkeypatch.setattr(video_ltx2, "checkpoint_variant", lambda _p: "dev")
for name in ("transformer", "connectors", "vae", "audio_vae_and_vocoder"):
monkeypatch.setattr(
video_ltx2,
f"load_ltx23_{name}",
(
lambda *_a, tag = name, **kwargs: (
seen.__setitem__(tag, kwargs),
(None, None) if tag == "audio_vae_and_vocoder" else None,
)[1]
),
)
video_ltx2.load_ltx23_pipeline(
"/models/ltx-2.3-dev-Q4_K_M.gguf",
# A repo id, which is what this branch always gets: the 2.3 snapshot lacks the base VAEs,
# so _run_load sets _base_local_dir to None for it deliberately.
base_repo = "Lightricks/LTX-Video-2",
torch_dtype = "bf16",
is_gguf = True,
local_files_only = local_files_only,
)
return seen
def test_the_ltx23_assembler_opens_every_component_from_the_cache_offline(monkeypatch):
"""The base model_index, the scheduler, the tokenizer, the dense Gemma3 encoder (~50 GB) and
every companion loader: none of them sees the guarded pipe_kwargs, and there is no staged
snapshot to fall back on, so each resolves the hub id itself."""
seen = _drive_ltx23(monkeypatch, local_files_only = True)
assert seen["load_config"]["local_files_only"] is True
for name in ("scheduler", "tokenizer", "text_encoder"):
assert seen[name]["local_files_only"] is True, name
for name in ("transformer", "connectors", "vae", "audio_vae_and_vocoder"):
assert seen[name]["local_files_only"] is True, name
def test_a_user_initiated_ltx23_load_still_fetches_every_component(monkeypatch):
seen = _drive_ltx23(monkeypatch, local_files_only = False)
assert seen["load_config"]["local_files_only"] is False
for name in ("scheduler", "tokenizer", "text_encoder", "transformer", "connectors"):
assert seen[name]["local_files_only"] is False, name
def test_the_video_assembly_hands_the_flag_to_the_ltx23_assembler():
for keywords in _call_keyword_sets(
"core/inference/video.py", "load_pipeline", "load_ltx23_pipeline"
):
assert "local_files_only" in keywords
# ── the live cache root ──────────────────────────────────────────────────────────
# Studio's HF cache folder is a SETTING (PUT /settings/hugging-face-cache), and changing it only
# rewrites the DB: os.environ and huggingface_hub's import-time constant keep the startup value.
# So after a change the live root and the import-time root differ, and an unset cache_dir resolves
# through the stale one. That mismatch predates this PR and used to be survivable, because a miss
# in the stale root just downloaded again. It is not survivable with local_files_only: the switch's
# locality gate reads the LIVE root (media_locality passes cache_dir = hub_cache_dir()), so it
# clears a model that is fully present, the resident pipeline is evicted, and the assembler then
# raises LocalEntryNotFoundError against the other root. Pinning is what keeps the gate's verdict
# and the load looking in the same place.
LIVE_ROOT = "/live-hub"
@pytest.fixture
def live_cache_root(monkeypatch):
"""Point the live root somewhere unmistakable, so a stale-root read cannot pass by accident."""
import utils.hf_cache_settings as cache_settings
monkeypatch.setattr(cache_settings, "active_hf_hub_cache", lambda: LIVE_ROOT)
return LIVE_ROOT
def test_the_krea_assembler_pins_every_component_to_the_live_cache(
monkeypatch, tmp_path, live_cache_root
):
seen = _drive_krea(monkeypatch, tmp_path, local_files_only = True)
# The direct loader calls. "tokenizer" and "text_encoder" are absent by design: those two tags
# record the kwargs handed to load_krea2_tokenizer / load_krea2_text_encoder, which are Studio
# helpers rather than hub calls, so they take no cache_dir and pin internally instead. The test
# below drives them for real.
for tag in ("scheduler", "vae", "transformer", "model_index"):
assert seen[tag].get("cache_dir") == live_cache_root, tag
def test_the_krea_tokenizer_and_encoder_helpers_pin_internally(monkeypatch, live_cache_root):
"""Both build their own kwargs dict, so the pin has to be inside each one."""
transformers = pytest.importorskip("transformers")
from core.inference import diffusion_krea2
seen: dict = {}
class _Component:
def __init__(self, tag):
self.tag = tag
def from_pretrained(self, repo_id, **kwargs):
seen[self.tag] = kwargs
return SimpleNamespace(tag = self.tag, text_config = SimpleNamespace())
monkeypatch.setattr(transformers, "AutoTokenizer", _Component("tokenizer"), raising = False)
monkeypatch.setattr(transformers, "AutoConfig", _Component("config"), raising = False)
monkeypatch.setattr(transformers, "Qwen3VLModel", _Component("text_encoder"), raising = False)
diffusion_krea2.load_krea2_tokenizer("krea/Krea-2-Turbo", local_files_only = True)
diffusion_krea2.load_krea2_text_encoder("krea/Krea-2-Turbo", "bf16", local_files_only = True)
assert set(seen) == {"tokenizer", "config", "text_encoder"}
for tag, kwargs in seen.items():
assert kwargs.get("cache_dir") == live_cache_root, tag
assert kwargs.get("local_files_only") is True, tag
def test_the_ltx23_assembler_pins_the_base_reads_to_the_live_cache(monkeypatch, live_cache_root):
seen = _drive_ltx23(monkeypatch, local_files_only = True)
# The base-repo reads only: the companion loaders take a checkpoint path, not a hub id.
assert seen["load_config"]["cache_dir"] == live_cache_root
for name in ("scheduler", "tokenizer", "text_encoder"):
assert seen[name]["cache_dir"] == live_cache_root, name
def test_the_hidream_external_encoder_is_pinned_to_the_live_cache(monkeypatch, live_cache_root):
"""TE4 lives in its own standalone repo, so it never rides the pipeline's pipe_kwargs (which
do carry cache_dir); it is the one 16 GB component that resolves its hub id unaided."""
transformers = pytest.importorskip("transformers")
from core.inference import diffusion_hidream
seen: dict = {}
class _Component:
def __init__(self, tag):
self.tag = tag
def from_pretrained(self, repo_id, **kwargs):
seen[self.tag] = kwargs
return SimpleNamespace(tag = self.tag)
monkeypatch.setattr(transformers, "AutoTokenizer", _Component("tokenizer_4"), raising = False)
monkeypatch.setattr(
transformers, "LlamaForCausalLM", _Component("text_encoder_4"), raising = False
)
diffusion_hidream.hidream_te4_kwargs(
dtype = "bf16",
hf_token = None,
local_files_only = True,
)
assert set(seen) == {"tokenizer_4", "text_encoder_4"}
for tag, kwargs in seen.items():
assert kwargs.get("cache_dir") == live_cache_root, tag
assert kwargs.get("local_files_only") is True, tag
# ── [E] the transformer-only single-file build ───────────────────────────────
# from_single_file(config = <repo id>, subfolder = "transformer") is not a local read: diffusers
# forwards local_files_only into the load_config() that resolves that id (single_file_model.py in
# 0.39 pops the kwarg and passes it on), and an unset flag is None, which is falsy, which permits
# the network. The pipeline assembly after it was already guarded, so this one call was the last
# unguarded Hub read on the GGUF/safetensors path -- and it runs AFTER eviction.
#
# The flag alone would not have been enough. transformer/config.json was deliberately excluded from
# the staged base file set on both paths (the shards come from the checkpoint), so the locality gate
# cleared picks that had never cached it and local_files_only would have turned a silent ~1 KB fetch
# into a hard failure on essentially every API-initiated GGUF load. Admitting the config -- and only
# the config -- is what makes the promise keepable, and lets the gate refuse up front instead.
def test_the_image_base_file_set_stages_the_transformer_config_but_not_its_shards():
from core.inference.diffusion import _base_file_downloaded as keep
assert keep("transformer/config.json", include_transformer = False)
assert not keep(
"transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
include_transformer = False,
)
def _sf_kwargs_keys(module_path: str) -> list[set[str]]:
"""The literal keys of every ``sf_kwargs = {...}`` in *module_path*, one set per assignment.
Read from the source, like the call-site checks above: reaching this branch needs a real
multi-GB GGUF plus its base repo, which no unit test can stage. The call itself is
``from_single_file(path, **sf_kwargs)``, so the keyword lives in the dict, not the call.
"""
backend_root = pathlib.Path(diffusion_mod.__file__).resolve().parents[2]
tree = ast.parse((backend_root / module_path).read_text(encoding = "utf-8"))
found: list[set[str]] = []
for node in ast.walk(tree):
if not isinstance(node, (ast.Assign, ast.AnnAssign)):
continue
targets = node.targets if isinstance(node, ast.Assign) else [node.target]
if not any(getattr(t, "id", None) == "sf_kwargs" for t in targets):
continue
if isinstance(node.value, ast.Dict):
found.append({k.value for k in node.value.keys if isinstance(k, ast.Constant)})
if not found:
raise AssertionError(f"no sf_kwargs dict literal in {module_path}")
return found
@pytest.mark.parametrize("module_path", ["core/inference/diffusion.py", "core/inference/video.py"])
def test_every_single_file_build_hands_over_the_flag(module_path):
for keys in _sf_kwargs_keys(module_path):
assert "local_files_only" in keys
# cache_dir is set here too, but diffusers does NOT forward it to the config lookup, so it
# pins only the checkpoint read. The flag is what keeps the config resolution off the Hub.
assert "config" in keys

View file

@ -79,6 +79,7 @@ def media(monkeypatch):
monkeypatch.setattr(tracker, "seen", None)
monkeypatch.setattr(tracker, "was_busy", False, raising = False)
monkeypatch.setattr(tracker, "completed", None, raising = False)
monkeypatch.setattr(mk, "_LOAD_ORIGINS", {})
return engines
@ -158,21 +159,15 @@ def test_media_ttl_env_behaves_like_the_chat_env(store, monkeypatch):
assert settings.get_media_auto_unload_idle_seconds() == 0
def test_api_only_disables_the_media_ttl(store, monkeypatch):
# "Only unload models loaded by the API" promises a model the user loaded from
# Studio stays resident, and nothing but the user ever loads an image or video
# model: /images/load and /video/load are the only entry points, and
# /v1/images/generations 503s rather than loading one. So with the setting on
# there is nothing here the idle unload is allowed to free.
def test_api_only_does_not_veto_the_media_ttl(store, monkeypatch):
# Media auto-switch gives an API request its own way to load a pipeline, so "only
# unload models loaded by the API" is a per-model rule here (see the tick tests
# below) rather than something that holds the whole TTL off.
store[settings.MEDIA_AUTO_UNLOAD_IDLE_SETTING_KEY] = 600
store[settings.AUTO_UNLOAD_API_ONLY_SETTING_KEY] = True
assert settings.get_media_auto_unload_idle_seconds() == 0
# The env-backed TTL is vetoed too, exactly as residency vetoes it.
assert settings.get_media_auto_unload_idle_seconds() == 600
del store[settings.MEDIA_AUTO_UNLOAD_IDLE_SETTING_KEY]
monkeypatch.setenv(settings.MEDIA_IDLE_TTL_ENV_VAR, "900")
assert settings.get_media_auto_unload_idle_seconds() == 0
# The stored seconds survive it, so turning the veto off brings them back.
store[settings.AUTO_UNLOAD_API_ONLY_SETTING_KEY] = False
assert settings.get_media_auto_unload_idle_seconds() == 900
@ -378,16 +373,16 @@ def test_a_different_model_restarts_the_ttl(media, monkeypatch):
def test_api_only_spares_a_model_the_user_loaded(media, store):
# The whole feature is off while the setting is on, and off means today's behaviour:
# nothing resolved, nothing unloaded.
# Unknown provenance reads as user-loaded, so an install that never recorded one is
# spared exactly as it was before media auto-switch existed.
store[settings.MEDIA_AUTO_UNLOAD_IDLE_SETTING_KEY] = 60
store[settings.AUTO_UNLOAD_API_ONLY_SETTING_KEY] = True
mk.note_load_origin(arb.DIFFUSION, "unsloth/FLUX.1-dev", None, user_action = True)
_step()
_step(*_BOTH)
assert media[arb.DIFFUSION].unloads == 0
assert media[arb.VIDEO].unloads == 0
# Turned off again, the same idle models are collectable: the setting was the only
# thing sparing them, so this does not cost the feature anything else.
# Turned off again, the same idle models are collectable.
store[settings.AUTO_UNLOAD_API_ONLY_SETTING_KEY] = False
_step()
_step(*_BOTH)
@ -395,6 +390,45 @@ def test_api_only_spares_a_model_the_user_loaded(media, store):
assert media[arb.VIDEO].unloads == 1
def test_api_only_still_frees_a_model_the_api_loaded(media, store):
# The other half of the per-model rule: auto-switch marks its own load, and that one
# is what the setting exists to collect.
store[settings.MEDIA_AUTO_UNLOAD_IDLE_SETTING_KEY] = 60
store[settings.AUTO_UNLOAD_API_ONLY_SETTING_KEY] = True
mk.note_load_origin(arb.DIFFUSION, "unsloth/FLUX.1-dev", None, user_action = False)
mk.note_load_origin(arb.VIDEO, "unsloth/Wan2.2", None, user_action = True)
_step()
_step(*_BOTH)
assert media[arb.DIFFUSION].unloads == 1
assert media[arb.VIDEO].unloads == 0
def test_a_failed_api_load_does_not_unpin_the_resident_user_model(media, store):
# A load is recorded when it is accepted, and it can still fail with the previous model
# resident. Reading that failed load's origin off the surviving model would evict a
# pipeline the setting promises to keep.
store[settings.MEDIA_AUTO_UNLOAD_IDLE_SETTING_KEY] = 60
store[settings.AUTO_UNLOAD_API_ONLY_SETTING_KEY] = True
mk.note_load_origin(arb.DIFFUSION, "unsloth/FLUX.1-dev", None, user_action = True)
mk.note_load_origin(arb.DIFFUSION, "unsloth/Z-Image-Turbo", None, user_action = False)
_step()
_step(*_BOTH)
assert media[arb.DIFFUSION].unloads == 0
def test_a_failed_api_load_of_another_quant_does_not_unpin_the_user_build(media, store):
# Same repo, different quant: the path alone is not the build, so a failed API load of Q8
# would otherwise mark the user's resident Q4 as API-loaded and free it.
store[settings.MEDIA_AUTO_UNLOAD_IDLE_SETTING_KEY] = 60
store[settings.AUTO_UNLOAD_API_ONLY_SETTING_KEY] = True
media[arb.DIFFUSION].build["gguf_variant"] = "Q4_K_M"
mk.note_load_origin(arb.DIFFUSION, "unsloth/FLUX.1-dev", "Q4_K_M", user_action = True)
mk.note_load_origin(arb.DIFFUSION, "unsloth/FLUX.1-dev", "Q8_0", user_action = False)
_step()
_step(*_BOTH)
assert media[arb.DIFFUSION].unloads == 0
def test_a_cached_reload_of_another_h3_partition_is_not_unloaded(media, monkeypatch):
# MiniMax-H3 keeps its identity in more than the repo id: fl2va and ref2va are
# different denoiser partitions, and the quants are part of the build too. A cached

View file

@ -0,0 +1,110 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""``media_locality`` reads the same verdict off both Hugging Face cache materialisations.
On Linux and macOS a snapshot entry is a relative symlink into ``blobs/``; on Windows without
developer mode ``huggingface_hub`` copies the blob into the snapshot instead. The completeness
check has to answer identically for the two, because the switch evicts the resident pipeline on
the strength of that answer: a component that reads as downloaded and then fails in
``from_pretrained`` leaves the user with nothing loaded.
The interesting state exists only in the symlink layout. A blob deleted by a cache sweep, or an
aborted pull, leaves the link behind, so a listing that matches on NAMES sees a complete
component where the copy layout would simply see an absent file.
"""
from __future__ import annotations
import os
from pathlib import Path
import core.inference.media_locality as locality
import pytest
def _cached_file(repo_dir: Path, snapshot: Path, name: str, payload: str, *, layout: str) -> Path:
"""One cached file, materialised the way *layout* says. Returns the blob path."""
blob = repo_dir / "blobs" / f"blob-{name.replace('/', '-')}"
blob.parent.mkdir(parents = True, exist_ok = True)
blob.write_bytes(payload.encode("utf-8"))
target = snapshot / name
target.parent.mkdir(parents = True, exist_ok = True)
if layout == "symlink":
# the relative spelling huggingface_hub writes, so the tree survives being moved
target.symlink_to(os.path.relpath(blob, target.parent))
else:
target.write_bytes(payload.encode("utf-8"))
return blob
def _snapshot(tmp_path: Path, files: dict, *, layout: str) -> tuple[Path, Path]:
repo_dir = tmp_path / "models--unsloth--z-image"
snapshot = repo_dir / "snapshots" / ("a" * 40)
snapshot.mkdir(parents = True)
for name, payload in files.items():
_cached_file(repo_dir, snapshot, name, payload, layout = layout)
return repo_dir, snapshot
@pytest.mark.parametrize("layout", ["symlink", "copy"])
def test_a_complete_component_reads_the_same_in_both_cache_layouts(tmp_path, layout):
_, snapshot = _snapshot(
tmp_path,
{
"transformer/config.json": "{}",
"transformer/diffusion_pytorch_model.safetensors": "weights",
},
layout = layout,
)
assert locality._component_present(snapshot / "transformer") is True
def test_a_dangling_weight_symlink_is_not_a_downloaded_component(tmp_path):
"""The blob is gone and the link remains: the copy layout cannot even express this."""
repo_dir, snapshot = _snapshot(
tmp_path,
{
"transformer/config.json": "{}",
"transformer/diffusion_pytorch_model.safetensors": "weights",
},
layout = "symlink",
)
weight = snapshot / "transformer" / "diffusion_pytorch_model.safetensors"
(repo_dir / "blobs" / "blob-transformer-diffusion_pytorch_model.safetensors").unlink()
assert weight.is_symlink() and not weight.exists()
assert locality._component_present(snapshot / "transformer") is False
def test_a_dangling_metadata_symlink_is_not_a_downloaded_component(tmp_path):
"""Same hole one branch over: a scheduler whose only config is a broken link."""
repo_dir, snapshot = _snapshot(
tmp_path, {"scheduler/scheduler_config.json": "{}"}, layout = "symlink"
)
(repo_dir / "blobs" / "blob-scheduler-scheduler_config.json").unlink()
assert locality._component_present(snapshot / "scheduler") is False
def test_a_directory_named_like_a_weight_file_is_not_a_weight(tmp_path):
"""A name test alone would take the directory for the checkpoint it is named after."""
_, snapshot = _snapshot(tmp_path, {"transformer/config.json": "{}"}, layout = "copy")
(snapshot / "transformer" / "diffusion_pytorch_model.safetensors").mkdir()
assert locality._component_present(snapshot / "transformer") is False
def test_a_pinned_variant_is_not_satisfied_by_a_dangling_link(tmp_path):
"""``variant`` names a weight set from_pretrained requires by name, so it needs a real file."""
repo_dir, snapshot = _snapshot(
tmp_path,
{
"transformer/config.json": "{}",
"transformer/diffusion_pytorch_model.fp16.safetensors": "weights",
},
layout = "symlink",
)
(repo_dir / "blobs" / "blob-transformer-diffusion_pytorch_model.fp16.safetensors").unlink()
assert locality._component_present(snapshot / "transformer", "fp16") is False

View file

@ -1172,7 +1172,7 @@ def test_setter_round_trips_auto_download_in_one_transaction(monkeypatch):
monkeypatch.setattr(settings, "_cached_setting", lambda k, d = None: store.get(k, d))
result = settings.set_openai_auto_switch(True, 120, None, True)
assert result == (True, 120, True, True, False, 0)
assert result == (True, 120, True, True, False, 0, False)
assert len(calls) == 1
assert calls[0][settings.OPENAI_AUTO_DOWNLOAD_SETTING_KEY] is True

View file

@ -788,7 +788,9 @@ def test_disabling_idle_unload_purges_saved_kv(monkeypatch, tmp_path):
"slots": [{"id": 0, "filename": saved.name}],
}
monkeypatch.setattr(
settings_route, "set_openai_auto_switch", lambda *a: (False, 300, True, False, False, 0)
settings_route,
"set_openai_auto_switch",
lambda *a: (False, 300, True, False, False, 0, False),
)
monkeypatch.setattr(settings_route, "get_auto_unload_idle_seconds", lambda: 0)
@ -813,7 +815,9 @@ def test_residency_does_not_purge_saved_kv(monkeypatch, tmp_path):
}
kw._kv_resume = manifest
monkeypatch.setattr(
settings_route, "set_openai_auto_switch", lambda *a: (True, 300, True, False, False, 0)
settings_route,
"set_openai_auto_switch",
lambda *a: (True, 300, True, False, False, 0, False),
)
monkeypatch.setattr(settings_route, "get_auto_unload_idle_seconds", lambda: 0)
monkeypatch.setattr(settings_route, "idle_unload_is_configured", lambda: True)
@ -5554,20 +5558,28 @@ def test_keep_kv_only_update_leaves_env_idle_ttl_active(monkeypatch):
monkeypatch.setenv(settings.MODEL_IDLE_TTL_ENV_VAR, "600")
assert settings_route.OpenAIAutoSwitchPayload(enabled = False).auto_unload_idle_seconds is None
enabled, idle, keep_kv, auto_dl, api_only, media_idle = settings.set_openai_auto_switch(
False, None, False
)
(
enabled,
idle,
keep_kv,
auto_dl,
api_only,
media_idle,
media_switch,
) = settings.set_openai_auto_switch(False, None, False)
assert settings.AUTO_UNLOAD_IDLE_SETTING_KEY not in store # idle untouched
assert settings.OPENAI_AUTO_DOWNLOAD_SETTING_KEY not in store # nor auto-download
assert settings.MEDIA_AUTO_UNLOAD_IDLE_SETTING_KEY not in store # nor the media TTL
assert settings.MEDIA_AUTO_SWITCH_SETTING_KEY not in store # nor media auto-switch
assert settings.get_auto_unload_idle_seconds() == 600 # env TTL still active
assert (enabled, idle, keep_kv, auto_dl, api_only, media_idle) == (
assert (enabled, idle, keep_kv, auto_dl, api_only, media_idle, media_switch) == (
False,
600,
False,
False,
False,
0,
False,
)

View file

@ -1263,7 +1263,21 @@ def _ltx23_assembly_stubs(monkeypatch, tmp_path):
_FakeLTX2Pipeline.last = kwargs
@staticmethod
def load_config(base_repo, token = None):
# An exact hand-written signature, so it follows the production one: the assembly bypasses
# the caller's guarded pipe_kwargs, and this base-repo read is the first of the calls that
# a load nobody asked for must not turn into a fetch.
def load_config(
base_repo,
token = None,
local_files_only = False,
cache_dir = None,
):
_FakeLTX2Pipeline.last_config_kwargs = {
"local_files_only": local_files_only,
# Pinned to Studio's LIVE root: unset, this resolves through huggingface_hub's
# import-time constant, which a mid-session cache-folder change leaves stale.
"cache_dir": cache_dir,
}
return {
"scheduler": ["diffusers", "_Loaded"],
"tokenizer": ["transformers", "_Loaded"],
@ -2078,10 +2092,15 @@ def test_base_download_files_scopes_pipeline_pull():
def test_base_download_files_gguf_drops_transformer():
# A GGUF/single-file checkpoint replaces the DiT: the base transformer never pulls.
# A GGUF/single-file checkpoint replaces the DiT: the base transformer WEIGHTS never pull.
info = types.SimpleNamespace(siblings = _LTX2_SIBLINGS)
names = [n for n, _ in VideoBackend._base_download_files(info, "gguf")]
assert not any(n.startswith("transformer/") for n in names)
transformer = [n for n in names if n.startswith("transformer/")]
# config.json is the one exception, and it is not an oversight: from_single_file resolves
# config = <repo id> through the Hub, so an API load that promised to download nothing needs
# this ~1 KB file staged, and the locality gate needs to count it. Everything else under
# transformer/ is supplied by the checkpoint itself.
assert transformer == ["transformer/config.json"]
assert "text_encoder/model-00001-of-00002.safetensors" in names

View file

@ -0,0 +1,321 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""An API-initiated video load downloads NOTHING.
``local_files_only=True`` is the contract the OpenAI-compatible routes load under: the model was
already staged, and the request may only open what is on disk. Every network-capable helper the
load reaches is replaced with a sentinel that RAISES when it is asked to fetch, so a load that
regains the network is a failing test rather than a multi-GB surprise on a user's connection. The
mirror test proves the user-initiated (UI) path still calls exactly those helpers, which is the
pre-PR behaviour nothing here is allowed to change.
"""
from __future__ import annotations
import inspect
import types
import pytest
import utils.hf_xet_fallback as xet
from core.inference import video as video_mod
from core.inference.video import VideoBackend
from core.inference.video_families import detect_video_family
# A plain (non-modular) family: the H3 conditioner / denoiser substitutions are exercised by their
# own suites, and this one keeps the load on the shared path every video pick walks.
WAN_GGUF = "unsloth/Wan2.2-TI2V-5B-GGUF"
WAN_BASE = "Wan-AI/Wan2.2-TI2V-5B-Diffusers"
WAN_FILE = "wan2.2-ti2v-5b-Q4_K_M.gguf"
class _Calls:
"""Every Hub call the load made, and how it made it."""
def __init__(self):
self.model_info: list[str] = []
self.downloads: list[tuple[str, str, bool]] = []
def _install_sentinels(monkeypatch, calls, tmp_path, *, offline):
"""Replace every network helper the load can reach.
``offline`` is the assertion: a metadata probe is refused outright (there is no offline form of
``model_info``), and a download is refused unless it carries ``local_files_only=True``, which is
what makes it a cache lookup rather than a fetch. Online they only record, so the same fake
serves both directions and the two tests differ by one flag.
"""
import huggingface_hub
def _model_info(self, repo_id, **_kwargs):
calls.model_info.append(repo_id)
if offline:
raise AssertionError(f"model_info({repo_id!r}) reached the Hub on an offline load")
return types.SimpleNamespace(siblings = [], sha = "deadbeef")
def _download(
repo_id,
filename,
token = None,
**kwargs,
):
local_files_only = bool(kwargs.get("local_files_only"))
calls.downloads.append((repo_id, filename, local_files_only))
if offline and not local_files_only:
raise AssertionError(
f"{repo_id}/{filename} was fetched without local_files_only on an offline load"
)
path = tmp_path / filename
path.parent.mkdir(parents = True, exist_ok = True)
path.write_bytes(b"")
return str(path)
monkeypatch.setattr(huggingface_hub.HfApi, "model_info", _model_info, raising = False)
monkeypatch.setattr(xet, "hf_hub_download_with_xet_fallback", _download)
# The wrapper's own offline branch calls this directly; a sentinel here catches a bypass.
monkeypatch.setattr(
huggingface_hub,
"hf_hub_download",
lambda **kwargs: _download(
kwargs.get("repo_id"), kwargs.get("filename"), kwargs.get("token"), **kwargs
),
raising = False,
)
def _backend(monkeypatch, calls_seen):
"""A backend whose family detection is pinned and whose pipeline build is a capture."""
backend = VideoBackend()
backend._load_token = 1
backend._loading = video_mod._VideoLoadingState(repo_id = WAN_GGUF, base_repo = WAN_BASE)
fam = detect_video_family(WAN_BASE)
assert fam is not None and not fam.modular_workflow
monkeypatch.setattr(video_mod, "_detect_load_family", lambda *_a, **_k: fam)
monkeypatch.setattr(backend, "load_pipeline", lambda **kwargs: calls_seen.update(kwargs))
return backend
def test_an_api_initiated_load_opens_the_cache_and_downloads_nothing(monkeypatch, tmp_path):
"""The whole promise, end to end: every helper on the load path either stays off the Hub or
asks it for a cached file only."""
calls = _Calls()
_install_sentinels(monkeypatch, calls, tmp_path, offline = True)
seen: dict = {}
backend = _backend(monkeypatch, seen)
backend._run_load(
repo_id = WAN_GGUF,
gguf_filename = WAN_FILE,
local_files_only = True,
_load_token = 1,
)
# _run_load swallows failures onto load_progress rather than raising, so the state IS the
# result: cleared means the load ran through, an error string means a sentinel fired.
assert backend._loading is None, getattr(backend._loading, "error", None)
assert seen.get("local_files_only") is True
# Not one metadata probe: the byte estimate and the base prefetch both stand down offline.
assert calls.model_info == []
# The checkpoint is still resolved -- as a cache lookup.
assert calls.downloads == [(WAN_GGUF, WAN_FILE, True)]
# And nothing was staged for from_pretrained, which resolves the cached snapshot itself.
assert seen.get("_base_local_dir") is None
def test_a_user_initiated_load_still_calls_every_one_of_them(monkeypatch, tmp_path):
"""The pre-PR path, unchanged: the UI load asks the Hub for sizes and pulls the checkpoint."""
calls = _Calls()
_install_sentinels(monkeypatch, calls, tmp_path, offline = False)
seen: dict = {}
backend = _backend(monkeypatch, seen)
backend._run_load(
repo_id = WAN_GGUF,
gguf_filename = WAN_FILE,
_load_token = 1,
)
assert backend._loading is None, getattr(backend._loading, "error", None)
assert seen.get("local_files_only") in (False, None)
# The byte estimate probes the checkpoint repo and the base; the base prefetch probes it again.
assert WAN_GGUF in calls.model_info and WAN_BASE in calls.model_info
# And the checkpoint is FETCHED, not looked up.
assert calls.downloads == [(WAN_GGUF, WAN_FILE, False)]
def test_the_native_h3_path_binds_the_flag_instead_of_swallowing_it(monkeypatch):
"""``_run_load_h3_native`` used to take ``**_``, so the flag arrived and vanished -- and its
four-file bundle, its sizing metadata and its sd-cli install all downloaded anyway."""
assert (
inspect.signature(VideoBackend._run_load_h3_native).parameters["local_files_only"].default
is False
)
from core.inference.video_minimax_h3 import H3_GGUF_REPO
fam = detect_video_family("MiniMaxAI/MiniMax-H3")
monkeypatch.setattr(video_mod, "_detect_load_family", lambda *_a, **_k: fam)
backend = VideoBackend()
backend._load_token = 1
backend._loading = video_mod._VideoLoadingState(repo_id = H3_GGUF_REPO, base_repo = fam.base_repo)
seen: dict = {}
monkeypatch.setattr(backend, "_run_load_h3_native", lambda **kwargs: seen.update(kwargs))
backend._run_load(
repo_id = H3_GGUF_REPO,
gguf_filename = "MiniMax-H3-Q4_K_M.gguf",
local_files_only = True,
_load_token = 1,
)
assert seen.get("local_files_only") is True
def test_load_pipeline_carries_the_flag_into_the_native_path(monkeypatch):
"""load_pipeline is reachable without _run_load (tests, and the keep-warm path), so the
dispatch has to pass the flag rather than let the parameter default re-enable downloads."""
from core.inference.video_minimax_h3 import H3_GGUF_REPO
backend = VideoBackend()
seen: dict = {}
monkeypatch.setattr(backend, "_run_load_h3_native", lambda **kwargs: seen.update(kwargs))
monkeypatch.setattr(
backend,
"validate_load_request",
lambda *a, **k: detect_video_family("MiniMaxAI/MiniMax-H3"),
)
backend.load_pipeline(
H3_GGUF_REPO,
gguf_filename = "MiniMax-H3-Q4_K_M.gguf",
model_kind = "gguf",
local_files_only = True,
)
assert seen.get("local_files_only") is True
def test_an_offline_denoiser_probe_reads_the_cache_instead_of_the_hub(monkeypatch):
"""Offline the "is there a replacement denoiser?" question is answered from disk. Refusing the
load instead would break the one case the flag exists for; answering yes without checking would
drop the dense shards a cache that lacks the artifact still needs."""
from core.inference.diffusion import DiffusionBackend
fam = detect_video_family("MiniMaxAI/MiniMax-H3")
monkeypatch.setattr(DiffusionBackend, "_hub_file_is_cached", staticmethod(lambda *a, **k: True))
assert (
VideoBackend._denoiser_prequant_cached_repo(fam, "int8", "MiniMaxAI/MiniMax-H3", "fl2va")
== "unsloth/MiniMax-H3-FP8"
)
monkeypatch.setattr(
DiffusionBackend, "_hub_file_is_cached", staticmethod(lambda *a, **k: False)
)
assert (
VideoBackend._denoiser_prequant_cached_repo(fam, "int8", "MiniMaxAI/MiniMax-H3", "fl2va")
is None
)
def test_the_estimate_and_the_base_prefetch_stand_down_offline(monkeypatch):
"""Both are pure Hub metadata, and both already have a "could not tell" answer the callers
handle, so offline they take it rather than inventing a probe."""
class _Boom:
def __init__(self, *_a, **_k):
pass
def model_info(self, *_a, **_k):
raise AssertionError("the Hub was asked about an offline load")
import huggingface_hub
monkeypatch.setattr(huggingface_hub, "HfApi", _Boom)
backend = VideoBackend()
assert (
backend._estimate_download_bytes(
WAN_GGUF, WAN_FILE, WAN_BASE, None, "gguf", local_files_only = True
)
is None
)
assert backend._predownload_base(WAN_BASE, None, "gguf", local_files_only = True) is None
def test_the_xet_wrapper_resolves_offline_without_the_shared_backend(monkeypatch, tmp_path):
"""``local_files_only`` must not depend on which unsloth_zoo is installed: the degraded stub
drops unknown keywords, so a forwarded flag would silently become a download."""
monkeypatch.setattr(
xet,
"_shared_hf_hub_download_with_xet_fallback",
lambda *a, **k: pytest.fail("the shared transport ran for an offline request"),
)
seen: dict = {}
def _hf_hub_download(**kwargs):
seen.update(kwargs)
path = tmp_path / "file.bin"
path.write_bytes(b"")
return str(path)
import huggingface_hub
monkeypatch.setattr(huggingface_hub, "hf_hub_download", _hf_hub_download)
out = xet.hf_hub_download_with_xet_fallback(
"org/repo", "file.bin", None, cache_dir = str(tmp_path), local_files_only = True
)
assert out == str(tmp_path / "file.bin")
assert seen["local_files_only"] is True
def test_the_xet_wrapper_is_unchanged_for_every_existing_caller(monkeypatch, tmp_path):
"""Default False: the online path still goes through the shared transport, and the flag is not
even forwarded, so an older shared backend cannot trip over it."""
seen: dict = {}
def _shared(*args, **kwargs):
seen["args"] = args
seen["kwargs"] = kwargs
return str(tmp_path / "file.bin")
monkeypatch.setattr(xet, "_shared_hf_hub_download_with_xet_fallback", _shared)
xet.hf_hub_download_with_xet_fallback("org/repo", "file.bin", None, cache_dir = str(tmp_path))
assert seen["args"] == ("org/repo", "file.bin", None)
assert "local_files_only" not in seen["kwargs"]
def _keywords_of(module_path: str, function: str, callee: str) -> set[str]:
"""The keyword names a call to *callee* inside *function* actually spells out.
Read from the source rather than driven, because the branch that reaches this call needs a
Modular Diffusers H3 pipeline and a card whose free memory has moved since the plan: the
condition is real but not one a unit test can stage, and the keyword either is there or is not.
"""
import ast
import pathlib
# Anchored on the package, not on the process CWD: CI runs pytest from the repo root with the
# backend merely on PYTHONPATH, where a relative open raises FileNotFoundError.
backend_root = pathlib.Path(video_mod.__file__).resolve().parents[2]
tree = ast.parse((backend_root / module_path).read_text(encoding = "utf-8"))
for node in ast.walk(tree):
if not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) or node.name != function:
continue
for call in ast.walk(node):
if isinstance(call, ast.Call) and getattr(call.func, "id", None) == callee:
return {kw.arg for kw in call.keywords if kw.arg}
raise AssertionError(f"no call to {callee}() inside {function}()")
def test_the_hosted_prequantized_denoiser_is_not_fetched_by_a_load_nobody_asked_for():
# `auto` is settled twice: the download plan decides against the card's CAPACITY, and this
# branch re-decides against LIVE free memory once the previous pipeline is gone. So a pick the
# plan sized as "the released bfloat16 denoiser fits" -- nothing hosted staged, the locality
# gate reporting zero missing bytes, the switch starting the load -- can be re-decided here as
# "take the hosted int8 checkpoint", and that is a multi-GB pull on a load that promised none.
# The image twin has passed the flag here all along; the video path did not.
assert "local_files_only" in _keywords_of(
"core/inference/video.py", "_load_h3_modular_pipeline", "load_prequantized_transformer"
)
assert "local_files_only" in _keywords_of(
"core/inference/diffusion.py", "_load_dense_quant_pipeline", "load_prequantized_transformer"
)

View file

@ -616,10 +616,19 @@ def hf_hub_download_with_xet_fallback(
force_download: bool = False,
cache_dir: Optional[str] = None,
reuse_other_cache_root: bool = False,
local_files_only: bool = False,
) -> str:
"""Single-file download via the shared fallback with Unsloth's marker-aware HTTP-retry prep.
``force_download`` re-fetches a newer blob over a cached one (Unsloth's model-update path).
``local_files_only`` resolves from the cache and never from the network, raising
huggingface_hub's ``LocalEntryNotFoundError`` on a miss. It deliberately BYPASSES the shared
fallback rather than forwarding the kwarg: that ladder exists only to recover a wedged
network transfer, so with no transfer permitted there is nothing to watch, and -- decisively
-- ``start_watchdog``-style version skew means an older installed ``unsloth_zoo`` could drop
an unrecognised kwarg on the floor. A dropped ``local_files_only`` DOWNLOADS, which is the one
outcome this parameter exists to prevent, so it must not depend on the installed zoo.
``reuse_other_cache_root`` (opt-in) resolves a file cached ONLY under huggingface_hub's
import-time root through that root. Studio's cache folder is a setting, so after it changes every
cached asset is invisible to a call pinned to the new root: GBs re-download, and a gated base with
@ -647,6 +656,27 @@ def hf_hub_download_with_xet_fallback(
cache_dir = None
except Exception: # noqa: BLE001 — a cache we cannot read just keeps the live root
pass
if local_files_only:
# Straight to huggingface_hub, after the root switch above (which is pure cache lookups and
# is exactly what lets an offline caller reach a file left under the import-time root).
# Cancellation is still honoured either side, as the fallback path does it. ``force_download``
# is not forwarded: there is nothing to re-fetch offline, and huggingface_hub rejects the pair.
from huggingface_hub import hf_hub_download
if cancel_event is not None and cancel_event.is_set():
raise RuntimeError("Cancelled")
path = hf_hub_download(
repo_id = repo_id,
filename = filename,
token = token,
repo_type = repo_type,
revision = revision,
cache_dir = cache_dir,
local_files_only = True,
)
if cancel_event is not None and cancel_event.is_set():
raise RuntimeError("Cancelled")
return path
# Omit rather than forward None: an older unsloth_zoo hands `interval` straight to Event.wait(),
# where None blocks forever and a hung Xet download never falls back. Omitting also lets the
# shared layer pick its per-transport defaults.

View file

@ -14,6 +14,11 @@ All off by default so existing API behavior is unchanged:
unloaded after this many idle seconds to free VRAM. Enabled values have a
60s floor (0 stays "off"): a tiny TTL tears the model down between turns of
an active chat, forcing a full weight reload + prompt re-prefill per turn.
- ``media_api_auto_switch_model``: the image/video twin of the first setting.
A media request naming a downloaded image or video model loads it before
generating, unloading the resident one once the work in flight has drained.
Its own setting for the same reason the media TTL is: the chat toggle says
nothing about pipelines the user loaded on the Image or Video page.
- ``media_auto_unload_idle_seconds``: the same for the image and video
pipelines. Its own setting, not a share of the chat one: this section is
about the OpenAI API and nothing here says it frees a model the user loaded
@ -41,6 +46,7 @@ from typing import Any, Optional
OPENAI_AUTO_SWITCH_SETTING_KEY = "openai_api_auto_switch_model"
OPENAI_AUTO_DOWNLOAD_SETTING_KEY = "openai_api_auto_download_model"
AUTO_UNLOAD_IDLE_SETTING_KEY = "openai_api_auto_unload_idle_seconds"
MEDIA_AUTO_SWITCH_SETTING_KEY = "media_api_auto_switch_model"
MEDIA_AUTO_UNLOAD_IDLE_SETTING_KEY = "media_auto_unload_idle_seconds"
AUTO_UNLOAD_KEEP_KV_SETTING_KEY = "openai_api_auto_unload_keep_kv"
AUTO_UNLOAD_API_ONLY_SETTING_KEY = "openai_api_auto_unload_api_only"
@ -50,6 +56,7 @@ MEDIA_IDLE_TTL_ENV_VAR = "UNSLOTH_MEDIA_IDLE_TTL"
DEFAULT_OPENAI_AUTO_SWITCH_ENABLED = False
DEFAULT_OPENAI_AUTO_DOWNLOAD_ENABLED = False
DEFAULT_MEDIA_AUTO_SWITCH_ENABLED = False
DEFAULT_AUTO_UNLOAD_IDLE_SECONDS = 0
DEFAULT_MEDIA_AUTO_UNLOAD_IDLE_SECONDS = 0
DEFAULT_AUTO_UNLOAD_KEEP_KV = True
@ -112,6 +119,12 @@ def get_openai_auto_switch_enabled() -> bool:
return parsed if parsed is not None else DEFAULT_OPENAI_AUTO_SWITCH_ENABLED
def get_media_auto_switch_enabled() -> bool:
"""Whether a media request may load the image or video model it names."""
parsed = _coerce_bool(_cached_setting(MEDIA_AUTO_SWITCH_SETTING_KEY, None))
return parsed if parsed is not None else DEFAULT_MEDIA_AUTO_SWITCH_ENABLED
def get_stored_openai_auto_download_enabled() -> bool:
"""The persisted auto-download flag, independent of auto-switch, so the UI
round-trips the saved value across an auto-switch toggle instead of erasing it."""
@ -240,14 +253,14 @@ def get_media_auto_unload_idle_seconds() -> int:
UNSLOTH_MEDIA_IDLE_TTL is the startup default when nothing is stored, exactly
as UNSLOTH_MODEL_IDLE_TTL is for chat.
Residency vetoes it like the chat reader, and so does "only unload models
loaded by the API": /images/load and /video/load are the only way a pipeline
is ever loaded (the OpenAI images route 503s instead of loading one), so every
resident image or video model is one the user loaded from Studio and the
setting promises to leave it alone. Chat can tell its two origins apart per
model and still frees the API-loaded ones; here there is nothing to free.
Residency vetoes it like the chat reader. "Only unload models loaded by the
API" does not veto it here: media auto-switch gives a request its own way to
load a pipeline, so the two origins now have to be told apart per model, which
media_keepwarm does with the provenance the load routes record. With
auto-switch off nothing but the user ever loads one, so that per-model rule
spares every resident model and the outcome is unchanged.
"""
if get_auto_unload_api_only() or _residency_vetoes_unload():
if _residency_vetoes_unload():
return 0
return get_stored_media_auto_unload_idle_seconds()
@ -284,7 +297,8 @@ def set_openai_auto_switch(
auto_download: Any = None,
api_only: Any = None,
media_idle_seconds: Any = None,
) -> tuple[bool, int, bool, bool, bool, int]:
media_auto_switch: Any = None,
) -> tuple[bool, int, bool, bool, bool, int, bool]:
"""One-transaction write; ``None`` leaves a stored value untouched."""
parsed_enabled = _coerce_bool(enabled)
if parsed_enabled is None:
@ -324,6 +338,11 @@ def set_openai_auto_switch(
parsed_api_only = _coerce_bool(api_only)
if parsed_api_only is None:
raise ValueError("Auto-unload API-loaded only must be true or false.")
parsed_media_auto_switch = None
if media_auto_switch is not None:
parsed_media_auto_switch = _coerce_bool(media_auto_switch)
if parsed_media_auto_switch is None:
raise ValueError("Media auto-switch must be true or false.")
from storage.studio_db import upsert_app_settings
updates: dict[str, Any] = {OPENAI_AUTO_SWITCH_SETTING_KEY: parsed_enabled}
@ -337,6 +356,8 @@ def set_openai_auto_switch(
updates[OPENAI_AUTO_DOWNLOAD_SETTING_KEY] = parsed_auto_download
if parsed_api_only is not None:
updates[AUTO_UNLOAD_API_ONLY_SETTING_KEY] = parsed_api_only
if parsed_media_auto_switch is not None:
updates[MEDIA_AUTO_SWITCH_SETTING_KEY] = parsed_media_auto_switch
upsert_app_settings(updates)
_invalidate(OPENAI_AUTO_SWITCH_SETTING_KEY)
if parsed_idle is not None:
@ -349,6 +370,8 @@ def set_openai_auto_switch(
_invalidate(OPENAI_AUTO_DOWNLOAD_SETTING_KEY)
if parsed_api_only is not None:
_invalidate(AUTO_UNLOAD_API_ONLY_SETTING_KEY)
if parsed_media_auto_switch is not None:
_invalidate(MEDIA_AUTO_SWITCH_SETTING_KEY)
return (
parsed_enabled,
parsed_idle if parsed_idle is not None else get_stored_auto_unload_idle_seconds(),
@ -364,6 +387,11 @@ def set_openai_auto_switch(
if parsed_media_idle is not None
else get_stored_media_auto_unload_idle_seconds()
),
(
parsed_media_auto_switch
if parsed_media_auto_switch is not None
else get_media_auto_switch_enabled()
),
)

View file

@ -21,8 +21,11 @@ export type OpenAIAutoSwitchSettings = {
// the chat TTL above is about the OpenAI API and never implied these.
mediaAutoUnloadIdleSeconds: number;
// True when the media idle unload will actually run, so the UI can say a veto
// (residency, or API-loaded only) is holding a saved TTL off.
// (residency) is holding a saved TTL off.
mediaIdleUnloadActive: boolean;
// Load the image or video model a media request names. Its own setting: the chat
// toggle above says nothing about pipelines loaded on the Images or Video page.
mediaAutoSwitchModel: boolean;
};
type ApiOpenAIAutoSwitchSettings = {
@ -43,6 +46,8 @@ type ApiOpenAIAutoSwitchSettings = {
media_auto_unload_idle_seconds?: number;
// biome-ignore lint/style/useNamingConvention: API schema
media_idle_unload_active?: boolean;
// biome-ignore lint/style/useNamingConvention: API schema
media_auto_switch_model?: boolean;
};
let cachedSettings: OpenAIAutoSwitchSettings | null = null;
@ -65,6 +70,7 @@ function fromApi(
autoUnloadApiOnly: settings.auto_unload_api_only ?? false,
mediaAutoUnloadIdleSeconds: settings.media_auto_unload_idle_seconds ?? 0,
mediaIdleUnloadActive: settings.media_idle_unload_active ?? false,
mediaAutoSwitchModel: settings.media_auto_switch_model ?? false,
};
}
@ -137,14 +143,37 @@ export async function loadOpenAIAutoSwitchSettings() {
return settings as OpenAIAutoSwitchSettings;
}
/** A partial write: `enabled` is always sent, and an omitted field keeps its stored value. */
export type OpenAIAutoSwitchUpdate = {
enabled: boolean;
autoUnloadIdleSeconds?: number;
autoUnloadKeepKv?: boolean;
autoDownloadModel?: boolean;
autoUnloadApiOnly?: boolean;
mediaAutoUnloadIdleSeconds?: number;
mediaAutoSwitchModel?: boolean;
};
// Camel-cased update field -> the API schema key it is sent as.
const UPDATE_KEYS = {
autoUnloadIdleSeconds: "auto_unload_idle_seconds",
autoUnloadKeepKv: "auto_unload_keep_kv",
autoDownloadModel: "auto_download_model",
autoUnloadApiOnly: "auto_unload_api_only",
mediaAutoUnloadIdleSeconds: "media_auto_unload_idle_seconds",
mediaAutoSwitchModel: "media_auto_switch_model",
} as const;
export async function updateOpenAIAutoSwitchSettings(
enabled: boolean,
autoUnloadIdleSeconds?: number,
autoUnloadKeepKv?: boolean,
autoDownloadModel?: boolean,
autoUnloadApiOnly?: boolean,
mediaAutoUnloadIdleSeconds?: number,
update: OpenAIAutoSwitchUpdate,
): Promise<OpenAIAutoSwitchSettings> {
const body: Record<string, unknown> = { enabled: update.enabled };
for (const [field, key] of Object.entries(UPDATE_KEYS)) {
const value = update[field as keyof typeof UPDATE_KEYS];
if (value !== undefined) {
body[key] = value;
}
}
// Read BEFORE the request: idleUnloadActive depends on the Model Memory
// setting, so a residency write landing mid-flight makes this response stale
// even though it is our own write's reply.
@ -152,30 +181,7 @@ export async function updateOpenAIAutoSwitchSettings(
const res = await authFetch("/api/settings/openai-auto-switch", {
method: "PUT",
headers: { "Content-Type": "application/json" },
body: JSON.stringify({
enabled,
// Omitted fields keep their stored value.
...(autoUnloadIdleSeconds === undefined
? {}
: // biome-ignore lint/style/useNamingConvention: API schema
{ auto_unload_idle_seconds: autoUnloadIdleSeconds }),
...(autoUnloadKeepKv === undefined
? {}
: // biome-ignore lint/style/useNamingConvention: API schema
{ auto_unload_keep_kv: autoUnloadKeepKv }),
...(autoDownloadModel === undefined
? {}
: // biome-ignore lint/style/useNamingConvention: API schema
{ auto_download_model: autoDownloadModel }),
...(autoUnloadApiOnly === undefined
? {}
: // biome-ignore lint/style/useNamingConvention: API schema
{ auto_unload_api_only: autoUnloadApiOnly }),
...(mediaAutoUnloadIdleSeconds === undefined
? {}
: // biome-ignore lint/style/useNamingConvention: API schema
{ media_auto_unload_idle_seconds: mediaAutoUnloadIdleSeconds }),
}),
body: JSON.stringify(body),
});
if (!res.ok) {
throw new Error(

View file

@ -8,6 +8,7 @@ import { useT } from "@/i18n";
import { useEffect, useState } from "react";
import {
type OpenAIAutoSwitchSettings,
type OpenAIAutoSwitchUpdate,
loadOpenAIAutoSwitchSettings,
updateOpenAIAutoSwitchSettings,
} from "../api/openai-auto-switch";
@ -36,8 +37,9 @@ function MediaIdleUnloadRow({
}) {
const t = useT();
const disabled = !settings || isSaving;
// A saved TTL that cannot run: residency or "only unload models loaded by the
// API" is vetoing it, and the number alone would not say so.
// A saved TTL that cannot run because residency is vetoing it; the number alone
// would not say so. "Only unload models loaded by the API" is per-model now, so
// it spares individual pipelines rather than holding the whole TTL off.
const paused =
settings !== null &&
settings.mediaAutoUnloadIdleSeconds > 0 &&
@ -136,28 +138,18 @@ export function ModelAutoSwitchSection() {
return parsed === 0 || parsed >= MIN_IDLE_SECONDS ? parsed : null;
};
// syncDraft only for a write the chat idle-seconds input owns; every other row
// leaves that draft alone so a save elsewhere cannot discard what is typed there.
const persist = async (
enabled: boolean,
idleSeconds: number | undefined,
syncDraft = true,
keepKv?: boolean,
autoDownload?: boolean,
apiOnly?: boolean,
mediaIdleSeconds?: number,
update: OpenAIAutoSwitchUpdate,
{ syncDraft = false }: { syncDraft?: boolean } = {},
) => {
setIsSaving(true);
setError(null);
try {
const saved = await updateOpenAIAutoSwitchSettings(
enabled,
idleSeconds,
keepKv,
autoDownload,
apiOnly,
mediaIdleSeconds,
);
const saved = await updateOpenAIAutoSwitchSettings(update);
setSettings(saved);
if (mediaIdleSeconds !== undefined) {
if (update.mediaAutoUnloadIdleSeconds !== undefined) {
setDraftMediaIdleSeconds(String(saved.mediaAutoUnloadIdleSeconds));
}
if (syncDraft) {
@ -182,10 +174,17 @@ export function ModelAutoSwitchSection() {
const handleToggle = (enabled: boolean) => {
const savedIdleSeconds = settings?.autoUnloadIdleSeconds ?? 0;
if (!enabled) {
void persist(false, savedIdleSeconds, false);
void persist({ enabled: false, autoUnloadIdleSeconds: savedIdleSeconds });
return;
}
void persist(true, parseIdleSeconds(draftIdleSeconds) ?? savedIdleSeconds);
void persist(
{
enabled: true,
autoUnloadIdleSeconds:
parseIdleSeconds(draftIdleSeconds) ?? savedIdleSeconds,
},
{ syncDraft: true },
);
};
const handleSaveIdle = () => {
@ -194,7 +193,10 @@ export function ModelAutoSwitchSection() {
setError(t("settings.general.modelAutoSwitch.idleError"));
return;
}
void persist(true, idleSeconds);
void persist(
{ enabled: true, autoUnloadIdleSeconds: idleSeconds },
{ syncDraft: true },
);
};
// The image/video TTL is its own setting, so it saves on its own: no enable
@ -208,37 +210,37 @@ export function ModelAutoSwitchSection() {
: null,
);
if (mediaIdleSeconds === null) return;
void persist(
settings.enabled,
undefined,
false,
undefined,
undefined,
undefined,
mediaIdleSeconds,
);
void persist({
enabled: settings.enabled,
mediaAutoUnloadIdleSeconds: mediaIdleSeconds,
});
};
const handleKeepKvToggle = (keepKv: boolean) => {
if (!settings) return;
void persist(settings.enabled, undefined, false, keepKv);
void persist({ enabled: settings.enabled, autoUnloadKeepKv: keepKv });
};
const handleAutoDownloadToggle = (autoDownload: boolean) => {
if (!settings) return;
void persist(settings.enabled, undefined, false, undefined, autoDownload);
void persist({
enabled: settings.enabled,
autoDownloadModel: autoDownload,
});
};
// Its own setting, so it saves alone: the chat toggle above is left untouched.
const handleMediaAutoSwitchToggle = (mediaAutoSwitch: boolean) => {
if (!settings) return;
void persist({
enabled: settings.enabled,
mediaAutoSwitchModel: mediaAutoSwitch,
});
};
const handleApiOnlyToggle = (apiOnly: boolean) => {
if (!settings) return;
void persist(
settings.enabled,
undefined,
false,
undefined,
undefined,
apiOnly,
);
void persist({ enabled: settings.enabled, autoUnloadApiOnly: apiOnly });
};
return (
@ -314,6 +316,18 @@ export function ModelAutoSwitchSection() {
) : null}
</div>
</SettingsRow>
<SettingsRow
label={t("settings.general.modelAutoSwitch.mediaEnable")}
description={t(
"settings.general.modelAutoSwitch.mediaEnableDescription",
)}
>
<Switch
checked={settings?.mediaAutoSwitchModel ?? false}
disabled={!settings || isSaving}
onCheckedChange={handleMediaAutoSwitchToggle}
/>
</SettingsRow>
<MediaIdleUnloadRow
draftSeconds={draftMediaIdleSeconds}
onDraftChange={setDraftMediaIdleSeconds}
@ -334,9 +348,8 @@ export function ModelAutoSwitchSection() {
/>
</SettingsRow>
) : null}
{/* Also whenever a media TTL is saved: this switch vetoes that TTL, and hiding it
behind the chat one leaves the media row stuck on "paused" with the only way out
being to re-enable chat unloading first. */}
{/* Also whenever a media TTL is saved: this switch decides which media models that
TTL may free, so it has to be reachable without re-enabling chat unloading first. */}
{settings &&
(settings.idleUnloadActive || settings.mediaAutoUnloadIdleSeconds > 0) ? (
<SettingsRow

View file

@ -394,13 +394,16 @@ export const ar = {
"حرّر ذاكرة VRAM بعد هذا العدد من ثواني الخمول. تُبقي القيمة 0 النموذج محمّلًا، والحد الأدنى 60 ثانية.",
idleSecondsAriaLabel:
"عدد الثواني قبل التفريغ التلقائي عند الخمول",
mediaEnable: "تبديل نموذج الصور والفيديو حسب الطلب",
mediaEnableDescription:
"تحميل نموذج صور أو فيديو منزّل ومذكور في طلب API قبل التوليد. إعداد مستقل: الإعداد أعلاه يخص نموذج المحادثة فقط. مُعطّل افتراضيًا.",
mediaIdleUnload: "التفريغ التلقائي عند الخمول للصور والفيديو",
mediaIdleUnloadDescription:
"حرّر ذاكرة VRAM بتفريغ نموذجَي الصور والفيديو بعد هذا العدد من ثواني الخمول. إنه إعداد مستقل: الإعداد أعلاه يخصّ نموذج المحادثة فقط. تُبقي القيمة 0 النموذجين محمَّلين، والحد الأدنى 60 ثانية.",
mediaIdleSecondsAriaLabel:
"عدد الثواني قبل التفريغ التلقائي عند الخمول للصور والفيديو",
mediaIdlePaused:
"متوقف مؤقتًا ما دام «إبقاء النموذج في ذاكرة كرت الرسوميات» أو «تفريغ النماذج التي حمّلتها واجهة API فقط» مفعّلًا.",
"متوقف مؤقتًا أثناء تفعيل إبقاء النموذج في ذاكرة وحدة معالجة الرسوميات.",
idleNeedsEnable: "فعّل «تبديل النموذج حسب الطلب» أولًا.",
idleActiveViaEnv:
"التفريغ التلقائي عند الخمول مُفعَّل عبر متغير البيئة UNSLOTH_MODEL_IDLE_TTL.",

View file

@ -407,6 +407,9 @@ export const de = {
"Gibt VRAM nach der angegebenen Anzahl von Sekunden ohne Aktivität frei. Bei 0 bleibt das Modell geladen; der Mindestwert ist 60.",
idleSecondsAriaLabel:
"Inaktivitätsdauer bis zum automatischen Entladen in Sekunden",
mediaEnable: "Bild- und Videomodell je Anfrage wechseln",
mediaEnableDescription:
"Lädt vor der Generierung ein darin angegebenes, bereits heruntergeladenes Bild- oder Videomodell aus einer API-Anfrage. Eine eigene Einstellung: Die Einstellung darüber gilt nur für das Chat-Modell. Standardmäßig deaktiviert.",
mediaIdleUnload:
"Automatisches Entladen bei Inaktivität für Bild und Video",
mediaIdleUnloadDescription:
@ -414,7 +417,7 @@ export const de = {
mediaIdleSecondsAriaLabel:
"Inaktivitätsdauer bis zum automatischen Entladen von Bild und Video in Sekunden",
mediaIdlePaused:
"Pausiert, solange „Modell im GPU-Speicher behalten“ oder „Nur über die API geladene Modelle entladen“ aktiv ist.",
"Pausiert, solange „Modell im GPU-Speicher behalten“ aktiv ist.",
idleNeedsEnable: "Aktivieren Sie zuerst „Modell je Anfrage wechseln“.",
idleActiveViaEnv:
"Automatisches Entladen bei Inaktivität ist über die Umgebungsvariable UNSLOTH_MODEL_IDLE_TTL aktiv.",

View file

@ -394,12 +394,14 @@ export const en = {
idleUnloadDescription:
"Free VRAM after this many idle seconds. 0 keeps it loaded, minimum 60.",
idleSecondsAriaLabel: "Idle auto-unload seconds",
mediaEnable: "Switch image and video model by request",
mediaEnableDescription:
"Load a downloaded image or video model named in an API request before generating. Its own setting: the one above covers the chat model only. Off by default.",
mediaIdleUnload: "Idle auto-unload for image and video",
mediaIdleUnloadDescription:
"Free VRAM by unloading the image and video models after this many idle seconds. Its own setting: the one above covers the chat model only. 0 keeps them loaded, minimum 60.",
mediaIdleSecondsAriaLabel: "Image and video idle auto-unload seconds",
mediaIdlePaused:
"Paused while Keep model in GPU memory or Only unload models loaded by the API is on.",
mediaIdlePaused: "Paused while Keep model in GPU memory is on.",
idleNeedsEnable: "Turn on Switch model by request first.",
idleActiveViaEnv: "Active via UNSLOTH_MODEL_IDLE_TTL.",
loadError: "Failed to load model auto-switch settings.",

View file

@ -403,6 +403,9 @@ export const es = {
"Libera la VRAM después de este número de segundos de inactividad. El valor 0 mantiene el modelo cargado; el mínimo es 60.",
idleSecondsAriaLabel:
"Segundos de inactividad antes de liberar el modelo",
mediaEnable: "Cambiar de modelo de imagen y vídeo según la solicitud",
mediaEnableDescription:
"Si una solicitud de la API especifica un modelo de imagen o vídeo ya descargado, lo carga antes de generar. Es una opción independiente: la de arriba solo se aplica al modelo de chat. Desactivado por defecto.",
mediaIdleUnload:
"Liberar imagen y vídeo automáticamente por inactividad",
mediaIdleUnloadDescription:
@ -410,7 +413,7 @@ export const es = {
mediaIdleSecondsAriaLabel:
"Segundos de inactividad antes de liberar los modelos de imagen y vídeo",
mediaIdlePaused:
"En pausa mientras «Mantener el modelo en la memoria de la GPU» o «Liberar solo los modelos cargados por la API» esté activado.",
"En pausa mientras «Mantener el modelo en la memoria de la GPU» está activado.",
idleNeedsEnable: "Activa primero «Cambiar de modelo según la solicitud».",
idleActiveViaEnv:
"La descarga automática por inactividad está activa mediante la variable de entorno UNSLOTH_MODEL_IDLE_TTL.",

View file

@ -405,6 +405,9 @@ export const fr = {
"Libérer la VRAM après ce nombre de secondes dinactivité. 0 maintient le modèle chargé ; le minimum est 60.",
idleSecondsAriaLabel:
"Délai dinactivité avant le déchargement automatique, en secondes",
mediaEnable: "Changer de modèle dimage et de vidéo par requête",
mediaEnableDescription:
"Charger, avant la génération, un modèle dimage ou de vidéo téléchargé indiqué dans une requête API. Réglage distinct : celui ci-dessus ne concerne que le modèle de discussion. Désactivé par défaut.",
mediaIdleUnload:
"Déchargement automatique en cas dinactivité pour limage et la vidéo",
mediaIdleUnloadDescription:
@ -412,7 +415,7 @@ export const fr = {
mediaIdleSecondsAriaLabel:
"Délai dinactivité avant le déchargement automatique de limage et de la vidéo, en secondes",
mediaIdlePaused:
"En pause tant que « Conserver le modèle dans la mémoire du GPU » ou « Décharger uniquement les modèles chargés par l'API » est activé.",
"En pause tant que « Conserver le modèle en mémoire GPU » est activé.",
idleNeedsEnable:
"Activez dabord « Changer de modèle par requête ».",
idleActiveViaEnv: "Actif via UNSLOTH_MODEL_IDLE_TTL.",

View file

@ -397,13 +397,16 @@ export const hi = {
idleUnloadDescription:
"इतने सेकंड तक निष्क्रिय रहने के बाद VRAM खाली करें। 0 पर मॉडल लोड रहता है; न्यूनतम 60 सेकंड।",
idleSecondsAriaLabel: "निष्क्रिय ऑटो-अनलोड की अवधि (सेकंड में)",
mediaEnable: "अनुरोध के अनुसार इमेज और वीडियो मॉडल बदलें",
mediaEnableDescription:
"जनरेट करने से पहले API अनुरोध में दिए गए डाउनलोड किए गए इमेज या वीडियो मॉडल को लोड करें। यह अलग सेटिंग है: ऊपर वाली केवल चैट मॉडल पर लागू होती है। डिफ़ॉल्ट रूप से बंद।",
mediaIdleUnload: "इमेज और वीडियो के लिए निष्क्रिय ऑटो-अनलोड",
mediaIdleUnloadDescription:
"इतने सेकंड तक निष्क्रिय रहने के बाद इमेज और वीडियो मॉडल अनलोड करके VRAM खाली करें। यह अपनी अलग सेटिंग है: ऊपर वाली सेटिंग केवल चैट मॉडल पर लागू होती है। 0 पर वे लोड रहते हैं; न्यूनतम 60 सेकंड।",
mediaIdleSecondsAriaLabel:
"इमेज और वीडियो के निष्क्रिय ऑटो-अनलोड की अवधि (सेकंड में)",
mediaIdlePaused:
"‘मॉडल को GPU मेमोरी में रखें’ या ‘केवल API द्वारा लोड किए गए मॉडल अनलोड करें’ चालू रहने तक रुका हुआ है।",
"जब तक “मॉडल को GPU मेमोरी में रखें” चालू है, तब तक रुका हुआ है।",
idleNeedsEnable: "पहले ‘अनुरोध के अनुसार मॉडल बदलें’ चालू करें।",
idleActiveViaEnv:
"निष्क्रिय ऑटो-अनलोड UNSLOTH_MODEL_IDLE_TTL एनवायरनमेंट वेरिएबल के माध्यम से सक्रिय है।",

View file

@ -373,6 +373,9 @@ export const it = {
"Libera la VRAM dopo il numero indicato di secondi di inattività. 0 mantiene il modello in memoria; il minimo è 60.",
idleSecondsAriaLabel:
"Secondi di inattività prima dello scaricamento automatico",
mediaEnable: "Cambia modello di immagini e video in base alla richiesta",
mediaEnableDescription:
"Carica un modello di immagini o video già scaricato indicato in una richiesta API prima di generare. È un'impostazione a sé: quella sopra riguarda solo il modello di chat. Disattivato per impostazione predefinita.",
mediaIdleUnload:
"Scaricamento automatico dalla memoria per inattività di immagini e video",
mediaIdleUnloadDescription:
@ -380,7 +383,7 @@ export const it = {
mediaIdleSecondsAriaLabel:
"Secondi di inattività prima dello scaricamento automatico di immagini e video",
mediaIdlePaused:
"In pausa finché «Mantieni il modello nella memoria della GPU» o «Scarica solo i modelli caricati dall'API» è attivo.",
"In pausa finché «Mantieni il modello nella memoria della GPU» è attivo.",
idleNeedsEnable:
"Attiva prima «Cambia modello in base alla richiesta».",
idleActiveViaEnv: "Attivo tramite UNSLOTH_MODEL_IDLE_TTL.",

View file

@ -388,13 +388,16 @@ export const ja = {
idleUnloadDescription:
"指定した秒数だけアイドル状態が続くと、モデルをアンロードして VRAM を解放します。0 にすると読み込んだままになります。最小値は 60 秒です。",
idleSecondsAriaLabel: "アイドル時の自動アンロードまでの秒数",
mediaEnable: "リクエストごとに画像・動画モデルを切り替え",
mediaEnableDescription:
"API リクエストで指定されたダウンロード済みの画像・動画モデルを、生成前に読み込みます。上の設定はチャットモデルのみが対象で、これは独立した設定です。デフォルトではオフです。",
mediaIdleUnload: "画像と動画のアイドル時の自動アンロード",
mediaIdleUnloadDescription:
"指定した秒数だけアイドル状態が続くと、画像モデルと動画モデルをアンロードして VRAM を解放します。これは独立した設定です。上の設定はチャットモデルのみが対象です。0 にすると読み込んだままになります。最小値は 60 秒です。",
mediaIdleSecondsAriaLabel:
"画像と動画のアイドル時の自動アンロードまでの秒数",
mediaIdlePaused:
"「モデルを GPU メモリに保持」または「API が読み込んだモデルのみアンロード」がオンの間は一時停止します。",
"「モデルを GPU メモリに保持」がオンの間は一時停止します。",
idleNeedsEnable: "アンロードされたモデルが次回使用時に再読み込みされるように、「リクエストごとにモデルを切り替え」をオンにしてください。",
idleActiveViaEnv: "アイドル時の自動アンロードは UNSLOTH_MODEL_IDLE_TTL 環境変数によって有効になっています。",
loadError: "モデル自動切り替え設定の読み込みに失敗しました。",

View file

@ -391,13 +391,16 @@ export const ko = {
idleUnloadDescription:
"지정한 유휴 시간(초)이 지나면 모델을 해제하여 VRAM을 확보합니다. 다음 요청 시 다시 불러옵니다. 0으로 설정하면 계속 로드된 상태로 유지됩니다. 최소 60초입니다.",
idleSecondsAriaLabel: "유휴 시 자동 해제까지의 시간(초)",
mediaEnable: "요청에 따라 이미지·동영상 모델 전환",
mediaEnableDescription:
"API 요청에 지정된 이미지 또는 동영상 모델이 다운로드되어 있으면 생성 전에 불러옵니다. 위 설정은 채팅 모델에만 적용되는 별개의 설정입니다. 기본값은 꺼짐입니다.",
mediaIdleUnload: "이미지와 동영상의 유휴 시 자동 해제",
mediaIdleUnloadDescription:
"지정한 유휴 시간(초)이 지나면 이미지와 동영상 모델을 해제하여 VRAM을 확보합니다. 별도의 설정입니다. 위 설정은 채팅 모델에만 적용됩니다. 0으로 설정하면 계속 로드된 상태로 유지됩니다. 최소 60초입니다.",
mediaIdleSecondsAriaLabel:
"이미지와 동영상의 유휴 시 자동 해제까지의 시간(초)",
mediaIdlePaused:
"‘모델을 GPU 메모리에 유지’ 또는 API가 불러온 모델만 해제’가 켜져 있는 동안 일시 중지됩니다.",
"‘모델을 GPU 메모리에 유지’가 켜져 있는 동안 일시 중지됩니다.",
idleNeedsEnable: "먼저 ‘요청에 따라 모델 전환’을 켜세요.",
idleActiveViaEnv:
"유휴 시 자동 해제가 UNSLOTH_MODEL_IDLE_TTL 환경 변수를 통해 활성화되어 있습니다.",

View file

@ -399,6 +399,9 @@ export const ptBR = {
"Libera a VRAM após esta quantidade de segundos de inatividade. 0 mantém o modelo carregado; o mínimo é 60.",
idleSecondsAriaLabel:
"Segundos até o descarregamento automático por inatividade",
mediaEnable: "Trocar de modelo de imagem e vídeo por requisição",
mediaEnableDescription:
"Carrega, antes de gerar, um modelo de imagem ou vídeo baixado indicado em uma requisição da API. É uma configuração própria: a de cima vale apenas para o modelo de chat. Desativado por padrão.",
mediaIdleUnload:
"Descarregamento automático por inatividade de imagem e vídeo",
mediaIdleUnloadDescription:
@ -406,7 +409,7 @@ export const ptBR = {
mediaIdleSecondsAriaLabel:
"Segundos até o descarregamento automático por inatividade de imagem e vídeo",
mediaIdlePaused:
"Pausado enquanto Manter o modelo na memória da GPU ou Descarregar apenas modelos carregados pela API estiver ativado.",
"Pausado enquanto “Manter o modelo na memória da GPU” estiver ativado.",
idleNeedsEnable: "Primeiro, ative Trocar de modelo por requisição.",
idleActiveViaEnv: "Ativo por meio de UNSLOTH_MODEL_IDLE_TTL.",
loadError: "Falha ao carregar as configurações de troca automática de modelo.",

View file

@ -397,13 +397,16 @@ export const ru = {
idleUnloadDescription:
"Освобождать VRAM после указанного числа секунд простоя. 0 оставляет модель загруженной; минимальное значение: 60.",
idleSecondsAriaLabel: "Число секунд до автовыгрузки при простое",
mediaEnable: "Переключать модель изображений и видео по запросу",
mediaEnableDescription:
"Загружать скачанную модель изображений или видео, указанную в запросе к API, перед генерацией. Отдельная настройка: та, что выше, относится только к модели чата. По умолчанию выключено.",
mediaIdleUnload: "Автовыгрузка при простое для изображений и видео",
mediaIdleUnloadDescription:
"Освобождать VRAM, выгружая модели изображений и видео после указанного числа секунд простоя. Это отдельная настройка: та, что выше, относится только к модели чата. 0 оставляет их загруженными; минимальное значение: 60.",
mediaIdleSecondsAriaLabel:
"Число секунд до автовыгрузки при простое для изображений и видео",
mediaIdlePaused:
"Приостановлено, пока включено «Держать модель в памяти GPU» или «Выгружать только модели, загруженные через API».",
"Приостановлено, пока включено «Держать модель в памяти GPU».",
idleNeedsEnable: "Сначала включите «Переключать модель по запросу».",
idleActiveViaEnv: "Активно через UNSLOTH_MODEL_IDLE_TTL.",
loadError: "Не удалось загрузить настройки автопереключения модели.",

View file

@ -379,12 +379,15 @@ export const zhCN = {
idleUnloadDescription:
"空闲达到该秒数后释放 VRAM。设为 0 则保持加载,最小值为 60 秒。",
idleSecondsAriaLabel: "空闲自动卸载秒数",
mediaEnable: "按请求切换图像和视频模型",
mediaEnableDescription:
"在生成前加载 API 请求中指定的已下载图像或视频模型。这是独立的设置:上面那项仅适用于聊天模型。默认关闭。",
mediaIdleUnload: "图像和视频的空闲自动卸载",
mediaIdleUnloadDescription:
"空闲达到该秒数后卸载图像和视频模型以释放 VRAM。这是独立的设置上面那项仅适用于聊天模型。设为 0 则保持加载,最小值为 60 秒。",
mediaIdleSecondsAriaLabel: "图像和视频空闲自动卸载秒数",
mediaIdlePaused:
"当“将模型保留在显存中”或“仅卸载由 API 加载的模型”开启时暂停。",
"当“将模型保留在 GPU 内存中”开启时暂停。",
idleNeedsEnable: "请先开启“按请求切换模型”。",
idleActiveViaEnv: "已通过 UNSLOTH_MODEL_IDLE_TTL 启用。",
loadError: "加载模型自动切换设置失败。",

View file

@ -61,13 +61,10 @@ test("a backend without the field reads as off", async () => {
test("the toggle round-trips", async () => {
invalidateOpenAIAutoSwitchSettings();
nextBody = { ...API, auto_unload_api_only: true };
const saved = await updateOpenAIAutoSwitchSettings(
true,
undefined,
undefined,
undefined,
true,
);
const saved = await updateOpenAIAutoSwitchSettings({
enabled: true,
autoUnloadApiOnly: true,
});
assert.equal(saved.autoUnloadApiOnly, true);
nextBody = { ...API };
});
@ -75,13 +72,10 @@ test("the toggle round-trips", async () => {
test("saving it alone leaves the other switches untouched", async () => {
invalidateOpenAIAutoSwitchSettings();
bodies.length = 0;
await updateOpenAIAutoSwitchSettings(
true,
undefined,
undefined,
undefined,
true,
);
await updateOpenAIAutoSwitchSettings({
enabled: true,
autoUnloadApiOnly: true,
});
assert.deepEqual(JSON.parse(bodies[0] ?? "{}"), {
enabled: true,
// biome-ignore lint/style/useNamingConvention: API schema

View file

@ -0,0 +1,107 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
// Image/video auto-switch rides the shared auto-switch PUT but is its own
// setting: saving it must not carry any other field along, and a backend that
// predates it must read as off rather than inheriting the chat toggle.
import assert from "node:assert/strict";
import { register } from "node:module";
import test from "node:test";
import { installLocalStorageFake } from "./helpers/kit.ts";
// The settings API modules reach authFetch through the auth barrel, which
// re-exports login-page.tsx. See helpers/auth-stub.mjs.
register("./helpers/settings-api-resolver.mjs", import.meta.url);
installLocalStorageFake();
const API = {
enabled: true,
// biome-ignore lint/style/useNamingConvention: API schema
auto_unload_idle_seconds: 300,
// biome-ignore lint/style/useNamingConvention: API schema
default_enabled: false,
// biome-ignore lint/style/useNamingConvention: API schema
idle_unload_active: true,
// biome-ignore lint/style/useNamingConvention: API schema
auto_unload_keep_kv: true,
// biome-ignore lint/style/useNamingConvention: API schema
auto_download_model: false,
// biome-ignore lint/style/useNamingConvention: API schema
auto_unload_api_only: false,
// biome-ignore lint/style/useNamingConvention: API schema
media_auto_unload_idle_seconds: 0,
// biome-ignore lint/style/useNamingConvention: API schema
media_idle_unload_active: false,
// biome-ignore lint/style/useNamingConvention: API schema
media_auto_switch_model: false,
};
let nextBody: Record<string, unknown> = { ...API };
const bodies: string[] = [];
globalThis.fetch = (async (_input: RequestInfo | URL, init?: RequestInit) => {
if (init?.body) bodies.push(String(init.body));
return new Response(JSON.stringify(nextBody), {
status: 200,
headers: { "Content-Type": "application/json" },
});
}) as typeof fetch;
const {
invalidateOpenAIAutoSwitchSettings,
loadOpenAIAutoSwitchSettings,
updateOpenAIAutoSwitchSettings,
} = await import("../src/features/settings/api/openai-auto-switch.ts");
test("a backend without the field reads as off, not as the chat toggle", async () => {
invalidateOpenAIAutoSwitchSettings();
const { media_auto_switch_model: _switch, ...older } = API;
nextBody = older;
const settings = await loadOpenAIAutoSwitchSettings();
assert.equal(settings.enabled, true);
assert.equal(settings.mediaAutoSwitchModel, false);
nextBody = { ...API };
});
test("the toggle round-trips", async () => {
invalidateOpenAIAutoSwitchSettings();
// biome-ignore lint/style/useNamingConvention: API schema
nextBody = { ...API, media_auto_switch_model: true };
const saved = await updateOpenAIAutoSwitchSettings({
enabled: true,
mediaAutoSwitchModel: true,
});
assert.equal(saved.mediaAutoSwitchModel, true);
nextBody = { ...API };
});
test("saving it alone leaves the other switches untouched", async () => {
invalidateOpenAIAutoSwitchSettings();
bodies.length = 0;
await updateOpenAIAutoSwitchSettings({
enabled: true,
mediaAutoSwitchModel: true,
});
assert.deepEqual(JSON.parse(bodies[0] ?? "{}"), {
enabled: true,
// biome-ignore lint/style/useNamingConvention: API schema
media_auto_switch_model: true,
});
});
test("a false toggle is sent, not dropped as absent", async () => {
// Only `undefined` means "leave stored"; turning the switch OFF has to reach the server.
invalidateOpenAIAutoSwitchSettings();
bodies.length = 0;
await updateOpenAIAutoSwitchSettings({
enabled: true,
mediaAutoSwitchModel: false,
});
assert.deepEqual(JSON.parse(bodies[0] ?? "{}"), {
enabled: true,
// biome-ignore lint/style/useNamingConvention: API schema
media_auto_switch_model: false,
});
});

View file

@ -77,14 +77,10 @@ test("the seconds round-trip", async () => {
// biome-ignore lint/style/useNamingConvention: API schema
media_idle_unload_active: true,
};
const saved = await updateOpenAIAutoSwitchSettings(
true,
undefined,
undefined,
undefined,
undefined,
600,
);
const saved = await updateOpenAIAutoSwitchSettings({
enabled: true,
mediaAutoUnloadIdleSeconds: 600,
});
assert.equal(saved.mediaAutoUnloadIdleSeconds, 600);
assert.equal(saved.mediaIdleUnloadActive, true);
nextBody = { ...API };
@ -93,14 +89,10 @@ test("the seconds round-trip", async () => {
test("saving it alone leaves the chat TTL untouched", async () => {
invalidateOpenAIAutoSwitchSettings();
bodies.length = 0;
await updateOpenAIAutoSwitchSettings(
true,
undefined,
undefined,
undefined,
undefined,
600,
);
await updateOpenAIAutoSwitchSettings({
enabled: true,
mediaAutoUnloadIdleSeconds: 600,
});
assert.deepEqual(JSON.parse(bodies[0] ?? "{}"), {
enabled: true,
// biome-ignore lint/style/useNamingConvention: API schema
@ -111,7 +103,7 @@ test("saving it alone leaves the chat TTL untouched", async () => {
test("saving the chat TTL does not send a media TTL", async () => {
invalidateOpenAIAutoSwitchSettings();
bodies.length = 0;
await updateOpenAIAutoSwitchSettings(true, 300);
await updateOpenAIAutoSwitchSettings({ enabled: true, autoUnloadIdleSeconds: 300 });
assert.deepEqual(JSON.parse(bodies[0] ?? "{}"), {
enabled: true,
// biome-ignore lint/style/useNamingConvention: API schema