Commit graph

2,867 commits

Author SHA1 Message Date
Daniel Han
cffee73135 Add Wan2.2-I2V-A14B image-to-video support to the video backend
Wan-AI/Wan2.2-I2V-A14B-Diffusers is the image-to-video sibling of the already
supported T2V-A14B: the same dual-expert WanTransformer3DModel pair (boundary_ratio
0.9 in the pipeline config) behind WanImageToVideoPipeline, conditioning through the
VAE latent (no CLIP-vision image encoder in this 2.2 variant).

- New wan2.2-i2v-a14b family: image_conditioned flag, card recipe defaults (40 steps,
  CFG 3.5, 81 frames at 16 fps), the T2V memory table (57.2 GB both experts bf16),
  fp32-pinned VAE, and a wan2.2-i2v generation-defaults key ahead of the generic wan
  50/5.0 entry.
- Source-image plumbing: /video/generate takes init_image (base64/data URL);
  begin_generate 400s synchronously when an image-to-video family has no image or a
  text-only family is given one, and generate() decodes and resizes the image to the
  snapped output size before threading it as the pipeline's image kwarg. status()
  reports image_input so the UI can gate the control.
- Trust the official -Diffusers repo for pipeline loads and transfer the measured wan
  quant recipes: fp8 keeps condition_embedder in bf16 (zero padding-row scale), mxfp8
  and nvfp4 stay denied, the UMT5 auto TE quant resolves dense, and the balanced
  FBCache pin carries over. All tables share the T2V DiT profile.
- Tests: family detection/aliases/defaults, the dual-DiT image pipeline load, the
  image gates on both begin_generate and generate, init_image route pass-through, and
  the quant exclude/deny/auto entries.

GPU-verified on a B200: bf16 resident load (offload none, 72.8 GB peak) animating a
conditioning image at 832x480/33f in 28.7 s with first-frame MAE 5.7 vs the source,
and an int8 load with both experts quantized (43.1 GB peak, clean output).
2026-07-17 11:18:01 +00:00
Daniel Han
63cbf7878e Add the Lumina Image 2.0 family to the image catalog
Alpha-VLLM/Lumina-Image-2.0 is a 2.6B single-stream DiT with a Gemma2-2B
encoder and a standard 16-channel VAE, all transformers-4.x-compatible, so the
generic from_pretrained pipeline path loads it as a new lumina-2 family:

- Family entry (Lumina2Pipeline / Lumina2Transformer2DModel), aliased to
  lumina-image-2.0 / lumina-image-2 / lumina2. No bare lumina alias: Lumina-Next
  checkpoints are a different arch and must stay unknown rather than crash
  mid-load. bf16-only upstream, so the fp16 fallback stays off like z-image.
- Trust the official repo for non-GGUF loads; bf16 component table entry
  (ships fp32, ~5.2 GB transformer + 5.2 GB encoder bf16-resident).
- Generation defaults 50 steps / guidance 4.0 per the model card, and the
  generate call passes the card's cfg_trunc_ratio=0.25 itself (family-gated,
  signature-gated): the pipeline default (1.0) runs the CFG double-forward on
  every step and oversaturates output.
- Catalog group with the single ungated bf16 pipeline artifact (11 GB resident)
  plus routing assertions; images page defaults row.
- No GGUF artifact: none exists upstream (only finetune/LLM quants), so the
  dense transformer_quant fast path (GGUF-kind-only) stays unreachable for now.
  Offline probes of the future prequant campaign: int8 and fp8 both engage and
  render cleanly (fp8 LPIPS 0.11 vs bf16, int8 0.33 from 50-step trajectory
  drift with intact quality), so neither scheme is family-denied.
2026-07-17 11:15:42 +00:00
Daniel Han
e1896eeaad Resolve pre-quantized checkpoints per base variant
One family entry covers several published variants whose weights differ
(flux.1: schnell, dev, Krea-dev), but prequant resolution was keyed on
(family, scheme) alone, so only the default base could ever be served: the
loader's baked base_model_id validation correctly refused the schnell
checkpoint for dev and Krea-dev bases and every such load paid the dense
download plus on-the-fly quantise.

Add an optional prequant_variant_repos table on DiffusionFamily as
(base_repo, scheme, repo_id) triples and thread the resolved base repo
through resolve_prequant_source / usable_prequant_source and their three
call sites (load fast path, memory-plan probe, auto-policy candidate). A
base without its own entry keeps returning the family default, preserving
the existing refuse-then-dense behavior exactly.

