Commit graph

7,228 commits

Author SHA1 Message Date
Daniel Han
bb78cca820 Correct the ltx-2 resident TE estimate to the bf16 cast size
The memory plan's bf16_components_gb held 50.4 GB for the LTX text
encoder, which is the fp32 hub store of Gemma3-12B (~49 GB download),
not what sits on device: the pipeline loads it torch_dtype=bf16, ~24.4
GB resident. The 26 GB over-estimate pushed the auto plan toward offload
on cards that fit the real footprint. Comments and the size-table test
now pin the resident semantics.
2026-07-18 08:00:25 +00:00
Daniel Han
a440ee748a Extend the fp8 TE quant to HiDream's Llama text_encoder_4
The generic quantize_text_encoders pass only covers text_encoder.._3, so
HiDream's HEAVIEST encoder (Llama-3.1-8B TE4, 16.1 GB bf16) always stayed
dense. TE4 is assembled separately (hidream_te4_kwargs), so the fp8 path
now lives there: when the requested TE quant is layerwise fp8 and the
device/family qualify, TE4 prefers the hosted pre-cast checkpoint
(unsloth/HiDream-I1-Full-FP8, 8.6 GB) and falls back to dense-load-then-
cast; a mid-pass cast failure reloads a fresh dense encoder instead of
shipping partial state. The pre-cast loader and builder gain
config_subfolder/config_overrides for standalone encoder repos whose
config sits at the root and whose pipeline needs forward flags
(output_hidden_states/attentions).

Verified on B200: bit-identity 291 tensors (225 fp8, 0 mismatches),
hosted checkpoint engages through the real backend (marker + status fp8),
load 24.3 s vs 48.0 s dense, LPIPS 0.133 mean over 3 same-seed pairs vs
the dense-TE render (gate 0.25), non-black frames.
2026-07-18 07:55:24 +00:00
Daniel Han
569220fc9a Key the fp8 cast idempotency on an explicit completion marker
Hook presence alone cannot distinguish a legitimately pre-cast text
encoder from leftover hooks after a cast that failed mid-pass, so the
early return now requires the completion marker _cast_fp8 sets once the
hooks are fully installed. Leftover partial state keeps failing closed.
Also tolerates non-Module encoder doubles in the hook probe and the
dtype override.
2026-07-18 07:27:05 +00:00
Daniel Han
bdf676cbe5 Report the compute dtype on fp8-cast encoders and inject the pre-cast TE on the dense fast path
Two more findings from the hosted-TE GPU smokes:

- Module.dtype reports the first floating parameter, which after the
  layerwise fp8 cast is the fp8 STORAGE dtype. Flux2 derives its prompt
  embed and latent dtypes from encoder.dtype and feeds them to
  randn_tensor, which has no fp8 kernel, so ANY flux.2 load with
  text_encoder_quant=fp8 crashed at generation (pre-existing, runtime
  cast included). The cast now swaps in a subclass whose dtype property
  reports the compute dtype; forward behaviour is unchanged.
- The dense transformer_quant fast path assembles companions through
  _assemble_pipe, which never received the pre-cast TE injection, so the
  hosted encoder engaged on full-pipeline and GGUF builds but not on the
  fast path. Threaded through like the other two branches.

Verified live on B200: qwen-image (full pipeline), flux.2-dev (GGUF picker
with int8 DiT prequant), ltx-2 (video backend) all engage the hosted TE,
render non-black, and report text_encoder_quant=fp8 truthfully.
2026-07-18 07:19:03 +00:00
Daniel Han
99486d8f6e Wire the hosted pre-cast fp8 text encoders
qwen-image and flux.2-dev (diffusion) and ltx-2 (video) now resolve a
hosted pre-cast fp8 text encoder from their unsloth -FP8 repos:

- unsloth/Qwen-Image-FP8: Qwen2.5-VL-7B, 16.6 GB dense -> 8.8 GB
- unsloth/FLUX.2-dev-FP8: Mistral-Small-24B, 48.0 GB dense -> 24.7 GB
- unsloth/LTX-2-FP8: Gemma3-12B, 48.7 GB fp32 store -> 13.2 GB

