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

Merge remote-tracking branch 'origin/video-wan' into video-hunyuan-gate

Browse source
This commit is contained in:
Daniel Han 2026-07-05 12:11:34 +00:00
commit 26fb733f25
6 changed files with 129 additions and 22 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -45,7 +45,7 @@ from .diffusion_device import (
diffusion_device_target_from_torch_device,
resolve_diffusion_device_target,
)
from .diffusion_ideogram4 import load_ideogram4_pipeline
from .diffusion_ideogram4 import ideogram4_repo_is_fp8, load_ideogram4_pipeline
from .diffusion_krea2 import KREA2_FAMILY_NAME, load_krea2_pipeline
from .diffusion_memory import (
OFFLOAD_NONE,
@ -556,6 +556,17 @@ class DiffusionBackend:
f"'{fam.name}' checkpoints are whole-pipeline single files and have no GGUF "
f"transformer variant; load the .safetensors pipeline instead of a GGUF."
)
# A family that assembles MULTIPLE denoisers per-component (Ideogram 4's dual
# DiTs) has no transformer-only single-file or GGUF path: those kinds build one
# transformer and would assemble a pipeline missing its second DiT (or fail deep
# in from_pretrained). Reject them here -- before the route evicts the current
# model -- so only a full pipeline load reaches the per-component loader.
if kind in ("gguf", "single_file") and fam.pipeline_only:
raise ValueError(
f"'{fam.name}' loads only as a full diffusers pipeline (it assembles "
f"multiple transformers), not from a single-file or GGUF checkpoint; "
f"select the pipeline repo."
)
# Non-GGUF loads (a single-file safetensors transformer, or a full pipeline)
# are gated to the unsloth org or a local path -- they fetch + deserialise
# weights, so an arbitrary remote repo is rejected here, before any work.
@ -1689,7 +1700,19 @@ class DiffusionBackend:
# (the family base -- prequant repos like the bnb-4bit exports have
# different ids and really do stay compressed), plan against the larger
# of the two estimates.
if repo_id and repo_id.strip().lower() == fam.base_repo.lower():
is_narrow_base = bool(repo_id) and repo_id.strip().lower() == fam.base_repo.lower()
if (
not is_narrow_base
and fam.name == IDEOGRAM4_FAMILY_NAME
and local_repo is not None
and local_repo.is_dir()
):
# A LOCAL directory mirror of the fp8 base never string-matches base_repo,
# so detect the fp8 layout from its transformer shard headers and reserve
# the bf16 footprint too (a local nf4 mirror has no fp8 scales and stays
# compressed). Header-only read, so this stays cheap and network-free.
is_narrow_base = ideogram4_repo_is_fp8(repo_id)
if is_narrow_base:
table = family_bf16_components_gb(fam, fam.base_repo)
if table is not None and model_dense_mib is not None:
table_mib = int(sum(table) * (1000.0**3) / (1024.0 * 1024.0))

10
studio/backend/core/inference/diffusion_families.py
View file

@ -49,6 +49,13 @@ class DiffusionFamily:
# rather than ``transformer_class.from_single_file`` + a companion base repo.
# DiT families leave this False (their single file is transformer-only).
single_file_is_pipeline: bool = False
# True for families whose full pipeline assembles MULTIPLE denoiser modules that a
# transformer-only file cannot supply (Ideogram 4 pairs a conditional ``transformer``
# with a separate ``unconditional_transformer``): there is no single-file or GGUF
# artifact carrying both, so only a full ``pipeline`` load is valid. The single-file /
# GGUF branches build just one transformer and would assemble a pipeline missing its
# second DiT, so validate_load_request rejects those kinds for such a family up front.
pipeline_only: bool = False
# Optional diffusers pipeline classes for image-conditioned workflows. The backend
# builds these around the ALREADY-loaded transformer/VAE/text-encoder via
# ``Pipeline.from_pipe`` (no extra weights, no reload), so a family only needs the
@ -322,6 +329,9 @@ _FAMILIES: tuple[DiffusionFamily, ...] = (
transformer_class = "Ideogram4Transformer2DModel",
base_repo = "ideogram-ai/ideogram-4-fp8",
aliases = ("ideogram4", "ideogram-v4", "ideogram"),
# Two DiTs assembled per-component (conditional + unconditional_transformer), so
# there is no transformer-only single-file / GGUF load for this family.
pipeline_only = True,
),
# SDXL is the one U-Net family here: the denoiser is ``pipe.unet``
# (UNet2DConditionModel), not a DiT ``pipe.transformer``, and a single-file

