mirror of
https://github.com/unslothai/unsloth.git
synced 2026-08-17 04:43:52 +00:00
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:
parent
321c3ee019
commit
cfee13795e
55 changed files with 7903 additions and 301 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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}
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
522
studio/backend/core/inference/media_auto_switch.py
Normal file
522
studio/backend/core/inference/media_auto_switch.py
Normal 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",
|
||||
]
|
||||
|
|
@ -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()
|
||||
|
|
|
|||
539
studio/backend/core/inference/media_locality.py
Normal file
539
studio/backend/core/inference/media_locality.py
Normal 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",
|
||||
]
|
||||
526
studio/backend/core/inference/media_model_index.py
Normal file
526
studio/backend/core/inference/media_model_index.py
Normal 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",
|
||||
]
|
||||
189
studio/backend/core/inference/media_switch_backends.py
Normal file
189
studio/backend/core/inference/media_switch_backends.py
Normal 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",
|
||||
]
|
||||
171
studio/backend/core/inference/media_switch_errors.py
Normal file
171
studio/backend/core/inference/media_switch_errors.py
Normal 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",
|
||||
]
|
||||
113
studio/backend/core/inference/media_switch_locks.py
Normal file
113
studio/backend/core/inference/media_switch_locks.py
Normal 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",
|
||||
]
|
||||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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))
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
360
studio/backend/tests/test_diffusion_engine_load_contract.py
Normal file
360
studio/backend/tests/test_diffusion_engine_load_contract.py
Normal 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"
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
216
studio/backend/tests/test_diffusion_offline_load.py
Normal file
216
studio/backend/tests/test_diffusion_offline_load.py
Normal 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
|
||||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
2599
studio/backend/tests/test_media_auto_switch.py
Normal file
2599
studio/backend/tests/test_media_auto_switch.py
Normal file
File diff suppressed because it is too large
Load diff
595
studio/backend/tests/test_media_family_assembler_offline.py
Normal file
595
studio/backend/tests/test_media_family_assembler_offline.py
Normal 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
|
||||
|
|
@ -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
|
||||
|
|
|
|||
110
studio/backend/tests/test_media_locality_cache_layout.py
Normal file
110
studio/backend/tests/test_media_locality_cache_layout.py
Normal 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
|
||||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
321
studio/backend/tests/test_video_offline_load.py
Normal file
321
studio/backend/tests/test_video_offline_load.py
Normal 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"
|
||||
)
|
||||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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.",
|
||||
|
|
|
|||
|
|
@ -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.",
|
||||
|
|
|
|||
|
|
@ -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.",
|
||||
|
|
|
|||
|
|
@ -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.",
|
||||
|
|
|
|||
|
|
@ -405,6 +405,9 @@ export const fr = {
|
|||
"Libérer la VRAM après ce nombre de secondes d’inactivité. 0 maintient le modèle chargé ; le minimum est 60.",
|
||||
idleSecondsAriaLabel:
|
||||
"Délai d’inactivité avant le déchargement automatique, en secondes",
|
||||
mediaEnable: "Changer de modèle d’image et de vidéo par requête",
|
||||
mediaEnableDescription:
|
||||
"Charger, avant la génération, un modèle d’image 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 d’inactivité pour l’image et la vidéo",
|
||||
mediaIdleUnloadDescription:
|
||||
|
|
@ -412,7 +415,7 @@ export const fr = {
|
|||
mediaIdleSecondsAriaLabel:
|
||||
"Délai d’inactivité avant le déchargement automatique de l’image 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 d’abord « Changer de modèle par requête ».",
|
||||
idleActiveViaEnv: "Actif via UNSLOTH_MODEL_IDLE_TTL.",
|
||||
|
|
|
|||
|
|
@ -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 एनवायरनमेंट वेरिएबल के माध्यम से सक्रिय है।",
|
||||
|
|
|
|||
|
|
@ -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.",
|
||||
|
|
|
|||
|
|
@ -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: "モデル自動切り替え設定の読み込みに失敗しました。",
|
||||
|
|
|
|||
|
|
@ -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 환경 변수를 통해 활성화되어 있습니다.",
|
||||
|
|
|
|||
|
|
@ -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.",
|
||||
|
|
|
|||
|
|
@ -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: "Не удалось загрузить настройки автопереключения модели.",
|
||||
|
|
|
|||
|
|
@ -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: "加载模型自动切换设置失败。",
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
107
studio/frontend/tests/media-auto-switch-setting.test.ts
Normal file
107
studio/frontend/tests/media-auto-switch-setting.test.ts
Normal 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,
|
||||
});
|
||||
});
|
||||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue