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

Load Ideogram 4 fp8 repo by dequantizing and remapping its DiTs and text encoder

Browse source 
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.
This commit is contained in:
Daniel Han 2026-07-04 14:30:58 +00:00
commit a5195517cf
3 changed files with 477 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

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

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

392
studio/backend/core/inference/diffusion_ideogram4.py Normal file
View file

@ -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,
)

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

@ -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