mirror of
https://github.com/unslothai/unsloth.git
synced 2026-08-25 08:42:25 +00:00
cfee13795e
* 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>
211 lines
9.7 KiB
Python
211 lines
9.7 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
|
|
|
"""CPU-only unit tests for the SDXL diffusion family.
|
|
|
|
SDXL is the one U-Net family: the denoiser is ``pipe.unet`` (not ``pipe.transformer``)
|
|
and a single-file ``.safetensors`` is the whole pipeline (not a transformer-only file).
|
|
These tests cover the pure helpers that encode those differences -- family detection,
|
|
the ``denoiser_attr`` / ``single_file_is_pipeline`` flags, the non-GGUF trust allowlist,
|
|
the VAE-dtype alignment reading the U-Net denoiser, and the LoRA-support gate -- with no
|
|
torch/diffusers/GPU needed.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import types
|
|
|
|
import pytest
|
|
|
|
from core.inference import diffusion_lora
|
|
from core.inference.diffusion import (
|
|
DiffusionBackend,
|
|
_is_trusted_diffusion_repo,
|
|
resolve_model_kind,
|
|
)
|
|
from core.inference.diffusion_families import detect_family, family_sd_cpp_supported
|
|
|
|
|
|
def test_sdxl_family_shape():
|
|
fam = detect_family("stabilityai/stable-diffusion-xl-base-1.0")
|
|
assert fam is not None and fam.name == "sdxl"
|
|
assert fam.pipeline_class == "StableDiffusionXLPipeline"
|
|
# The denoiser is a U-Net, addressed via pipe.unet (DiT families use pipe.transformer).
|
|
assert fam.denoiser_attr == "unet"
|
|
assert fam.transformer_class == "UNet2DConditionModel"
|
|
# A single-file SDXL checkpoint is the whole pipeline, loaded via the pipeline class.
|
|
assert fam.single_file_is_pipeline is True
|
|
# Image-conditioned + ControlNet workflows are the standard SDXL pipelines.
|
|
assert fam.img2img_pipeline_class == "StableDiffusionXLImg2ImgPipeline"
|
|
assert fam.inpaint_pipeline_class == "StableDiffusionXLInpaintPipeline"
|
|
assert fam.controlnet_pipeline_class == "StableDiffusionXLControlNetPipeline"
|
|
assert fam.controlnet_model_class == "ControlNetModel"
|
|
# Real CFG; SDXL uses guidance_scale, not a distilled true_cfg_scale.
|
|
assert fam.cfg_kwarg == "guidance_scale"
|
|
|
|
|
|
def test_sdxl_detection_by_repo_and_override():
|
|
assert detect_family("stabilityai/sdxl-turbo").name == "sdxl"
|
|
assert detect_family("some-org/My-Cool-SDXL-Merge").name == "sdxl"
|
|
assert detect_family("some-org/stable-diffusion-xl-anime").name == "sdxl"
|
|
assert detect_family("x", override = "sdxl").name == "sdxl"
|
|
# A GGUF DiT family must NOT be swallowed by the SDXL match.
|
|
assert detect_family("unsloth/FLUX.1-schnell-GGUF").name == "flux.1"
|
|
|
|
|
|
def test_dit_families_keep_transformer_denoiser():
|
|
# The generalisation must not change existing DiT families: they stay on pipe.transformer and their single file is transformer-only.
|
|
for rid in ("unsloth/FLUX.1-schnell-GGUF", "unsloth/Qwen-Image-GGUF", "unsloth/Z-Image-GGUF"):
|
|
fam = detect_family(rid)
|
|
assert fam.denoiser_attr == "transformer"
|
|
assert fam.single_file_is_pipeline is False
|
|
|
|
|
|
def test_sdxl_has_no_native_sd_cpp_mapping():
|
|
# No single-file VAE/TE mapping yet, so the no-GPU route falls back to diffusers rather than driving sd-cli.
|
|
assert family_sd_cpp_supported(detect_family("stabilityai/sdxl-turbo")) is False
|
|
|
|
|
|
def test_sdxl_base_repos_are_trusted_non_gguf():
|
|
# Official safetensors-only base repos are allowlisted so their catalog entries load.
|
|
assert _is_trusted_diffusion_repo("stabilityai/stable-diffusion-xl-base-1.0")
|
|
assert _is_trusted_diffusion_repo("stabilityai/sdxl-turbo")
|
|
# The refiner is img2img-only and intentionally NOT allowlisted (see test_sdxl_refiner_not_trusted). Case-insensitive match.
|
|
assert _is_trusted_diffusion_repo("StabilityAI/SDXL-Turbo")
|
|
# A random repo (even one that detects as SDXL) is NOT trusted for a non-GGUF load.
|
|
assert not _is_trusted_diffusion_repo("randomorg/my-sdxl-merge")
|
|
assert not _is_trusted_diffusion_repo("stabilityai/sdxl-turbo-evil")
|
|
|
|
|
|
def test_sdxl_model_kind_resolution():
|
|
# A full-pipeline load (no single-file name) is "pipeline"; a single .safetensors is "single_file".
|
|
assert resolve_model_kind(None) == "pipeline"
|
|
assert resolve_model_kind("sdxl.safetensors") == "single_file"
|
|
|
|
|
|
class _FakeVae:
|
|
def __init__(self, dtype):
|
|
self._dtype = dtype
|
|
self.moved_to = None
|
|
|
|
def parameters(self):
|
|
yield types.SimpleNamespace(dtype = self._dtype)
|
|
|
|
def to(self, dtype = None):
|
|
self.moved_to = dtype
|
|
self._dtype = dtype
|
|
|
|
|
|
def test_align_vae_dtype_uses_unet_denoiser():
|
|
# For SDXL the denoiser lives at pipe.unet, so _align_vae_dtype must read it and cast the VAE to the U-Net's dtype, which comes from a parameter (hence the _FakeVae).
|
|
import torch
|
|
|
|
vae = _FakeVae(dtype = torch.float32)
|
|
unet = _FakeVae(dtype = torch.bfloat16)
|
|
pipe = types.SimpleNamespace(unet = unet, vae = vae)
|
|
DiffusionBackend._align_vae_dtype(pipe, "unet")
|
|
assert vae.moved_to == torch.bfloat16
|
|
|
|
|
|
def test_align_vae_dtype_transformer_default_unchanged():
|
|
# DiT default: reads pipe.transformer; a pipe with no transformer is a safe no-op.
|
|
import torch
|
|
|
|
vae = _FakeVae(dtype = torch.float32)
|
|
transformer = _FakeVae(dtype = torch.bfloat16)
|
|
pipe = types.SimpleNamespace(transformer = transformer, vae = vae)
|
|
DiffusionBackend._align_vae_dtype(pipe)
|
|
assert vae.moved_to == torch.bfloat16
|
|
# No denoiser attribute -> no-op (does not raise, does not move the VAE).
|
|
vae2 = _FakeVae(dtype = torch.float32)
|
|
DiffusionBackend._align_vae_dtype(types.SimpleNamespace(vae = vae2), "unet")
|
|
assert vae2.moved_to is None
|
|
|
|
|
|
def test_align_vae_dtype_skips_gguf_packed_uint8_params():
|
|
# A GGUF-quantized transformer's leading parameters are packed uint8, so the dtype probe must skip them and use the first
|
|
# FLOATING dtype, else nn.Module.to() rejects the integer dtype and an Edit/img2img call 500s. All-integer is a no-op.
|
|
import torch
|
|
|
|
class _GgufDenoiser:
|
|
def parameters(self):
|
|
yield types.SimpleNamespace(dtype = torch.uint8) # packed GGUF block
|
|
yield types.SimpleNamespace(dtype = torch.bfloat16) # compute dtype
|
|
|
|
vae = _FakeVae(dtype = torch.float32)
|
|
pipe = types.SimpleNamespace(transformer = _GgufDenoiser(), vae = vae)
|
|
DiffusionBackend._align_vae_dtype(pipe)
|
|
assert vae.moved_to == torch.bfloat16
|
|
|
|
class _AllPacked:
|
|
def parameters(self):
|
|
yield types.SimpleNamespace(dtype = torch.uint8)
|
|
|
|
vae2 = _FakeVae(dtype = torch.float32)
|
|
DiffusionBackend._align_vae_dtype(types.SimpleNamespace(transformer = _AllPacked(), vae = vae2))
|
|
assert vae2.moved_to is None
|
|
|
|
|
|
def test_sdxl_lora_supported_on_diffusers():
|
|
# SDXL is bf16/bnb-4bit on diffusers, so LoRA is allowed (unlike GGUF-via-diffusers).
|
|
assert diffusion_lora.supports_lora(
|
|
engine = "diffusers", family = "sdxl", model_kind = "pipeline", transformer_quant = None
|
|
)
|
|
assert diffusion_lora.supports_lora(
|
|
engine = "diffusers", family = "sdxl", model_kind = "single_file", transformer_quant = None
|
|
)
|
|
|
|
|
|
def test_pipeline_prefetch_skips_non_torch_artifacts():
|
|
# The SDXL Base repo ships fp16 variants, ONNX, OpenVINO and Flax exports beside the default safetensors, and
|
|
# from_pretrained loads only the default torch weights, so the prefetch filter must skip the rest or pull tens of GB.
|
|
from core.inference.diffusion import _pipeline_file_downloaded as keep
|
|
|
|
assert keep("model_index.json")
|
|
assert keep("unet/diffusion_pytorch_model.safetensors")
|
|
assert keep("text_encoder/model.safetensors")
|
|
assert keep("scheduler/scheduler_config.json")
|
|
assert not keep("sd_xl_base_1.0.safetensors") # top-level single-file twin
|
|
assert not keep("unet/diffusion_pytorch_model.fp16.safetensors")
|
|
assert not keep("text_encoder/model.onnx")
|
|
assert not keep("text_encoder/openvino_model.bin")
|
|
assert not keep("unet/flax_model.msgpack")
|
|
assert not keep("vae_decoder/model.onnx_data")
|
|
assert not keep("assets/preview.png")
|
|
|
|
|
|
def test_sdxl_refiner_not_trusted():
|
|
# The refiner is img2img-only and the sdxl family loads every repo as the base txt2img pipeline, so it must NOT be allowlisted for a non-GGUF load.
|
|
assert not _is_trusted_diffusion_repo("stabilityai/stable-diffusion-xl-refiner-1.0")
|
|
# The base and turbo remain trusted.
|
|
assert _is_trusted_diffusion_repo("stabilityai/stable-diffusion-xl-base-1.0")
|
|
assert _is_trusted_diffusion_repo("stabilityai/sdxl-turbo")
|
|
|
|
|
|
def test_sdxl_gguf_load_rejected_up_front():
|
|
# SDXL has no transformer-only GGUF variant (its single file is the whole pipeline), so a GGUF request fails cheap validation before the GPU handoff.
|
|
backend = DiffusionBackend()
|
|
with pytest.raises(ValueError, match = "no GGUF"):
|
|
backend.validate_load_request(
|
|
"some-org/my-sdxl.gguf", gguf_filename = "my-sdxl.gguf", family_override = "sdxl"
|
|
)
|
|
|
|
|
|
def test_base_config_filter_skips_weights():
|
|
# For a whole-pipeline single file the base repo supplies only config/tokenizer, not its unused weight tensors.
|
|
from core.inference.diffusion import _base_config_file_downloaded as keep
|
|
|
|
assert keep("model_index.json")
|
|
assert keep("text_encoder/config.json")
|
|
assert keep("tokenizer/vocab.json")
|
|
assert keep("scheduler/scheduler_config.json")
|
|
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), 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")
|