Merge remote-tracking branch 'origin/diffusion-more-families' into fold-integration

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,
@@ -43,6 +45,7 @@ from .diffusion_device import (
     diffusion_device_target_from_torch_device,
     resolve_diffusion_device_target,
 )
+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 (
     MEMORY_MODE_BALANCED,
@@ -91,7 +94,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,
@@ -228,6 +235,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",
     }
 )
 
@@ -554,6 +568,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, "
@@ -570,6 +590,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.
@@ -1197,6 +1228,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:
@@ -1707,6 +1744,39 @@ 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.
+            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:
+                    # 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 = (
+                        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:
@@ -2261,6 +2331,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.

studio/backend/core/inference/diffusion_attention.py

@@ -276,13 +276,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
@@ -293,28 +307,35 @@ 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
 
 

studio/backend/core/inference/diffusion_auto_policy.py

@@ -59,6 +59,13 @@ _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
+    # 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

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
@@ -314,6 +321,28 @@ _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"),
+        # 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
     # ``.safetensors`` is the WHOLE pipeline rather than a transformer-only file.
@@ -353,6 +382,38 @@ 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], ...] = (
+    (
+        # "-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.",
+    ),
+)
+
+
+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(token, 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
@@ -478,6 +539,10 @@ _GENERATION_DEFAULTS: tuple[tuple[str, int, float], ...] = (
     ("flux.2-dev", 28, 4.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),
     # 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.
     ("sdxl-turbo", 3, 0.0),

studio/backend/core/inference/diffusion_ideogram4.py

@@ -0,0 +1,469 @@
+# 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)
+        # 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():
+        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]
+            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]
+                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)
+            # 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)
+
+    # Construct normally (so __init__ computes the non-persistent rotary inv_freq
+    # buffers the checkpoint omits) then copy the dequantized weights in with
+    # 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:
+        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 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,
+    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
+    import torch
+
+    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 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.
+        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"])
+    # 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
+    # (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,
+    )

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:

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,10 @@ 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 +334,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:
@@ -404,16 +427,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:

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")

studio/backend/tests/test_diffusion_backend.py

@@ -383,6 +383,11 @@ 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. 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
     # pipeline class carries from_single_file; UNet2DConditionModel is the denoiser
     # class (fetched but unused on the pipeline/single-file-pipeline paths).
@@ -391,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
@@ -1223,6 +1233,58 @@ 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):
+    # 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):
+    # 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

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):

studio/backend/tests/test_diffusion_more_families.py