Every checkpoint verified bit-identical to dense-load-then-cast
(729 / 585 / 1066 tensors, zero mismatches) and smoke-tested through the
real backends with the repo engagement marker. Tests cover the wired
entries, the resolver filenames, builder metadata weights_only survival,
and the idempotent re-cast.
2026-07-18 06:52:46 +00:00
Daniel Han
645f3858d1 Fix pre-cast TE checkpoint loading and engagement reporting
Two bugs found while building the hosted checkpoints:

- The builder recorded torch.__version__ (a TorchVersion object) in the
  checkpoint metadata, so torch.load(weights_only=True) rejected every
  artifact and the loader silently fell back to the dense download.
  Record plain strings.
- Re-applying the layerwise fp8 cast to an injected pre-cast encoder
  raised on the duplicate hook registration, making quantize_text_encoders
  report the engaged cast as failed (status showed no TE quant while the
  encoder ran fp8). _cast_fp8 now returns early when the hooks are
  already installed.

Also corrects the LTX TE size note: Gemma3-12B stored fp32 (~49 GB), not 27B.
2026-07-18 06:52:46 +00:00
Daniel Han
52af944bdc Test the pre-cast text-encoder load path
Hermetic CPU coverage for diffusion_te_prequant: the checkpoint
filename convention, family-table resolution by scheme and component
with malformed entries skipped, resolution priority (path override,
hosted repo, none) and the fp8-only scheme gate, the checkpoint
validation matrix (wrong format, missing state_dict, wrong scheme,
wrong component, wrong or missing base_model_id) with base case
folding, the local-path allowlist refusal and missing-file fallback,
and the assembly injection gating (mode, hosted entry, device support,
family deny, load failure, successful injection). Also pins the
te_prequant_repos field on both family dataclasses and that no family
ships a hosted TE checkpoint until the campaign wires one.
2026-07-18 06:24:48 +00:00
Daniel Han
4e9aab60ae Add the pre-cast text-encoder checkpoint builder
Applies the runtime layerwise fp8 storage cast to a model's dense text
encoder once and saves the cast state dict with baked metadata (format
tag, base_model_id, family, scheme, component, te_class, versions) in
the layout diffusion_te_prequant.py validates. Resolves the encoder
class from the checkpoint's config.architectures so the recorded
te_class matches what the pipeline instantiates. CPU-runnable: the cast
touches storage dtypes only.
2026-07-18 06:24:48 +00:00
Daniel Han
bbc8c232c3 Inject hosted pre-cast text encoders during pipeline assembly
Wire te_prequant_pipe_kwargs into the three pipeline assembly sites:
the diffusion full-pipeline branch, the diffusion transformer-only and
GGUF branch (where the companion TE is the big remaining download), and
the shared video assembly path before the pipeline/component split.

Injection is gated exactly like the runtime cast (mode normalized to
fp8, device supported, family not denied), so it can never engage where
quantize_text_encoders would not; the later quantize_text_encoders call
re-applies the cast idempotently and keeps status reporting truthful.
With no hosted checkpoint configured the call returns {} and assembly
loads the dense encoder as before.
2026-07-18 06:24:40 +00:00
Daniel Han
79b4bfff1d Add a pre-cast text-encoder loader for the layerwise fp8 scheme
The runtime text_encoder_quant=fp8 path downloads the full bf16 text
encoder and layerwise-casts it in place on every fresh load. For the
heavyweight encoders (LTX's Gemma3-27B ~50 GB, FLUX.2-dev's Mistral-24B
~48 GB, Qwen-Image's Qwen2.5-VL ~16.6 GB) that download dominates load
time on a fresh machine.

diffusion_te_prequant.py loads a pre-cast fp8-storage state dict
instead: meta-init the encoder skeleton from the checkpoint's te_class,
load_state_dict(assign=True), rebuild on CPU if non-persistent buffers
stay on meta, then re-apply the same layerwise cast to install the
upcast hooks. The cast is a deterministic storage transform, so the
loaded encoder is bit-identical to dense-load-then-cast by construction.

v1 hosts the layerwise fp8 storage scheme only: its state dict is plain
tensors (torch.load(weights_only=True), no pickle execution). The
dynamic-compute schemes (fp8_dynamic, int8, nvfp4) build torchao
subclass wrappers at runtime and are deliberately not hosted.

Checkpoints validate format, scheme, component and base_model_id before
use and any problem falls back to the dense download and cast. Local
path overrides reuse the DiT prequant allowlist env var. Families opt in
via a new te_prequant_repos (scheme, component, repo_id) field on both
DiffusionFamily and VideoFamily; the field defaults empty so nothing
changes until a gate-validated artifact is wired.
2026-07-18 06:24:34 +00:00
Daniel Han
930d29be98 Wire the hosted HiDream I1 int8/fp8 checkpoints
Gate-validated: all 28 per-case pairs pass per scheme (LPIPS suite means 0.291
int8 / 0.278 fp8, in the 50-step trajectory-divergence band; CLIP delta means
0.007-0.008), and the int8 checkpoint is verified bit-identical to on-the-fly
quantize across all 1615 state dict tensors (1073 quantized, max abs diff 0.0).
Uploaded to unsloth/HiDream-I1-Full-FP8.
2026-07-18 05:19:36 +00:00
Daniel Han
400c950eee Fix video progress under-reporting during load and generate
Two live-test findings on the video progress endpoints:

- load-progress downloaded_bytes froze mid-download: the counter used
  scan_cache_dir, which skips in-flight *.incomplete blobs, so it sat at the
  last completed blob for the whole multi-GB shard pull while the disk kept
  filling. Count the repo's cache directory directly (completed plus incomplete
  blobs, snapshot symlinks skipped so nothing is double-counted).
