ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Inject hosted pre-cast text encoders during pipeline assembly

Browse source 
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.
This commit is contained in:
Daniel Han 2026-07-18 06:24:40 +00:00
commit 15b37e129b
2 changed files with 44 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

28
studio/backend/core/inference/diffusion.py
View file

@ -95,6 +95,7 @@ from .diffusion_cache import (
normalize_transformer_cache,
)
from .diffusion_precision import TE_QUANT_AUTO, normalize_te_quant, quantize_text_encoders
from .diffusion_te_prequant import te_prequant_pipe_kwargs
from .diffusion_vae_quant import VAE_QUANT_AUTO, normalize_vae_quant, quantize_vae
from .diffusion_prequant import (
load_prequantized_transformer,
@ -1462,6 +1463,20 @@ class DiffusionBackend:
# The repo names a Llama text_encoder_4 it does not ship;
# supply it from the open mirror (diffusion_hidream.py).
pipe_kwargs.update(hidream_te4_kwargs(dtype, hf_token))
# A hosted pre-cast fp8 text encoder (when the family ships one and
# the runtime cast would engage) skips the dense TE download; the
# later quantize_text_encoders re-applies the cast idempotently.
pipe_kwargs.update(
te_prequant_pipe_kwargs(
fam,
repo_id,
te_quant_mode = text_encoder_quant,
target = target,
dtype = dtype,
hf_token = hf_token,
logger = logger,
)
)
# The prefetched snapshot dir keeps from_pretrained off the hub (its
# sweep re-pulls files the scoped prefetch skipped: 24 GB per FLUX.1).
pipe = pipeline_cls.from_pretrained(
@ -1505,6 +1520,19 @@ class DiffusionBackend:
if fam.name == HIDREAM_FAMILY_NAME:
# Same Llama TE4 assembly as the full-pipeline branch above.
pipe_kwargs.update(hidream_te4_kwargs(dtype, hf_token))
# Same pre-cast TE injection as the full-pipeline branch: the GGUF
# supplies the transformer, so the companion TE is the big download.
pipe_kwargs.update(
te_prequant_pipe_kwargs(
fam,
base,
te_quant_mode = text_encoder_quant,
target = target,
dtype = dtype,
hf_token = hf_token,
logger = logger,
)
)
pipe = pipeline_cls.from_pretrained(
_base_local_dir or base, **pipe_kwargs
)

16
studio/backend/core/inference/video.py
View file

@ -1376,6 +1376,22 @@ class VideoBackend:
pipe_kwargs["torch_dtype"] = {"vae": torch.float32, "default": dtype}
if hf_token:
pipe_kwargs["token"] = hf_token
# A hosted pre-cast fp8 text encoder (when the family ships one and the runtime cast
# would engage) skips the dense TE download -- for LTX's Gemma3-27B that is the ~50 GB
# heavyweight of the load. quantize_text_encoders below re-applies the cast idempotently.
from .diffusion_te_prequant import te_prequant_pipe_kwargs
pipe_kwargs.update(
te_prequant_pipe_kwargs(
fam,
repo_id if kind == "pipeline" else base,
te_quant_mode = text_encoder_quant,
target = target,
dtype = dtype,
hf_token = hf_token,
logger = logger,
)
)
if kind == "pipeline":
# The pre-downloaded snapshot dir keeps from_pretrained off the hub (its sweep would
# also pull root checkpoints + duplicate shards); hub id when pre-download was skipped.