From cbfc43215d1ae5c187751e57da8e8b97a99fa8fd Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 4 Jul 2026 12:46:24 +0000 Subject: [PATCH 01/12] Add Ideogram 4 family, structured HunyuanImage exclusion, curated Krea 2 LoRAs Ideogram 4 (diffusers 0.39 Ideogram4Pipeline) as a new image family. The vendor publishes no bf16 checkpoint, so ideogram-ai/ideogram-4-fp8 (raw float8 DiTs, upcast by from_pretrained) is the family base and ideogram-4-nf4-diffusers is the bnb-4bit pipeline artifact (ideogram-4-nf4 is byte-identical and detects to the same family). All three repos join the trusted non-GGUF allowlist and the frontend safetensors catalog. Family specifics handled: - Dual-branch CFG runs through a SEPARATE unconditional_transformer, so the auto-policy size table entry counts two ~9.3B DiTs (37.2 GB bf16), and the pipeline-kind memory plan now takes max(cached bytes, family table) for the family base repo: the fp8 repo's cached bytes undershoot the bf16-resident footprint by ~2x, which would let auto planning pick a resident placement that OOMs. - The pipeline accepts EITHER guidance_scale OR a per-step guidance_schedule (its default: the recommended 45x7.0 + 3x3.0 taper, valid only at 48 steps) and raises when both are set. At the advertised defaults (48 steps, guidance 7) generate() drops the constant so the recommended taper engages; any other request nulls the schedule so the constant broadcasts legally. - Generation defaults per the model card: 48 steps, guidance 7 (both tables). tencent/HunyuanImage-3.0 is deliberately excluded: it has no diffusers pipeline (an 80B autoregressive MoE behind trust_remote_code). A structured exclusion map now surfaces that reason verbatim from validate_load_request instead of the generic unknown-family error. The curated diffusion LoRA catalog gains the nine official krea/Krea-2-LoRA-* style adapters (family-tagged krea-2, explicit weight filenames), so they show up in the picker instead of requiring a typed repo id. Tests: new test_diffusion_more_families.py (detection, trust, defaults, size table, exclusion reason, curated catalog + family filter), two generate() tests for the guidance_scale/guidance_schedule pairing, and the local-scan LoRA test updated for a non-empty curated list. Backend suite + CI-sim (block_diffusers/block_torchao) green; frontend builds. --- studio/backend/core/inference/diffusion.py | 46 ++++++- .../core/inference/diffusion_auto_policy.py | 6 + .../core/inference/diffusion_families.py | 53 +++++++++ .../backend/core/inference/diffusion_lora.py | 33 +++++- .../backend/tests/test_diffusion_backend.py | 38 ++++++ studio/backend/tests/test_diffusion_lora.py | 9 +- .../tests/test_diffusion_more_families.py | 112 ++++++++++++++++++ .../src/features/images/images-page.tsx | 15 +++ 8 files changed, 304 insertions(+), 8 deletions(-) create mode 100644 studio/backend/tests/test_diffusion_more_families.py diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 9555b3c38a..c0bc560255 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -31,9 +31,11 @@ from utils.hardware import clear_gpu_cache from .diffusion_families import ( DIFFUSION_CANCELLED_MSG, DIFFUSION_NOT_LOADED_MSG, + IDEOGRAM4_FAMILY_NAME, DiffusionFamily, default_generation_params, detect_family_for_pick, + excluded_model_reason, resolve_base_repo, resolve_local_gguf_child, supported_family_names, @@ -84,7 +86,11 @@ from .diffusion_prequant import ( load_prequantized_transformer, resolve_prequant_source, ) -from .diffusion_auto_policy import build_resolved_record, resolve_dense_quant_candidate +from .diffusion_auto_policy import ( + build_resolved_record, + family_bf16_components_gb, + resolve_dense_quant_candidate, +) from .diffusion_transformer_quant import ( TQ_AUTO, DEFAULT_MIN_LINEAR_FEATURES, @@ -221,6 +227,13 @@ _TRUSTED_NON_GGUF_REPOS = frozenset( # training LoRAs on (train on Raw, run adapters on Turbo). "krea/krea-2-turbo", "krea/krea-2-raw", + # Ideogram 4: official vendor repos, safetensors-only diffusers pipelines, no + # remote code. The vendor ships no bf16 checkpoint: -fp8 stores the two DiTs + # as raw float8 (highest precision available, the family base); the two nf4 + # repos are identical bnb-4bit exports (both listed so either id loads). + "ideogram-ai/ideogram-4-fp8", + "ideogram-ai/ideogram-4-nf4", + "ideogram-ai/ideogram-4-nf4-diffusers", } ) @@ -508,6 +521,12 @@ class DiffusionBackend: kind = resolve_model_kind(gguf_filename, model_kind) fam = detect_family_for_pick(repo_id, gguf_filename, family_override) if fam is None: + # A deliberately-excluded model gets its stated reason, not the generic + # unknown-family message (which reads like a detection gap and invites a + # family_override retry that would fail deeper and less clearly). + excluded = excluded_model_reason(repo_id) + if excluded: + raise ValueError(f"'{repo_id}' cannot be loaded: {excluded}") raise ValueError( f"'{repo_id}' is not a supported diffusion image model. Supported families: " f"{', '.join(supported_family_names())}. If this is a variant of one of them, " @@ -1590,6 +1609,19 @@ class DiffusionBackend: cached = self._cache_bytes(repo_id) if repo_id else 0 cached_mib = int(cached // (1024 * 1024)) if cached else None model_dense_mib = estimate_safetensors_dense_mib(cached_mib) + # A repo can store weights in a NARROWER dtype than they occupy after the + # loader's torch_dtype cast: ideogram-4's base repo ships its two DiTs as + # raw float8, so the cached bytes undershoot the bf16-resident footprint + # by ~2x and auto planning would pick a resident placement that OOMs. + # When the family size table knows the bf16-resident total for THIS repo + # (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(): + 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)) + model_dense_mib = max(model_dense_mib, table_mib) companion_mib = None else: if transformer_resident_override_mib is not None: @@ -2118,6 +2150,18 @@ class DiffusionBackend: # share this call's seed, drawn sequentially from one generator. "num_images_per_prompt": batch_size, } + if state.family.name == IDEOGRAM4_FAMILY_NAME: + # Ideogram 4 drives CFG through EITHER a constant guidance_scale OR + # a per-step guidance_schedule; its check_inputs rejects the call + # when both are set, and the schedule DEFAULTS to the recommended + # 45x7.0 + 3x3.0 polish taper (valid only at exactly 48 steps). At + # the family's advertised defaults, drop the constant so the + # recommended taper engages; any other request nulls the schedule + # so the constant broadcasts legally to the chosen step count. + if steps == 48 and abs(float(guidance) - 7.0) < 1e-6: + kwargs.pop(state.family.cfg_kwarg, None) + else: + kwargs["guidance_schedule"] = None if init_pil is not None: # Reference with extra images passes the whole list (FLUX.2 combines them); # every other workflow takes the single image. diff --git a/studio/backend/core/inference/diffusion_auto_policy.py b/studio/backend/core/inference/diffusion_auto_policy.py index 0aa0a51371..8433e4e2a2 100644 --- a/studio/backend/core/inference/diffusion_auto_policy.py +++ b/studio/backend/core/inference/diffusion_auto_policy.py @@ -59,6 +59,12 @@ _FAMILY_BF16_GB: dict[str, tuple[float, float, float]] = { "qwen-image-edit": (40.9, 16.6, 0.3), "z-image": (12.3, 8.0, 0.2), "krea-2": (26.3, 8.9, 0.5), + # Two ~9.3B DiTs (the conditional transformer PLUS the separate + # unconditional_transformer driving Ideogram's dual-branch CFG), both resident + # for every generation, and a Qwen3-VL text encoder. The vendor repo stores the + # DiTs as raw float8 (9.29 GB each); these are the bf16-resident sizes after the + # loader's dtype cast, per this table's contract. + "ideogram-4": (37.2, 8.8, 0.2), } # Base-repo overrides for families whose picker offers multiple sizes under one family diff --git a/studio/backend/core/inference/diffusion_families.py b/studio/backend/core/inference/diffusion_families.py index ed2dc1e9c1..a6e47f2388 100644 --- a/studio/backend/core/inference/diffusion_families.py +++ b/studio/backend/core/inference/diffusion_families.py @@ -304,6 +304,25 @@ _FAMILIES: tuple[DiffusionFamily, ...] = ( # unvalidated upstream, so keep the fp16 fallback off like z-image. fp16_incompatible = True, ), + # Ideogram 4 (diffusers >= 0.39): a 34-layer single-stream flow-matching DiT PAIR -- + # the conditional transformer plus a separate ``unconditional_transformer`` driving + # its dual-branch CFG (both ~9B params, so memory planning must count two DiTs) -- + # with a Qwen3-VL text encoder. The vendor publishes no bf16 checkpoint: + # ideogram-4-fp8 stores the DiTs as raw float8 tensors (from_pretrained upcasts + # them to the compute dtype) and is the highest-precision artifact, so it is the + # family base; ideogram-4-nf4-diffusers / ideogram-4-nf4 (identical contents) + # carry bnb-4bit quantization_configs the pipeline kind re-applies automatically. + # All three repos are gated="auto" on the Hub, so a load may need the user's HF + # token. No GGUF variant and no sd.cpp mapping, so the no-GPU route falls back to + # diffusers. CFG quirk: the pipeline takes EITHER guidance_scale OR a per-step + # guidance_schedule (see the loader's IDEOGRAM4 branch in diffusion.py). + DiffusionFamily( + name = "ideogram-4", + pipeline_class = "Ideogram4Pipeline", + transformer_class = "Ideogram4Transformer2DModel", + base_repo = "ideogram-ai/ideogram-4-fp8", + aliases = ("ideogram4", "ideogram-v4", "ideogram"), + ), # SDXL is the one U-Net family here: the denoiser is ``pipe.unet`` # (UNet2DConditionModel), not a DiT ``pipe.transformer``, and a single-file # ``.safetensors`` is the WHOLE pipeline rather than a transformer-only file. @@ -343,6 +362,36 @@ def trainable_family_names() -> tuple[str, ...]: return tuple(fam.name for fam in _FAMILIES if fam.trainable) +# The family whose CFG runs through a guidance_scale/guidance_schedule pair rather +# than a plain guidance_scale (the loader special-cases the call, like krea-2's +# per-component assembly). Named here so the two modules cannot drift apart. +IDEOGRAM4_FAMILY_NAME = "ideogram-4" + + +# Models Studio deliberately does NOT support, with the reason surfaced verbatim in +# the load error (instead of the generic unknown-family message, which reads like a +# detection gap). Keyed by a lowercase substring of the repo id. The bar for support +# is a diffusers pipeline: HunyuanImage-3.0 is an 80B autoregressive MoE loaded via +# AutoModelForCausalLM + trust_remote_code -- there is nothing for this backend to +# assemble, and remote-code execution is out of the question for a load path. +_EXCLUDED_MODELS: tuple[tuple[str, str], ...] = ( + ( + "hunyuanimage", + "HunyuanImage-3.0 has no diffusers pipeline (it is an 80B autoregressive MoE " + "that requires trust_remote_code), so Studio does not support it.", + ), +) + + +def excluded_model_reason(repo_id: str) -> Optional[str]: + """The stated reason ``repo_id`` is unsupported, or None when it is simply unknown.""" + needle = (repo_id or "").lower() + for token, reason in _EXCLUDED_MODELS: + if token in needle: + return reason + return None + + # Editing / inpaint checkpoints share an arch keyword but need a different # pipeline and an input image, which this text-to-image backend doesn't drive. # "layered" rejects Qwen-Image-Layered: its transformer sets additional_t_cond=True @@ -456,6 +505,10 @@ _GENERATION_DEFAULTS: tuple[tuple[str, int, float], ...] = ( ("flux.2-klein", 4, 0.0), ("qwen-image", 20, 4.0), ("z-image", 20, 4.0), + # Ideogram 4's model-card settings: 48 steps, guidance 7 (its recommended + # schedule tapers the last 3 steps to 3.0 -- the loader keeps that taper when + # the request matches these defaults exactly; see the IDEOGRAM4 branch). + ("ideogram", 48, 7.0), ) # Unrecognised model: distilled few-step / no-CFG shape, matching the UI fallback. _GENERATION_DEFAULT_FALLBACK = (9, 0.0) diff --git a/studio/backend/core/inference/diffusion_lora.py b/studio/backend/core/inference/diffusion_lora.py index cad6d3eae7..a971ffb8b3 100644 --- a/studio/backend/core/inference/diffusion_lora.py +++ b/studio/backend/core/inference/diffusion_lora.py @@ -66,9 +66,36 @@ class ResolvedLora: # Curated, family-tagged catalog of known-good diffusion LoRAs. Kept intentionally small # and data-driven; extend as unsloth hosts/curates more. Entries are HF repos with a -# single-file weight. (Left minimal on purpose -- local discovery is the primary source, -# and users can also reference any public HF LoRA repo id directly.) -_CURATED: tuple[LoraCatalogEntry, ...] = () +# single-file weight. (Local discovery remains a primary source, and users can also +# reference any public HF LoRA repo id directly.) + + +def _krea2_lora(style: str, display_name: str) -> LoraCatalogEntry: + """One official krea/Krea-2-LoRA-* style adapter. All nine follow the same repo + shape (a single ``{style}.safetensors`` at the root) and are trained on Krea-2-Raw + for use on Krea-2-Turbo, per Krea's release guidance.""" + return LoraCatalogEntry( + id = f"krea/Krea-2-LoRA-{style}", + display_name = display_name, + source = "hub", + fmt = "safetensors", + families = ("krea-2",), + repo_id = f"krea/Krea-2-LoRA-{style}", + weight_name = f"{style}.safetensors", + ) + + +_CURATED: tuple[LoraCatalogEntry, ...] = ( + _krea2_lora("retroanime", "Krea 2 Retro Anime"), + _krea2_lora("neondrip", "Krea 2 Neon Drip"), + _krea2_lora("darkbrush", "Krea 2 Dark Brush"), + _krea2_lora("softwatercolor", "Krea 2 Soft Watercolor"), + _krea2_lora("dotmatrix", "Krea 2 Dot Matrix"), + _krea2_lora("rainywindow", "Krea 2 Rainy Window"), + _krea2_lora("vintagetarot", "Krea 2 Vintage Tarot"), + _krea2_lora("sunsetblur", "Krea 2 Sunset Blur"), + _krea2_lora("kidsdrawing", "Krea 2 Kids Drawing"), +) def loras_dir() -> Path: diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index 78c320c0be..b0890c4966 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -383,6 +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. + diffusers.Ideogram4Pipeline = _FakePipeline + diffusers.Ideogram4Transformer2DModel = _FakeTransformer # SDXL: a U-Net family. Its single-file checkpoint is the whole pipeline, so the # pipeline class carries from_single_file; UNet2DConditionModel is the denoiser # class (fetched but unused on the pipeline/single-file-pipeline paths). @@ -1223,6 +1226,41 @@ def test_generate_qwen_uses_true_cfg_scale(fake_runtime, tmp_path): assert call["true_cfg_scale"] == 4.0 and call["guidance_scale"] is None +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", + ) + + +def test_generate_ideogram_defaults_keep_recommended_schedule(fake_runtime, tmp_path): + # Ideogram 4's pipeline defaults to its recommended tapered guidance_schedule + # (45x7.0 + 3x3.0, valid only at 48 steps) and REJECTS guidance_scale while the + # schedule is set. At the family's advertised defaults the backend must drop the + # constant so the recommended taper engages. + backend = DiffusionBackend() + _load_ideogram(backend, tmp_path) + backend.generate(prompt = "a sloth", steps = 48, guidance = 7.0) + call = backend._state.pipe.last_kwargs + assert call["guidance_scale"] is None # not passed: the pipe default engages + assert "guidance_schedule" not in call + + +def test_generate_ideogram_custom_guidance_nulls_schedule(fake_runtime, tmp_path): + # Any non-default request must broadcast the constant legally: guidance_scale set + # AND guidance_schedule explicitly nulled (the pipeline raises when both are set, + # and its default schedule is non-None). + backend = DiffusionBackend() + _load_ideogram(backend, tmp_path) + backend.generate(prompt = "a sloth", steps = 20, guidance = 5.0) + call = backend._state.pipe.last_kwargs + assert call["guidance_scale"] == 5.0 + assert "guidance_schedule" in call and call["guidance_schedule"] is None + + def test_begin_load_rejects_concurrent(monkeypatch): backend = DiffusionBackend() # The worker resolves the base + downloads, both over the network; stub them diff --git a/studio/backend/tests/test_diffusion_lora.py b/studio/backend/tests/test_diffusion_lora.py index 2743546255..a60140c7fc 100644 --- a/studio/backend/tests/test_diffusion_lora.py +++ b/studio/backend/tests/test_diffusion_lora.py @@ -153,10 +153,11 @@ def test_list_loras_scans_local(tmp_path, monkeypatch): (d / "other.gguf").write_bytes(b"y") (d / "ignore.txt").write_bytes(b"z") monkeypatch.setattr(dl, "loras_dir", lambda: d) - ids = {e.id for e in dl.list_loras()} - assert ids == {"mystyle", "other"} - fmts = {e.id: e.fmt for e in dl.list_loras()} - assert fmts["other"] == "gguf" and fmts["mystyle"] == "safetensors" + # The merged catalog also carries the curated hub entries; the local scan is + # exactly the weight files dropped in the directory. + local = {e.id: e for e in dl.list_loras() if e.source == "local"} + assert set(local) == {"mystyle", "other"} + assert local["other"].fmt == "gguf" and local["mystyle"].fmt == "safetensors" def test_resolve_one_local_and_unknown(tmp_path, monkeypatch): diff --git a/studio/backend/tests/test_diffusion_more_families.py b/studio/backend/tests/test_diffusion_more_families.py new file mode 100644 index 0000000000..ec4b241057 --- /dev/null +++ b/studio/backend/tests/test_diffusion_more_families.py @@ -0,0 +1,112 @@ +# 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") + + +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_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 + + +# ── 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) diff --git a/studio/frontend/src/features/images/images-page.tsx b/studio/frontend/src/features/images/images-page.tsx index ff75c0cc9b..648f8a0c15 100644 --- a/studio/frontend/src/features/images/images-page.tsx +++ b/studio/frontend/src/features/images/images-page.tsx @@ -100,6 +100,11 @@ const SAFETENSORS_MODELS: Record = { "unsloth/Z-Image-Turbo-unsloth-bnb-4bit": { kind: "pipeline" }, // Krea 2 Turbo: official vendor repo (bf16 pipeline), on the backend allowlist. "krea/Krea-2-Turbo": { kind: "pipeline" }, + // Ideogram 4: official vendor pipelines, on the backend allowlist. No bf16 repo + // exists: -fp8 stores its two DiTs as raw float8 (highest precision; ~46 GB + // resident after the bf16 cast); -nf4-diffusers is the bnb-4bit export (~11 GB). + "ideogram-ai/ideogram-4-fp8": { kind: "pipeline" }, + "ideogram-ai/ideogram-4-nf4-diffusers": { kind: "pipeline" }, "unsloth/Qwen-Image-2512-unsloth-bnb-4bit": { kind: "pipeline" }, "unsloth/Qwen-Image-2512-FP8": { kind: "single_file", @@ -135,6 +140,12 @@ const MODELS: ModelOption[] = [ "Safetensors · bnb-4bit", ), safetensors("krea/Krea-2-Turbo", "Krea 2 Turbo", "Safetensors · bf16"), + safetensors("ideogram-ai/ideogram-4-fp8", "Ideogram 4 (FP8)", "Safetensors · fp8"), + safetensors( + "ideogram-ai/ideogram-4-nf4-diffusers", + "Ideogram 4 (bnb-4bit)", + "Safetensors · bnb-4bit", + ), safetensors( "unsloth/Qwen-Image-2512-unsloth-bnb-4bit", "Qwen-Image 2512 (bnb-4bit)", @@ -223,6 +234,10 @@ const MODEL_DEFAULTS: Array<{ match: string; steps: number; guidance: number }> { match: "flux.2-dev", steps: 28, guidance: 4 }, { match: "qwen-image", steps: 20, guidance: 4 }, { match: "z-image", steps: 20, guidance: 4 }, + // Ideogram 4's model-card settings (48 steps, guidance 7). At exactly these + // defaults the backend keeps the pipeline's recommended tapered guidance schedule + // instead of a flat constant. + { match: "ideogram", steps: 48, guidance: 7 }, // SDXL: Turbo is distilled (few steps, no CFG); base/full SDXL wants ~30 steps and // real CFG (~7). "sdxl-turbo" must precede the generic "sdxl" substring match. { match: "sdxl-turbo", steps: 3, guidance: 0 }, From a5195517cf2a3735bd2322df7bbc577b599c5174 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 4 Jul 2026 14:30:58 +0000 Subject: [PATCH 02/12] Load Ideogram 4 fp8 repo by dequantizing and remapping its DiTs and text encoder The ideogram-ai/ideogram-4-fp8 repo stores its two DiTs and the Qwen3-VL text encoder in a vendor float8 layout that diffusers 0.39.0 (and diffusers main) cannot read, so a stock Ideogram4Pipeline.from_pretrained produced a pipeline with randomly initialized attention weights left on the meta device: the load then died at pipe.to(device) with "Cannot copy out of meta tensor", and any load that got past that would have generated noise. Two things broke: - The DiT attention is stored FUSED as attention.qkv.weight ([3*hidden, hidden], Q/K/V rows stacked) plus attention.o.weight, while the diffusers transformer has split to_q/to_k/to_v/to_out.0. from_pretrained mapped neither name and left them meta + random. - Every quantized weight is float8_e4m3 with a per-output-channel weight_scale; the real weight is fp8.float() * weight_scale[:, None]. diffusers dropped the scales and loaded the raw fp8 values (range +-448) as the weights, so even the weights that did map were wrong. load_ideogram4_transformer now reads the shards, dequantizes every scaled weight, splits the fused qkv into to_q/to_k/to_v and renames o to to_out.0, then loads the result into a config-constructed model. It fails loudly if any key stays unmatched so a partly random model can never ship. The dequantized fp8 projections match the byte-identical -nf4 export (already in the diffusers split layout with a bnb quantization_config) to cosine ~0.997, so the split order and scale axis are confirmed. The conversion is gated on the fp8 marker (a *.weight_scale key) read from the shard header only, so the -nf4 repos skip it and load through the stock from_pretrained path without a wasteful full-shard read. The fp8 text encoder needed the same float8 dequant (its keys already match the transformers Qwen3-VL module, so no rename). load_ideogram4_text_encoder handles the fp8 repo and delegates the bnb-4bit and dense repos to the shared krea shim. One more incompatibility was in the diffusers pipeline itself: it calls transformers create_causal_mask(inputs_embeds = ...) with no cache_position, but on transformers 4.57.6 the parameter is spelled input_embeds and cache_position is required. _patch_create_causal_mask installs a signature-aware wrapper that renames the kwarg and supplies cache_position, and is self-disabling on a matching signature. Adds unit tests for the fp8 dequant/split conversion and the causal-mask patch. Verified live on a B200: ideogram-4-fp8 (both CFG paths), ideogram-4-nf4-diffusers, and krea-2 with the retroanime LoRA all load and generate coherent images. --- studio/backend/core/inference/diffusion.py | 7 + .../core/inference/diffusion_ideogram4.py | 392 ++++++++++++++++++ .../tests/test_diffusion_more_families.py | 78 ++++ 3 files changed, 477 insertions(+) create mode 100644 studio/backend/core/inference/diffusion_ideogram4.py diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index c0bc560255..9629198532 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -45,6 +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_krea2 import KREA2_FAMILY_NAME, load_krea2_pipeline from .diffusion_memory import ( OFFLOAD_NONE, @@ -1115,6 +1116,12 @@ class DiffusionBackend: # line cannot parse; assemble the pipeline per-component # (see diffusion_krea2.py for the exact compat story). pipe = load_krea2_pipeline(repo_id, dtype, hf_token = hf_token) + elif fam.name == IDEOGRAM4_FAMILY_NAME: + # The ideogram repos ship the same transformers-5.x style Qwen + # text stack as krea (rope under rope_parameters, a slow-only + # tokenizer pin without its vocab files), so this family is + # assembled per-component too (see diffusion_ideogram4.py). + pipe = load_ideogram4_pipeline(repo_id, dtype, hf_token = hf_token) else: pipe_kwargs: dict[str, Any] = {"torch_dtype": dtype} if hf_token: diff --git a/studio/backend/core/inference/diffusion_ideogram4.py b/studio/backend/core/inference/diffusion_ideogram4.py new file mode 100644 index 0000000000..999a38e37c --- /dev/null +++ b/studio/backend/core/inference/diffusion_ideogram4.py @@ -0,0 +1,392 @@ +# 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 pipeline assembly for a transformers-4.x runtime. + +The ideogram-ai repos ship the same transformers-5.x style Qwen text stack as the +krea repos, which breaks ``Ideogram4Pipeline.from_pretrained`` twice on the 4.x +line: + +- ``text_encoder/config.json`` keeps rope settings under ``rope_parameters`` (the + 5.x name); 4.x's Qwen3-VL rotary embedding reads ``config.rope_scaling`` and + crashes on None. Fixed by ``diffusion_krea2.load_krea2_text_encoder`` (the shared + remap shim). +- ``model_index.json`` pins the SLOW ``Qwen2Tokenizer`` while the repo ships only + ``tokenizer.json`` (no vocab.json/merges.txt), so the slow class cannot even + construct -- and diffusers' passed-component type gate rejects the fast class + against the slow pin, so the fast tokenizer cannot be handed to from_pretrained + either. + +So the pipeline is assembled per-component (the constructor registers modules +without from_pretrained's type gate), mirroring ``diffusion_krea2``. + +The two DiTs need one more fix on the ``-fp8`` base repo. Its transformer shards +store the vendor's OWN float8 layout, which diffusers 0.39.0 cannot read: + +- attention is stored FUSED as ``attention.qkv.weight`` (shape ``[3*hidden, hidden]``, + the Q/K/V rows stacked in that order) plus ``attention.o.weight``, whereas the + diffusers ``Ideogram4Transformer2DModel`` has SPLIT ``to_q`` / ``to_k`` / ``to_v`` + and ``to_out.0`` projections. from_pretrained can map neither name, so it leaves + every attention projection randomly initialized (garbage images) AND on the meta + device (a later ``pipe.to(device)`` then dies with "Cannot copy out of meta tensor"). +- each quantized ``*.weight`` is float8_e4m3 with a companion per-output-channel + ``*.weight_scale`` (float32); the real weight is ``fp8.float() * weight_scale[:, None]``. + diffusers 0.39.0 has no float8 dequant path here, so it drops the scales entirely + and loads the raw fp8 values (range +-448) as if they were the weights. + +diffusers ``main`` still ships neither the fused->split rename nor the float8 dequant +(the attention module is split-only and there is no ideogram single-file converter), +so ``load_ideogram4_transformer`` does the conversion here: it reads the shards, +dequantizes every scaled weight, splits the fused ``qkv`` into ``to_q``/``to_k``/``to_v`` +and renames ``o`` -> ``to_out.0``, then loads the result into a config-constructed model +(verified against the byte-identical ``-nf4`` repo, whose transformer is ALREADY exported +in the diffusers split layout: the dequantized fp8 projections match its bnb-4bit weights +to cosine ~0.997, i.e. only quant noise apart). The already-split ``-nf4`` repos carry a +``quantization_config`` and load through the stock diffusers path, so the conversion is +gated on the fp8 marker (a ``*.weight_scale`` key) and is a no-op for them. + +The VAE loads through ``AutoencoderKLFlux2`` and the scheduler is stock +``FlowMatchEulerDiscreteScheduler``. + +One last 4.x incompatibility is in the diffusers pipeline itself, not the repo: +``Ideogram4Pipeline._get_text_encoder_hidden_states`` calls transformers' +``create_causal_mask(inputs_embeds = ...)`` with no ``cache_position``, but on the +4.x line (and even transformers 5.0) the parameter is spelled ``input_embeds`` and +``cache_position`` is required. ``_patch_create_causal_mask`` installs a signature-aware +wrapper over the name the pipeline module imported, which renames the kwarg and derives +``cache_position`` when the installed function needs one. It is self-disabling: on a +transformers whose ``create_causal_mask`` already accepts the pipeline's exact kwargs the +wrapper forwards them unchanged. +""" + +from __future__ import annotations + +import inspect +import json +from pathlib import Path +from typing import Any, Optional + +from loggers import get_logger + +from .diffusion_krea2 import load_krea2_text_encoder, load_krea2_tokenizer + +logger = get_logger(__name__) + +_CAUSAL_MASK_PATCHED = False + + +def _patch_create_causal_mask() -> None: + """Adapt the diffusers Ideogram4 pipeline's ``create_causal_mask`` call to the + installed transformers signature (see module doc). Idempotent and self-disabling. + """ + global _CAUSAL_MASK_PATCHED + if _CAUSAL_MASK_PATCHED: + return + import torch + from diffusers.pipelines.ideogram4 import pipeline_ideogram4 as pipe_mod + + original = pipe_mod.create_causal_mask + params = inspect.signature(original).parameters + + def create_causal_mask_compat(*args, **kwargs): + # The pipeline always calls this by keyword. Rename inputs_embeds -> input_embeds + # when the installed function uses the (older/5.x) spelling. + if "inputs_embeds" in kwargs and "inputs_embeds" not in params and "input_embeds" in params: + kwargs["input_embeds"] = kwargs.pop("inputs_embeds") + # Supply a cache_position when the function requires one and the caller omitted it: + # past_key_values is None here, so positions run 0..seq_len-1 over the text region. + if "cache_position" in params and "cache_position" not in kwargs: + embeds = kwargs.get("input_embeds", kwargs.get("inputs_embeds")) + if embeds is not None: + kwargs["cache_position"] = torch.arange(embeds.shape[1], device = embeds.device) + return original(*args, **kwargs) + + pipe_mod.create_causal_mask = create_causal_mask_compat + _CAUSAL_MASK_PATCHED = True + +# The fp8 attention is stored as a single fused ``qkv`` matrix with the Q, K and V +# rows stacked in that order; each block is ``hidden_size`` rows tall. hidden_size = +# attention_head_dim * num_attention_heads, read from the transformer config so a +# future config change cannot silently mis-split the matrix. +_QKV_SPLIT = ("to_q", "to_k", "to_v") + + +def _transformer_shard_paths(repo_id: str, subfolder: str, token: Optional[str]) -> list[str]: + """The local safetensors shard paths for ``repo_id/subfolder``. + + Prefers the sharded index; falls back to the single-file name when the subfolder + ships one file. Resolves through a local dir when ``repo_id`` is a path, else the + Hub cache. + """ + from huggingface_hub import hf_hub_download + + local_root = Path(repo_id).expanduser() + if local_root.is_dir(): + sub = local_root / subfolder + index = sub / "diffusion_pytorch_model.safetensors.index.json" + if index.is_file(): + weight_map = json.loads(index.read_text())["weight_map"] + return [str(sub / name) for name in sorted(set(weight_map.values()))] + single = sub / "diffusion_pytorch_model.safetensors" + if single.is_file(): + return [str(single)] + raise FileNotFoundError(f"no transformer safetensors under {sub}") + + index_name = f"{subfolder}/diffusion_pytorch_model.safetensors.index.json" + try: + index_path = hf_hub_download(repo_id, index_name, token = token) + weight_map = json.loads(Path(index_path).read_text())["weight_map"] + shards = sorted(set(weight_map.values())) + except Exception: # noqa: BLE001 -- single-file subfolder has no index + shards = ["diffusion_pytorch_model.safetensors"] + return [hf_hub_download(repo_id, f"{subfolder}/{name}", token = token) for name in shards] + + +def _read_transformer_config(repo_id: str, subfolder: str, token: Optional[str]) -> dict[str, Any]: + """``subfolder/config.json`` as a dict, from a local path or the Hub cache.""" + local = Path(repo_id).expanduser() / subfolder / "config.json" + if local.is_file(): + return json.loads(local.read_text()) + from huggingface_hub import hf_hub_download + + path = hf_hub_download(repo_id, f"{subfolder}/config.json", token = token) + return json.loads(Path(path).read_text()) + + +def _convert_fp8_state_dict(raw: dict, hidden_size: int, dtype) -> dict: + """Dequantize + rename the vendor fp8 shards into the diffusers split layout. + + A ``*.weight`` with a companion ``*.weight_scale`` is float8 stored per-output-channel: + the real weight is ``fp8.float() * weight_scale[:, None]``. The fused ``attention.qkv`` + is split into ``to_q``/``to_k``/``to_v`` (``hidden_size`` rows each, Q/K/V order) and + ``attention.o`` is renamed ``to_out.0``. Everything else (norms, biases, embeddings) is + stored dense and passes through cast to ``dtype``. + """ + import torch + + def dequantize(name: str): + weight = raw[name].to(torch.float32) + scale = raw[name + "_scale"].to(torch.float32) + return (weight * scale[:, None]).to(dtype) + + converted: dict = {} + for key, value in raw.items(): + if key.endswith("_scale"): + continue + if key + "_scale" not in raw: + # Dense (non-fp8) tensor: norms, biases, embeddings -- load as-is. + converted[key] = value.to(dtype) + continue + if key.endswith("attention.qkv.weight"): + fused = dequantize(key) # [3 * hidden_size, hidden_size] + base = key[: -len("qkv.weight")] + for index, proj in enumerate(_QKV_SPLIT): + block = fused[index * hidden_size : (index + 1) * hidden_size] + converted[f"{base}{proj}.weight"] = block.clone() + elif key.endswith("attention.o.weight"): + converted[key[: -len("o.weight")] + "to_out.0.weight"] = dequantize(key) + else: + converted[key] = dequantize(key) + return converted + + +def _text_encoder_shard_paths(repo_id: str, token: Optional[str]) -> list[str]: + """The local safetensors shard paths for ``repo_id/text_encoder`` (index or single file).""" + from huggingface_hub import hf_hub_download + + local_root = Path(repo_id).expanduser() + if local_root.is_dir(): + sub = local_root / "text_encoder" + index = sub / "model.safetensors.index.json" + if index.is_file(): + weight_map = json.loads(index.read_text())["weight_map"] + return [str(sub / name) for name in sorted(set(weight_map.values()))] + single = sub / "model.safetensors" + if single.is_file(): + return [str(single)] + raise FileNotFoundError(f"no text_encoder safetensors under {sub}") + + try: + index_path = hf_hub_download(repo_id, "text_encoder/model.safetensors.index.json", token = token) + weight_map = json.loads(Path(index_path).read_text())["weight_map"] + shards = sorted(set(weight_map.values())) + except Exception: # noqa: BLE001 -- single-file text encoder has no index + shards = ["model.safetensors"] + return [hf_hub_download(repo_id, f"text_encoder/{name}", token = token) for name in shards] + + +def _text_encoder_is_fp8(repo_id: str, token: Optional[str]) -> bool: + """True when the text_encoder ships the vendor fp8 layout (a ``*.weight_scale`` key).""" + from huggingface_hub import hf_hub_download + + local_root = Path(repo_id).expanduser() + if local_root.is_dir(): + index = local_root / "text_encoder" / "model.safetensors.index.json" + if index.is_file(): + return any(k.endswith("_scale") for k in json.loads(index.read_text())["weight_map"]) + else: + try: + index_path = hf_hub_download( + repo_id, "text_encoder/model.safetensors.index.json", token = token + ) + weight_map = json.loads(Path(index_path).read_text())["weight_map"] + return any(k.endswith("_scale") for k in weight_map) + except Exception: # noqa: BLE001 -- single-file (nf4) text encoder, not fp8 + return False + # Single-file local text encoder: peek the header keys. + import safetensors + + single = local_root / "text_encoder" / "model.safetensors" + if single.is_file(): + with safetensors.safe_open(str(single), "pt") as handle: + return any(k.endswith("_scale") for k in handle.keys()) + return False + + +def load_ideogram4_text_encoder(repo_id: str, dtype, hf_token: Optional[str] = None): + """The Qwen3-VL text encoder for ``repo_id``. + + The ``-fp8`` repo stores this encoder in the SAME float8-plus-per-channel-scale + layout as its DiTs, and its keys already match the transformers Qwen3-VL module + (only the DiTs used the fused ``qkv``; Qwen3-VL's own attention is already split + and its visual tower's fused ``qkv`` matches transformers), so it needs no rename + -- only the float8 dequant diffusers/transformers skip. So the fp8 encoder is + dequantized and loaded into a config-constructed model; the ``-nf4`` (bnb-4bit) + and any dense repo fall through to the shared krea shim (which also applies the + rope_parameters remap). + """ + token = hf_token or None + if not _text_encoder_is_fp8(repo_id, token): + return load_krea2_text_encoder(repo_id, dtype, hf_token = token) + + import safetensors + import torch + from transformers import AutoConfig, Qwen3VLModel + + from .diffusion_krea2 import remap_rope_parameters + + config_kwargs: dict[str, Any] = {"subfolder": "text_encoder"} + if token: + config_kwargs["token"] = token + config = AutoConfig.from_pretrained(repo_id, **config_kwargs) + remap_rope_parameters(getattr(config, "text_config", config)) + + raw: dict = {} + for path in _text_encoder_shard_paths(repo_id, token): + with safetensors.safe_open(path, "pt") as handle: + for key in handle.keys(): + raw[key] = handle.get_tensor(key) + + state_dict: dict = {} + for key, value in raw.items(): + if key.endswith("_scale"): + continue + if key + "_scale" in raw: + weight = value.to(torch.float32) + scale = raw[key + "_scale"].to(torch.float32) + state_dict[key] = (weight * scale[:, None]).to(dtype) + else: + state_dict[key] = value.to(dtype) + + # Construct normally (so __init__ computes the non-persistent rotary inv_freq + # buffers the checkpoint omits) then copy the dequantized weights in with + # assign=False. Host RAM is ample, so the transient dense init is fine. + model = Qwen3VLModel(config).to(dtype) + missing, unexpected = model.load_state_dict(state_dict, strict = False) + real_missing = [k for k in missing if not k.endswith("inv_freq")] + if real_missing or unexpected: + raise RuntimeError( + f"ideogram4 fp8 text_encoder remap left keys unmatched for {repo_id}: " + f"missing={real_missing[:8]} unexpected={unexpected[:8]}" + ) + return model + + +def load_ideogram4_transformer(repo_id: str, subfolder: str, dtype, hf_token: Optional[str] = None): + """An ``Ideogram4Transformer2DModel`` for ``repo_id/subfolder`` (still on CPU). + + Reads the transformer config, and if the shards carry the vendor fp8 layout + (a ``*.weight_scale`` key), dequantizes + renames them into the diffusers split + layout and loads that into a config-constructed model. When the shards are already + in the diffusers layout (the ``-nf4`` repos, which carry a ``quantization_config``), + delegates to the stock ``from_pretrained`` so bnb re-applies the 4-bit weights. + """ + import diffusers + import safetensors + + token = hf_token or None + config = _read_transformer_config(repo_id, subfolder, token) + shard_paths = _transformer_shard_paths(repo_id, subfolder, token) + + # Detect the fp8 layout from the shard HEADER (safe_open.keys() reads metadata only, + # not the multi-GB tensor bodies). Only the fp8 path then materializes the tensors; + # the -nf4 path goes straight to from_pretrained without a wasteful full-shard read. + with safetensors.safe_open(shard_paths[0], "pt") as handle: + is_fp8 = any(key.endswith("_scale") for key in handle.keys()) + if not is_fp8: + # Already the diffusers split layout (the quantized -nf4 exports). Let + # from_pretrained re-apply the embedded quantization_config unchanged. + model_kwargs: dict[str, Any] = {"subfolder": subfolder, "torch_dtype": dtype} + if token: + model_kwargs["token"] = token + return diffusers.Ideogram4Transformer2DModel.from_pretrained(repo_id, **model_kwargs) + + raw: dict = {} + for path in shard_paths: + with safetensors.safe_open(path, "pt") as handle: + for key in handle.keys(): + raw[key] = handle.get_tensor(key) + + config.pop("quantization_config", None) + hidden_size = int(config["attention_head_dim"]) * int(config["num_attention_heads"]) + model = diffusers.Ideogram4Transformer2DModel.from_config(config) + state_dict = _convert_fp8_state_dict(raw, hidden_size, dtype) + missing, unexpected = model.load_state_dict(state_dict, strict = False) + # rotary_emb.inv_freq is a non-persistent buffer built in __init__, so it is + # (correctly) absent from the checkpoint and the only expected "missing" key; a + # real gap (an unmapped weight) or any leftover checkpoint key must fail loudly + # rather than ship a partly random model. + real_missing = [k for k in missing if not k.endswith("rotary_emb.inv_freq")] + if real_missing or unexpected: + raise RuntimeError( + f"ideogram4 fp8 remap left keys unmatched for {repo_id}/{subfolder}: " + f"missing={real_missing[:8]} unexpected={unexpected[:8]}" + ) + model.to(dtype) + return model + + +def load_ideogram4_pipeline(repo_id: str, dtype, hf_token: Optional[str] = None): + """Assemble Ideogram4Pipeline from ``repo_id`` per-component (see module doc).""" + import diffusers + + # The pipeline's text-encoder call uses a transformers-5.x create_causal_mask + # signature; adapt it to the installed one before any generate runs. + _patch_create_causal_mask() + + token = hf_token or None + model_kwargs: dict[str, Any] = {"torch_dtype": dtype} + if token: + model_kwargs["token"] = token + + text_encoder = load_ideogram4_text_encoder(repo_id, dtype, hf_token = token) + tokenizer = load_krea2_tokenizer(repo_id, hf_token = token) + transformer = load_ideogram4_transformer(repo_id, "transformer", dtype, hf_token = token) + # The second DiT drives the unconditional branch of Ideogram's dual-branch CFG; + # it is the same class and size as the conditional one and always required. + unconditional_transformer = load_ideogram4_transformer( + repo_id, "unconditional_transformer", dtype, hf_token = token + ) + vae = diffusers.AutoencoderKLFlux2.from_pretrained(repo_id, subfolder = "vae", **model_kwargs) + scheduler = diffusers.FlowMatchEulerDiscreteScheduler.from_pretrained( + repo_id, subfolder = "scheduler", token = token + ) + logger.info("diffusion.ideogram4: assembled pipeline from %s per-component", repo_id) + return diffusers.Ideogram4Pipeline( + scheduler = scheduler, + vae = vae, + text_encoder = text_encoder, + tokenizer = tokenizer, + transformer = transformer, + unconditional_transformer = unconditional_transformer, + ) diff --git a/studio/backend/tests/test_diffusion_more_families.py b/studio/backend/tests/test_diffusion_more_families.py index ec4b241057..738c8b90a5 100644 --- a/studio/backend/tests/test_diffusion_more_families.py +++ b/studio/backend/tests/test_diffusion_more_families.py @@ -110,3 +110,81 @@ def test_list_loras_family_filter_gates_krea_entries(): 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_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 From 65d233836400c99504b38aa5b6bc150c5b70654f Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 4 Jul 2026 14:44:49 +0000 Subject: [PATCH 03/12] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../core/inference/diffusion_ideogram4.py | 24 +++++++++++++++---- .../tests/test_diffusion_more_families.py | 1 - 2 files changed, 20 insertions(+), 5 deletions(-) diff --git a/studio/backend/core/inference/diffusion_ideogram4.py b/studio/backend/core/inference/diffusion_ideogram4.py index 999a38e37c..448240eecf 100644 --- a/studio/backend/core/inference/diffusion_ideogram4.py +++ b/studio/backend/core/inference/diffusion_ideogram4.py @@ -104,6 +104,7 @@ def _patch_create_causal_mask() -> None: pipe_mod.create_causal_mask = create_causal_mask_compat _CAUSAL_MASK_PATCHED = True + # The fp8 attention is stored as a single fused ``qkv`` matrix with the Q, K and V # rows stacked in that order; each block is ``hidden_size`` rows tall. hidden_size = # attention_head_dim * num_attention_heads, read from the transformer config so a @@ -207,7 +208,9 @@ def _text_encoder_shard_paths(repo_id: str, token: Optional[str]) -> list[str]: raise FileNotFoundError(f"no text_encoder safetensors under {sub}") try: - index_path = hf_hub_download(repo_id, "text_encoder/model.safetensors.index.json", token = token) + index_path = hf_hub_download( + repo_id, "text_encoder/model.safetensors.index.json", token = token + ) weight_map = json.loads(Path(index_path).read_text())["weight_map"] shards = sorted(set(weight_map.values())) except Exception: # noqa: BLE001 -- single-file text encoder has no index @@ -243,7 +246,11 @@ def _text_encoder_is_fp8(repo_id: str, token: Optional[str]) -> bool: return False -def load_ideogram4_text_encoder(repo_id: str, dtype, hf_token: Optional[str] = None): +def load_ideogram4_text_encoder( + repo_id: str, + dtype, + hf_token: Optional[str] = None, +): """The Qwen3-VL text encoder for ``repo_id``. The ``-fp8`` repo stores this encoder in the SAME float8-plus-per-channel-scale @@ -302,7 +309,12 @@ def load_ideogram4_text_encoder(repo_id: str, dtype, hf_token: Optional[str] = N return model -def load_ideogram4_transformer(repo_id: str, subfolder: str, dtype, hf_token: Optional[str] = None): +def load_ideogram4_transformer( + repo_id: str, + subfolder: str, + dtype, + hf_token: Optional[str] = None, +): """An ``Ideogram4Transformer2DModel`` for ``repo_id/subfolder`` (still on CPU). Reads the transformer config, and if the shards carry the vendor fp8 layout @@ -356,7 +368,11 @@ def load_ideogram4_transformer(repo_id: str, subfolder: str, dtype, hf_token: Op return model -def load_ideogram4_pipeline(repo_id: str, dtype, hf_token: Optional[str] = None): +def load_ideogram4_pipeline( + repo_id: str, + dtype, + hf_token: Optional[str] = None, +): """Assemble Ideogram4Pipeline from ``repo_id`` per-component (see module doc).""" import diffusers diff --git a/studio/backend/tests/test_diffusion_more_families.py b/studio/backend/tests/test_diffusion_more_families.py index 738c8b90a5..2e21e3023f 100644 --- a/studio/backend/tests/test_diffusion_more_families.py +++ b/studio/backend/tests/test_diffusion_more_families.py @@ -85,7 +85,6 @@ def test_excluded_model_reason_none_for_supported_and_unknown(): 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") From 85395e3b94465c779947c55f3325e36224d4485e Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 5 Jul 2026 00:16:14 +0000 Subject: [PATCH 04/12] Harden ideogram fp8 dequant and scope the HunyuanImage exclusion Review follow ups on the more-families branch: the per channel scale now broadcasts rank aware instead of assuming 2D (all shipped tensors are 2D, verified across all three fp8 components, but a future non 2D quantized tensor would have mis broadcast silently), the fused qkv split asserts the expected 3x hidden row count so a GQA style export fails loudly, fp8 detection scans every shard header rather than the first, and the excluded model match uses the segment aware token helper with a hunyuanimage-3 token so a future HunyuanImage 2.x is not blocked with a 3.0 reason. --- .../core/inference/diffusion_families.py | 6 ++-- .../core/inference/diffusion_ideogram4.py | 31 ++++++++++++++----- 2 files changed, 28 insertions(+), 9 deletions(-) diff --git a/studio/backend/core/inference/diffusion_families.py b/studio/backend/core/inference/diffusion_families.py index a6e47f2388..38afb0b210 100644 --- a/studio/backend/core/inference/diffusion_families.py +++ b/studio/backend/core/inference/diffusion_families.py @@ -376,7 +376,9 @@ IDEOGRAM4_FAMILY_NAME = "ideogram-4" # assemble, and remote-code execution is out of the question for a load path. _EXCLUDED_MODELS: tuple[tuple[str, str], ...] = ( ( - "hunyuanimage", + # "-3" scoped: the reason is 3.0-specific, and a future HunyuanImage 2.x with a + # diffusers pipeline must fall through to normal (unknown-family) handling. + "hunyuanimage-3", "HunyuanImage-3.0 has no diffusers pipeline (it is an 80B autoregressive MoE " "that requires trust_remote_code), so Studio does not support it.", ), @@ -387,7 +389,7 @@ def excluded_model_reason(repo_id: str) -> Optional[str]: """The stated reason ``repo_id`` is unsupported, or None when it is simply unknown.""" needle = (repo_id or "").lower() for token, reason in _EXCLUDED_MODELS: - if token in needle: + if _token_in_needle(token, needle): return reason return None diff --git a/studio/backend/core/inference/diffusion_ideogram4.py b/studio/backend/core/inference/diffusion_ideogram4.py index 448240eecf..0bbaf63bda 100644 --- a/studio/backend/core/inference/diffusion_ideogram4.py +++ b/studio/backend/core/inference/diffusion_ideogram4.py @@ -168,7 +168,10 @@ def _convert_fp8_state_dict(raw: dict, hidden_size: int, dtype) -> dict: def dequantize(name: str): weight = raw[name].to(torch.float32) scale = raw[name + "_scale"].to(torch.float32) - return (weight * scale[:, None]).to(dtype) + # Per-output-channel scale, broadcast over the remaining dims. Every scaled + # tensor in the shipped repos is 2D; the rank-aware view keeps a future + # non-2D quantized tensor correct instead of silently mis-broadcasting. + return (weight * scale.view(-1, *([1] * (weight.ndim - 1)))).to(dtype) converted: dict = {} for key, value in raw.items(): @@ -180,6 +183,13 @@ def _convert_fp8_state_dict(raw: dict, hidden_size: int, dtype) -> dict: continue if key.endswith("attention.qkv.weight"): fused = dequantize(key) # [3 * hidden_size, hidden_size] + if fused.shape[0] != 3 * hidden_size: + # Equal-thirds is only correct for full multi-head attention; a GQA + # export (fewer K/V rows) must fail loudly, not split into garbage. + raise RuntimeError( + f"fused qkv at {key} has {fused.shape[0]} rows, expected " + f"{3 * hidden_size}; cannot split into equal Q/K/V blocks" + ) base = key[: -len("qkv.weight")] for index, proj in enumerate(_QKV_SPLIT): block = fused[index * hidden_size : (index + 1) * hidden_size] @@ -291,7 +301,8 @@ def load_ideogram4_text_encoder( if key + "_scale" in raw: weight = value.to(torch.float32) scale = raw[key + "_scale"].to(torch.float32) - state_dict[key] = (weight * scale[:, None]).to(dtype) + # Rank-aware broadcast, matching _convert_fp8_state_dict. + state_dict[key] = (weight * scale.view(-1, *([1] * (weight.ndim - 1)))).to(dtype) else: state_dict[key] = value.to(dtype) @@ -330,11 +341,17 @@ def load_ideogram4_transformer( config = _read_transformer_config(repo_id, subfolder, token) shard_paths = _transformer_shard_paths(repo_id, subfolder, token) - # Detect the fp8 layout from the shard HEADER (safe_open.keys() reads metadata only, - # not the multi-GB tensor bodies). Only the fp8 path then materializes the tensors; - # the -nf4 path goes straight to from_pretrained without a wasteful full-shard read. - with safetensors.safe_open(shard_paths[0], "pt") as handle: - is_fp8 = any(key.endswith("_scale") for key in handle.keys()) + # Detect the fp8 layout from the shard HEADERS (safe_open.keys() reads metadata only, + # not the multi-GB tensor bodies). All shards are checked so a multi-shard export + # whose first shard happens to hold only dense tensors still routes to the dequant + # path. Only the fp8 path then materializes the tensors; the -nf4 path goes straight + # to from_pretrained without a wasteful full-shard read. + is_fp8 = False + for path in shard_paths: + with safetensors.safe_open(path, "pt") as handle: + if any(key.endswith("_scale") for key in handle.keys()): + is_fp8 = True + break if not is_fp8: # Already the diffusers split layout (the quantized -nf4 exports). Let # from_pretrained re-apply the embedded quantization_config unchanged. From 330897bb62b1b3adfa5c9f628f80111ff7f53d65 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 5 Jul 2026 04:48:37 +0000 Subject: [PATCH 05/12] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/tests/test_diffusion_backend.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index 623c210b7a..97ea60c6da 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -2209,8 +2209,7 @@ def test_prefetch_returns_snapshot_dir_for_manifest(monkeypatch): ) assert root == "/cache/snap" assert ( - backend._prefetch_files("base/repo", None, "base/repo", ["vae/x.safetensors"], None) - is None + backend._prefetch_files("base/repo", None, "base/repo", ["vae/x.safetensors"], None) is None ) From c23c0d6047f3afba610927e0bfac7e4e62527d07 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 5 Jul 2026 08:02:42 +0000 Subject: [PATCH 06/12] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/core/training/diffusion_train_common.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/studio/backend/core/training/diffusion_train_common.py b/studio/backend/core/training/diffusion_train_common.py index 30b6ec727a..bdf53d5b6b 100644 --- a/studio/backend/core/training/diffusion_train_common.py +++ b/studio/backend/core/training/diffusion_train_common.py @@ -597,7 +597,7 @@ def _plan_cache_variants( # fallback. Over this budget the default falls back to per-step VAE encoding. A fixed # constant (rather than a psutil RAM fraction) keeps the gate dependency-free and identical # across hosts; it is deliberately conservative, well under a typical training host's RAM. -_LATENT_CACHE_BUDGET_BYTES = 4 * 1024 ** 3 # 4 GiB +_LATENT_CACHE_BUDGET_BYTES = 4 * 1024**3 # 4 GiB # Returned by the cache builders when the estimated cache exceeds the budget: the caller # keeps the VAE resident and encodes each step's latents in-loop. A distinct sentinel from @@ -614,7 +614,9 @@ def _latent_cache_forced() -> bool: def _latent_cache_over_budget( - per_variant_bytes: int, total_variants: int, budget_bytes: Optional[int] = None + per_variant_bytes: int, + total_variants: int, + budget_bytes: Optional[int] = None, ) -> bool: """True when a cache of ``total_variants`` entries, each two fp32 tensors totalling ``per_variant_bytes``, is estimated to exceed ``budget_bytes``. ``per_variant_bytes`` is From 31f04c026ecca57092d1e3669b33f9f2bf93f4ee Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 5 Jul 2026 11:46:44 +0000 Subject: [PATCH 07/12] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/core/inference/diffusion.py | 1 - studio/backend/tests/test_diffusion_controlnet.py | 8 ++++++-- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index fef4cad433..b0712af2d0 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -1801,7 +1801,6 @@ class DiffusionBackend: # local dir the user picked has no Hub scan and is exempt (fail-open there). if not getattr(resolved_cn, "is_local", False): from utils.security import evaluate_file_security - _cn_fs = evaluate_file_security(resolved_cn.path, hf_token = state.hf_token or None) if _cn_fs.blocked: raise ValueError(_cn_fs.reason) diff --git a/studio/backend/tests/test_diffusion_controlnet.py b/studio/backend/tests/test_diffusion_controlnet.py index b7ae037197..bcb7393457 100644 --- a/studio/backend/tests/test_diffusion_controlnet.py +++ b/studio/backend/tests/test_diffusion_controlnet.py @@ -234,7 +234,6 @@ def _state(): def _allow_cn_security(monkeypatch): """Stub the Hub malware preflight to allow the load (hermetic, no network).""" import utils.security - monkeypatch.setattr( utils.security, "evaluate_file_security", @@ -277,7 +276,12 @@ def test_controlnet_pipe_blocks_flagged_remote_repo(monkeypatch): class _TrapModel(_FakeCNModel): @classmethod - def from_pretrained(cls, path, torch_dtype = None, token = None): + def from_pretrained( + cls, + path, + torch_dtype = None, + token = None, + ): loaded["called"] = True return super().from_pretrained(path, torch_dtype = torch_dtype, token = token) From f30713181984111dc39b1547c4072341008ecfae Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 5 Jul 2026 12:08:45 +0000 Subject: [PATCH 08/12] Studio: reject transformer-only loads for Ideogram 4 and size local fp8 mirrors Ideogram 4 assembles two DiTs per-component (a conditional transformer plus a separate unconditional_transformer), so there is no transformer-only single-file or GGUF artifact that could supply both. Add a pipeline_only family flag and reject the gguf/single_file kinds in validate_load_request, before a load evicts the current model, instead of assembling a pipeline missing its second DiT. Extend the fp8 bf16-resident size override to a LOCAL directory mirror of the ideogram-4-fp8 base: such a path never string-matches base_repo, so detect the fp8 layout from the transformer shard headers (a *.weight_scale marker) and reserve the bf16 footprint, matching the remote-base behaviour. A local nf4 mirror has no fp8 scales and correctly stays planned against its compressed bytes. --- studio/backend/core/inference/diffusion.py | 27 ++++++++++++- .../core/inference/diffusion_families.py | 10 +++++ .../core/inference/diffusion_ideogram4.py | 23 +++++++++++ .../backend/tests/test_diffusion_backend.py | 40 +++++++++++++++---- .../tests/test_diffusion_more_families.py | 33 +++++++++++++++ 5 files changed, 123 insertions(+), 10 deletions(-) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index b0712af2d0..91a2acb699 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -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)) diff --git a/studio/backend/core/inference/diffusion_families.py b/studio/backend/core/inference/diffusion_families.py index f6ccd23431..2fbefeacc0 100644 --- a/studio/backend/core/inference/diffusion_families.py +++ b/studio/backend/core/inference/diffusion_families.py @@ -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 diff --git a/studio/backend/core/inference/diffusion_ideogram4.py b/studio/backend/core/inference/diffusion_ideogram4.py index 0bbaf63bda..ddc35d46ec 100644 --- a/studio/backend/core/inference/diffusion_ideogram4.py +++ b/studio/backend/core/inference/diffusion_ideogram4.py @@ -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, diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index 97ea60c6da..bfe93e5b58 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -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): diff --git a/studio/backend/tests/test_diffusion_more_families.py b/studio/backend/tests/test_diffusion_more_families.py index 2e21e3023f..9d89eb5ab3 100644 --- a/studio/backend/tests/test_diffusion_more_families.py +++ b/studio/backend/tests/test_diffusion_more_families.py @@ -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, From 70af9fa75a6833da7145f10b9dcc72620eb6cfbc Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 6 Jul 2026 08:38:03 +0000 Subject: [PATCH 09/12] ideogram4: build fp8 DiT at target dtype to halve host-RAM transient The fp8 loader builds Ideogram4Transformer2DModel via from_config, which materializes the full ~9B-parameter module at the process default dtype (fp32) before the dequantized bf16 weights are copied in and cast at the end. That fp32 scaffold is ~2x the bf16 model (~37 GB vs ~18 GB) on host RAM, and the second (unconditional) DiT builds while the first DiT and the text encoder are already resident, so it can OOM smaller hosts. Wrap from_config in set_default_dtype(dtype) so the module is built at the target dtype directly. rotary_emb.inv_freq (the only __init__ state absent from the checkpoint) is computed in explicit fp32, so a bf16 default leaves it correct. --- .../backend/core/inference/diffusion_ideogram4.py | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/studio/backend/core/inference/diffusion_ideogram4.py b/studio/backend/core/inference/diffusion_ideogram4.py index ddc35d46ec..a4ac1dfec2 100644 --- a/studio/backend/core/inference/diffusion_ideogram4.py +++ b/studio/backend/core/inference/diffusion_ideogram4.py @@ -359,6 +359,7 @@ def load_ideogram4_transformer( """ import diffusers import safetensors + import torch token = hf_token or None config = _read_transformer_config(repo_id, subfolder, token) @@ -391,7 +392,19 @@ def load_ideogram4_transformer( config.pop("quantization_config", None) hidden_size = int(config["attention_head_dim"]) * int(config["num_attention_heads"]) - model = diffusers.Ideogram4Transformer2DModel.from_config(config) + # from_config materializes the full ~9B-param module before the dequantized weights + # are copied in. At the process default (fp32) that scaffold is ~2x the bf16 model + # (~37 GB vs ~18 GB) on host RAM, and the second (unconditional) DiT builds while the + # first DiT and the text encoder are already resident, so the fp32 transient can OOM + # smaller hosts. Build at the target dtype instead; the only __init__ state absent from + # the checkpoint is rotary_emb.inv_freq (computed in explicit fp32), so a bf16 default + # leaves it correct while halving each DiT's transient peak. + default_dtype = torch.get_default_dtype() + torch.set_default_dtype(dtype) + try: + model = diffusers.Ideogram4Transformer2DModel.from_config(config) + finally: + torch.set_default_dtype(default_dtype) state_dict = _convert_fp8_state_dict(raw, hidden_size, dtype) missing, unexpected = model.load_state_dict(state_dict, strict = False) # rotary_emb.inv_freq is a non-persistent buffer built in __init__, so it is From bb2b14db9795ec135e1c1b3d1744eab9ca1b1317 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 6 Jul 2026 10:21:10 +0000 Subject: [PATCH 10/12] ideogram-4: build the FP8 text encoder at target dtype, size it as bf16-resident, optimize both DiTs - The FP8 Qwen3-VL text encoder was constructed at the process fp32 default before the dequantized bf16 weights are copied in. That ~8B-param fp32 scaffold peaks ~2x on host RAM (loading FIRST, before the DiTs), so a 64 GB host can OOM. Build it at the target dtype under set_default_dtype, mirroring the DiT loader; rotary inv_freq is still computed in explicit fp32. - The auto-policy memory table listed the text encoder at 8.8 GB, its FP8 on-disk size, while the DiTs were doubled to their bf16-resident sizes. The loader dequantizes the encoder to bf16 too (~16.3 GB), so the entry understated the resident footprint by ~7.5 GB and could let the planner pick a resident placement that OOMs. Size it as bf16-resident. - Speed (regional compile, QKV fuse) and the attention backend only touched pipe.transformer, so ideogram-4's second denoiser (unconditional_transformer, run every step for dual-branch CFG) stayed eager/native while status reported the optimization as engaged. Iterate every denoiser DiT (mirroring the offload path) so both experts are optimized. Guarded on attr presence, so single-DiT families are unchanged. --- .../core/inference/diffusion_attention.py | 50 +++++++++++------ .../core/inference/diffusion_auto_policy.py | 9 +-- .../core/inference/diffusion_ideogram4.py | 12 +++- .../backend/core/inference/diffusion_speed.py | 55 +++++++++++++++---- .../backend/tests/test_diffusion_attention.py | 11 ++++ studio/backend/tests/test_diffusion_speed.py | 24 ++++++++ 6 files changed, 127 insertions(+), 34 deletions(-) diff --git a/studio/backend/core/inference/diffusion_attention.py b/studio/backend/core/inference/diffusion_attention.py index 7377f8d919..6873026cd1 100644 --- a/studio/backend/core/inference/diffusion_attention.py +++ b/studio/backend/core/inference/diffusion_attention.py @@ -244,13 +244,27 @@ def _ensure_attention_backend_installed(backend: str, logger: Any = None) -> Non ) +def _attention_dits(pipe: Any) -> list: + """Every DiT the denoise loop runs each step: the primary ``transformer`` plus a second + expert some families carry (Ideogram's ``unconditional_transformer`` for its dual-branch + CFG, an MoE ``transformer_2``). The attention backend must be set on ALL of them, else the + second DiT keeps the native default while status reports the requested kernel as engaged.""" + dits: list = [] + for attr in ("transformer", "transformer_2", "unconditional_transformer"): + m = getattr(pipe, attr, None) + if m is not None and m not in dits: + dits.append(m) + return dits + + def apply_attention_backend( pipe: Any, backend: Optional[str], *, logger: Any = None, ) -> Optional[str]: - """Set ``backend`` on ``pipe.transformer`` via the diffusers dispatcher. + """Set ``backend`` on EVERY denoiser DiT (``pipe.transformer`` plus a second expert such as + Ideogram's ``unconditional_transformer``) via the diffusers dispatcher. Returns the backend actually engaged, or None when left at the native default (either because ``backend`` was None or because the requested kernel was unavailable -> graceful @@ -261,28 +275,32 @@ def apply_attention_backend( defaults to None). So a load that wants native must restore it explicitly: otherwise it silently inherits a backend an earlier load pinned (e.g. cuDNN under a speed profile), breaking the bit-identical/``off`` guarantee. Best-effort throughout.""" - transformer = getattr(pipe, "transformer", None) - fn = getattr(transformer, "set_attention_backend", None) - if not callable(fn): + setters = [s for s in (getattr(t, "set_attention_backend", None) for t in _attention_dits(pipe)) + if callable(s)] + if not setters: return None if backend is not None: _ensure_attention_backend_installed(backend, logger) - try: - fn(backend) - # set_attention_backend also pins the backend in diffusers' process-wide - # registry. This transformer's own processors keep it locally (their - # _attention_backend is now explicit), so reset the global default back to - # native -- otherwise a later component whose processors are unconfigured - # (backend None) silently inherits this kernel. + engaged = False + for fn in setters: + try: + fn(backend) + engaged = True + except Exception as exc: # noqa: BLE001 — unavailable kernel -> restore native below + _warn(logger, backend, exc) + if engaged: + # set_attention_backend also pins the backend in diffusers' process-wide registry. + # Each DiT's own processors now keep it locally (their _attention_backend is now + # explicit), so reset the global default back to native ONCE -- otherwise a later + # component whose processors are unconfigured (backend None) inherits this kernel. _reset_global_backend_to_native(logger) if logger is not None: logger.info("diffusion.attention: backend=%s", backend) return backend - except Exception as exc: # noqa: BLE001 — unavailable kernel -> restore native below - _warn(logger, backend, exc) - # No backend requested, or the requested one failed: pin the native default so a stale - # process-wide backend from a previous load can't leak into this one. - _restore_native_backend(fn, logger) + # No backend requested, or every set failed: pin the native default so a stale process-wide + # backend from a previous load can't leak into this one. Fresh DiTs follow the process-wide + # backend, so one reset via any DiT's setter covers them all. + _restore_native_backend(setters[0], logger) return None diff --git a/studio/backend/core/inference/diffusion_auto_policy.py b/studio/backend/core/inference/diffusion_auto_policy.py index 8433e4e2a2..85653364dd 100644 --- a/studio/backend/core/inference/diffusion_auto_policy.py +++ b/studio/backend/core/inference/diffusion_auto_policy.py @@ -61,10 +61,11 @@ _FAMILY_BF16_GB: dict[str, tuple[float, float, float]] = { "krea-2": (26.3, 8.9, 0.5), # Two ~9.3B DiTs (the conditional transformer PLUS the separate # unconditional_transformer driving Ideogram's dual-branch CFG), both resident - # for every generation, and a Qwen3-VL text encoder. The vendor repo stores the - # DiTs as raw float8 (9.29 GB each); these are the bf16-resident sizes after the - # loader's dtype cast, per this table's contract. - "ideogram-4": (37.2, 8.8, 0.2), + # for every generation, and a Qwen3-VL text encoder. The vendor repo stores the DiTs + # AND the text encoder as raw float8 (9.29 GB per DiT, 8.8 GB encoder); these are the + # bf16-resident sizes after the loader's dtype cast, per this table's contract, so the + # encoder doubles to ~16.3 GB just like each DiT (37.2 = 2 x 18.6). + "ideogram-4": (37.2, 16.3, 0.2), } # Base-repo overrides for families whose picker offers multiple sizes under one family diff --git a/studio/backend/core/inference/diffusion_ideogram4.py b/studio/backend/core/inference/diffusion_ideogram4.py index a4ac1dfec2..4e3e321604 100644 --- a/studio/backend/core/inference/diffusion_ideogram4.py +++ b/studio/backend/core/inference/diffusion_ideogram4.py @@ -308,8 +308,16 @@ def load_ideogram4_text_encoder( # Construct normally (so __init__ computes the non-persistent rotary inv_freq # buffers the checkpoint omits) then copy the dequantized weights in with - # assign=False. Host RAM is ample, so the transient dense init is fine. - model = Qwen3VLModel(config).to(dtype) + # assign=False. Build at the target dtype (mirrors the DiT loader below): this ~8B-param + # Qwen3-VL scaffold is ~2x at the process fp32 default (~33 GB vs ~16 GB) and loads FIRST + # on host RAM, so the fp32 transient can OOM a 64 GB host. rotary inv_freq is computed in + # explicit fp32 in __init__, so a bf16 default leaves it correct. + default_dtype = torch.get_default_dtype() + torch.set_default_dtype(dtype) + try: + model = Qwen3VLModel(config).to(dtype) + finally: + torch.set_default_dtype(default_dtype) missing, unexpected = model.load_state_dict(state_dict, strict = False) real_missing = [k for k in missing if not k.endswith("inv_freq")] if real_missing or unexpected: diff --git a/studio/backend/core/inference/diffusion_speed.py b/studio/backend/core/inference/diffusion_speed.py index 8c2524dfda..1f9a736f2f 100644 --- a/studio/backend/core/inference/diffusion_speed.py +++ b/studio/backend/core/inference/diffusion_speed.py @@ -272,6 +272,20 @@ def _vae_channels_last(pipe: Any, logger: Any) -> bool: return False +def _denoiser_dits(pipe: Any) -> list: + """Every DiT the denoise loop runs each step: the primary ``transformer`` plus a second + expert some families carry (Ideogram's ``unconditional_transformer`` for its dual-branch + CFG, an MoE ``transformer_2``). Speed / attention optims must reach ALL of them -- mirroring + the offload path (diffusion_memory streams the same set) -- else the second DiT runs + eager / native for every generation while status over-reports the optimisation as engaged.""" + dits: list = [] + for attr in ("transformer", "transformer_2", "unconditional_transformer"): + m = getattr(pipe, attr, None) + if m is not None and m not in dits: + dits.append(m) + return dits + + def _compile_repeated_blocks( pipe: Any, logger: Any, @@ -280,9 +294,8 @@ def _compile_repeated_blocks( cache_active: bool = False, offload_active: bool = False, ) -> bool: - transformer = getattr(pipe, "transformer", None) - fn = getattr(transformer, "compile_repeated_blocks", None) - if not callable(fn): + dits = [t for t in _denoiser_dits(pipe) if callable(getattr(t, "compile_repeated_blocks", None))] + if not dits: return False # default: mode="default" + dynamic=True -- fast cold start, robust to resolution # changes (no recompile). max: mode="max-autotune-no-cudagraphs" + dynamic=False -- @@ -319,11 +332,19 @@ def _compile_repeated_blocks( for _limit_attr in ("recompile_limit", "cache_size_limit"): # name varies by torch ver if hasattr(dynamo_cfg, _limit_attr): setattr(dynamo_cfg, _limit_attr, max(getattr(dynamo_cfg, _limit_attr) or 0, 64)) - fn(**kwargs) - return True except Exception as exc: # noqa: BLE001 — optimisation only _warn(logger, "compile_repeated_blocks", exc) return False + # Compile every denoiser DiT (a dual-DiT family such as Ideogram runs both each step); a + # per-DiT failure degrades that one to eager without dropping the others. + engaged = False + for transformer in dits: + try: + transformer.compile_repeated_blocks(**kwargs) + engaged = True + except Exception as exc: # noqa: BLE001 — optimisation only + _warn(logger, "compile_repeated_blocks", exc) + return engaged def _enable_cudnn_benchmark(logger: Any) -> bool: @@ -399,16 +420,26 @@ def _enable_fp16_accumulation( def _fuse_qkv(pipe: Any, logger: Any) -> bool: - for owner in (pipe, getattr(pipe, "transformer", None)): - fn = getattr(owner, "fuse_qkv_projections", None) - if callable(fn): + # Prefer the pipe-level fuse (it covers every component the pipe knows about); else fuse each + # denoiser DiT directly so a dual-DiT family (Ideogram) fuses BOTH experts, not just the first. + fn = getattr(pipe, "fuse_qkv_projections", None) + if callable(fn): + try: + fn() + return True + except Exception as exc: # noqa: BLE001 — optimisation only + _warn(logger, "fuse_qkv_projections", exc) + return False + engaged = False + for transformer in _denoiser_dits(pipe): + tfn = getattr(transformer, "fuse_qkv_projections", None) + if callable(tfn): try: - fn() - return True + tfn() + engaged = True except Exception as exc: # noqa: BLE001 — optimisation only _warn(logger, "fuse_qkv_projections", exc) - return False - return False + return engaged def _warn(logger: Any, what: str, exc: Exception) -> None: diff --git a/studio/backend/tests/test_diffusion_attention.py b/studio/backend/tests/test_diffusion_attention.py index 5abaa86fa7..aa8ccd8229 100644 --- a/studio/backend/tests/test_diffusion_attention.py +++ b/studio/backend/tests/test_diffusion_attention.py @@ -177,6 +177,17 @@ def test_apply_sets_backend(): assert engaged == "_native_cudnn" and t.set_to == "_native_cudnn" +def test_apply_sets_backend_on_both_dits(): + # A dual-DiT family (Ideogram) runs transformer + unconditional_transformer each step, so the + # backend must be set on BOTH; otherwise the second DiT keeps the native default while status + # reports the requested kernel as engaged. + t1, t2 = _FakeTransformer(), _FakeTransformer() + pipe = types.SimpleNamespace(transformer = t1, unconditional_transformer = t2) + engaged = apply_attention_backend(pipe, "_native_cudnn") + assert engaged == "_native_cudnn" + assert t1.set_to == "_native_cudnn" and t2.set_to == "_native_cudnn" + + def test_apply_falls_back_on_unavailable_kernel(monkeypatch): # an unavailable kernel must not fail the load -> returns None (diffusers default). monkeypatch.setattr(att, "_active_attention_backend", lambda: "native") diff --git a/studio/backend/tests/test_diffusion_speed.py b/studio/backend/tests/test_diffusion_speed.py index 9a18090e36..28cd4e0416 100644 --- a/studio/backend/tests/test_diffusion_speed.py +++ b/studio/backend/tests/test_diffusion_speed.py @@ -177,6 +177,7 @@ class _Pipe: *, with_compile = False, with_fuse = False, + with_second_dit = False, ) -> None: self.vae = types.SimpleNamespace(mem_format = None, to = self._vae_to) self.transformer = types.SimpleNamespace() @@ -186,6 +187,12 @@ class _Pipe: self.fuse_qkv_projections = self._fuse self.compiled = False self.fused = False + # A dual-DiT family (Ideogram) carries a second denoiser expert that runs every step. + self.second_compiled = False + if with_second_dit: + self.unconditional_transformer = types.SimpleNamespace() + if with_compile: + self.unconditional_transformer.compile_repeated_blocks = self._compile2 def _vae_to(self, *, memory_format): self.vae.mem_format = memory_format @@ -194,6 +201,9 @@ class _Pipe: self.compiled = True self.compile_kwargs = kwargs + def _compile2(self, **kwargs): + self.second_compiled = True + def _fuse(self): self.fused = True @@ -218,6 +228,20 @@ def test_speed_off_applies_nothing(monkeypatch): assert torch.backends.cudnn.benchmark is False +def test_speed_compiles_both_dits_for_dual_dit_family(monkeypatch): + # A dual-DiT family (Ideogram: transformer + unconditional_transformer) runs BOTH DiTs each + # denoise step, so the regional block compile must engage on both, not just the first -- + # otherwise the second DiT runs eager while status reports compile as engaged. + _stub_torch(monkeypatch) + _stub_gguf_accel(monkeypatch) + pipe = _Pipe(with_compile = True, with_second_dit = True) + applied = apply_speed_optims( + pipe, _target(), is_gguf = False, family = _family(), speed_mode = SPEED_DEFAULT + ) + assert applied["compiled"] is True + assert pipe.compiled is True and pipe.second_compiled is True + + def test_speed_default_dense_falls_back_to_regional_compile(monkeypatch): # A DENSE model has no GGUF dequant to compile, so `default` falls back to the # regional block compile (its only compile lever) -- and no GGUF accelerators. From 6e14e0373554ac31c2d1a9456286c90785ca4d43 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 6 Jul 2026 10:21:41 +0000 Subject: [PATCH 11/12] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/core/inference/diffusion_attention.py | 7 +++++-- studio/backend/core/inference/diffusion_speed.py | 4 +++- 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/studio/backend/core/inference/diffusion_attention.py b/studio/backend/core/inference/diffusion_attention.py index 6873026cd1..b14fea508c 100644 --- a/studio/backend/core/inference/diffusion_attention.py +++ b/studio/backend/core/inference/diffusion_attention.py @@ -275,8 +275,11 @@ def apply_attention_backend( defaults to None). So a load that wants native must restore it explicitly: otherwise it silently inherits a backend an earlier load pinned (e.g. cuDNN under a speed profile), breaking the bit-identical/``off`` guarantee. Best-effort throughout.""" - setters = [s for s in (getattr(t, "set_attention_backend", None) for t in _attention_dits(pipe)) - if callable(s)] + setters = [ + s + for s in (getattr(t, "set_attention_backend", None) for t in _attention_dits(pipe)) + if callable(s) + ] if not setters: return None if backend is not None: diff --git a/studio/backend/core/inference/diffusion_speed.py b/studio/backend/core/inference/diffusion_speed.py index 1f9a736f2f..7dd8147ce6 100644 --- a/studio/backend/core/inference/diffusion_speed.py +++ b/studio/backend/core/inference/diffusion_speed.py @@ -294,7 +294,9 @@ def _compile_repeated_blocks( cache_active: bool = False, offload_active: bool = False, ) -> bool: - dits = [t for t in _denoiser_dits(pipe) if callable(getattr(t, "compile_repeated_blocks", None))] + dits = [ + t for t in _denoiser_dits(pipe) if callable(getattr(t, "compile_repeated_blocks", None)) + ] if not dits: return False # default: mode="default" + dynamic=True -- fast cold start, robust to resolution From 4f00222c2ed27db0ffae2e977a672e1c86efeda2 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 6 Jul 2026 11:48:11 +0000 Subject: [PATCH 12/12] Reserve the fp8 bf16 footprint even when the cache estimate is absent For a narrow (fp8) Ideogram-4 base the planner reserves the family's known bf16 component total, but the reservation was gated on model_dense_mib being non-None. On a first-time load an empty blob cache (or a best-effort download probe that swallowed a transient HF error) leaves model_dense_mib None, so the guard skipped the reservation exactly when it was needed: the planner then read 'size unknown -> stay resident' and the ~54 GB pipeline OOMed a card that offload would have fit. family_bf16_components_gb is a network-free constant, so reserve it whenever the cache signal is absent (use it directly when None, else take the max). --- studio/backend/core/inference/diffusion.py | 12 ++++++++++-- studio/backend/tests/test_diffusion_more_families.py | 12 ++++++++++++ 2 files changed, 22 insertions(+), 2 deletions(-) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 91a2acb699..25dc36bee2 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -1714,9 +1714,17 @@ class DiffusionBackend: 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: + if table is not None: + # family_bf16_components_gb is a network-free constant, so reserve the bf16 + # footprint even when the cache-derived estimate is absent (empty blob cache, + # or a best-effort download probe that swallowed a transient HF error and + # returned nothing). Otherwise model_dense_mib stays None and the planner + # reads "size unknown -> stay resident", so the ~54 GB fp8 pipeline plans a + # resident placement and OOMs a card that offload would have fit. table_mib = int(sum(table) * (1000.0**3) / (1024.0 * 1024.0)) - model_dense_mib = max(model_dense_mib, table_mib) + model_dense_mib = ( + table_mib if model_dense_mib is None else max(model_dense_mib, table_mib) + ) companion_mib = None else: if transformer_resident_override_mib is not None: diff --git a/studio/backend/tests/test_diffusion_more_families.py b/studio/backend/tests/test_diffusion_more_families.py index 9d89eb5ab3..b2ffcd093e 100644 --- a/studio/backend/tests/test_diffusion_more_families.py +++ b/studio/backend/tests/test_diffusion_more_families.py @@ -59,6 +59,18 @@ def test_ideogram4_generation_defaults(): 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)