- generate-progress reported total_steps=null / fraction=0 while step advanced:
  the video API only carried the native total field while the image API exposes
  total_steps and fraction, so one poller could not work against both. Derive
  the image-compatible aliases in generate_progress and declare them on the
  response model; the native total stays for back-compat.
2026-07-18 03:49:10 +00:00
Daniel Han
0d866998f0 Fix silent LoRA drop and wasted transformer prefetch on GGUF quant loads
Two live-test findings on the images load path:

- transformer_quant with baked LoRAs, when the dense quantized build is
  declined for memory or fails: the load completed as a plain GGUF with the
  adapters silently dropped (HTTP success, supports_lora=false after the
  fact) -- wrong output with no signal. The load now fails with the recovery
  options (drop the adapters, free VRAM, or pick a smaller model). Weight-0
  adapters still count as no bake request, and the plain no-LoRA decline
  keeps its silent GGUF fallback.
- A fresh GGUF load on a small GPU prefetched the base repo's full bf16
  transformer shards (~47 GB on Qwen-Image) because the dense-quant prefetch
  widening only checked scheme viability, not whether the device could ever
  hold the candidate resident. Gate the widening on total device capacity
  (reserve + 0.85 margin, the plan_fits_total_capacity bar) so a card that is
  certain to decline the dense build never pays the download; capable devices
  keep the prefetch.
2026-07-18 03:49:05 +00:00
Daniel Han
361cbbb287 Wire the hosted HunyuanImage 2.1 int8/fp8 checkpoints
Verified bit-identical to on-the-fly quantize: all 1264 state dict tensors
(456 quantized) dequantize equal between the loaded checkpoint and a fresh
quantize_ pass, so quality matches the runtime Dtype path exactly. Same-seed
LPIPS suite means (0.35 int8 / 0.28 fp8) blend trajectory divergence with this
family's own run-to-run nondeterminism (identical weights and seed reproduce a
17/255 mean pixel delta through the 50-step guider pipeline); per-case hard
checks pass and the drift is compositional, reviewed visually. Uploaded to
unsloth/HunyuanImage-2.1-FP8.
2026-07-17 23:59:07 +00:00
Daniel Han
958645bb6b Pin the measured HiDream quant verdict in tests
int8 and fp8 both engage and render cleanly on this family, including short
prompts on int8: the routed MoE expert Linears only ever see the concatenated
image+text stream (M far above the torch._int_mm minimum), so no deny entry
and no family exclude tokens are warranted. Assert that so a future table edit
cannot silently regress the measured behavior.
2026-07-17 23:13:58 +00:00
Daniel Han
9174535616 Surface HiDream I1 in the image model catalog
One catalog group with the three official bf16 pipelines (Full, plus the Dev
and Fast distillations as labeled artifacts) at their ~63 GB resident size, so
auto-routing keeps this a datacenter-GPU pick. city96's GGUF is deliberately
not wired: the GGUF path would need the same Llama TE4 assembly for very small
demand. Images-page defaults mirror the backend table with the variant keys
ahead of the generic hidream key.
2026-07-17 13:24:57 +00:00
Daniel Han
e098ee52ea Add the HiDream-I1 family to the image backend
A 17B MoE DiT (16 double + 32 single layers, 4 routed experts) with four text
encoders, on HiDreamImagePipeline (diffusers 0.39). One family covers the open
Full / Dev / Fast repos (same arch); per-variant generation defaults follow the
upstream inference recipes (Full 50 steps at guidance 5, the distilled Dev 28
and Fast 16 guidance-free).

