From 1d193f58bb30eb42cdf8bdac6ff1b0f2e08e90c9 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 5 Jul 2026 12:07:31 +0000 Subject: [PATCH 1/2] Drop the now-redundant second-expert group offload on the A14B load After the upstream cascade generalised _apply_group_offload to stream every DiT it finds on the pipe (transformer AND transformer_2), apply_memory_plan's group tier already block-streams both A14B experts, and model/sequential offload hook every top-level module. The explicit per-expert pass this branch added is therefore redundant: on a real group-tier load it re-registers the group-offload hooks transformer_2 already carries, which diffusers rejects with a duplicate-hook ValueError. That was swallowed by the helper's try/except and surfaced as a misleading "second expert stays resident" warning even though the expert was in fact streamed. Remove the pass and document why MoE needs none. --- studio/backend/core/inference/video.py | 18 ++++++------------ 1 file changed, 6 insertions(+), 12 deletions(-) diff --git a/studio/backend/core/inference/video.py b/studio/backend/core/inference/video.py index 40210c0fb7..9a3a1d82e0 100644 --- a/studio/backend/core/inference/video.py +++ b/studio/backend/core/inference/video.py @@ -1004,18 +1004,12 @@ class VideoBackend: clear_gpu_cache() raise RuntimeError("Video load was cancelled or superseded.") offload_policy, vae_tiling = apply_memory_plan(pipe, plan, device = device, logger = logger) - if offload_policy == "group" and len(views) > 1: - # Group offload streams only ``pipe.transformer``; the second expert would - # otherwise sit resident (~57 GB bf16 on the A14B) and defeat the tier. - # model/sequential offload hook every top-level module, so only group needs - # this. Applied through the view so the helper streams transformer_2. - from .diffusion_memory import _apply_group_offload - for view in views[1:]: - if not _apply_group_offload(view, device, logger): - logger.warning( - "video.memory: group offload did not engage on the second " - "expert; it stays resident" - ) + # A dual-DiT MoE pipe (Wan2.2-A14B) needs no extra per-expert offload pass here: + # apply_memory_plan's group tier (_apply_group_offload) already block-streams every + # DiT it finds on the pipe -- transformer AND transformer_2 -- and model/sequential + # offload hook every top-level module, so the second expert is covered under all tiers. + # A second _apply_group_offload on transformer_2 would re-register the group-offload + # hooks it already carries, which diffusers rejects with a duplicate-hook ValueError. if not vae_tiling: # Decode of a whole clip is the video memory peak; tiling is near-free # in quality and keeps the decode bounded, so it is always on. From f30713181984111dc39b1547c4072341008ecfae Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 5 Jul 2026 12:08:45 +0000 Subject: [PATCH 2/2] 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,