unsloth/studio/backend/core
Daniel Han aea138cc85 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:29:31 +00:00
..
data_recipe Studio: harden background consumer loops and streaming paths against silent UI freezes (#6653) 2026-06-26 03:31:33 -07:00
export Studio: multi-select export formats, portable FP8/INT8, GGUF LoRA, and source parity (#6767) 2026-07-03 08:25:10 -07:00
inference Add a pre-cast text-encoder loader for the layerwise fp8 scheme 2026-07-18 06:29:31 +00:00
rag Run the malware gate on the RAG embedding model before it loads (#6887) 2026-07-07 04:30:21 -07:00
training Add FLUX.2 Klein and FLUX.2-dev DiT LoRA training 2026-07-17 09:34:14 +00:00
__init__.py Reduce and tighten code comments and docstrings repo-wide (#6095) 2026-06-08 23:09:51 -07:00
_torchao_stub.py Reduce and tighten code comments and docstrings repo-wide (#6095) 2026-06-08 23:09:51 -07:00
import_guards.py Studio: self-heal unsloth namespace shadows; clearer failed-load messages (#6532) 2026-06-21 22:43:31 -07:00
tool_healing.py Studio: parse Mistral [TOOL_CALLS] and rehearsal tool-call shapes (#5704) 2026-07-06 18:52:13 -07:00