The repos name a Llama-3.1-8B text_encoder_4 in their model_index but do not
ship its weights; the official example passes the gated meta-llama repo in by
hand. The loader instead assembles the component from the open unsloth mirror
(byte-identical weights, already inside the non-GGUF trust gate), injected at
the three pipeline from_pretrained sites, with output_hidden_states matching
the official example. Memory planning counts the assembled TE4: 34.2 GB DiT +
28.8 GB encoders, ~63 GB bf16-resident.
2026-07-17 13:24:57 +00:00
Daniel Han
5887d2489c Surface HunyuanImage 2.1 in the image model catalog
Catalog group with the open bf16 mirror pipeline (~50 GB resident, so a
bare click on a consumer card routes to the QuantStack GGUF quants, which
load and render through the generic GGUF path, verified live) plus the
images page defaults (50 steps, guidance 3.25 feeding
distilled_guidance_scale).
2026-07-17 12:52:45 +00:00
Daniel Han
c2743eb697 Add the HunyuanImage 2.1 family to the image backend
The hunyuanvideo-community diffusers mirror carries the full stack in
standard layout: a 17B dual-stream DiT (32.5 GB bf16), a Qwen2.5-VL text
encoder, a ByT5 glyph encoder, the 32x HunyuanImage VAE, and
guider/ocr_guider components (AdaptiveProjectedMixGuidance) that diffusers
0.39 loads natively, so the generic from_pretrained pipeline path covers
everything with no per-component assembly.

Family notes:
- The call's guidance knob is distilled_guidance_scale (there is no
  guidance_scale kwarg), so cfg_kwarg routes the UI value there; real CFG
  runs inside the repo's guider at its baked scale. Defaults follow the
  card recipe: 50 steps, 3.25.
- 2K-native: verified live at both 1024 and 2048.
- Coexists with the HunyuanImage-3.0 structured exclusion (3.0 has no
  diffusers pipeline and stays excluded with its stated reason).
- int8/fp8 dense quantization verified live (LPIPS 0.186 both vs same-seed
  bf16); a short prompt does not trip the int8 torch._int_mm minimum on
  this arch, so no family exclude entry is needed.
- bf16 component table for the memory planner: (32.5, 16.3, 0.8) GB.
2026-07-17 12:52:45 +00:00
Daniel Han
98673456c8 Wire the hosted Lumina Image 2.0 int8/fp8 checkpoints
Gate-validated against same-seed bf16 renders (28/28 pairs per scheme, zero
failures): int8 LPIPS mean 0.146 / SSIM 0.937, fp8 LPIPS mean 0.116 /
SSIM 0.946. Uploaded to unsloth/Lumina-Image-2.0-FP8 following the existing
checkpoint repo conventions.
2026-07-17 12:20:30 +00:00
Daniel Han
350e46bf2e 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:07:17 +00:00
Daniel Han
570cef6c79 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:07:17 +00:00
Daniel Han
08b853df15 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:49:04 +00:00
Daniel Han
362cacc448 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: 296 passed.
2026-07-17 09:37:17 +00:00
Daniel Han
07076430b3 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:01 +00:00
Daniel Han
b35985162f 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:12 +00:00
Daniel Han
5d8bf5d094 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 per-family int8 exclusions (txt_in, add_q/k/v_proj, to_add_out, txt_mlp) for
qwen-image and qwen-image-edit, threaded through exclude_tokens_for_scheme(scheme,
family) and 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:23 +00:00
Daniel Han
91be795a9b 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:24 +00:00
Daniel Han
5e1614259c 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:47:45 +00:00
Daniel Han
57c40125ad 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). 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:05 +00:00
Daniel Han
117abb320a Studio: restore Reapply target on async image/video load errors; recheck training state before dataset commit; reclaim partial sd.cpp installs on retry
Images/Video: a background model load that fails AFTER starting (error/eviction during download) leaves the previous pipeline resident, but handleLoad had already overwritten lastLoad.current with the failed pick, so "Reapply to loaded model" reloaded the failed model. Carry the prior Reapply target into the poller and restore it on the async error/null paths, mirroring the quant rollback.