23
studio/backend/core/inference/diffusion_ideogram4.py
View file

@ -320,6 +320,29 @@ def load_ideogram4_text_encoder(
return model
def ideogram4_repo_is_fp8(repo_id: str, hf_token: Optional[str] = None) -> bool:
"""True when ``repo_id``'s transformer ships the vendor fp8 layout (a ``*.weight_scale``
shard key).
Those weights dequantize to a WIDER resident dtype, so the on-disk bytes undershoot
the bf16 footprint -- memory planning uses this to reserve the real size for a LOCAL
mirror of the fp8 base (whose path cannot string-match ``base_repo``; the bnb-4bit
``-nf4`` mirrors carry no ``_scale`` marker and correctly stay compressed). Reads shard
HEADERS only (metadata, not tensor bodies). Any failure (no transformer shards, no
reader) resolves to False so the caller falls back to the file-size estimate.
"""
try:
shard_paths = _transformer_shard_paths(repo_id, "transformer", hf_token or None)
import safetensors
except Exception: # noqa: BLE001 -- treat an unreadable / absent transformer as not fp8
return False
for path in shard_paths:
with safetensors.safe_open(path, "pt") as handle:
if any(key.endswith("_scale") for key in handle.keys()):
return True
return False
def load_ideogram4_transformer(
repo_id: str,
subfolder: str,

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

@ -1062,18 +1062,12 @@ class VideoBackend:
clear_gpu_cache()
raise RuntimeError("Video load was cancelled or superseded.")
offload_policy, vae_tiling = apply_memory_plan(pipe, plan, device = device, logger = logger)
if offload_policy == "group" and len(views) > 1:
# Group offload streams only ``pipe.transformer``; the second expert would
# otherwise sit resident (~57 GB bf16 on the A14B) and defeat the tier.
# model/sequential offload hook every top-level module, so only group needs
# this. Applied through the view so the helper streams transformer_2.
from .diffusion_memory import _apply_group_offload
for view in views[1:]:
if not _apply_group_offload(view, device, logger):
logger.warning(
"video.memory: group offload did not engage on the second "
"expert; it stays resident"
)
# A dual-DiT MoE pipe (Wan2.2-A14B) needs no extra per-expert offload pass here:
# apply_memory_plan's group tier (_apply_group_offload) already block-streams every
# DiT it finds on the pipe -- transformer AND transformer_2 -- and model/sequential
# offload hook every top-level module, so the second expert is covered under all tiers.
# A second _apply_group_offload on transformer_2 would re-register the group-offload
# hooks it already carries, which diffusers rejects with a duplicate-hook ValueError.
if not vae_tiling:
# Decode of a whole clip is the video memory peak; tiling is near-free
# in quality and keeps the decode bounded, so it is always on.

40
studio/backend/tests/test_diffusion_backend.py
View file

@ -383,7 +383,9 @@ def fake_runtime(monkeypatch):
diffusers.QwenImageInpaintPipeline = _FakeInpaintPipeline
# Instruction-editing pipeline (Qwen-Image-Edit): its own pipeline IS the loaded one.
diffusers.QwenImageEditPlusPipeline = _FakePipeline
# Ideogram 4, so its guidance_scale/guidance_schedule pairing is exercisable.
# Ideogram 4, so its guidance_scale/guidance_schedule pairing is exercisable. It loads
# only as a full pipeline (two DiTs), assembled per-component by load_ideogram4_pipeline
# -- stub that to a fake pipe so the guidance path is reachable without real weights.
diffusers.Ideogram4Pipeline = _FakePipeline
diffusers.Ideogram4Transformer2DModel = _FakeTransformer
# SDXL: a U-Net family. Its single-file checkpoint is the whole pipeline, so the
@ -394,6 +396,11 @@ def fake_runtime(monkeypatch):
diffusers.StableDiffusionXLImg2ImgPipeline = _FakeImg2ImgPipeline
diffusers.StableDiffusionXLInpaintPipeline = _FakeInpaintPipeline
monkeypatch.setattr(
"core.inference.diffusion.load_ideogram4_pipeline",
lambda repo_id, dtype, hf_token = None: _FakePipe(),
)
monkeypatch.setitem(sys.modules, "torch", torch)
monkeypatch.setitem(sys.modules, "diffusers", diffusers)
# The backend imports clear_gpu_cache by reference; no-op it so unload doesn't
@ -1227,13 +1234,30 @@ def test_generate_qwen_uses_true_cfg_scale(fake_runtime, tmp_path):
def _load_ideogram(backend, tmp_path):
(tmp_path / "model.gguf").write_bytes(b"weights")
backend.load_pipeline(
str(tmp_path),
gguf_filename = "model.gguf",
base_repo = "ideogram-ai/ideogram-4-fp8",
family_override = "ideogram-4",
)
# Ideogram 4 loads only as a full pipeline (its two DiTs are assembled per-component
# by the stubbed load_ideogram4_pipeline); a local pipeline dir is enough here.
(tmp_path / "model_index.json").write_text("{}")
backend.load_pipeline(str(tmp_path), family_override = "ideogram-4")
def test_ideogram_rejects_single_file_and_gguf_kinds(fake_runtime, tmp_path):
# Ideogram 4 needs two DiTs assembled per-component, so there is no transformer-only
# single-file or GGUF load: the explicit kinds must be rejected up front (before a
# load evicts a working model), not assembled into a pipeline missing its second DiT.
backend = DiffusionBackend()
(tmp_path / "model.gguf").write_bytes(b"x")
with pytest.raises(ValueError, match = "full diffusers pipeline"):
backend.load_pipeline(
str(tmp_path), gguf_filename = "model.gguf", family_override = "ideogram-4"
)
(tmp_path / "model.safetensors").write_bytes(b"x")
with pytest.raises(ValueError, match = "full diffusers pipeline"):
backend.load_pipeline(
str(tmp_path),
gguf_filename = "model.safetensors",
model_kind = "single_file",
family_override = "ideogram-4",
)
def test_generate_ideogram_defaults_keep_recommended_schedule(fake_runtime, tmp_path):

33
studio/backend/tests/test_diffusion_more_families.py
View file

@ -166,6 +166,39 @@ def test_convert_fp8_state_dict_dequantizes_and_splits_qkv():
torch.testing.assert_close(out["layers.0.attention_norm1.weight"], norm)
def test_ideogram4_repo_is_fp8_detects_local_layout(tmp_path):
# A local mirror of the fp8 base never string-matches base_repo, so memory planning
# relies on this shard-header probe to reserve the bf16 footprint. The fp8 layout is
# marked by a companion ``*.weight_scale``; the bnb-4bit (nf4) mirror carries none and
# must read as not-fp8 so it stays (correctly) planned against its compressed bytes.
torch = pytest.importorskip("torch")
st = pytest.importorskip("safetensors.torch")
from core.inference.diffusion_ideogram4 import ideogram4_repo_is_fp8
fp8 = tmp_path / "fp8"
(fp8 / "transformer").mkdir(parents = True)
st.save_file(
{
"layers.0.attention.o.weight": torch.zeros(2, 2),
"layers.0.attention.o.weight_scale": torch.ones(2),
},
str(fp8 / "transformer" / "diffusion_pytorch_model.safetensors"),
)
assert ideogram4_repo_is_fp8(str(fp8)) is True
nf4 = tmp_path / "nf4"
(nf4 / "transformer").mkdir(parents = True)
st.save_file(
{"layers.0.attention.to_q.weight": torch.zeros(2, 2)},
str(nf4 / "transformer" / "diffusion_pytorch_model.safetensors"),
)
assert ideogram4_repo_is_fp8(str(nf4)) is False
# A directory with no transformer shards at all resolves to False, not an error.
assert ideogram4_repo_is_fp8(str(tmp_path / "missing")) is False
def test_create_causal_mask_patch_is_self_disabling_and_idempotent():
# The patch adapts the pipeline's inputs_embeds kwarg to the installed transformers
# create_causal_mask signature; on a matching signature it must forward unchanged,