@@ -0,0 +1,234 @@
+# 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_bf16_reservation_table_present():
+    # The memory planner reserves this bf16 footprint for a narrow (fp8) ideogram-4 base even
+    # when the blob-cache estimate is absent (empty cache / a best-effort download probe that
+    # swallowed a transient HF error), so the ~54 GB pipeline never plans a resident placement
+    # it cannot fit. If this constant table ever went None, that fp8 OOM safeguard would
+    # silently disable, so pin that it is present and sums to the expected ~54 GB.
+    fam = detect_family("ideogram-ai/ideogram-4-fp8")
+    table = family_bf16_components_gb(fam, fam.base_repo)
+    assert table is not None
+    assert sum(table) > 50.0  # transformer (37.2) + bf16 text encoder (16.3) + VAE (0.2)
+
+
+def test_ideogram4_memory_table_counts_both_dits():
+    fam = detect_family("ideogram-ai/ideogram-4-fp8")
+    components = family_bf16_components_gb(fam)
+    assert components is not None
+    transformer_gb, text_encoders_gb, _vae_gb = components
+    # Two ~9.3B DiTs (conditional + unconditional) at bf16: well above one DiT's
+    # ~18.6 GB. A single-DiT entry here would let auto planning under-reserve and OOM.
+    assert transformer_gb > 30.0
+    assert text_encoders_gb > 5.0
+
+
+# ── structured exclusions ────────────────────────────────────────────────────
+def test_hunyuanimage_is_excluded_with_reason():
+    reason = excluded_model_reason("tencent/HunyuanImage-3.0")
+    assert reason is not None and "diffusers" in reason
+    # Not detectable as any family: the exclusion reason is the load error surface.
+    assert detect_family("tencent/HunyuanImage-3.0") is None
+
+
+def test_excluded_model_reason_none_for_supported_and_unknown():
+    assert excluded_model_reason("unsloth/Z-Image-Turbo-GGUF") is None
+    assert excluded_model_reason("someorg/some-model") is None
+
+
+def test_validate_load_request_surfaces_exclusion_reason():
+    from core.inference.diffusion import DiffusionBackend
+    backend = DiffusionBackend()
+    with pytest.raises(ValueError, match = "trust_remote_code"):
+        backend.validate_load_request("tencent/HunyuanImage-3.0")
+
+
+# ── curated krea LoRA catalog ────────────────────────────────────────────────
+def test_curated_krea2_loras_present_and_well_formed():
+    krea = [e for e in _CURATED if e.repo_id and e.repo_id.startswith("krea/Krea-2-LoRA-")]
+    assert len(krea) == 9
+    for entry in krea:
+        assert entry.source == "hub" and entry.fmt == "safetensors"
+        assert entry.families == ("krea-2",)
+        # Every official style repo carries a single "{style}.safetensors" at the root.
+        style = entry.repo_id.split("Krea-2-LoRA-")[-1]
+        assert entry.weight_name == f"{style}.safetensors"
+
+
+def test_list_loras_family_filter_gates_krea_entries():
+    krea_ids = {e.id for e in _CURATED if e.families == ("krea-2",)}
+    assert krea_ids  # curated entries exist
+    listed_for_krea = {e.id for e in list_loras(family = "krea-2")}
+    assert krea_ids <= listed_for_krea
+    listed_for_flux = {e.id for e in list_loras(family = "flux.1")}
+    assert not (krea_ids & listed_for_flux)
+
+
+# ── ideogram-4 fp8 transformer remap ─────────────────────────────────────────
+def test_convert_fp8_state_dict_dequantizes_and_splits_qkv():
+    # The vendor fp8 transformer stores fused attention.qkv (Q/K/V rows stacked) +
+    # attention.o, each with a per-output-channel weight_scale; diffusers expects split
+    # to_q/to_k/to_v/to_out.0 with the scale already applied. The converter must undo
+    # both, or every attention weight loads wrong (garbage) and on meta (a load crash).
+    torch = pytest.importorskip("torch")
+
+    from core.inference.diffusion_ideogram4 import _convert_fp8_state_dict
+
+    hidden = 4  # tiny stand-in for attention_head_dim * num_attention_heads
+    # Reference (real) weights, then a fake per-channel fp8 encoding: value / scale.
+    q = torch.randn(hidden, hidden)
+    k = torch.randn(hidden, hidden)
+    v = torch.randn(hidden, hidden)
+    o = torch.randn(hidden, hidden)
+    ff = torch.randn(hidden, hidden)
+    fused = torch.cat([q, k, v], dim = 0)  # [3 * hidden, hidden]
+    qkv_scale = torch.rand(3 * hidden) + 0.5
+    o_scale = torch.rand(hidden) + 0.5
+    ff_scale = torch.rand(hidden) + 0.5
+    norm = torch.randn(hidden)  # dense (unscaled) weight passes through
+    raw = {
+        "layers.0.attention.qkv.weight": fused / qkv_scale[:, None],
+        "layers.0.attention.qkv.weight_scale": qkv_scale,
+        "layers.0.attention.o.weight": o / o_scale[:, None],
+        "layers.0.attention.o.weight_scale": o_scale,
+        "layers.0.feed_forward.w1.weight": ff / ff_scale[:, None],
+        "layers.0.feed_forward.w1.weight_scale": ff_scale,
+        "layers.0.attention_norm1.weight": norm,
+    }
+    out = _convert_fp8_state_dict(raw, hidden, torch.bfloat16)
+
+    # Every converted tensor is cast to the requested compute dtype (the load_state_dict
+    # copy would silently up/down-cast otherwise).
+    assert all(t.dtype == torch.bfloat16 for t in out.values())
+    # Re-run in float32 for the exact value checks below (bf16 loses precision).
+    out = _convert_fp8_state_dict(raw, hidden, torch.float32)
+
+    # No scale keys leak through; fused/renamed keys are gone.
+    assert not any(key.endswith("_scale") for key in out)
+    assert "layers.0.attention.qkv.weight" not in out
+    assert "layers.0.attention.o.weight" not in out
+    # QKV split back to the reference weights in Q/K/V order.
+    torch.testing.assert_close(out["layers.0.attention.to_q.weight"], q)
+    torch.testing.assert_close(out["layers.0.attention.to_k.weight"], k)
+    torch.testing.assert_close(out["layers.0.attention.to_v.weight"], v)
+    # o renamed to to_out.0 with the scale applied.
+    torch.testing.assert_close(out["layers.0.attention.to_out.0.weight"], o)
+    # A non-attention fp8 weight keeps its name, scale applied.
+    torch.testing.assert_close(out["layers.0.feed_forward.w1.weight"], ff)
+    # A dense weight passes through unchanged.
+    torch.testing.assert_close(out["layers.0.attention_norm1.weight"], norm)
+
+
+def test_ideogram4_repo_is_fp8_detects_local_layout(tmp_path):
+    # A local mirror of the fp8 base never string-matches base_repo, so memory planning
+    # relies on this shard-header probe to reserve the bf16 footprint. The fp8 layout is
+    # marked by a companion ``*.weight_scale``; the bnb-4bit (nf4) mirror carries none and
+    # must read as not-fp8 so it stays (correctly) planned against its compressed bytes.
+    torch = pytest.importorskip("torch")
+    st = pytest.importorskip("safetensors.torch")
+
+    from core.inference.diffusion_ideogram4 import ideogram4_repo_is_fp8
+
+    fp8 = tmp_path / "fp8"
+    (fp8 / "transformer").mkdir(parents = True)
+    st.save_file(
+        {
+            "layers.0.attention.o.weight": torch.zeros(2, 2),
+            "layers.0.attention.o.weight_scale": torch.ones(2),
+        },
+        str(fp8 / "transformer" / "diffusion_pytorch_model.safetensors"),
+    )
+    assert ideogram4_repo_is_fp8(str(fp8)) is True
+
+    nf4 = tmp_path / "nf4"
+    (nf4 / "transformer").mkdir(parents = True)
+    st.save_file(
+        {"layers.0.attention.to_q.weight": torch.zeros(2, 2)},
+        str(nf4 / "transformer" / "diffusion_pytorch_model.safetensors"),
+    )
+    assert ideogram4_repo_is_fp8(str(nf4)) is False
+
+    # A directory with no transformer shards at all resolves to False, not an error.
+    assert ideogram4_repo_is_fp8(str(tmp_path / "missing")) is False
+
+
+def test_create_causal_mask_patch_is_self_disabling_and_idempotent():
+    # The patch adapts the pipeline's inputs_embeds kwarg to the installed transformers
+    # create_causal_mask signature; on a matching signature it must forward unchanged,
+    # and a second apply must not double-wrap.
+    pytest.importorskip("torch")
+    pytest.importorskip("diffusers")
+
+    import core.inference.diffusion_ideogram4 as ig4
+    from diffusers.pipelines.ideogram4 import pipeline_ideogram4 as pipe_mod
+
+    original = pipe_mod.create_causal_mask
+    try:
+        ig4._CAUSAL_MASK_PATCHED = False
+        ig4._patch_create_causal_mask()
+        wrapped = pipe_mod.create_causal_mask
+        assert wrapped is not original  # the patch installed a wrapper
+        ig4._patch_create_causal_mask()  # idempotent: no re-wrap
+        assert pipe_mod.create_causal_mask is wrapped
+    finally:
+        pipe_mod.create_causal_mask = original
+        ig4._CAUSAL_MASK_PATCHED = False

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.

studio/frontend/src/features/images/images-page.tsx

@@ -100,6 +100,11 @@ const SAFETENSORS_MODELS: Record<string, SafetensorsSpec> = {
   "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)",
@@ -227,6 +238,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 },