Training: an in-flight diffusion dataset upload passed _require_diffusion_dataset_mutable() at entry but could still commit files after a concurrent /diffusion/start reserved the training slot, mutating the dataset underneath the trainer. Re-check the interlock immediately before the commit phase; a 409 there leaves the staged temps for the finally to clean.

sd.cpp install: an interrupted extraction (disk full, killed process, a raising post-extract cudart fetch) left the target non-empty with no owner marker, so the next lazy install tripped the "not a Studio-managed directory" refusal and wedged native install. Write the ownership marker before the partial writes when the target is reclaimable, so a retry recognises the debris as ours and re-extracts.
2026-07-13 15:27:46 +00:00
Daniel Han
90827bc05c Studio: gate gallery serve/export on ownership; keep image progress active until persisted; reserve diffusion training before the dataset scan 2026-07-13 14:40:18 +00:00
Daniel Han
53912b2f99 Studio: tighten image-generation fix comments and docstrings 2026-07-13 13:32:28 +00:00
Daniel Han
be0bd00064 Studio: gate local-pipeline image tagging on a real family; validate video sidecars before delete/clear 2026-07-13 12:28:52 +00:00
pre-commit-ci[bot]
d0f00f2f5f [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2026-07-13 11:22:23 +00:00
Daniel Han
06b543b880 Studio: track the loaded GGUF filename so native companion resolution reproduces the load identity 2026-07-13 11:21:34 +00:00
pre-commit-ci[bot]
3fa9fdbc0d [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2026-07-13 10:04:07 +00:00
Daniel Han
5eef2f4003 Studio: preserve foreign gallery files, force safetensors on remote ControlNets, and close dataset/seed/GPU gaps
Gallery clear/delete now scope to Studio-owned files: image_gallery and
video_gallery skip PNGs / MP4s without a readable recipe (a hand-dropped or
orphan file the listing already hides), so clear() and a guessed-id delete no
longer destroy files the gallery never surfaced.

Remote ControlNets now force use_safetensors: a bare owner/name reaches
from_pretrained without the base trust gate, and the Hub scan fails open when
unavailable, so requiring safetensors closes the pickle deserialization vector.

POSIX uninstall now stops resident sd-server / sd-cli under an owned sd.cpp root
before removing the tree (marker-gated), mirroring the Windows stop-before-delete
scan; a live native server no longer survives unlinking its binary.

Diffusion dataset containment: the training-start read path and the discovery
picker route bare names through the protected resolver, so a symlinked dataset
is rejected / not advertised like the caption/delete routes already do. Uploads
gain the inference decode guard (oversized real images 400 before OOMing the
trainer) and dataset upload/caption/delete/import are blocked with 409 while a
diffusion run is active.

JSONL readers (trainer + routes) tolerate non-object JSON and invalid UTF-8
instead of raising AttributeError / 500.

LoRA family compatibility is enforced in the shared resolver, not only the
picker, so a direct API client cannot apply a mismatched-family adapter.

GPU arbiter gains release_if so the image/video unload idle-check and release
are atomic against a concurrent same-owner load's registration. Native batch
recipes persist the base batch_seed and restore replays from it, so a native
batch_index>0 image no longer advances its seed twice.

FLUX.2-klein selects its sd.cpp text encoder by variant (4B -> Qwen3-4B,
9B -> Qwen3-8B) instead of the single family default.
2026-07-13 10:02:42 +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
Daniel Han
53668c9f66 Studio: include marker-owned custom sd.cpp roots in the uninstall stop scan; harden Pester install against the nuget.exe PSGallery bootstrap 2026-07-13 07:26:18 +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
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
5a17614b51 Reject native batch seeds outside the JSON-safe range 2026-07-13 02:09:15 +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
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