Wire the flux.1 variants: the gate-validated unsloth/FLUX.1-dev-FP8
checkpoints (built in the earlier campaign but never reachable) and the
new unsloth/FLUX.1-Krea-dev-FP8.
2026-07-17 11:15:29 +00:00
Daniel Han
5a58e9df13 Add FLUX.1 Krea dev to the image model catalog
Krea's guidance-distilled FLUX.1-dev finetune keeps the exact dev layout, so it
runs under the existing flux.1 family unchanged. Wire it up end to end:

- Catalog group with the gated official bf16 pipeline and the open QuantStack
  GGUF quants; the gated artifact is skipped on auto-routing when undownloaded.
- Trust the official repo for non-GGUF from_pretrained loads, next to the other
  black-forest-labs bases.
- Generation defaults: 28 steps at guidance 4.5 per the model card. The generic
  "krea" defaults key (Krea-2-Turbo's 8-step no-CFG recipe) used to swallow the
  id, which would have produced garbage output; the new flux.1-krea key precedes
  it on both the backend table and the images page table.
- The flux.1 prequant checkpoints are schnell-based; the loader's baked
  base_model_id validation refuses them for the Krea-dev base, so int8/fp8
  requests dense-quantize instead (covered by existing prequant tests).
2026-07-17 10:51:21 +00:00
Daniel Han
44f100af97 Support LoRA adapters on torchao int8/fp8 quantized image pipelines
Adapters are baked at load time: they attach to the dense transformer,
then quantize_ converts only the frozen base linears (the lora_ side
path is excluded by name), then the loader compiles. Post-quant PEFT
injection is not possible on a manually quantized module, so the
prequant shortcut is skipped for a baked load and the memory plan is
sized for the dense build (force_dense on the quant candidate).

At generation time the baked topology is frozen: weight tweaks and
disabling (scale 0 reproduces the quantized base exactly) go through
set_adapters, while adding or removing adapters returns a clean 400
telling the client to reload with the new selection.

supports_lora now returns True for int8/fp8 diffusers loads (checked
before the gguf-kind early return, since the quant fast path keeps the
picker kind); nvfp4/mxfp8 and GGUF-via-diffusers stay blocked. The
load request model takes an optional loras list, threaded through
begin_load on both engines (native ignores it and keeps applying LoRA
at generation).

Verified end to end on GPU: Z-Image GGUF picker + int8 + trained
adapter loads through the API, bake marker logged, weight 1.0 vs 0
renders differ visibly, weight 0.5 accepted live, unknown adapter
rejected as 400. Affected suites: 304 passed.
2026-07-17 09:37:39 +00:00
Daniel Han
d5be6fb5a3 Add FLUX.2 Klein and FLUX.2-dev DiT LoRA training
Register flux.2-klein and flux.2-dev in the DiT trainer following the
upstream DreamBooth references: latents train patchified and batch-norm
normalized from the VAE posterior mode, the packed forward reuses
step-invariant position ids, and the guidance vector (3.5) is gated on
the variant's guidance_embeds config. Conditioning stacks load per
variant (Mistral via Flux2Pipeline for dev, Qwen3 via Flux2KleinPipeline
for Klein) and are encoded and freed before the transformer lands on the
device. The fused single-stream to_qkv_mlp_proj joins the attention
projections in the LoRA targets; the single-stream out projection stays
dense because its to_out suffix would also match the double-stream
ModuleList container.

Wire both families through the training registry (family set, labels,
VRAM notes, rank 16 / lr 1e-4 defaults, bf16-only preflight), mark them
trainable with train base repos in the family registry, add FLUX.2-dev
to the gated-repo token check, and trust both official bases for
training downloads.

Verified on B200: 30-step klein int8 (19.6s) and nf4 (20.9s) and dev
int8 (52.0s) runs train with finite decreasing loss and the saved
adapters apply on the bf16 base pipeline (weight 0 reproduces the base
image exactly, weight 1 visibly restyles it).
2026-07-17 09:34:14 +00:00
Daniel Han
cf1910c92d Harden the diffusion memory plan against transient free-VRAM undercounts
A cold FLUX.2-dev int8 load on an idle 183 GB B200 planned offload=model
(companions exceed budget) and silently served the GGUF as-is; the identical
retry went resident and engaged the hosted prequant. The plan arithmetic was
byte-identical across both loads (required 90,228 MiB, resident needs free of
about 124 GB); the only divergent input was torch.cuda.mem_get_info, which is
device-wide and instantaneous: a transient foreign CUDA context briefly held
about 100 GB at the first snapshot, and the planner trusted that single read.

Three changes:
- settled_snapshot_device_memory: on cuda, synchronize + empty_cache
  (best-effort) and take the MAX free over up to 3 spaced reads. A transient
  can only shrink free, so the max rejects transient undercounts while a
  persistent tenant still caps every read. _plan_memory now uses it.
- plan_fits_total_capacity + one replan retry: when the dense/prequant
  candidate fits TOTAL device capacity under the standard reserve and the 0.85
  resident margin, an offload verdict can only stem from the free reading, so
  the loader re-snapshots and replans once before declining the fast path.
  Explicit balanced/low_vram modes skip the retry (they offload by mode).
- diffusion.transformer_quant_declined log line with required/budget/free and
  the plan reasons, so the next decline is diagnosable from the server log
  (previously silent).

Verified: cold FLUX.2-dev int8 first load in a fresh server now engages the
hosted prequant resident (offload=none).
2026-07-17 08:46:24 +00:00
Daniel Han
6d33582b88 Keep Qwen-Image's text-stream linears bf16 on int8 (short prompts break torch._int_mm)
Qwen-Image's MMDiT runs every text-stream Linear at M = actual prompt tokens: the
Qwen2.5-VL embeds are not padded to a fixed length like FLUX's 512-token T5. A short
prompt (13 tokens) or the near-empty negative prompt drives torch._int_mm below its
M > 16 floor and the first denoise step raises 'self.size(0) needs to be greater than
16, but got 13' (measured on B200 through the Studio images tab).

Add qwen-image / qwen-image-edit to the per-family int8 exclusions (txt_in,
add_q/k/v_proj, to_add_out, txt_mlp), the same recipe HunyuanVideo-1.5 already uses
for its trimmed text streams. The exclude list feeds the prequant checkpoint
validation, so a checkpoint baked under the old token list is rejected and
re-quantised instead of loaded crashing. The text stream runs at M = tens vs the
image stream's M ~ 4k, so the exclusion costs nothing; the rebuilt hosted checkpoint
gates 28/28 PASS with LPIPS mean 0.057 (was 0.069).
2026-07-17 07:16:27 +00:00
Daniel Han
ee0e059007 Route krea-2 through its per-component loader on the transformer-quant fast path
_assemble_pipe used Pipeline.from_pretrained for every family, but the krea repo
ships transformers-5.x configs and no top-level tokenizer files, so the tokenizer
dies with vocab_file=None. The pre-quantized checkpoint loaded fine and then the
assembly crashed, dropping the load to the GGUF build, which krea-2 cannot take
(Krea2Transformer2DModel has no from_single_file). Assemble per-component via
load_krea2_pipeline like the pipeline-kind and single-file paths already do.

Verified live: Krea-2-Turbo int8 and fp8 hosted prequant loads now assemble and
render through the Studio images tab.
2026-07-17 06:36:27 +00:00
Daniel Han
fbcf070fce Wire hosted pre-quantized DiT checkpoints into the image families
Point prequant_repos for flux.1, flux.2-klein, flux.2-dev, qwen-image
(int8 only there; fp8 is family-denied), z-image and krea-2 at the
unsloth/<Model>-FP8 Hub repos carrying gate-validated int8 and fp8
transformer checkpoints, so the fast quant path loads the small
pre-quantized file instead of materialising the dense bf16 transformer
and quantising on device. Measured on FLUX.2-dev int8: build peak drops
from 60.7 GB (dense + quantize) to 30.7 GB (hosted prequant), identical
30.7 GB resident after either path since loading a checkpoint is
bit-identical to on-the-fly quantisation.

The hosted repos name files <Model>-<SCHEME>.pt, so resolve_prequant_source
now derives that model-name filename from the repo id (scheme suffix
stripped case-insensitively) and carries the legacy transformer_<scheme>.pt
as a fallback the resolver tries when the primary 404s, keeping older
repos loadable.

Wiring a repo also exposed a fallback hazard: with a prequant source
present, the dense-fit preflight used to be skipped entirely, so a failed
prequant download would fall through to the dense bf16 load the memory
plan never budgeted, OOMing after eviction. The preflight now always runs
and gates an allow_dense_fallback flag through _load_dense_quant_pipeline:
a dense misfit still skips the fast path when no prequant exists, but with
one it proceeds and a prequant failure raises to the GGUF build instead of
loading dense. The same flag is set when the auto-policy replans an
offloaded GGUF against a prequant-sized transient.

Tests updated to the new filename convention plus new coverage for the
derivation and the legacy-name fallback; the prequant-skips-refit test now
asserts the re-check runs and forbids the dense fallback. Verified end to
end on GPU: z-image int8 resolves the hosted repo, downloads the
model-name file and renders (6.8s load, 5.9 GB peak).
2026-07-17 05:48:23 +00:00
pre-commit-ci[bot]
9ac6bbe67d [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2026-07-16 09:29:04 +00:00
Daniel Han
62264c8b11 fp8 DiT quant: floor the dynamic activation scale with activation_value_lb
An all-zero activation token row makes the dynamic per-row fp8 scale 0,
which turns the quantized data to NaN and the render to black frames on
torchao's plain-torch kernel path. The fused fbgemm/mslk quantize kernels
clamp zero rows internally, so the bug only reproduces on machines without
them, which is most user environments. Zero rows are real inputs, not a
corner case: Wan 2.2 zero-pads its text conditioning, and Hunyuan-1.5 and
Qwen-Image regenerate zero rows inside their transformer blocks every step.

Pass activation_value_lb=1e-12 to Float8DynamicActivationFloat8WeightConfig
whenever the installed torchao supports the kwarg (Float8Tensor rework,
0.13+), checked via inspect.signature so older torchao keeps exactly the
current behaviour; the existing Float8MMConfig fallback chain is unchanged.
Verified on GPU: with the forced plain-torch kernel path a zero-row input
NaNs without the floor and stays finite with it, and end to end on
HunyuanVideo-1.5 fp8 goes from a solid black frame (LPIPS 1.00) to a normal
render (LPIPS 0.225); on Wan the floor matches the condition_embedder
exclusion (LPIPS 0.211 vs 0.206). This is defense in depth on top of the
family excludes and deny list, which stay as-is: it changes the failure
mode of any future zero-row family from black frames to graceful
degradation. Same-seed renders with fused kernels present are unaffected,
and pre-quantized fp8 checkpoints stay valid since weight scales are
untouched.
2026-07-16 09:28:11 +00:00
Daniel Han
7e1bc5cbde Tighten fault-path comments added by the video/diffusion hardening pass 2026-07-13 17:33:01 +00:00
Daniel Han
133f6fecf7 Studio: harden video/diffusion cache, attention, and CFG-parallel fault paths
- diffusion_attention: arch-gate FlashAttention 2 to Ampere (SM80)+ in both the
  primary selector and the heterogeneous-replica guard (it crashed on pre-Ampere).
- diffusion_cfg_parallel: convert boolean attn masks to additive bias before the direct
  cuDNN op so partial masks match F.scaled_dot_product_attention; make proxy disable_cache
  transactional (clean both branches, mark broken, surface a reload-required error).
- diffusion_cache: fail closed when a magcache step-count resize or below-threshold
  disable cannot remove the old cache; surface a failed enable+cleanup instead of a false
  uncached None.
- video: roll back earlier experts when a later expert raises in the all-or-none step-cache
  loop; fail the load when the primary-only cache cannot be re-engaged through the
  CFG-parallel proxy; validate transformer_cache_quality and cfg_parallel before the worker.
- scripts: place the fp8 ablation pipeline on CUDA; fail closed on a failed magcache resize
  in the speedmem bench; label OOM distinctly in the SDPA mask probe.
- tests: regressions for the FA2 arch gate, transactional proxy disable, all-or-none
  exception rollback, magcache fail-closed transitions, and enable+cleanup failure.
2026-07-13 09:46:17 +00:00
Daniel Han
9eadfc822d Merge remote-tracking branch 'origin/image-generation' into r7021
# Conflicts:
#	studio/backend/routes/inference.py
#	studio/backend/routes/video.py
2026-07-13 09:28:04 +00:00
Daniel Han
b127256eb4 Studio: surface local pipeline scan roots, tag single-file checkpoints by filename, mark companion-only pipelines partial
- _scan_models_dir: admit a scan folder that is itself a diffusers pipeline
  (root model_index.json, weights in transformer/ vae/ subdirs). _is_model_directory
  rejects such a root, so the child scan would list the component subdirs as bogus
  models and hide the real pipeline; treat the root as one model via _local_pipeline_index.

- _local_is_diffusers / _local_model_task: include the sole checkpoint filename in the
  family-detection needles (_local_family_needles, resolved via resolve_local_single_file).
  A generically named folder holding one loadable qwen-image-*.safetensors / ltx-*.safetensors
  identifies its family only from the filename; the load route already resolves that file, so
  tag it or the task-scoped picker (which rejects task=null) hides the on-device model.

- list_cached_models: mark a companion-only base snapshot partial. A GGUF image load prefetches
  the base repo's VAE / text-encoder / model_index.json but skips the transformer (the GGUF
  supplies it); the snapshot has a pipeline manifest yet is not a loadable BF16 pipeline, and
  _cached_repo_partial misses it. _repo_pipeline_missing_denoiser flags a pipeline snapshot whose
  transformer/ or unet/ component carries no weight, so the picker drops it instead of advertising
  it as fully on-device.
2026-07-13 08:26:43 +00:00
pre-commit-ci[bot]
11330b8de5 [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2026-07-13 06:31:51 +00:00
Daniel Han
899465ed80 Studio: close arbiter load-registration race and surface native progress + local pipeline folders
Publish native sd.cpp generate progress (_gen) before LoRA resolution so a reload probe reads active during setup, matching the diffusers path.

Register the diffusion/video GPU load under the arbiter lock (acquire_for now takes a register callback) so a competing acquire cannot evict an owner before its load is marked in-flight and let two loaders allocate VRAM at once.

Admit local diffusers pipeline folders (root model_index.json, weights in component subdirs) in the local model scan so they reach task tagging and the On Device picker.
2026-07-13 06:31:00 +00:00
Daniel Han
e0ef488f47 Tighten comments and docstrings added by the image-generation fixes 2026-07-13 05:29:09 +00:00
Daniel Han
dabb12e198 Studio: gate dataset uploads on the symlink check and surface local video single-file checkpoints 2026-07-13 04:21:36 +00:00
Daniel Han
a3388820a2 Reflow verbose diffusion comments to fewer lines 2026-07-13 03:30:48 +00:00
Daniel Han
da1770bb44 Refuse sd.cpp install into unowned non-empty target dir
When the install target already exists, is non-empty and lacks the
.unsloth-studio-owned marker (a user's own stable-diffusion.cpp checkout,
or unrelated files beside a custom Studio root), install() previously still
extracted the release into it. Skipping the ownership marker only stopped the
uninstaller from deleting the directory; extraction still merged binaries into
the user's working tree and could overwrite same-named files.

Fail up front with a clear message pointing the user at a fresh/empty location
before any download or extraction, leaving their directory untouched. Update
the ownership test suite to assert the refusal.
2026-07-13 03:17:49 +00:00
Daniel Han
31c8ab776a Register static video compile-cache shapes before save
Under speed=max the video DiT compiles dynamic=False, so inductor
artifacts are per (width, height, frames). The save path never called
compile_cache.register_shape, so after a bundle hit ctx.saved stayed
true and later resolutions/frame counts never re-dirtied the bundle,
leaving those shapes to recompile on every restart. Register the actual
generation shape before saving, gated on the static tier, mirroring the
image backend.
2026-07-13 02:12:13 +00:00
Daniel Han
5a17614b51 Reject native batch seeds outside the JSON-safe range 2026-07-13 02:09:15 +00:00
pre-commit-ci[bot]
30d3cd91dc [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2026-07-13 01:31:35 +00:00
Daniel Han
6e2e8c846c Harden video diffusion cache, CFG-parallel replica, and layerwise-fp8 rollback
- diffusion_attention: clear the HunyuanVideo-1.5 null-mask flag with an always_call
  post-hook so it is scoped to one hooked forward and never latches across an
  exception; add attention_backend_supported_on_device to arch-gate an
  already-resolved backend on a specific (heterogeneous) CUDA device.
- video: make the explicit MagCache resize transactional via _step_cache_all_or_none
  (refuse to stack a fresh cache over one that could not be disabled; roll a mixed
  resize back and report the true state); raise on a failed all-or-none rollback
  instead of falsely reporting an uncached pipeline.
- diffusion_cfg_parallel: re-validate the attention backend on the replica device
  and pin native there when unsupported; mirror the primary's max tier on the
  replica (max-autotune compile + direct QKV fusion) via a new speed_mode arg;
  prefer a viable heterogeneous secondary GPU over an unusable identical one; clear
  the const cache at each plan_generation.
- diffusion_vae_quant / diffusion_precision: detect a partial diffusers
  layerwise-fp8 mutation (leftover casting hooks the torchao detector cannot see)
  and fail the load closed, while a clean failure still falls back to dense.
- video_speedmem_bench: engage the dual-expert cache all-or-none like the loader.
- frontend video api: add text_encoder_quant / vae_quant and the auto/off literals
  to VideoLoadRequest so typed callers match the backend contract.
2026-07-13 01:30:22 +00:00
pre-commit-ci[bot]
499bada598 [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2026-07-13 01:23:46 +00:00
Daniel Han
7ec9e77a5d Studio: fix diffusion install ownership, dataset upload atomicity, gallery pagination, and teardown races
install_sd_cpp_prebuilt: only write the .unsloth-studio-owned marker when the
install created the target directory or it was empty. Adopting a pre-existing,
unowned, non-empty directory (a user's own stable-diffusion.cpp checkout) made
it eligible for the uninstaller's recursive delete.

routes/training upload: make the multi-file promotion transactional. Back up
each displaced original and roll every destination back on any failure, so a
mid-loop rename error can no longer partially overwrite the live dataset.

routes/training _resolve_dataset_folder: reject a symlinked dataset directory
and prove the resolved folder stays under the datasets root, so image
read/caption/delete cannot escape the root through a link.

routes/training delete: escape glob metacharacters in the thumbnail filename so
deleting an image named like [ab].png removes only its own thumbnails.

image_gallery / video_gallery listing: filter records against the response
schema inside the pager via a valid callback, so offset/limit/has_more all count
over accepted records. A leading schema-invalid record no longer returns an
empty page with has_more=true and stalls infinite scroll at offset 0.

image_gallery / video_gallery save: publish via a temp file plus atomic rename
(the sidecar is the video pair's commit marker) and clean up on failure, so a
partial write never surfaces a truncated PNG or strands an orphan MP4.

diffusion_train_common discovery: treat an empty caption sidecar as a metadata
tombstone that still falls through to the dreambooth instance prompt, so
clearing every metadata caption no longer fails with no captioned images found.

diffusion backend unload: wait for an in-flight denoise to exit before tearing
down process-wide patches and state, mirroring the load path.

diffusion_engine_router: serialize the whole check/unload/publish transition so
a concurrent selection cannot return the engine being unloaded.

uninstall.ps1: gate the default sd.cpp process stop on the owner marker so a
user's own sd-server is not terminated for a directory we then keep.
2026-07-13 01:22:54 +00:00
Daniel Han
e727b9e82f Merge branch 'image-generation' into video-diffusion-improvements 2026-07-12 13:59:40 +00:00
Daniel Han
21052db120 Publish image generation active state before pre-denoise setup
generate() assigned self._gen only at the pipe() call, after deferred
compile, LoRA resolution/application, and ControlNet download/build had
run. Across that setup window generate_progress() reported inactive even
though _generate_lock was held, so a reloaded page's mount probe showed
idle and let a second generate queue behind the first.

Publish an active step-0 _GenState the moment the generation lock is
acquired, before the setup work, and clear it in the outer finally so a
setup-time error cannot leave the UI stuck active. Mirrors the video
backend's queued phase and the training start guard.
2026-07-12 13:42:24 +00:00
pre-commit-ci[bot]
7776148463 [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2026-07-12 12:44:19 +00:00
Daniel Han
23a71b1c2c Close video single-file, training reservation, and image mount-resume gaps
Route on-device single-checkpoint video folders through the single_file loader:
a bare local .safetensors directory (no model_index.json) is advertised as a
pipeline with no filename, so validation rejected it before it could load.
Reinterpret the pick as a single_file load of the sole checkpoint, mirroring the
image load route.

Treat a reserved-but-not-yet-spawned LLM training start as active in
is_training_active() so /images/load, /video/load, and /diffusion/start cannot
race the reserved run for VRAM during the pre-spawn free window. Mirrors the
diffusion training service reservation.

Resume an in-flight image generation on the Images page mount: probe
generate-progress, re-enter the poll loop, and refresh the gallery on completion
so a run started elsewhere is reflected and its saved image appears without a
manual refresh. Seed resident image defaults from the resolved base_repo rather
than a possibly path-shaped repo_id so the first resident generation uses the
right recipe.
2026-07-12 12:41:21 +00:00
Daniel Han
4cc35aa87a Tighten comments in the image stack tests and scripts 2026-07-12 12:21:14 +00:00
Daniel Han
0017674664 Tighten comments in the image and video frontend 2026-07-12 12:15:36 +00:00
Daniel Han
ecae46cbfb Tighten comments in the diffusion training core and API models 2026-07-12 12:06:44 +00:00
Daniel Han
3a6a007038 Tighten comments in the image stack routes 2026-07-12 12:00:44 +00:00
Daniel Han
1762dab12f Tighten comments across the remaining image stack files 2026-07-12 11:46:23 +00:00
Daniel Han
9fde4b9991 Tighten comments across the remaining image stack files 2026-07-12 11:40:05 +00:00
Daniel Han
226363c7f0 Tighten comments across the video speed stack 2026-07-12 10:59:41 +00:00
Daniel Han
892696733e Tighten comments across the image generation stack 2026-07-12 10:55:39 +00:00
Daniel Han
dcd6666f8d Thread the task-scoped GGUF fit budget through every picker expander
The previous commit introduced the single-device budget but wired it
into only the downloaded-group and Hub search sites. The LM Studio,
custom-folder, local-dir, live-search and exported-GGUF expanders
reachable from the Images/Video pickers still measured against the
summed multi-GPU total, as did the size-based GGUF row badge, so those
paths could still recommend a quant that OOMs on a single device. All
GgufVariantExpander call sites in HubModelPicker now share
expanderGpuGb, and the row badge derives the same task-scoped budget.
2026-07-11 18:36:31 +00:00
Daniel Han
cece544fd9 Fix stale LoRA closure, variant fit budget, and hidden dataset remove button
The generate callback omitted loraCapable from its dependencies, so when
an auto-compile flips supports_lora off mid-session the memoized handler
still sent the previously selected adapters and the next generation
failed with the backend's LoRA-not-supported error instead of omitting
adapters the UI had already hidden.

HubModelPicker's GGUF variant expanders, format lists and Hub row fit
hints measured against the summed multi-GPU total. When the picker is
task-scoped (Images/Video) the loaders place the whole pipeline on one
device, so a variant could be recommended as fitting and then OOM at
load; those sites now share the single-device budget the group fit gate
already uses, while chat pickers keep the summed total since llama.cpp
splits layers across devices.

The dataset labeling grid's Remove button relied on group-hover with no
group parent, leaving it permanently invisible to mouse users; the image
wrapper now carries the group class.
2026-07-11 17:32:05 +00:00
pre-commit-ci[bot]
668088436e [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2026-07-11 16:51:12 +00:00
Daniel Han
e1fa4fec04 Studio diffusion: fix static compile shape registration and prequant path validation
Register the dims the forward actually compiled with: image-conditioned
workflows (img2img, inpaint, upscale, edit) run at the input image's size,
not the slider's, so recording the slider values marked never-compiled
shapes as covered and warm restarts kept paying compile for the real one.

Validate a request-supplied transformer_prequant_path (existence plus the
UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH allowlist) before treating prequant as
available at the resident-fit re-check: an unusable path skipped the dense
fit check up front and then fell back to materializing dense bf16 after
the previous pipeline was evicted, recreating the post-eviction OOM path.
Shared as usable_prequant_source, also used by the auto-policy planner.
2026-07-11 16:50:19 +00:00
Daniel Han
316c8b5a81 Studio video: resume a background generation on page mount
A generation runs on a backend daemon thread and survives a page reload,
but the mount effect only probed the load progress, so reloading during a
generate showed an idle page that never picked up the finished clip until
a manual refresh. Hoist the generate poll loop out of handleGenerate and
re-enter it on mount when generate-progress reports an active job; merge a
terminal completed record into the gallery to cover the race with the
mount gallery fetch.
2026-07-11 16:50:19 +00:00
Daniel Han
57f08ebfb2 Merge remote-tracking branch 'origin/main' into ig_merge 2026-07-11 15:15:54 +00:00
Daniel Han
6412efd7d9
Studio: auto-detect completion masking markers, stop silent full-sequence training (#7054)
* Auto-detect completion masking markers with template table fallback

Studio's train_on_completions previously relied only on the hardcoded
MODEL_TO_TEMPLATE_MAPPER / TEMPLATE_TO_RESPONSES_MAPPER tables and
silently disabled masking when a model was not in the table, so unmapped
models (LFM2-8B-A1B, DeepSeek, and others) trained on full sequences
without telling the user. Several mapped templates (glm, mistral, llama,
starling, zephyr, qwen3-thinking) also carried markers that mask every
assistant token, which made every row drop in the post-masking filter.

Both training callsites (CUDA trainer.py and MLX worker.py) now share
utils.datasets.completion_masking.apply_completion_masking:

- Try unsloth_zoo chat template auto-detection first; it raises loudly
  when the template cannot be parsed and never masks the EOS token.
- gpt-oss models keep their manual markers so non-final assistant
  <|end|> tokens stay trained, matching current behavior.
- If auto-detection raises, fall back to the template table exactly as
  before.
- If the table also misses, emit an explicit user-visible warning that
  completion masking could not be applied and full-sequence training
  will occur, instead of a quiet log line.

The >30 percent dropped-rows safety net in trainer.py now guards the
auto path as well. Table consumers for inference and chat templates are
unchanged. Validated against one representative tokenizer for every
template in TEMPLATE_TO_RESPONSES_MAPPER plus the unmapped models:
no template regresses; unit tests cover the four decision paths.

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

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

* Restrict masking fallback to marker detection failures

The auto branch wrapped the whole train_on_responses_only call, so a real
failure while applying the masking (dataset map, tokenization) was treated
as a detection miss and training silently proceeded on full sequences.
Detect markers separately via get_chat_template_parts (test seam via
detect_fn), then apply them with errors propagating, matching the manual
path. Tokenizers with preset unsloth marker attrs skip detection and call
bare so zoo reuses the stored parts.

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

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

* Fail the run when applying completion masking raises

The helper already falls back internally on detection failures and returns
applied=False on a double miss, so an exception reaching the callsites is a
real failure applying the masking. Remove the callsite catches that
downgraded it to full-sequence training; the run now fails visibly instead.

Also use the explicit re-export alias form in utils/datasets/__init__.py for
the two new names, satisfying the import-hoist source lint.

* Import completion masking from its submodule

The import-hoist source lint counts only real name loads, so package-level
re-exports of the two new names cannot satisfy it. Import
apply_completion_masking from utils.datasets.completion_masking directly at
both callsites and leave utils/datasets/__init__.py untouched.

* Completion masking: gpt-oss renames and MLX raw/alpaca parity

Renamed or private gpt-oss checkpoints are name-detected as gpt-oss but miss
the exact-name table; default them to the gpt-oss template markers instead of
falling through to full-sequence training.

Gate the MLX masking call on not raw_text_mode and format_type != alpaca,
mirroring the CUDA path: raw/CPT text has no chat turns to mask and
Alpaca-rendered text lacks the tokenizer's chat markers.

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

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

* Define raw_text_mode outside the MLX feature-detect block

With an older zoo lacking the append_eos config field, the masking
gate referenced raw_text_mode before assignment. Hoist the assignment
above the feature detection so both consumers see it.

* Gate MLX masking on the formatter's resolved format

format_type auto can resolve to alpaca or raw text; the masking skip
checked only the requested value, so auto-detected Alpaca data got
chat-template markers applied to rendered prompt text. Track the
final_format returned by format_and_template_dataset and gate on it,
matching the CUDA path.

* Unwrap the mlx-lm TokenizerWrapper before marker checks

The wrapper delegates plain reads to the wrapped HF tokenizer but hides
underscore attrs, so preset unsloth markers were invisible and detection
relied on the loader's call patch. Unwrap to the real tokenizer first,
as the zoo MLX resolver does.

* Tighten masking comments

* gpt-oss: auto-detect markers first like every other template

The quantized and BF16 gpt-oss checkpoints ship a chat template without
the channel final header, so the pinned manual markers match nothing
there and masking trained zero tokens. Auto-detection derives markers
from whichever template the checkpoint ships and keeps the final
terminator trained; the manual gpt-oss markers remain the detection
failure fallback, including for renamed checkpoints.

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

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

* Tighten comments

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-07-11 05:13:45 -07:00
Daniel Han
97161c89d6
Studio: route models by CONFIG_MAPPING_NAMES instead of hardcoded tables (#7043)
* Studio: route models by CONFIG_MAPPING_NAMES instead of hardcoded tables

A model whose model_type is absent from an overlay's transformers cannot load
there, so a new MoE arch not yet in the tier tables gets routed to default and
fails (e.g. lfm2_moe, deepseek_v4). Add a static resolver that parses each
overlay's CONFIG_MAPPING_NAMES straight from source (AST only, no import, no
network, no trust_remote_code) and picks the lowest tier that ships the
model_type. Runs after the existing checks and only ever upgrades default, so
no existing routing changes and new archs no longer need a table edit.

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

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

* Studio router: harden the CONFIG_MAPPING_NAMES resolver

- Resolve the default tier map from the base install, skipping any .venv_t5_*
  sidecar on sys.path, so an in-process 5.x activation cannot make a 5.x-only
  model look loadable by 4.x.
- Do not cache an overlay whose sidecar dir is absent, so a later call re-reads
  it once provisioned instead of serving a stale empty map.
- Also collect model types added via CONFIG_MAPPING_NAMES.update({...}) and
  **{...} unpacking, not just the literal assignment (5.10 uses both).
- Wrap the AST walk in the try/except so a malformed source can never crash tier
  resolution.
- Feed the mapping fallback from _load_config_json so a config served from the
  hub cache during a transient outage still routes new architectures.

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

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

* Tighten comments

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-07-11 05:12:37 -07:00
Daniel Han
c3feac6160
Studio: route lfm2_moe (LFM2-8B-A1B) to transformers 5.3.0 (#7040)
LFM2-8B-A1B and any other lfm2_moe checkpoint were missing from the
transformers tier tables, so they fell through to the default 4.57.x
sidecar, which does not register lfm2_moe and errors with
"not supported yet in transformers==4.57.6". Only lfm2_vl was listed.

Add Lfm2MoeForCausalLM / lfm2_moe to the 5.3.0 tier (lfm2_moe is
registered in transformers 5.3.0). get_transformers_tier now returns
530 for LFM2-8B-A1B and the model loads and trains as expected.
2026-07-11 05:08:07 -07:00
pre-commit-ci[bot]
a3593ba78a [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2026-07-11 10:09:03 +00:00