ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions
unsloth/studio/backend/tests/test_diffusion_more_families.py
pre-commit-ci[bot] 7c9521810e [pre-commit.ci] auto fixes from pre-commit.com hooks 
for more information, see https://pre-commit.ci
2026-07-25 09:00:00 +00:00

472 lines
22 KiB
Python
Raw Blame History

# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Ideogram 4 family registration, the HunyuanImage structured exclusion, and the
curated krea/Krea-2-LoRA-* catalog entries. Pure-module tests: no torch, no network."""
import pytest
from core.inference.diffusion import _is_trusted_diffusion_repo
from core.inference.diffusion_auto_policy import family_bf16_components_gb
from core.inference.diffusion_families import (
IDEOGRAM4_FAMILY_NAME,
default_generation_params,
detect_family,
excluded_model_reason,
)
from core.inference.diffusion_lora import _CURATED, list_loras
# ── ideogram-4 family detection ──────────────────────────────────────────────
@pytest.mark.parametrize(
"repo_id",
[
"ideogram-ai/ideogram-4-fp8",
"ideogram-ai/ideogram-4-nf4",
"ideogram-ai/ideogram-4-nf4-diffusers",
],
)
def test_detect_family_ideogram4_repos(repo_id):
fam = detect_family(repo_id)
assert fam is not None and fam.name == IDEOGRAM4_FAMILY_NAME
assert fam.pipeline_class == "Ideogram4Pipeline"
assert fam.transformer_class == "Ideogram4Transformer2DModel"
# The vendor ships no bf16 repo: the raw-float8 export is the family base.
assert fam.base_repo == "ideogram-ai/ideogram-4-fp8"
def test_detect_family_ideogram4_override():
fam = detect_family("some/local-path", override = "ideogram-4")
assert fam is not None and fam.name == IDEOGRAM4_FAMILY_NAME
assert detect_family("x", override = "ideogram4").name == IDEOGRAM4_FAMILY_NAME
def test_ideogram4_repos_are_trusted_non_gguf():
# The three official vendor pipelines load via from_pretrained, which is gated
# to the unsloth org + the explicit allowlist.
for rid in (
"ideogram-ai/ideogram-4-fp8",
"ideogram-ai/ideogram-4-nf4",
"ideogram-ai/ideogram-4-nf4-diffusers",
):
assert _is_trusted_diffusion_repo(rid)
assert not _is_trusted_diffusion_repo("ideogram-ai/some-future-repo")
# ── FLUX.1 Krea dev (flux.1 family variant) ──────────────────────────────────
@pytest.mark.parametrize(
"repo_id",
[
"black-forest-labs/FLUX.1-Krea-dev",
"QuantStack/FLUX.1-Krea-dev-GGUF",
# A local GGUF pick where the family keyword lives in the filename.
"QuantStack/FLUX.1-Krea-dev-GGUF/flux1-krea-dev-Q4_K_M.gguf",
],
)
def test_detect_family_flux1_krea_dev(repo_id):
# Krea's FLUX.1-dev finetune keeps the exact dev layout, so it must resolve to the
# existing flux.1 family (FluxPipeline), never to krea-2 (a different arch).
fam = detect_family(repo_id)
assert fam is not None and fam.name == "flux.1"
assert fam.pipeline_class == "FluxPipeline"
def test_flux1_krea_dev_is_trusted_non_gguf():
# The gated official pipeline loads via from_pretrained -> needs the allowlist.
assert _is_trusted_diffusion_repo("black-forest-labs/FLUX.1-Krea-dev")
def test_flux1_krea_dev_generation_defaults():
# Model-card recipe: 28 steps at guidance 4.5. The generic "krea" key (Krea-2-Turbo's
# 8-step no-CFG shape) must NOT swallow it, and the krea-2 defaults must stay intact.
assert default_generation_params("black-forest-labs/FLUX.1-Krea-dev") == (28, 4.5)
assert default_generation_params("QuantStack/FLUX.1-Krea-dev-GGUF") == (28, 4.5)
assert default_generation_params("krea/Krea-2-Turbo") == (8, 0.0)
assert default_generation_params("krea/Krea-2-Raw") == (52, 3.5)
# ── lumina-2 family ──────────────────────────────────────────────────────────
@pytest.mark.parametrize(
"repo_id",
[
"Alpha-VLLM/Lumina-Image-2.0",
# A same-arch finetune must group here via the lumina-image-2.0 token.
"neta-art/NetaYume-Lumina-Image-2.0",
],
)
def test_detect_family_lumina2_repos(repo_id):
fam = detect_family(repo_id)
assert fam is not None and fam.name == "lumina-2"
assert fam.pipeline_class == "Lumina2Pipeline"
assert fam.transformer_class == "Lumina2Transformer2DModel"
assert fam.base_repo == "Alpha-VLLM/Lumina-Image-2.0"
# Published bf16-only upstream; the fp16 fallback stays off.
assert fam.fp16_incompatible is True
def test_detect_family_lumina2_override_and_next_rejected():
assert detect_family("x", override = "lumina-2").name == "lumina-2"
assert detect_family("x", override = "lumina2").name == "lumina-2"
# Lumina-Next is a DIFFERENT arch (LuminaText2ImgPipeline): it must stay unknown
# instead of resolving here and crashing mid-load.
assert detect_family("Alpha-VLLM/Lumina-Next-SFT-diffusers") is None
def test_lumina2_is_trusted_non_gguf():
# The official pipeline loads via from_pretrained -> needs the allowlist.
assert _is_trusted_diffusion_repo("Alpha-VLLM/Lumina-Image-2.0")
assert not _is_trusted_diffusion_repo("Alpha-VLLM/some-future-repo")
def test_lumina2_generation_defaults():
# Model-card recipe: 50 steps at guidance 4.0 (cfg_trunc_ratio is added by the
# backend generate call itself, not the defaults table).
assert default_generation_params("Alpha-VLLM/Lumina-Image-2.0") == (50, 4.0)
def test_lumina2_prequant_wiring():
# Hosted int8/fp8 checkpoints (gate-validated) serve the family default base.
from core.inference.diffusion_families import family_prequant_repo
fam = detect_family("Alpha-VLLM/Lumina-Image-2.0")
for scheme in ("int8", "fp8"):
assert family_prequant_repo(fam, scheme) == "unsloth/Lumina-Image-2.0-FP8"
def test_lumina2_bf16_component_table_present():
fam = detect_family("Alpha-VLLM/Lumina-Image-2.0")
sizes = family_bf16_components_gb(fam)
assert sizes is not None
transformer_gb, encoders_gb, vae_gb = sizes
# 2.6B DiT + Gemma2-2B, both fp32 on disk -> ~5.2 GB each bf16-resident.
assert 4.0 <= transformer_gb <= 7.0
assert 4.0 <= encoders_gb <= 7.0
assert vae_gb <= 0.5
# ── hunyuanimage-2.1 family ──────────────────────────────────────────────────
@pytest.mark.parametrize(
"repo_id",
[
"hunyuanvideo-community/HunyuanImage-2.1-Diffusers",
"QuantStack/HunyuanImage-2.1-GGUF",
# A local GGUF pick where the family keyword lives in the filename (QuantStack's
# actual naming drops the dash: the hunyuanimage2.1 alias covers it).
"QuantStack/HunyuanImage-2.1-GGUF/HunyuanImage2.1-Q4_K_M.gguf",
],
)
def test_detect_family_hunyuanimage21_repos(repo_id):
fam = detect_family(repo_id)
assert fam is not None and fam.name == "hunyuanimage-2.1"
assert fam.pipeline_class == "HunyuanImagePipeline"
assert fam.transformer_class == "HunyuanImageTransformer2DModel"
assert fam.base_repo == "hunyuanvideo-community/HunyuanImage-2.1-Diffusers"
# The call's guidance knob is distilled_guidance_scale; there is no guidance_scale kwarg.
assert fam.cfg_kwarg == "distilled_guidance_scale"
# Published bf16-only upstream; the fp16 fallback stays off.
assert fam.fp16_incompatible is True
def test_detect_family_hunyuanimage21_override_and_30_still_excluded():
assert detect_family("x", override = "hunyuanimage-2.1").name == "hunyuanimage-2.1"
assert detect_family("x", override = "hunyuanimage2.1").name == "hunyuanimage-2.1"
# The HunyuanImage-3.0 structured exclusion must survive the 2.1 family: 3.0 has
# no diffusers pipeline and must stay unknown with its stated reason.
assert detect_family("tencent/HunyuanImage-3.0") is None
assert excluded_model_reason("tencent/HunyuanImage-3.0") is not None
assert excluded_model_reason("hunyuanvideo-community/HunyuanImage-2.1-Diffusers") is None
def test_hunyuanimage21_is_trusted_non_gguf():
# The mirror pipeline loads via from_pretrained -> needs the allowlist.
assert _is_trusted_diffusion_repo("hunyuanvideo-community/HunyuanImage-2.1-Diffusers")
assert not _is_trusted_diffusion_repo("hunyuanvideo-community/some-future-repo")
def test_hunyuanimage21_generation_defaults():
# Card recipe: 50 steps; guidance feeds the call's distilled_guidance_scale (3.25
# default), while classifier-free guidance runs inside the repo's guider components.
assert default_generation_params("hunyuanvideo-community/HunyuanImage-2.1-Diffusers") == (
50,
3.25,
)
def test_hunyuanimage21_prequant_wiring():
# Hosted int8/fp8 checkpoints, verified bit-identical to on-the-fly quantize.
from core.inference.diffusion_families import family_prequant_repo
fam = detect_family("hunyuanvideo-community/HunyuanImage-2.1-Diffusers")
for scheme in ("int8", "fp8"):
assert family_prequant_repo(fam, scheme) == "unsloth/HunyuanImage-2.1-FP8"
def test_hunyuanimage21_bf16_component_table_present():
fam = detect_family("hunyuanvideo-community/HunyuanImage-2.1-Diffusers")
sizes = family_bf16_components_gb(fam)
assert sizes is not None
transformer_gb, encoders_gb, vae_gb = sizes
# 17B DiT (32.5 GB bf16 on disk) + Qwen2.5-VL 15.5 GB + ByT5 0.8 GB.
assert 30.0 <= transformer_gb <= 35.0
assert 15.0 <= encoders_gb <= 18.0
assert vae_gb <= 1.0
# ── hidream-i1 family ────────────────────────────────────────────────────────
@pytest.mark.parametrize(
"repo_id",
[
"HiDream-ai/HiDream-I1-Full",
"HiDream-ai/HiDream-I1-Dev",
"HiDream-ai/HiDream-I1-Fast",
],
)
def test_detect_family_hidream_repos(repo_id):
# One family covers all three variants (same 17B MoE arch + 4-TE stack).
fam = detect_family(repo_id)
assert fam is not None and fam.name == "hidream-i1"
assert fam.pipeline_class == "HiDreamImagePipeline"
assert fam.transformer_class == "HiDreamImageTransformer2DModel"
assert fam.base_repo == "HiDream-ai/HiDream-I1-Full"
# Published bf16-only upstream; the fp16 fallback stays off.
assert fam.fp16_incompatible is True
def test_hidream_override_and_trust():
assert detect_family("x", override = "hidream-i1").name == "hidream-i1"
assert detect_family("x", override = "hidream").name == "hidream-i1"
# The three official repos load via from_pretrained -> allowlisted; the Llama TE4
# comes from the unsloth mirror, which the org prefix already trusts.
for rid in (
"HiDream-ai/HiDream-I1-Full",
"HiDream-ai/HiDream-I1-Dev",
"HiDream-ai/HiDream-I1-Fast",
):
assert _is_trusted_diffusion_repo(rid)
assert not _is_trusted_diffusion_repo("HiDream-ai/some-future-repo")
assert _is_trusted_diffusion_repo("unsloth/Meta-Llama-3.1-8B-Instruct")
def test_hidream_generation_defaults():
# Upstream inference.py: Full 50 steps / guidance 5; Dev and Fast are distilled and
# run guidance-free at 28 / 16 steps. The specific keys must beat the generic
# "hidream" (which also appears in the owner segment of every variant id).
assert default_generation_params("HiDream-ai/HiDream-I1-Full") == (50, 5.0)
assert default_generation_params("HiDream-ai/HiDream-I1-Dev") == (28, 0.0)
assert default_generation_params("HiDream-ai/HiDream-I1-Fast") == (16, 0.0)
def test_hidream_bf16_component_table_present():
fam = detect_family("HiDream-ai/HiDream-I1-Full")
sizes = family_bf16_components_gb(fam)
assert sizes is not None
transformer_gb, encoders_gb, vae_gb = sizes
# 17B MoE DiT 34.2 GB; TEs = CLIP-L 0.5 + CLIP-G 2.8 + T5-XXL 9.5 from the repo plus
# the ~16 GB Llama TE4 assembled from the mirror -> ~28.8 GB.
assert 32.0 <= transformer_gb <= 37.0
assert 26.0 <= encoders_gb <= 32.0
assert vae_gb <= 0.5
def test_ideogram4_generation_defaults():
# Model-card settings: 48 steps, guidance 7 (the backend keeps the pipeline's
# recommended tapered schedule when the request matches exactly).
assert default_generation_params("ideogram-ai/ideogram-4-fp8") == (48, 7.0)
def test_ideogram4_bf16_reservation_table_present():
# The memory planner reserves this bf16 footprint for a narrow (fp8) ideogram-4 base even
# when the blob-cache estimate is absent (empty cache / a best-effort download probe that
# swallowed a transient HF error), so the ~54 GB pipeline never plans a resident placement
# it cannot fit. If this constant table ever went None, that fp8 OOM safeguard would
# silently disable, so pin that it is present and sums to the expected ~54 GB.
fam = detect_family("ideogram-ai/ideogram-4-fp8")
table = family_bf16_components_gb(fam, fam.base_repo)
assert table is not None
assert sum(table) > 50.0 # transformer (37.2) + bf16 text encoder (16.3) + VAE (0.2)
def test_ideogram4_memory_table_counts_both_dits():
fam = detect_family("ideogram-ai/ideogram-4-fp8")
components = family_bf16_components_gb(fam)
assert components is not None
transformer_gb, text_encoders_gb, _vae_gb = components
# Two ~9.3B DiTs (conditional + unconditional) at bf16: well above one DiT's
# ~18.6 GB. A single-DiT entry here would let auto planning under-reserve and OOM.
assert transformer_gb > 30.0
assert text_encoders_gb > 5.0
def test_hidream_prequant_wiring():
# Hosted int8/fp8 checkpoints (28/28 per-case gate pairs per scheme; int8 verified
# bit-identical to on-the-fly quantize) serve the family default base.
from core.inference.diffusion_families import family_prequant_repo
fam = detect_family("HiDream-ai/HiDream-I1-Full")
for scheme in ("int8", "fp8"):
assert family_prequant_repo(fam, scheme) == "unsloth/HiDream-I1-Full-FP8"
def test_hidream_quant_schemes_not_denied_and_no_extra_excludes():
# Measured on a B200 (outputs/hidream_smoke): int8 and fp8 both engage and render
# cleanly, including a 2-3 token prompt on int8 -- the routed MoE expert Linears
# only ever see the concatenated image+text token stream (M >> 16), so the
# torch._int_mm minimum never binds and no family exclude tokens are needed.
from core.inference.diffusion_transformer_quant import (
_FAMILY_SCHEME_DENY,
_INT8_EXCLUDE_NAME_TOKENS,
exclude_tokens_for_scheme,
)
assert "hidream-i1" not in _FAMILY_SCHEME_DENY
assert exclude_tokens_for_scheme("int8", "hidream-i1") == _INT8_EXCLUDE_NAME_TOKENS
assert exclude_tokens_for_scheme("fp8", "hidream-i1") == ()
# ── structured exclusions ────────────────────────────────────────────────────
def test_hunyuanimage_is_excluded_with_reason():
reason = excluded_model_reason("tencent/HunyuanImage-3.0")
assert reason is not None and "diffusers" in reason
# Not detectable as any family: the exclusion reason is the load error surface.
assert detect_family("tencent/HunyuanImage-3.0") is None
def test_excluded_model_reason_none_for_supported_and_unknown():
assert excluded_model_reason("unsloth/Z-Image-Turbo-GGUF") is None
assert excluded_model_reason("someorg/some-model") is None
def test_validate_load_request_surfaces_exclusion_reason():
from core.inference.diffusion import DiffusionBackend
backend = DiffusionBackend()
with pytest.raises(ValueError, match = "trust_remote_code"):
backend.validate_load_request("tencent/HunyuanImage-3.0")
# ── curated krea LoRA catalog ────────────────────────────────────────────────
def test_curated_krea2_loras_present_and_well_formed():
krea = [e for e in _CURATED if e.repo_id and e.repo_id.startswith("krea/Krea-2-LoRA-")]
assert len(krea) == 9
for entry in krea:
assert entry.source == "hub" and entry.fmt == "safetensors"
assert entry.families == ("krea-2",)
# Every official style repo carries a single "{style}.safetensors" at the root.
style = entry.repo_id.split("Krea-2-LoRA-")[-1]
assert entry.weight_name == f"{style}.safetensors"
def test_list_loras_family_filter_gates_krea_entries():
krea_ids = {e.id for e in _CURATED if e.families == ("krea-2",)}
assert krea_ids # curated entries exist
listed_for_krea = {e.id for e in list_loras(family = "krea-2")}
assert krea_ids <= listed_for_krea
listed_for_flux = {e.id for e in list_loras(family = "flux.1")}
assert not (krea_ids & listed_for_flux)
# ── ideogram-4 fp8 transformer remap ─────────────────────────────────────────
def test_convert_fp8_state_dict_dequantizes_and_splits_qkv():
# The vendor fp8 transformer stores fused attention.qkv (Q/K/V rows stacked) +
# attention.o, each with a per-output-channel weight_scale; diffusers expects split
# to_q/to_k/to_v/to_out.0 with the scale already applied. The converter must undo
# both, or every attention weight loads wrong (garbage) and on meta (a load crash).
torch = pytest.importorskip("torch")
from core.inference.diffusion_ideogram4 import _convert_fp8_state_dict
hidden = 4 # tiny stand-in for attention_head_dim * num_attention_heads
# Reference (real) weights, then a fake per-channel fp8 encoding: value / scale.
q = torch.randn(hidden, hidden)
k = torch.randn(hidden, hidden)
v = torch.randn(hidden, hidden)
o = torch.randn(hidden, hidden)
ff = torch.randn(hidden, hidden)
fused = torch.cat([q, k, v], dim = 0) # [3 * hidden, hidden]
qkv_scale = torch.rand(3 * hidden) + 0.5
o_scale = torch.rand(hidden) + 0.5
ff_scale = torch.rand(hidden) + 0.5
norm = torch.randn(hidden) # dense (unscaled) weight passes through
raw = {
"layers.0.attention.qkv.weight": fused / qkv_scale[:, None],
"layers.0.attention.qkv.weight_scale": qkv_scale,
"layers.0.attention.o.weight": o / o_scale[:, None],
"layers.0.attention.o.weight_scale": o_scale,
"layers.0.feed_forward.w1.weight": ff / ff_scale[:, None],
"layers.0.feed_forward.w1.weight_scale": ff_scale,
"layers.0.attention_norm1.weight": norm,
}
out = _convert_fp8_state_dict(raw, hidden, torch.bfloat16)
# Every converted tensor is cast to the requested compute dtype (the load_state_dict
# copy would silently up/down-cast otherwise).
assert all(t.dtype == torch.bfloat16 for t in out.values())
# Re-run in float32 for the exact value checks below (bf16 loses precision).
out = _convert_fp8_state_dict(raw, hidden, torch.float32)
# No scale keys leak through; fused/renamed keys are gone.
assert not any(key.endswith("_scale") for key in out)
assert "layers.0.attention.qkv.weight" not in out
assert "layers.0.attention.o.weight" not in out
# QKV split back to the reference weights in Q/K/V order.
torch.testing.assert_close(out["layers.0.attention.to_q.weight"], q)
torch.testing.assert_close(out["layers.0.attention.to_k.weight"], k)
torch.testing.assert_close(out["layers.0.attention.to_v.weight"], v)
# o renamed to to_out.0 with the scale applied.
torch.testing.assert_close(out["layers.0.attention.to_out.0.weight"], o)
# A non-attention fp8 weight keeps its name, scale applied.
torch.testing.assert_close(out["layers.0.feed_forward.w1.weight"], ff)
# A dense weight passes through unchanged.
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,
# and a second apply must not double-wrap.
pytest.importorskip("torch")
pytest.importorskip("diffusers")
import core.inference.diffusion_ideogram4 as ig4
from diffusers.pipelines.ideogram4 import pipeline_ideogram4 as pipe_mod
original = pipe_mod.create_causal_mask
try:
ig4._CAUSAL_MASK_PATCHED = False
ig4._patch_create_causal_mask()
wrapped = pipe_mod.create_causal_mask
assert wrapped is not original # the patch installed a wrapper
ig4._patch_create_causal_mask() # idempotent: no re-wrap
assert pipe_mod.create_causal_mask is wrapped
finally:
pipe_mod.create_causal_mask = original
ig4._CAUSAL_MASK_PATCHED = False
Reference in a new issue View git blame Copy permalink