[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
This commit is contained in:
pre-commit-ci[bot]
parent
b14e2f9be7
commit
a62dc39303
15 changed files with 467 additions and 239 deletions
|
|
@ -143,7 +143,7 @@ def _decode_b64_image(data: str, *, mode: str = "RGB") -> Any:
|
|||
# data:[<mime>][;base64],<payload>
|
||||
_, _, raw = raw.partition(",")
|
||||
try:
|
||||
blob = base64.b64decode(raw, validate=False)
|
||||
blob = base64.b64decode(raw, validate = False)
|
||||
except (binascii.Error, ValueError) as exc:
|
||||
raise ValueError(f"Invalid base64 image data: {exc}") from exc
|
||||
try:
|
||||
|
|
@ -418,9 +418,7 @@ class DiffusionBackend:
|
|||
path_shaped = repo_id.startswith(("/", "~", "./", "../")) or local_root.is_absolute()
|
||||
if kind in ("gguf", "single_file"):
|
||||
if not gguf_filename:
|
||||
raise ValueError(
|
||||
f"a single-file checkpoint name is required for a '{kind}' load."
|
||||
)
|
||||
raise ValueError(f"a single-file checkpoint name is required for a '{kind}' load.")
|
||||
if local_root.exists():
|
||||
resolve_local_gguf_child(local_root, gguf_filename)
|
||||
elif path_shaped:
|
||||
|
|
@ -525,7 +523,11 @@ class DiffusionBackend:
|
|||
)
|
||||
kwargs["base_repo"] = base
|
||||
expected, base_files = self._estimate_download_bytes(
|
||||
kwargs["repo_id"], kwargs.get("gguf_filename"), base, kwargs.get("hf_token"), kind = kind
|
||||
kwargs["repo_id"],
|
||||
kwargs.get("gguf_filename"),
|
||||
base,
|
||||
kwargs.get("hf_token"),
|
||||
kind = kind,
|
||||
)
|
||||
loading = self._loading
|
||||
if loading is not None:
|
||||
|
|
@ -669,7 +671,9 @@ class DiffusionBackend:
|
|||
kind = resolve_model_kind(gguf_filename, model_kind)
|
||||
# For a full pipeline the repo itself supplies every component, so it is its
|
||||
# own base; the single-file kinds resolve the companion base diffusers repo.
|
||||
base = repo_id if kind == "pipeline" else _resolve_base_repo(repo_id, base_repo, fam, hf_token)
|
||||
base = (
|
||||
repo_id if kind == "pipeline" else _resolve_base_repo(repo_id, base_repo, fam, hf_token)
|
||||
)
|
||||
target = self._resolve_device_target(fam)
|
||||
device, dtype = target.device, target.dtype
|
||||
|
||||
|
|
@ -786,7 +790,9 @@ class DiffusionBackend:
|
|||
)
|
||||
# A safetensors single-file (e.g. fp8) carries its own dtype, so no
|
||||
# GGUF dequant config is passed.
|
||||
transformer = transformer_cls.from_single_file(single_file_path, **sf_kwargs)
|
||||
transformer = transformer_cls.from_single_file(
|
||||
single_file_path, **sf_kwargs
|
||||
)
|
||||
|
||||
pipe_kwargs = {"torch_dtype": dtype, "transformer": transformer}
|
||||
if hf_token:
|
||||
|
|
@ -1138,7 +1144,7 @@ class DiffusionBackend:
|
|||
try:
|
||||
target_dtype = transformer.dtype
|
||||
if next(vae.parameters()).dtype != target_dtype:
|
||||
vae.to(dtype=target_dtype)
|
||||
vae.to(dtype = target_dtype)
|
||||
except (StopIteration, AttributeError, RuntimeError):
|
||||
pass
|
||||
|
||||
|
|
@ -1261,8 +1267,7 @@ class DiffusionBackend:
|
|||
# Additional references (FLUX.2 accepts a list): decode them so the
|
||||
# conditioning combines all of them. Capped to keep VRAM bounded.
|
||||
ref_extra = [
|
||||
_decode_b64_image(x, mode = "RGB")
|
||||
for x in (reference_images or [])[:3]
|
||||
_decode_b64_image(x, mode = "RGB") for x in (reference_images or [])[:3]
|
||||
]
|
||||
elif init_image is not None:
|
||||
workflow = "img2img"
|
||||
|
|
@ -1279,7 +1284,6 @@ class DiffusionBackend:
|
|||
init_pil = _snap_to_multiple(init_pil, 16)
|
||||
if mask_pil is not None and mask_pil.size != init_pil.size:
|
||||
from PIL import Image as _PILImage
|
||||
|
||||
mask_pil = mask_pil.resize(init_pil.size, _PILImage.NEAREST)
|
||||
if init_pil is not None:
|
||||
# Keep the VAE encode dtype consistent with the input image.
|
||||
|
|
|
|||
|
|
@ -69,11 +69,16 @@ def _body_has(fn: Callable, *needles: str) -> bool:
|
|||
# =====================================================================================
|
||||
# qwen-image: QwenImageTransformerBlock._modulate (modulation addcmul, all 4 call sites)
|
||||
# =====================================================================================
|
||||
def _qwen_modulate(self, x, mod_params, index=None):
|
||||
def _qwen_modulate(
|
||||
self,
|
||||
x,
|
||||
mod_params,
|
||||
index = None,
|
||||
):
|
||||
"""diffusers 0.38 ``QwenImageTransformerBlock._modulate`` with the final
|
||||
``x*(1+scale)+shift`` fused to ``torch.addcmul`` (covers both the global and the
|
||||
per-token ``index`` branches, since both end in that same expression)."""
|
||||
shift, scale, gate = mod_params.chunk(3, dim=-1)
|
||||
shift, scale, gate = mod_params.chunk(3, dim = -1)
|
||||
|
||||
if index is not None:
|
||||
actual_batch = shift.size(0) // 2
|
||||
|
|
@ -101,7 +106,9 @@ def _qwen_modulate(self, x, mod_params, index=None):
|
|||
|
||||
def _spec_qwen_modulate():
|
||||
try:
|
||||
from diffusers.models.transformers.transformer_qwenimage import QwenImageTransformerBlock as cls
|
||||
from diffusers.models.transformers.transformer_qwenimage import (
|
||||
QwenImageTransformerBlock as cls,
|
||||
)
|
||||
except Exception: # noqa: BLE001
|
||||
return None
|
||||
orig = getattr(cls, "_modulate", None)
|
||||
|
|
@ -135,8 +142,12 @@ def _zimage_forward(
|
|||
mod_noisy = self.adaLN_modulation(adaln_noisy)
|
||||
mod_clean = self.adaLN_modulation(adaln_clean)
|
||||
|
||||
scale_msa_noisy, gate_msa_noisy, scale_mlp_noisy, gate_mlp_noisy = mod_noisy.chunk(4, dim=1)
|
||||
scale_msa_clean, gate_msa_clean, scale_mlp_clean, gate_mlp_clean = mod_clean.chunk(4, dim=1)
|
||||
scale_msa_noisy, gate_msa_noisy, scale_mlp_noisy, gate_mlp_noisy = mod_noisy.chunk(
|
||||
4, dim = 1
|
||||
)
|
||||
scale_msa_clean, gate_msa_clean, scale_mlp_clean, gate_mlp_clean = mod_clean.chunk(
|
||||
4, dim = 1
|
||||
)
|
||||
|
||||
gate_msa_noisy, gate_mlp_noisy = gate_msa_noisy.tanh(), gate_mlp_noisy.tanh()
|
||||
gate_msa_clean, gate_mlp_clean = gate_msa_clean.tanh(), gate_mlp_clean.tanh()
|
||||
|
|
@ -150,13 +161,13 @@ def _zimage_forward(
|
|||
gate_mlp = select_per_token(gate_mlp_noisy, gate_mlp_clean, noise_mask, seq_len)
|
||||
else:
|
||||
mod = self.adaLN_modulation(adaln_input)
|
||||
scale_msa, gate_msa, scale_mlp, gate_mlp = mod.unsqueeze(1).chunk(4, dim=2)
|
||||
scale_msa, gate_msa, scale_mlp, gate_mlp = mod.unsqueeze(1).chunk(4, dim = 2)
|
||||
gate_msa, gate_mlp = gate_msa.tanh(), gate_mlp.tanh()
|
||||
scale_msa, scale_mlp = 1.0 + scale_msa, 1.0 + scale_mlp
|
||||
|
||||
# Attention block -- fused gated residual: x + gate_msa * attention_norm2(attn_out)
|
||||
attn_out = self.attention(
|
||||
self.attention_norm1(x) * scale_msa, attention_mask=attn_mask, freqs_cis=freqs_cis
|
||||
self.attention_norm1(x) * scale_msa, attention_mask = attn_mask, freqs_cis = freqs_cis
|
||||
)
|
||||
x = torch.addcmul(x, gate_msa, self.attention_norm2(attn_out))
|
||||
|
||||
|
|
@ -165,7 +176,9 @@ def _zimage_forward(
|
|||
x, gate_mlp, self.ffn_norm2(self.feed_forward(self.ffn_norm1(x) * scale_mlp))
|
||||
)
|
||||
else:
|
||||
attn_out = self.attention(self.attention_norm1(x), attention_mask=attn_mask, freqs_cis=freqs_cis)
|
||||
attn_out = self.attention(
|
||||
self.attention_norm1(x), attention_mask = attn_mask, freqs_cis = freqs_cis
|
||||
)
|
||||
x = x + self.attention_norm2(attn_out)
|
||||
x = x + self.ffn_norm2(self.feed_forward(self.ffn_norm1(x)))
|
||||
|
||||
|
|
@ -193,17 +206,24 @@ def _spec_zimage_forward():
|
|||
# here we fuse the inline norm2 modulation + the gated residual adds.)
|
||||
# =====================================================================================
|
||||
def _flux_double_forward(
|
||||
self, hidden_states, encoder_hidden_states, temb, image_rotary_emb=None, joint_attention_kwargs=None
|
||||
self,
|
||||
hidden_states,
|
||||
encoder_hidden_states,
|
||||
temb,
|
||||
image_rotary_emb = None,
|
||||
joint_attention_kwargs = None,
|
||||
):
|
||||
norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, emb=temb)
|
||||
norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.norm1_context(
|
||||
encoder_hidden_states, emb=temb
|
||||
norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(
|
||||
hidden_states, emb = temb
|
||||
)
|
||||
norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = (
|
||||
self.norm1_context(encoder_hidden_states, emb = temb)
|
||||
)
|
||||
joint_attention_kwargs = joint_attention_kwargs or {}
|
||||
attention_outputs = self.attn(
|
||||
hidden_states=norm_hidden_states,
|
||||
encoder_hidden_states=norm_encoder_hidden_states,
|
||||
image_rotary_emb=image_rotary_emb,
|
||||
hidden_states = norm_hidden_states,
|
||||
encoder_hidden_states = norm_encoder_hidden_states,
|
||||
image_rotary_emb = image_rotary_emb,
|
||||
**joint_attention_kwargs,
|
||||
)
|
||||
if len(attention_outputs) == 2:
|
||||
|
|
@ -216,14 +236,18 @@ def _flux_double_forward(
|
|||
|
||||
norm_hidden_states = self.norm2(hidden_states)
|
||||
# fused: norm * (1 + scale_mlp) + shift_mlp
|
||||
norm_hidden_states = torch.addcmul(shift_mlp[:, None], norm_hidden_states, 1 + scale_mlp[:, None])
|
||||
norm_hidden_states = torch.addcmul(
|
||||
shift_mlp[:, None], norm_hidden_states, 1 + scale_mlp[:, None]
|
||||
)
|
||||
|
||||
ff_output = self.ff(norm_hidden_states)
|
||||
hidden_states = torch.addcmul(hidden_states, gate_mlp.unsqueeze(1), ff_output)
|
||||
if len(attention_outputs) == 3:
|
||||
hidden_states = hidden_states + ip_attn_output
|
||||
|
||||
encoder_hidden_states = torch.addcmul(encoder_hidden_states, c_gate_msa.unsqueeze(1), context_attn_output)
|
||||
encoder_hidden_states = torch.addcmul(
|
||||
encoder_hidden_states, c_gate_msa.unsqueeze(1), context_attn_output
|
||||
)
|
||||
|
||||
norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states)
|
||||
norm_encoder_hidden_states = torch.addcmul(
|
||||
|
|
@ -231,7 +255,9 @@ def _flux_double_forward(
|
|||
)
|
||||
|
||||
context_ff_output = self.ff_context(norm_encoder_hidden_states)
|
||||
encoder_hidden_states = torch.addcmul(encoder_hidden_states, c_gate_mlp.unsqueeze(1), context_ff_output)
|
||||
encoder_hidden_states = torch.addcmul(
|
||||
encoder_hidden_states, c_gate_mlp.unsqueeze(1), context_ff_output
|
||||
)
|
||||
if encoder_hidden_states.dtype == torch.float16:
|
||||
encoder_hidden_states = encoder_hidden_states.clip(-65504, 65504)
|
||||
|
||||
|
|
@ -255,29 +281,37 @@ def _spec_flux_double():
|
|||
|
||||
|
||||
def _flux_single_forward(
|
||||
self, hidden_states, encoder_hidden_states, temb, image_rotary_emb=None, joint_attention_kwargs=None
|
||||
self,
|
||||
hidden_states,
|
||||
encoder_hidden_states,
|
||||
temb,
|
||||
image_rotary_emb = None,
|
||||
joint_attention_kwargs = None,
|
||||
):
|
||||
text_seq_len = encoder_hidden_states.shape[1]
|
||||
hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
|
||||
hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim = 1)
|
||||
|
||||
residual = hidden_states
|
||||
norm_hidden_states, gate = self.norm(hidden_states, emb=temb)
|
||||
norm_hidden_states, gate = self.norm(hidden_states, emb = temb)
|
||||
mlp_hidden_states = self.act_mlp(self.proj_mlp(norm_hidden_states))
|
||||
joint_attention_kwargs = joint_attention_kwargs or {}
|
||||
attn_output = self.attn(
|
||||
hidden_states=norm_hidden_states,
|
||||
image_rotary_emb=image_rotary_emb,
|
||||
hidden_states = norm_hidden_states,
|
||||
image_rotary_emb = image_rotary_emb,
|
||||
**joint_attention_kwargs,
|
||||
)
|
||||
|
||||
hidden_states = torch.cat([attn_output, mlp_hidden_states], dim=2)
|
||||
hidden_states = torch.cat([attn_output, mlp_hidden_states], dim = 2)
|
||||
gate = gate.unsqueeze(1)
|
||||
# fused: residual + gate * proj_out(hidden_states)
|
||||
hidden_states = torch.addcmul(residual, gate, self.proj_out(hidden_states))
|
||||
if hidden_states.dtype == torch.float16:
|
||||
hidden_states = hidden_states.clip(-65504, 65504)
|
||||
|
||||
encoder_hidden_states, hidden_states = hidden_states[:, :text_seq_len], hidden_states[:, text_seq_len:]
|
||||
encoder_hidden_states, hidden_states = (
|
||||
hidden_states[:, :text_seq_len],
|
||||
hidden_states[:, text_seq_len:],
|
||||
)
|
||||
return encoder_hidden_states, hidden_states
|
||||
|
||||
|
||||
|
|
@ -302,27 +336,36 @@ def _spec_flux_single():
|
|||
# the gated residuals; scale/shift/gate are [B,1,dim] so no [:, None] is needed.)
|
||||
# =====================================================================================
|
||||
def _flux2_double_forward(
|
||||
self, hidden_states, encoder_hidden_states, temb_mod_img, temb_mod_txt,
|
||||
image_rotary_emb=None, joint_attention_kwargs=None,
|
||||
self,
|
||||
hidden_states,
|
||||
encoder_hidden_states,
|
||||
temb_mod_img,
|
||||
temb_mod_txt,
|
||||
image_rotary_emb = None,
|
||||
joint_attention_kwargs = None,
|
||||
):
|
||||
from diffusers.models.transformers.transformer_flux2 import Flux2Modulation
|
||||
|
||||
joint_attention_kwargs = joint_attention_kwargs or {}
|
||||
(shift_msa, scale_msa, gate_msa), (shift_mlp, scale_mlp, gate_mlp) = Flux2Modulation.split(temb_mod_img, 2)
|
||||
(c_shift_msa, c_scale_msa, c_gate_msa), (c_shift_mlp, c_scale_mlp, c_gate_mlp) = Flux2Modulation.split(
|
||||
temb_mod_txt, 2
|
||||
(shift_msa, scale_msa, gate_msa), (shift_mlp, scale_mlp, gate_mlp) = Flux2Modulation.split(
|
||||
temb_mod_img, 2
|
||||
)
|
||||
(c_shift_msa, c_scale_msa, c_gate_msa), (c_shift_mlp, c_scale_mlp, c_gate_mlp) = (
|
||||
Flux2Modulation.split(temb_mod_txt, 2)
|
||||
)
|
||||
|
||||
norm_hidden_states = self.norm1(hidden_states)
|
||||
norm_hidden_states = torch.addcmul(shift_msa, norm_hidden_states, 1 + scale_msa)
|
||||
|
||||
norm_encoder_hidden_states = self.norm1_context(encoder_hidden_states)
|
||||
norm_encoder_hidden_states = torch.addcmul(c_shift_msa, norm_encoder_hidden_states, 1 + c_scale_msa)
|
||||
norm_encoder_hidden_states = torch.addcmul(
|
||||
c_shift_msa, norm_encoder_hidden_states, 1 + c_scale_msa
|
||||
)
|
||||
|
||||
attention_outputs = self.attn(
|
||||
hidden_states=norm_hidden_states,
|
||||
encoder_hidden_states=norm_encoder_hidden_states,
|
||||
image_rotary_emb=image_rotary_emb,
|
||||
hidden_states = norm_hidden_states,
|
||||
encoder_hidden_states = norm_encoder_hidden_states,
|
||||
image_rotary_emb = image_rotary_emb,
|
||||
**joint_attention_kwargs,
|
||||
)
|
||||
attn_output, context_attn_output = attention_outputs
|
||||
|
|
@ -338,7 +381,9 @@ def _flux2_double_forward(
|
|||
encoder_hidden_states = torch.addcmul(encoder_hidden_states, c_gate_msa, context_attn_output)
|
||||
|
||||
norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states)
|
||||
norm_encoder_hidden_states = torch.addcmul(c_shift_mlp, norm_encoder_hidden_states, 1 + c_scale_mlp)
|
||||
norm_encoder_hidden_states = torch.addcmul(
|
||||
c_shift_mlp, norm_encoder_hidden_states, 1 + c_scale_mlp
|
||||
)
|
||||
|
||||
context_ff_output = self.ff_context(norm_encoder_hidden_states)
|
||||
encoder_hidden_states = torch.addcmul(encoder_hidden_states, c_gate_mlp, context_ff_output)
|
||||
|
|
@ -365,14 +410,20 @@ def _spec_flux2_double():
|
|||
|
||||
|
||||
def _flux2_single_forward(
|
||||
self, hidden_states, encoder_hidden_states, temb_mod, image_rotary_emb=None,
|
||||
joint_attention_kwargs=None, split_hidden_states=False, text_seq_len=None,
|
||||
self,
|
||||
hidden_states,
|
||||
encoder_hidden_states,
|
||||
temb_mod,
|
||||
image_rotary_emb = None,
|
||||
joint_attention_kwargs = None,
|
||||
split_hidden_states = False,
|
||||
text_seq_len = None,
|
||||
):
|
||||
from diffusers.models.transformers.transformer_flux2 import Flux2Modulation
|
||||
|
||||
if encoder_hidden_states is not None:
|
||||
text_seq_len = encoder_hidden_states.shape[1]
|
||||
hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
|
||||
hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim = 1)
|
||||
|
||||
mod_shift, mod_scale, mod_gate = Flux2Modulation.split(temb_mod, 1)[0]
|
||||
|
||||
|
|
@ -381,8 +432,8 @@ def _flux2_single_forward(
|
|||
|
||||
joint_attention_kwargs = joint_attention_kwargs or {}
|
||||
attn_output = self.attn(
|
||||
hidden_states=norm_hidden_states,
|
||||
image_rotary_emb=image_rotary_emb,
|
||||
hidden_states = norm_hidden_states,
|
||||
image_rotary_emb = image_rotary_emb,
|
||||
**joint_attention_kwargs,
|
||||
)
|
||||
|
||||
|
|
@ -391,7 +442,10 @@ def _flux2_single_forward(
|
|||
hidden_states = hidden_states.clip(-65504, 65504)
|
||||
|
||||
if split_hidden_states:
|
||||
encoder_hidden_states, hidden_states = hidden_states[:, :text_seq_len], hidden_states[:, text_seq_len:]
|
||||
encoder_hidden_states, hidden_states = (
|
||||
hidden_states[:, :text_seq_len],
|
||||
hidden_states[:, text_seq_len:],
|
||||
)
|
||||
return encoder_hidden_states, hidden_states
|
||||
else:
|
||||
return hidden_states
|
||||
|
|
@ -399,7 +453,9 @@ def _flux2_single_forward(
|
|||
|
||||
def _spec_flux2_single():
|
||||
try:
|
||||
from diffusers.models.transformers.transformer_flux2 import Flux2SingleTransformerBlock as cls
|
||||
from diffusers.models.transformers.transformer_flux2 import (
|
||||
Flux2SingleTransformerBlock as cls,
|
||||
)
|
||||
except Exception: # noqa: BLE001
|
||||
return None
|
||||
orig = getattr(cls, "forward", None)
|
||||
|
|
@ -444,16 +500,21 @@ def install_arch_patches() -> int:
|
|||
try:
|
||||
spec = resolve()
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logger.warning("arch-patch: resolver %s failed: %s", getattr(resolve, "__name__", resolve), exc)
|
||||
logger.warning(
|
||||
"arch-patch: resolver %s failed: %s", getattr(resolve, "__name__", resolve), exc
|
||||
)
|
||||
spec = None
|
||||
if spec is None:
|
||||
continue
|
||||
cls, attr, new_fn = spec
|
||||
if apply_patch(cls, attr, new_fn, match_level="relaxed"):
|
||||
if apply_patch(cls, attr, new_fn, match_level = "relaxed"):
|
||||
_patched.append((cls, attr))
|
||||
else:
|
||||
logger.warning("arch-patch: skipping %s.%s (signature mismatch / unavailable)",
|
||||
getattr(cls, "__name__", cls), attr)
|
||||
logger.warning(
|
||||
"arch-patch: skipping %s.%s (signature mismatch / unavailable)",
|
||||
getattr(cls, "__name__", cls),
|
||||
attr,
|
||||
)
|
||||
logger.info("arch-patch: installed %d/%d per-arch fusions", len(_patched), len(_SPECS))
|
||||
return len(_patched)
|
||||
|
||||
|
|
|
|||
|
|
@ -117,6 +117,7 @@ def environment_fingerprint() -> dict[str, Any]:
|
|||
}
|
||||
try:
|
||||
import torch # noqa: PLC0415
|
||||
|
||||
fp["torch"] = str(torch.__version__)
|
||||
fp["torch_cuda"] = str(torch.version.cuda)
|
||||
if torch.cuda.is_available():
|
||||
|
|
@ -157,7 +158,7 @@ def model_fingerprint(
|
|||
|
||||
|
||||
def cache_key(env_fp: dict[str, Any], model_fp: dict[str, Any]) -> str:
|
||||
payload = json.dumps({"env": env_fp, "model": model_fp}, sort_keys=True, default=str)
|
||||
payload = json.dumps({"env": env_fp, "model": model_fp}, sort_keys = True, default = str)
|
||||
return hashlib.sha256(payload.encode("utf-8")).hexdigest()[:32]
|
||||
|
||||
|
||||
|
|
@ -200,8 +201,10 @@ def begin(
|
|||
return None
|
||||
try:
|
||||
import torch # noqa: PLC0415
|
||||
if not (hasattr(torch.compiler, "save_cache_artifacts")
|
||||
and hasattr(torch.compiler, "load_cache_artifacts")):
|
||||
if not (
|
||||
hasattr(torch.compiler, "save_cache_artifacts")
|
||||
and hasattr(torch.compiler, "load_cache_artifacts")
|
||||
):
|
||||
_warn(logger, "Mega-cache API unavailable (need torch >= 2.7); skipping")
|
||||
return None
|
||||
except Exception as exc: # noqa: BLE001
|
||||
|
|
@ -210,20 +213,29 @@ def begin(
|
|||
|
||||
env_fp = environment_fingerprint()
|
||||
model_fp = model_fingerprint(
|
||||
family=family, transformer=transformer, dtype=dtype, quant=quant,
|
||||
attention_backend=attention_backend, compile_kwargs=compile_kwargs,
|
||||
shape_bucket=shape_bucket,
|
||||
family = family,
|
||||
transformer = transformer,
|
||||
dtype = dtype,
|
||||
quant = quant,
|
||||
attention_backend = attention_backend,
|
||||
compile_kwargs = compile_kwargs,
|
||||
shape_bucket = shape_bucket,
|
||||
)
|
||||
key = cache_key(env_fp, model_fp)
|
||||
cdir = cache_root() / key
|
||||
ctx = CacheContext(
|
||||
key=key, dir=cdir, bundle=cdir / _BUNDLE_NAME,
|
||||
manifest_path=cdir / _MANIFEST_NAME, env_fp=env_fp, model_fp=model_fp, mode=mode,
|
||||
key = key,
|
||||
dir = cdir,
|
||||
bundle = cdir / _BUNDLE_NAME,
|
||||
manifest_path = cdir / _MANIFEST_NAME,
|
||||
env_fp = env_fp,
|
||||
model_fp = model_fp,
|
||||
mode = mode,
|
||||
)
|
||||
|
||||
# Isolate inductor's on-disk cache per key so bundles never cross-contaminate.
|
||||
try:
|
||||
cdir.mkdir(parents=True, exist_ok=True)
|
||||
cdir.mkdir(parents = True, exist_ok = True)
|
||||
ctx.prev_inductor_dir = os.environ.get("TORCHINDUCTOR_CACHE_DIR")
|
||||
ctx.prev_inductor_dir_set = True
|
||||
os.environ["TORCHINDUCTOR_CACHE_DIR"] = str(cdir / "inductor")
|
||||
|
|
@ -264,6 +276,7 @@ def _try_load(ctx: CacheContext, logger: Any) -> bool:
|
|||
|
||||
try:
|
||||
import torch # noqa: PLC0415
|
||||
|
||||
info = torch.compiler.load_cache_artifacts(data)
|
||||
if info is None:
|
||||
_warn(logger, "compile-cache: load_cache_artifacts returned None (no hit)")
|
||||
|
|
@ -295,7 +308,7 @@ def save(ctx: Optional[CacheContext], *, logger: Any = None) -> bool:
|
|||
|
||||
data = result[0]
|
||||
try:
|
||||
ctx.dir.mkdir(parents=True, exist_ok=True)
|
||||
ctx.dir.mkdir(parents = True, exist_ok = True)
|
||||
ctx.bundle.write_bytes(data)
|
||||
manifest = {
|
||||
"format": _FORMAT_VERSION,
|
||||
|
|
@ -306,7 +319,7 @@ def save(ctx: Optional[CacheContext], *, logger: Any = None) -> bool:
|
|||
"env": ctx.env_fp,
|
||||
"model": ctx.model_fp,
|
||||
}
|
||||
ctx.manifest_path.write_text(json.dumps(manifest, indent=2, sort_keys=True, default=str))
|
||||
ctx.manifest_path.write_text(json.dumps(manifest, indent = 2, sort_keys = True, default = str))
|
||||
ctx.saved = True
|
||||
_info(logger, f"compile-cache: saved bundle ({len(data)} bytes) for key {ctx.key}")
|
||||
return True
|
||||
|
|
|
|||
|
|
@ -59,6 +59,7 @@ _ENV_ENABLE = "UNSLOTH_DIFFUSION_EAGER_PATCHES"
|
|||
def _patches_enabled() -> bool:
|
||||
return (os.environ.get(_ENV_ENABLE) or "").strip().lower() not in ("0", "off", "false", "no")
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Resolve the diffusers classes we patch. Any import failure -> that patch is
|
||||
# simply unavailable (None) and is skipped at install time.
|
||||
|
|
@ -86,24 +87,35 @@ except Exception: # noqa: BLE001
|
|||
# --------------------------------------------------------------------------- #
|
||||
def _adaln_continuous_forward(self, x, conditioning_embedding):
|
||||
emb = self.linear(self.silu(conditioning_embedding).to(x.dtype))
|
||||
scale, shift = torch.chunk(emb, 2, dim=1)
|
||||
scale, shift = torch.chunk(emb, 2, dim = 1)
|
||||
# original: self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :]
|
||||
return torch.addcmul(shift[:, None, :], self.norm(x), 1 + scale[:, None, :])
|
||||
|
||||
|
||||
def _adaln_zero_forward(self, x, timestep=None, class_labels=None, hidden_dtype=None, emb=None):
|
||||
def _adaln_zero_forward(
|
||||
self,
|
||||
x,
|
||||
timestep = None,
|
||||
class_labels = None,
|
||||
hidden_dtype = None,
|
||||
emb = None,
|
||||
):
|
||||
if self.emb is not None:
|
||||
emb = self.emb(timestep, class_labels, hidden_dtype=hidden_dtype)
|
||||
emb = self.emb(timestep, class_labels, hidden_dtype = hidden_dtype)
|
||||
emb = self.linear(self.silu(emb))
|
||||
shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = emb.chunk(6, dim=1)
|
||||
shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = emb.chunk(6, dim = 1)
|
||||
# original: self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None]
|
||||
x = torch.addcmul(shift_msa[:, None], self.norm(x), 1 + scale_msa[:, None])
|
||||
return x, gate_msa, shift_mlp, scale_mlp, gate_mlp
|
||||
|
||||
|
||||
def _adaln_zero_single_forward(self, x, emb=None):
|
||||
def _adaln_zero_single_forward(
|
||||
self,
|
||||
x,
|
||||
emb = None,
|
||||
):
|
||||
emb = self.linear(self.silu(emb))
|
||||
shift_msa, scale_msa, gate_msa = emb.chunk(3, dim=1)
|
||||
shift_msa, scale_msa, gate_msa = emb.chunk(3, dim = 1)
|
||||
x = torch.addcmul(shift_msa[:, None], self.norm(x), 1 + scale_msa[:, None])
|
||||
return x, gate_msa
|
||||
|
||||
|
|
@ -122,12 +134,7 @@ def _rmsnorm_forward(self, hidden_states):
|
|||
# * dtype mismatch (e.g. fp32 activations into an fp16/bf16-weight norm) -> diffusers
|
||||
# computes the variance in fp32 from the ORIGINAL tensor and only casts before the
|
||||
# weight multiply, so casting first would change the variance.
|
||||
if (
|
||||
_NPU
|
||||
or self.bias is not None
|
||||
or _orig_rmsnorm_forward is None
|
||||
or len(tuple(self.dim)) != 1
|
||||
):
|
||||
if _NPU or self.bias is not None or _orig_rmsnorm_forward is None or len(tuple(self.dim)) != 1:
|
||||
return _orig_rmsnorm_forward(self, hidden_states) # type: ignore[misc]
|
||||
weight = self.weight
|
||||
if weight is None:
|
||||
|
|
@ -178,14 +185,18 @@ def install_compile_safe_patches() -> int:
|
|||
# to it for the uncommon (NPU / bias / fp32-weight / tuple-dim) cases.
|
||||
if cls is _RMSNorm:
|
||||
_orig_rmsnorm_forward = cls.forward
|
||||
if apply_patch(cls, "forward", new_fn, match_level="relaxed"):
|
||||
if apply_patch(cls, "forward", new_fn, match_level = "relaxed"):
|
||||
_patched.append(cls)
|
||||
else:
|
||||
logger.warning("eager-patch: skipping %s (signature mismatch / unavailable)",
|
||||
getattr(cls, "__name__", cls))
|
||||
logger.warning(
|
||||
"eager-patch: skipping %s (signature mismatch / unavailable)",
|
||||
getattr(cls, "__name__", cls),
|
||||
)
|
||||
if cls is _RMSNorm:
|
||||
_orig_rmsnorm_forward = None
|
||||
logger.info("eager-patch: installed %d/%d shared diffusion patches", len(_patched), len(_specs()))
|
||||
logger.info(
|
||||
"eager-patch: installed %d/%d shared diffusion patches", len(_patched), len(_specs())
|
||||
)
|
||||
return len(_patched)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -100,7 +100,10 @@ def _activate(name: str, reason: Optional[str]) -> Any:
|
|||
|
||||
|
||||
def select_and_activate_engine(
|
||||
fam: DiffusionFamily, *, hf_token: Optional[str] = None, model_kind: Optional[str] = None
|
||||
fam: DiffusionFamily,
|
||||
*,
|
||||
hf_token: Optional[str] = None,
|
||||
model_kind: Optional[str] = None,
|
||||
) -> Any:
|
||||
"""Pick + activate the engine for loading ``fam`` on this host; return the engine.
|
||||
|
||||
|
|
|
|||
|
|
@ -255,8 +255,7 @@ def detect_family(repo_id: str, override: Optional[str] = None) -> Optional[Diff
|
|||
# matched family does not itself declare, reject so the load fails fast + clearly.
|
||||
matched_tokens = (match.name, *match.aliases)
|
||||
if any(
|
||||
kw in needle and not any(kw in tok for tok in matched_tokens)
|
||||
for kw in _EDIT_KEYWORDS
|
||||
kw in needle and not any(kw in tok for tok in matched_tokens) for kw in _EDIT_KEYWORDS
|
||||
):
|
||||
return None
|
||||
return match
|
||||
|
|
|
|||
|
|
@ -56,7 +56,6 @@ def _gguf_utils():
|
|||
"""The diffusers GGUF utils module, or None if this diffusers build lacks it."""
|
||||
try:
|
||||
from diffusers.quantizers.gguf import utils as gguf_utils # noqa: PLC0415
|
||||
|
||||
return gguf_utils
|
||||
except Exception: # noqa: BLE001 — old/!GGUF diffusers -> accelerator is a no-op
|
||||
return None
|
||||
|
|
@ -91,10 +90,10 @@ def install_compiled_dequant(logger: Any = None) -> bool:
|
|||
try:
|
||||
import torch # noqa: PLC0415
|
||||
|
||||
compiled = torch.compile(gguf_utils.dequantize_gguf_tensor, dynamic=True)
|
||||
compiled = torch.compile(gguf_utils.dequantize_gguf_tensor, dynamic = True)
|
||||
# force=True: the new callable is the SAME function compiled, so its fingerprint
|
||||
# differs from the original and can_safely_patch would (correctly) reject it.
|
||||
if apply_patch(gguf_utils, _DEQUANT_ATTR, compiled, force=True):
|
||||
if apply_patch(gguf_utils, _DEQUANT_ATTR, compiled, force = True):
|
||||
_compiled_dequant_installed = True
|
||||
return True
|
||||
return False
|
||||
|
|
|
|||
|
|
@ -49,9 +49,7 @@ def apply_patch(
|
|||
except Exception: # noqa: BLE001 — no unsloth_zoo / no-GPU host -> optimisation skipped
|
||||
return False
|
||||
try:
|
||||
return bool(
|
||||
patch_function(target, attr, new_fn, match_level=match_level, force=force)
|
||||
)
|
||||
return bool(patch_function(target, attr, new_fn, match_level = match_level, force = force))
|
||||
except Exception: # noqa: BLE001 — best-effort; leave the original in place
|
||||
return False
|
||||
|
||||
|
|
|
|||
|
|
@ -21,7 +21,7 @@ pytest.importorskip("diffusers")
|
|||
from core.inference import diffusion_arch_patches as ap # noqa: E402
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
@pytest.fixture(autouse = True)
|
||||
def _clean():
|
||||
ap.uninstall_arch_patches()
|
||||
yield
|
||||
|
|
@ -40,7 +40,7 @@ def test_qwen_modulate_matches_stock_global_and_indexed():
|
|||
# _modulate uses no real `self` state, so call it unbound with self=None.
|
||||
ref_x, ref_g = Q._modulate(None, x, mod)
|
||||
got_x, got_g = ap._qwen_modulate(None, x, mod)
|
||||
torch.testing.assert_close(got_x, ref_x, atol=1e-5, rtol=1e-4)
|
||||
torch.testing.assert_close(got_x, ref_x, atol = 1e-5, rtol = 1e-4)
|
||||
assert torch.equal(got_g, ref_g)
|
||||
|
||||
# per-token `index` branch (mod batch is 2*B).
|
||||
|
|
@ -48,7 +48,7 @@ def test_qwen_modulate_matches_stock_global_and_indexed():
|
|||
mod2 = torch.randn(2 * B, 3 * D)
|
||||
ref2_x, ref2_g = Q._modulate(None, x, mod2, idx)
|
||||
got2_x, got2_g = ap._qwen_modulate(None, x, mod2, idx)
|
||||
torch.testing.assert_close(got2_x, ref2_x, atol=1e-5, rtol=1e-4)
|
||||
torch.testing.assert_close(got2_x, ref2_x, atol = 1e-5, rtol = 1e-4)
|
||||
assert torch.equal(got2_g, ref2_g)
|
||||
|
||||
|
||||
|
|
@ -61,18 +61,23 @@ class _AttnStub(torch.nn.Module):
|
|||
|
||||
def __init__(self, dim):
|
||||
super().__init__()
|
||||
self.proj = torch.nn.Linear(dim, dim, bias=False)
|
||||
self.proj = torch.nn.Linear(dim, dim, bias = False)
|
||||
|
||||
def forward(self, h, **kwargs):
|
||||
return self.proj(h)
|
||||
|
||||
|
||||
def _zimage_block(dim=64, heads=4):
|
||||
def _zimage_block(dim = 64, heads = 4):
|
||||
from diffusers.models.transformers.transformer_z_image import ZImageTransformerBlock
|
||||
|
||||
blk = ZImageTransformerBlock(
|
||||
layer_id=0, dim=dim, n_heads=heads, n_kv_heads=heads,
|
||||
norm_eps=1e-5, qk_norm=True, modulation=True,
|
||||
layer_id = 0,
|
||||
dim = dim,
|
||||
n_heads = heads,
|
||||
n_kv_heads = heads,
|
||||
norm_eps = 1e-5,
|
||||
qk_norm = True,
|
||||
modulation = True,
|
||||
).eval()
|
||||
blk.attention = _AttnStub(dim).eval()
|
||||
return blk
|
||||
|
|
@ -93,9 +98,9 @@ def test_zimage_forward_matches_stock_global_modulation():
|
|||
adaln = torch.randn(B, _adaln_dim(D))
|
||||
|
||||
with torch.inference_mode():
|
||||
ref = ZImageTransformerBlock.forward(blk, x, None, None, adaln_input=adaln).clone()
|
||||
got = ap._zimage_forward(blk, x, None, None, adaln_input=adaln)
|
||||
torch.testing.assert_close(got, ref, atol=1e-5, rtol=1e-4)
|
||||
ref = ZImageTransformerBlock.forward(blk, x, None, None, adaln_input = adaln).clone()
|
||||
got = ap._zimage_forward(blk, x, None, None, adaln_input = adaln)
|
||||
torch.testing.assert_close(got, ref, atol = 1e-5, rtol = 1e-4)
|
||||
|
||||
|
||||
def test_zimage_forward_matches_stock_per_token_modulation():
|
||||
|
|
@ -112,14 +117,24 @@ def test_zimage_forward_matches_stock_per_token_modulation():
|
|||
|
||||
with torch.inference_mode():
|
||||
ref = ZImageTransformerBlock.forward(
|
||||
blk, x, None, None, noise_mask=noise_mask,
|
||||
adaln_noisy=adaln_noisy, adaln_clean=adaln_clean,
|
||||
blk,
|
||||
x,
|
||||
None,
|
||||
None,
|
||||
noise_mask = noise_mask,
|
||||
adaln_noisy = adaln_noisy,
|
||||
adaln_clean = adaln_clean,
|
||||
).clone()
|
||||
got = ap._zimage_forward(
|
||||
blk, x, None, None, noise_mask=noise_mask,
|
||||
adaln_noisy=adaln_noisy, adaln_clean=adaln_clean,
|
||||
blk,
|
||||
x,
|
||||
None,
|
||||
None,
|
||||
noise_mask = noise_mask,
|
||||
adaln_noisy = adaln_noisy,
|
||||
adaln_clean = adaln_clean,
|
||||
)
|
||||
torch.testing.assert_close(got, ref, atol=1e-5, rtol=1e-4)
|
||||
torch.testing.assert_close(got, ref, atol = 1e-5, rtol = 1e-4)
|
||||
|
||||
|
||||
# ── flux.1 / flux.2 block forwards (modulation + gated-residual addcmul) ─────────
|
||||
|
|
@ -130,10 +145,15 @@ class _Tuple2AttnStub(torch.nn.Module):
|
|||
|
||||
def __init__(self, dim):
|
||||
super().__init__()
|
||||
self.pi = torch.nn.Linear(dim, dim, bias=False)
|
||||
self.pc = torch.nn.Linear(dim, dim, bias=False)
|
||||
self.pi = torch.nn.Linear(dim, dim, bias = False)
|
||||
self.pc = torch.nn.Linear(dim, dim, bias = False)
|
||||
|
||||
def forward(self, hidden_states, encoder_hidden_states=None, **kwargs):
|
||||
def forward(
|
||||
self,
|
||||
hidden_states,
|
||||
encoder_hidden_states = None,
|
||||
**kwargs,
|
||||
):
|
||||
return self.pi(hidden_states), self.pc(encoder_hidden_states)
|
||||
|
||||
|
||||
|
|
@ -142,7 +162,7 @@ class _SingleAttnStub(torch.nn.Module):
|
|||
|
||||
def __init__(self, dim):
|
||||
super().__init__()
|
||||
self.p = torch.nn.Linear(dim, dim, bias=False)
|
||||
self.p = torch.nn.Linear(dim, dim, bias = False)
|
||||
|
||||
def forward(self, hidden_states, **kwargs):
|
||||
return self.p(hidden_states)
|
||||
|
|
@ -152,9 +172,9 @@ def _close_any(got, ref):
|
|||
if isinstance(ref, tuple):
|
||||
assert len(got) == len(ref)
|
||||
for g, r in zip(got, ref):
|
||||
torch.testing.assert_close(g, r, atol=1e-5, rtol=1e-4)
|
||||
torch.testing.assert_close(g, r, atol = 1e-5, rtol = 1e-4)
|
||||
else:
|
||||
torch.testing.assert_close(got, ref, atol=1e-5, rtol=1e-4)
|
||||
torch.testing.assert_close(got, ref, atol = 1e-5, rtol = 1e-4)
|
||||
|
||||
|
||||
D, H = 64, 4
|
||||
|
|
@ -165,7 +185,7 @@ def test_flux_double_forward_matches_stock():
|
|||
from diffusers.models.transformers.transformer_flux import FluxTransformerBlock
|
||||
|
||||
torch.manual_seed(0)
|
||||
blk = FluxTransformerBlock(dim=D, num_attention_heads=H, attention_head_dim=D // H).eval()
|
||||
blk = FluxTransformerBlock(dim = D, num_attention_heads = H, attention_head_dim = D // H).eval()
|
||||
blk.attn = _Tuple2AttnStub(D).eval()
|
||||
hs, ehs, temb = torch.randn(B, L, D), torch.randn(B, LC, D), torch.randn(B, D)
|
||||
with torch.inference_mode():
|
||||
|
|
@ -178,7 +198,7 @@ def test_flux_single_forward_matches_stock():
|
|||
from diffusers.models.transformers.transformer_flux import FluxSingleTransformerBlock
|
||||
|
||||
torch.manual_seed(1)
|
||||
blk = FluxSingleTransformerBlock(dim=D, num_attention_heads=H, attention_head_dim=D // H).eval()
|
||||
blk = FluxSingleTransformerBlock(dim = D, num_attention_heads = H, attention_head_dim = D // H).eval()
|
||||
blk.attn = _SingleAttnStub(D).eval()
|
||||
hs, ehs, temb = torch.randn(B, L, D), torch.randn(B, LC, D), torch.randn(B, D)
|
||||
with torch.inference_mode():
|
||||
|
|
@ -191,7 +211,7 @@ def test_flux2_double_forward_matches_stock():
|
|||
from diffusers.models.transformers.transformer_flux2 import Flux2TransformerBlock
|
||||
|
||||
torch.manual_seed(2)
|
||||
blk = Flux2TransformerBlock(dim=D, num_attention_heads=H, attention_head_dim=D // H).eval()
|
||||
blk = Flux2TransformerBlock(dim = D, num_attention_heads = H, attention_head_dim = D // H).eval()
|
||||
blk.attn = _Tuple2AttnStub(D).eval()
|
||||
hs, ehs = torch.randn(B, L, D), torch.randn(B, LC, D)
|
||||
tmi, tmt = torch.randn(B, 6 * D), torch.randn(B, 6 * D)
|
||||
|
|
@ -205,7 +225,9 @@ def test_flux2_single_forward_matches_stock():
|
|||
from diffusers.models.transformers.transformer_flux2 import Flux2SingleTransformerBlock
|
||||
|
||||
torch.manual_seed(3)
|
||||
blk = Flux2SingleTransformerBlock(dim=D, num_attention_heads=H, attention_head_dim=D // H).eval()
|
||||
blk = Flux2SingleTransformerBlock(
|
||||
dim = D, num_attention_heads = H, attention_head_dim = D // H
|
||||
).eval()
|
||||
blk.attn = _SingleAttnStub(D).eval()
|
||||
hs, ehs, tm = torch.randn(B, L, D), torch.randn(B, LC, D), torch.randn(B, 3 * D)
|
||||
with torch.inference_mode():
|
||||
|
|
|
|||
|
|
@ -404,8 +404,12 @@ def test_generate_img2img_uses_from_pipe(fake_runtime, tmp_path):
|
|||
|
||||
loaded_pipe = backend._state.pipe
|
||||
out = backend.generate(
|
||||
prompt = "a car at sunset", steps = 4, guidance = 0.0, seed = 3,
|
||||
init_image = _tiny_png_b64(), strength = 0.5,
|
||||
prompt = "a car at sunset",
|
||||
steps = 4,
|
||||
guidance = 0.0,
|
||||
seed = 3,
|
||||
init_image = _tiny_png_b64(),
|
||||
strength = 0.5,
|
||||
)
|
||||
assert len(out["images"]) == 1
|
||||
# from_pipe was handed the loaded text-to-image pipe (component reuse, no reload).
|
||||
|
|
@ -414,7 +418,7 @@ def test_generate_img2img_uses_from_pipe(fake_runtime, tmp_path):
|
|||
# both upcasts the reused bf16 modules and crashes on torchao-quantized weights.
|
||||
assert _FakeImg2ImgPipeline.from_pipe_kwargs.get("torch_dtype", "MISSING") is None
|
||||
call = _FakeImg2ImgPipe.last_kwargs
|
||||
assert call["image"] is not None # decoded source image passed through
|
||||
assert call["image"] is not None # decoded source image passed through
|
||||
assert call["strength"] == 0.5
|
||||
assert "width" not in call and "height" not in call # img2img derives size from image
|
||||
|
||||
|
|
@ -461,8 +465,12 @@ def test_generate_upscale_enlarges_and_low_strength(fake_runtime, tmp_path):
|
|||
|
||||
loaded_pipe = backend._state.pipe
|
||||
out = backend.generate(
|
||||
prompt = "a crisp photo", steps = 4, guidance = 0.0, seed = 3,
|
||||
init_image = _tiny_png_b64(), upscale = 2.0, # 64 -> 128, no explicit strength
|
||||
prompt = "a crisp photo",
|
||||
steps = 4,
|
||||
guidance = 0.0,
|
||||
seed = 3,
|
||||
init_image = _tiny_png_b64(),
|
||||
upscale = 2.0, # 64 -> 128, no explicit strength
|
||||
)
|
||||
assert len(out["images"]) == 1
|
||||
# Reuses the resident modules via from_pipe (no reload, no extra VRAM).
|
||||
|
|
@ -475,14 +483,22 @@ def test_generate_upscale_enlarges_and_low_strength(fake_runtime, tmp_path):
|
|||
|
||||
# The factor is capped at 4x so a large request can't blow up the VAE/transformer.
|
||||
backend.generate(
|
||||
prompt = "x", steps = 4, seed = 1, init_image = _tiny_png_b64(), upscale = 99.0,
|
||||
prompt = "x",
|
||||
steps = 4,
|
||||
seed = 1,
|
||||
init_image = _tiny_png_b64(),
|
||||
upscale = 99.0,
|
||||
)
|
||||
assert _FakeImg2ImgPipe.last_kwargs["image"].size == (256, 256) # 64 * 4 (capped)
|
||||
|
||||
# An explicit strength overrides the hires-fix default.
|
||||
backend.generate(
|
||||
prompt = "x", steps = 4, seed = 1, init_image = _tiny_png_b64(),
|
||||
upscale = 1.5, strength = 0.2,
|
||||
prompt = "x",
|
||||
steps = 4,
|
||||
seed = 1,
|
||||
init_image = _tiny_png_b64(),
|
||||
upscale = 1.5,
|
||||
strength = 0.2,
|
||||
)
|
||||
assert _FakeImg2ImgPipe.last_kwargs["strength"] == 0.2
|
||||
# 64 * 1.5 = 96, already a multiple of 16.
|
||||
|
|
@ -562,8 +578,12 @@ def test_inpaint_snaps_image_and_mask_together(fake_runtime, tmp_path):
|
|||
str(tmp_path), gguf_filename = "model.gguf", base_repo = "base/repo", family_override = "z-image"
|
||||
)
|
||||
backend.generate(
|
||||
prompt = "x", steps = 4, seed = 1,
|
||||
init_image = _png_b64(186), mask_image = _mask_b64(186), strength = 0.5,
|
||||
prompt = "x",
|
||||
steps = 4,
|
||||
seed = 1,
|
||||
init_image = _png_b64(186),
|
||||
mask_image = _mask_b64(186),
|
||||
strength = 0.5,
|
||||
)
|
||||
assert _FakeInpaintPipe.last_kwargs["image"].size == (192, 192)
|
||||
assert _FakeInpaintPipe.last_kwargs["mask_image"].size == (192, 192)
|
||||
|
|
@ -581,7 +601,9 @@ def test_generate_reference_uses_loaded_pipe_at_slider_size(fake_runtime, tmp_pa
|
|||
(tmp_path / "model.gguf").write_bytes(b"x")
|
||||
backend = DiffusionBackend()
|
||||
backend.load_pipeline(
|
||||
str(tmp_path), gguf_filename = "model.gguf", base_repo = "base/repo",
|
||||
str(tmp_path),
|
||||
gguf_filename = "model.gguf",
|
||||
base_repo = "base/repo",
|
||||
family_override = "flux.2-klein",
|
||||
)
|
||||
# FLUX.2-klein: txt2img + reference (own pipe) + inpaint (dedicated pipe). No img2img class,
|
||||
|
|
@ -590,14 +612,20 @@ def test_generate_reference_uses_loaded_pipe_at_slider_size(fake_runtime, tmp_pa
|
|||
|
||||
loaded_pipe = backend._state.pipe
|
||||
out = backend.generate(
|
||||
prompt = "a portrait in this style", steps = 6, guidance = 4.0, seed = 5,
|
||||
width = 768, height = 512, init_image = _tiny_png_b64(), strength = 0.5,
|
||||
prompt = "a portrait in this style",
|
||||
steps = 6,
|
||||
guidance = 4.0,
|
||||
seed = 5,
|
||||
width = 768,
|
||||
height = 512,
|
||||
init_image = _tiny_png_b64(),
|
||||
strength = 0.5,
|
||||
)
|
||||
assert len(out["images"]) == 1
|
||||
call = loaded_pipe.last_kwargs
|
||||
assert call["image"] is not None # reference handed to the loaded pipe
|
||||
assert call["image"] is not None # reference handed to the loaded pipe
|
||||
assert call["width"] == 768 and call["height"] == 512 # OUTPUT size = sliders, not input
|
||||
assert "strength" not in call # reference conditioning has no strength
|
||||
assert "strength" not in call # reference conditioning has no strength
|
||||
assert "mask_image" not in call
|
||||
# Guidance flows via guidance_scale (FLUX.2 default behaviour).
|
||||
assert call["guidance_scale"] == 4.0
|
||||
|
|
@ -605,8 +633,13 @@ def test_generate_reference_uses_loaded_pipe_at_slider_size(fake_runtime, tmp_pa
|
|||
# Multi-reference: extra reference_images are combined with init_image into a LIST so the
|
||||
# model can blend several references (subject + style).
|
||||
backend.generate(
|
||||
prompt = "combine these", steps = 6, seed = 9, width = 1024, height = 1024,
|
||||
init_image = _tiny_png_b64(), reference_images = [_tiny_png_b64(), _tiny_png_b64()],
|
||||
prompt = "combine these",
|
||||
steps = 6,
|
||||
seed = 9,
|
||||
width = 1024,
|
||||
height = 1024,
|
||||
init_image = _tiny_png_b64(),
|
||||
reference_images = [_tiny_png_b64(), _tiny_png_b64()],
|
||||
)
|
||||
img_arg = loaded_pipe.last_kwargs["image"]
|
||||
assert isinstance(img_arg, list) and len(img_arg) == 3 # primary + 2 extras
|
||||
|
|
@ -614,10 +647,14 @@ def test_generate_reference_uses_loaded_pipe_at_slider_size(fake_runtime, tmp_pa
|
|||
# Branch ordering: an init image + MASK on a reference family must route to inpaint (the
|
||||
# dedicated pipeline), NOT be swallowed by the reference branch (which ignores the mask).
|
||||
backend.generate(
|
||||
prompt = "repaint here", steps = 6, seed = 2,
|
||||
init_image = _tiny_png_b64(), mask_image = _tiny_mask_b64(), strength = 0.8,
|
||||
prompt = "repaint here",
|
||||
steps = 6,
|
||||
seed = 2,
|
||||
init_image = _tiny_png_b64(),
|
||||
mask_image = _tiny_mask_b64(),
|
||||
strength = 0.8,
|
||||
)
|
||||
assert _FakeInpaintPipeline.built_from is loaded_pipe # built via from_pipe off the load
|
||||
assert _FakeInpaintPipeline.built_from is loaded_pipe # built via from_pipe off the load
|
||||
assert _FakeInpaintPipe.last_kwargs["mask_image"] is not None
|
||||
assert _FakeInpaintPipe.last_kwargs["strength"] == 0.8
|
||||
|
||||
|
|
@ -653,8 +690,13 @@ def test_generate_inpaint_uses_from_pipe(fake_runtime, tmp_path):
|
|||
)
|
||||
loaded_pipe = backend._state.pipe
|
||||
out = backend.generate(
|
||||
prompt = "a red door", steps = 4, guidance = 0.0, seed = 5,
|
||||
init_image = _tiny_png_b64(), mask_image = _tiny_mask_b64(), strength = 0.7,
|
||||
prompt = "a red door",
|
||||
steps = 4,
|
||||
guidance = 0.0,
|
||||
seed = 5,
|
||||
init_image = _tiny_png_b64(),
|
||||
mask_image = _tiny_mask_b64(),
|
||||
strength = 0.7,
|
||||
)
|
||||
assert len(out["images"]) == 1
|
||||
# The inpaint pipe (not img2img) was selected and built from the loaded pipe.
|
||||
|
|
@ -725,7 +767,9 @@ def test_edit_family_uses_own_pipeline_and_requires_image(fake_runtime, tmp_path
|
|||
(tmp_path / "model.gguf").write_bytes(b"x")
|
||||
backend = DiffusionBackend()
|
||||
backend.load_pipeline(
|
||||
str(tmp_path), gguf_filename = "model.gguf", base_repo = "Qwen/Qwen-Image-Edit-2511",
|
||||
str(tmp_path),
|
||||
gguf_filename = "model.gguf",
|
||||
base_repo = "Qwen/Qwen-Image-Edit-2511",
|
||||
family_override = "qwen-image-edit",
|
||||
)
|
||||
# Edit families advertise only the edit workflow (no txt2img / img2img / inpaint).
|
||||
|
|
@ -733,7 +777,10 @@ def test_edit_family_uses_own_pipeline_and_requires_image(fake_runtime, tmp_path
|
|||
loaded_pipe = backend._state.pipe
|
||||
|
||||
out = backend.generate(
|
||||
prompt = "make it night", steps = 8, guidance = 4.0, seed = 1,
|
||||
prompt = "make it night",
|
||||
steps = 8,
|
||||
guidance = 4.0,
|
||||
seed = 1,
|
||||
init_image = _tiny_png_b64(),
|
||||
)
|
||||
assert len(out["images"]) == 1
|
||||
|
|
@ -1240,9 +1287,7 @@ def test_callback_cancellation_interrupts_denoise(fake_runtime):
|
|||
def test_validate_load_request(tmp_path):
|
||||
backend = DiffusionBackend()
|
||||
# No filename + unsloth repo -> a full-pipeline load (allowed for unsloth/*).
|
||||
assert (
|
||||
backend.validate_load_request("unsloth/Z-Image-Turbo-unsloth-bnb-4bit").name == "z-image"
|
||||
)
|
||||
assert backend.validate_load_request("unsloth/Z-Image-Turbo-unsloth-bnb-4bit").name == "z-image"
|
||||
# No filename + non-unsloth repo -> a pipeline load, gated to unsloth/* -> rejected.
|
||||
with pytest.raises(ValueError, match = "unsloth"):
|
||||
backend.validate_load_request("some-org/Z-Image-bnb-4bit")
|
||||
|
|
|
|||
|
|
@ -19,17 +19,17 @@ import pytest
|
|||
from core.inference import diffusion_compile_cache as cc
|
||||
|
||||
|
||||
def _transformer(blocks=("FluxTransformerBlock", "FluxSingleTransformerBlock")):
|
||||
return types.SimpleNamespace(_repeated_blocks=list(blocks))
|
||||
def _transformer(blocks = ("FluxTransformerBlock", "FluxSingleTransformerBlock")):
|
||||
return types.SimpleNamespace(_repeated_blocks = list(blocks))
|
||||
|
||||
|
||||
_BEGIN_KW = dict(
|
||||
family="flux.1",
|
||||
dtype="torch.bfloat16",
|
||||
quant=None,
|
||||
attention_backend="_native_cudnn",
|
||||
compile_kwargs={"fullgraph": True, "dynamic": True},
|
||||
shape_bucket="1024x1024",
|
||||
family = "flux.1",
|
||||
dtype = "torch.bfloat16",
|
||||
quant = None,
|
||||
attention_backend = "_native_cudnn",
|
||||
compile_kwargs = {"fullgraph": True, "dynamic": True},
|
||||
shape_bucket = "1024x1024",
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -43,25 +43,47 @@ def test_environment_fingerprint_has_hard_dimensions():
|
|||
def test_cache_key_stable_across_kwarg_order():
|
||||
efp = cc.environment_fingerprint()
|
||||
t = _transformer()
|
||||
a = cc.model_fingerprint(family="flux.1", transformer=t, dtype="bf16", quant=None,
|
||||
attention_backend="x", compile_kwargs={"fullgraph": True, "dynamic": True})
|
||||
b = cc.model_fingerprint(family="flux.1", transformer=t, dtype="bf16", quant=None,
|
||||
attention_backend="x", compile_kwargs={"dynamic": True, "fullgraph": True})
|
||||
a = cc.model_fingerprint(
|
||||
family = "flux.1",
|
||||
transformer = t,
|
||||
dtype = "bf16",
|
||||
quant = None,
|
||||
attention_backend = "x",
|
||||
compile_kwargs = {"fullgraph": True, "dynamic": True},
|
||||
)
|
||||
b = cc.model_fingerprint(
|
||||
family = "flux.1",
|
||||
transformer = t,
|
||||
dtype = "bf16",
|
||||
quant = None,
|
||||
attention_backend = "x",
|
||||
compile_kwargs = {"dynamic": True, "fullgraph": True},
|
||||
)
|
||||
assert cc.cache_key(efp, a) == cc.cache_key(efp, b)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("field,value", [
|
||||
("family", "qwen-image"),
|
||||
("dtype", "torch.float16"),
|
||||
("quant", "int8"),
|
||||
("attention_backend", "native"),
|
||||
("shape_bucket", "512x512"),
|
||||
])
|
||||
@pytest.mark.parametrize(
|
||||
"field,value",
|
||||
[
|
||||
("family", "qwen-image"),
|
||||
("dtype", "torch.float16"),
|
||||
("quant", "int8"),
|
||||
("attention_backend", "native"),
|
||||
("shape_bucket", "512x512"),
|
||||
],
|
||||
)
|
||||
def test_cache_key_sensitive_to_model_dims(field, value):
|
||||
efp = cc.environment_fingerprint()
|
||||
t = _transformer()
|
||||
base = dict(family="flux.1", transformer=t, dtype="bf16", quant=None,
|
||||
attention_backend="x", compile_kwargs={"fullgraph": True}, shape_bucket="1024x1024")
|
||||
base = dict(
|
||||
family = "flux.1",
|
||||
transformer = t,
|
||||
dtype = "bf16",
|
||||
quant = None,
|
||||
attention_backend = "x",
|
||||
compile_kwargs = {"fullgraph": True},
|
||||
shape_bucket = "1024x1024",
|
||||
)
|
||||
k0 = cc.cache_key(efp, cc.model_fingerprint(**base))
|
||||
base[field] = value
|
||||
assert cc.cache_key(efp, cc.model_fingerprint(**base)) != k0
|
||||
|
|
@ -69,32 +91,58 @@ def test_cache_key_sensitive_to_model_dims(field, value):
|
|||
|
||||
def test_repeated_blocks_change_key():
|
||||
efp = cc.environment_fingerprint()
|
||||
k1 = cc.cache_key(efp, cc.model_fingerprint(family="f", transformer=_transformer(("A",)),
|
||||
dtype="bf16", quant=None, attention_backend="x", compile_kwargs={}))
|
||||
k2 = cc.cache_key(efp, cc.model_fingerprint(family="f", transformer=_transformer(("B",)),
|
||||
dtype="bf16", quant=None, attention_backend="x", compile_kwargs={}))
|
||||
k1 = cc.cache_key(
|
||||
efp,
|
||||
cc.model_fingerprint(
|
||||
family = "f",
|
||||
transformer = _transformer(("A",)),
|
||||
dtype = "bf16",
|
||||
quant = None,
|
||||
attention_backend = "x",
|
||||
compile_kwargs = {},
|
||||
),
|
||||
)
|
||||
k2 = cc.cache_key(
|
||||
efp,
|
||||
cc.model_fingerprint(
|
||||
family = "f",
|
||||
transformer = _transformer(("B",)),
|
||||
dtype = "bf16",
|
||||
quant = None,
|
||||
attention_backend = "x",
|
||||
compile_kwargs = {},
|
||||
),
|
||||
)
|
||||
assert k1 != k2
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------------- env knobs
|
||||
@pytest.mark.parametrize("raw,expected", [
|
||||
("0", "off"), ("off", "off"), ("1", "on"), ("on", "on"),
|
||||
("auto", "auto"), ("", "auto"), ("garbage", "auto"),
|
||||
])
|
||||
@pytest.mark.parametrize(
|
||||
"raw,expected",
|
||||
[
|
||||
("0", "off"),
|
||||
("off", "off"),
|
||||
("1", "on"),
|
||||
("on", "on"),
|
||||
("auto", "auto"),
|
||||
("", "auto"),
|
||||
("garbage", "auto"),
|
||||
],
|
||||
)
|
||||
def test_cache_mode(monkeypatch, raw, expected):
|
||||
monkeypatch.setenv(cc._ENV_MODE, raw)
|
||||
assert cc.cache_mode() == expected
|
||||
|
||||
|
||||
def test_cache_mode_default_auto(monkeypatch):
|
||||
monkeypatch.delenv(cc._ENV_MODE, raising=False)
|
||||
monkeypatch.delenv(cc._ENV_MODE, raising = False)
|
||||
assert cc.cache_mode() == "auto"
|
||||
|
||||
|
||||
# ------------------------------------------------------------------------------ disabled
|
||||
def test_begin_returns_none_when_disabled(monkeypatch):
|
||||
monkeypatch.setenv(cc._ENV_MODE, "0")
|
||||
assert cc.begin(transformer=_transformer(), **_BEGIN_KW) is None
|
||||
assert cc.begin(transformer = _transformer(), **_BEGIN_KW) is None
|
||||
|
||||
|
||||
def test_begin_returns_none_without_megacache_api(monkeypatch):
|
||||
|
|
@ -102,7 +150,7 @@ def test_begin_returns_none_without_megacache_api(monkeypatch):
|
|||
fake_torch = types.ModuleType("torch")
|
||||
fake_torch.compiler = types.SimpleNamespace() # no save/load attrs
|
||||
monkeypatch.setitem(__import__("sys").modules, "torch", fake_torch)
|
||||
assert cc.begin(transformer=_transformer(), **_BEGIN_KW) is None
|
||||
assert cc.begin(transformer = _transformer(), **_BEGIN_KW) is None
|
||||
|
||||
|
||||
# ----------------------------------------------------------------- megacache fake + flow
|
||||
|
|
@ -120,23 +168,23 @@ def fake_megacache(monkeypatch):
|
|||
state["loaded_with"] = data
|
||||
return object() if data == b"ARTIFACT-BYTES" else None
|
||||
|
||||
monkeypatch.setattr(torch.compiler, "save_cache_artifacts", fake_save, raising=False)
|
||||
monkeypatch.setattr(torch.compiler, "load_cache_artifacts", fake_load, raising=False)
|
||||
monkeypatch.setattr(torch.compiler, "save_cache_artifacts", fake_save, raising = False)
|
||||
monkeypatch.setattr(torch.compiler, "load_cache_artifacts", fake_load, raising = False)
|
||||
return state
|
||||
|
||||
|
||||
def test_save_then_load_roundtrip(monkeypatch, tmp_path, fake_megacache):
|
||||
monkeypatch.setenv(cc._ENV_MODE, "on") # load + save
|
||||
monkeypatch.setenv(cc._ENV_MODE, "on") # load + save
|
||||
monkeypatch.setenv(cc._ENV_DIR, str(tmp_path))
|
||||
|
||||
# First load: cold (no bundle yet).
|
||||
ctx = cc.begin(transformer=_transformer(), **_BEGIN_KW)
|
||||
ctx = cc.begin(transformer = _transformer(), **_BEGIN_KW)
|
||||
assert ctx is not None and ctx.hit is False
|
||||
assert cc.save(ctx) is True
|
||||
assert ctx.bundle.exists() and ctx.manifest_path.exists()
|
||||
|
||||
# Second load with the SAME fingerprint: warm hit.
|
||||
ctx2 = cc.begin(transformer=_transformer(), **_BEGIN_KW)
|
||||
ctx2 = cc.begin(transformer = _transformer(), **_BEGIN_KW)
|
||||
assert ctx2 is not None and ctx2.hit is True
|
||||
assert fake_megacache["loaded_with"] == b"ARTIFACT-BYTES"
|
||||
assert ctx2.key == ctx.key
|
||||
|
|
@ -144,17 +192,17 @@ def test_save_then_load_roundtrip(monkeypatch, tmp_path, fake_megacache):
|
|||
|
||||
def test_no_save_in_auto_mode(monkeypatch, tmp_path, fake_megacache):
|
||||
monkeypatch.setenv(cc._ENV_MODE, "auto")
|
||||
monkeypatch.delenv(cc._ENV_SAVE, raising=False)
|
||||
monkeypatch.delenv(cc._ENV_SAVE, raising = False)
|
||||
monkeypatch.setenv(cc._ENV_DIR, str(tmp_path))
|
||||
ctx = cc.begin(transformer=_transformer(), **_BEGIN_KW)
|
||||
assert cc.save(ctx) is False # auto without SAVE opt-in does not write
|
||||
ctx = cc.begin(transformer = _transformer(), **_BEGIN_KW)
|
||||
assert cc.save(ctx) is False # auto without SAVE opt-in does not write
|
||||
assert not ctx.bundle.exists()
|
||||
|
||||
|
||||
def test_fingerprint_mismatch_falls_back(monkeypatch, tmp_path, fake_megacache):
|
||||
monkeypatch.setenv(cc._ENV_MODE, "on")
|
||||
monkeypatch.setenv(cc._ENV_DIR, str(tmp_path))
|
||||
ctx = cc.begin(transformer=_transformer(), **_BEGIN_KW)
|
||||
ctx = cc.begin(transformer = _transformer(), **_BEGIN_KW)
|
||||
cc.save(ctx)
|
||||
|
||||
# Tamper the manifest's env fingerprint -> exact-match guard must reject the bundle.
|
||||
|
|
@ -162,28 +210,29 @@ def test_fingerprint_mismatch_falls_back(monkeypatch, tmp_path, fake_megacache):
|
|||
manifest["env"]["torch"] = "0.0.0-other"
|
||||
ctx.manifest_path.write_text(json.dumps(manifest))
|
||||
|
||||
ctx2 = cc.begin(transformer=_transformer(), **_BEGIN_KW)
|
||||
assert ctx2.hit is False # mismatch -> local compile, non-fatal
|
||||
ctx2 = cc.begin(transformer = _transformer(), **_BEGIN_KW)
|
||||
assert ctx2.hit is False # mismatch -> local compile, non-fatal
|
||||
|
||||
|
||||
def test_corrupt_bundle_rejected(monkeypatch, tmp_path, fake_megacache):
|
||||
monkeypatch.setenv(cc._ENV_MODE, "on")
|
||||
monkeypatch.setenv(cc._ENV_DIR, str(tmp_path))
|
||||
ctx = cc.begin(transformer=_transformer(), **_BEGIN_KW)
|
||||
ctx = cc.begin(transformer = _transformer(), **_BEGIN_KW)
|
||||
cc.save(ctx)
|
||||
ctx.bundle.write_bytes(b"CORRUPTED") # manifest sha256 no longer matches
|
||||
ctx.bundle.write_bytes(b"CORRUPTED") # manifest sha256 no longer matches
|
||||
|
||||
ctx2 = cc.begin(transformer=_transformer(), **_BEGIN_KW)
|
||||
ctx2 = cc.begin(transformer = _transformer(), **_BEGIN_KW)
|
||||
assert ctx2.hit is False
|
||||
|
||||
|
||||
# ------------------------------------------------------------------------------- restore
|
||||
def test_restore_inductor_dir(monkeypatch, tmp_path, fake_megacache):
|
||||
import os
|
||||
|
||||
monkeypatch.setenv(cc._ENV_MODE, "auto")
|
||||
monkeypatch.setenv(cc._ENV_DIR, str(tmp_path))
|
||||
monkeypatch.setenv("TORCHINDUCTOR_CACHE_DIR", "/tmp/prior-inductor")
|
||||
ctx = cc.begin(transformer=_transformer(), **_BEGIN_KW)
|
||||
assert os.environ["TORCHINDUCTOR_CACHE_DIR"] != "/tmp/prior-inductor" # redirected
|
||||
ctx = cc.begin(transformer = _transformer(), **_BEGIN_KW)
|
||||
assert os.environ["TORCHINDUCTOR_CACHE_DIR"] != "/tmp/prior-inductor" # redirected
|
||||
cc.restore(ctx)
|
||||
assert os.environ["TORCHINDUCTOR_CACHE_DIR"] == "/tmp/prior-inductor" # restored
|
||||
assert os.environ["TORCHINDUCTOR_CACHE_DIR"] == "/tmp/prior-inductor" # restored
|
||||
|
|
|
|||
|
|
@ -29,7 +29,7 @@ from diffusers.models.normalization import ( # noqa: E402
|
|||
B, S, D, COND = 2, 16, 64, 32
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
@pytest.fixture(autouse = True)
|
||||
def _clean_patches():
|
||||
ep.uninstall_patches()
|
||||
yield
|
||||
|
|
@ -46,30 +46,32 @@ def _devices_dtypes():
|
|||
def _build(cls, device, dtype):
|
||||
torch.manual_seed(0)
|
||||
if cls is RMSNorm:
|
||||
m = RMSNorm(D, eps=1e-6, elementwise_affine=True)
|
||||
m = RMSNorm(D, eps = 1e-6, elementwise_affine = True)
|
||||
elif cls is AdaLayerNormContinuous:
|
||||
m = AdaLayerNormContinuous(D, COND, elementwise_affine=False, eps=1e-6, norm_type="layer_norm")
|
||||
m = AdaLayerNormContinuous(
|
||||
D, COND, elementwise_affine = False, eps = 1e-6, norm_type = "layer_norm"
|
||||
)
|
||||
elif cls is AdaLayerNormZero:
|
||||
m = AdaLayerNormZero(D, num_embeddings=None, norm_type="layer_norm")
|
||||
m = AdaLayerNormZero(D, num_embeddings = None, norm_type = "layer_norm")
|
||||
elif cls is AdaLayerNormZeroSingle:
|
||||
m = AdaLayerNormZeroSingle(D, norm_type="layer_norm")
|
||||
return m.to(device=device, dtype=dtype).eval()
|
||||
m = AdaLayerNormZeroSingle(D, norm_type = "layer_norm")
|
||||
return m.to(device = device, dtype = dtype).eval()
|
||||
|
||||
|
||||
def _inputs(cls, device, dtype):
|
||||
torch.manual_seed(1)
|
||||
x = torch.randn(B, S, D, device=device, dtype=dtype)
|
||||
x = torch.randn(B, S, D, device = device, dtype = dtype)
|
||||
if cls is RMSNorm:
|
||||
return (x,)
|
||||
if cls is AdaLayerNormContinuous:
|
||||
return (x, torch.randn(B, COND, device=device, dtype=dtype))
|
||||
return (x, torch.randn(B, COND, device = device, dtype = dtype))
|
||||
# AdaLayerNormZero / Single take the conditioning emb of width D
|
||||
return (x, torch.randn(B, D, device=device, dtype=dtype))
|
||||
return (x, torch.randn(B, D, device = device, dtype = dtype))
|
||||
|
||||
|
||||
def _call(cls, m, args):
|
||||
if cls is AdaLayerNormZero:
|
||||
return m(args[0], emb=args[1])
|
||||
return m(args[0], emb = args[1])
|
||||
return m(*args)
|
||||
|
||||
|
||||
|
|
@ -77,7 +79,9 @@ def _first(out):
|
|||
return out[0] if isinstance(out, tuple) else out
|
||||
|
||||
|
||||
@pytest.mark.parametrize("cls", [RMSNorm, AdaLayerNormContinuous, AdaLayerNormZero, AdaLayerNormZeroSingle])
|
||||
@pytest.mark.parametrize(
|
||||
"cls", [RMSNorm, AdaLayerNormContinuous, AdaLayerNormZero, AdaLayerNormZeroSingle]
|
||||
)
|
||||
@pytest.mark.parametrize("device,dtype", _devices_dtypes())
|
||||
def test_patched_matches_original(cls, device, dtype):
|
||||
m = _build(cls, device, dtype)
|
||||
|
|
@ -93,62 +97,63 @@ def test_patched_matches_original(cls, device, dtype):
|
|||
# The fused ops are FMA-based (addcmul) / fused (F.rms_norm): within ~1 ULP of the
|
||||
# stock mul+add (and more accurate, single rounding), NOT bit-identical in fp32.
|
||||
atol, rtol = (1e-5, 1e-4) if dtype == torch.float32 else (8e-3, 8e-3)
|
||||
torch.testing.assert_close(got, ref, atol=atol, rtol=rtol)
|
||||
torch.testing.assert_close(got, ref, atol = atol, rtol = rtol)
|
||||
|
||||
|
||||
def test_rmsnorm_mixed_dtype_falls_back():
|
||||
"""fp32 activations into a bf16-weight RMSNorm: diffusers reduces variance in fp32 from
|
||||
the original tensor, so the fused path must FALL BACK (identical output, not divergent)."""
|
||||
m = RMSNorm(D, eps=1e-6, elementwise_affine=True).to(torch.bfloat16).eval()
|
||||
x = torch.randn(B, S, D, dtype=torch.float32)
|
||||
m = RMSNorm(D, eps = 1e-6, elementwise_affine = True).to(torch.bfloat16).eval()
|
||||
x = torch.randn(B, S, D, dtype = torch.float32)
|
||||
with torch.inference_mode():
|
||||
ref = m(x).clone()
|
||||
ep.install_compile_safe_patches()
|
||||
with torch.inference_mode():
|
||||
got = m(x)
|
||||
torch.testing.assert_close(got, ref, atol=0.0, rtol=0.0) # exact fallback
|
||||
torch.testing.assert_close(got, ref, atol = 0.0, rtol = 0.0) # exact fallback
|
||||
|
||||
|
||||
def test_rmsnorm_tuple_dim_falls_back():
|
||||
"""diffusers RMSNorm always reduces the LAST dim even for a tuple `dim`; F.rms_norm
|
||||
would reduce all of them, so a multi-dim `dim` must FALL BACK to the original."""
|
||||
m = RMSNorm((2, D), eps=1e-6, elementwise_affine=True).eval()
|
||||
m = RMSNorm((2, D), eps = 1e-6, elementwise_affine = True).eval()
|
||||
x = torch.randn(B, 2, D)
|
||||
with torch.inference_mode():
|
||||
ref = m(x).clone()
|
||||
ep.install_compile_safe_patches()
|
||||
with torch.inference_mode():
|
||||
got = m(x)
|
||||
torch.testing.assert_close(got, ref, atol=0.0, rtol=0.0) # exact fallback
|
||||
torch.testing.assert_close(got, ref, atol = 0.0, rtol = 0.0) # exact fallback
|
||||
|
||||
|
||||
def test_install_idempotent_and_reversible():
|
||||
rms = RMSNorm(D, eps=1e-6)
|
||||
rms = RMSNorm(D, eps = 1e-6)
|
||||
orig = RMSNorm.forward
|
||||
n1 = ep.install_compile_safe_patches()
|
||||
n2 = ep.install_compile_safe_patches() # second call is a no-op
|
||||
n2 = ep.install_compile_safe_patches() # second call is a no-op
|
||||
assert n1 >= 1 and n2 == n1
|
||||
assert RMSNorm.forward is not orig
|
||||
assert ep.is_installed()
|
||||
ep.uninstall_patches()
|
||||
assert RMSNorm.forward is orig # exact restore
|
||||
assert RMSNorm.forward is orig # exact restore
|
||||
assert not ep.is_installed()
|
||||
ep.uninstall_patches() # idempotent uninstall
|
||||
ep.uninstall_patches() # idempotent uninstall
|
||||
del rms
|
||||
|
||||
|
||||
def test_kill_switch_disables_patches(monkeypatch):
|
||||
monkeypatch.setenv("UNSLOTH_DIFFUSION_EAGER_PATCHES", "0")
|
||||
orig = RMSNorm.forward
|
||||
assert ep.install_compile_safe_patches() == 0 # no-op
|
||||
assert ep.install_compile_safe_patches() == 0 # no-op
|
||||
assert not ep.is_installed()
|
||||
assert RMSNorm.forward is orig # untouched
|
||||
assert RMSNorm.forward is orig # untouched
|
||||
|
||||
|
||||
def test_signature_guard_skips_changed_class(monkeypatch):
|
||||
"""A diffusers class whose forward signature differs must be left untouched."""
|
||||
|
||||
class WeirdRMS(nn.Module):
|
||||
def forward(self, x, extra): # not (self, hidden_states)
|
||||
def forward(self, x, extra): # not (self, hidden_states)
|
||||
return x
|
||||
|
||||
orig = WeirdRMS.forward
|
||||
|
|
@ -157,27 +162,29 @@ def test_signature_guard_skips_changed_class(monkeypatch):
|
|||
monkeypatch.setattr(ep, "_AdaLayerNormZero", None)
|
||||
monkeypatch.setattr(ep, "_AdaLayerNormZeroSingle", None)
|
||||
applied = ep.install_compile_safe_patches()
|
||||
assert applied == 0 # nothing matched -> nothing patched
|
||||
assert WeirdRMS.forward is orig # left untouched
|
||||
assert applied == 0 # nothing matched -> nothing patched
|
||||
assert WeirdRMS.forward is orig # left untouched
|
||||
|
||||
|
||||
@pytest.mark.skipif(not torch.cuda.is_available(), reason="compile graph-break check needs CUDA")
|
||||
@pytest.mark.skipif(not torch.cuda.is_available(), reason = "compile graph-break check needs CUDA")
|
||||
def test_no_graph_break_under_fullgraph():
|
||||
ep.install_compile_safe_patches()
|
||||
|
||||
class Block(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.rms = RMSNorm(D, eps=1e-6)
|
||||
self.ada = AdaLayerNormContinuous(D, COND, elementwise_affine=False, norm_type="layer_norm")
|
||||
self.rms = RMSNorm(D, eps = 1e-6)
|
||||
self.ada = AdaLayerNormContinuous(
|
||||
D, COND, elementwise_affine = False, norm_type = "layer_norm"
|
||||
)
|
||||
|
||||
def forward(self, x, cond):
|
||||
return self.ada(self.rms(x), cond)
|
||||
|
||||
m = Block().to("cuda", torch.bfloat16).eval()
|
||||
x = torch.randn(B, S, D, device="cuda", dtype=torch.bfloat16)
|
||||
cond = torch.randn(B, COND, device="cuda", dtype=torch.bfloat16)
|
||||
compiled = torch.compile(m, fullgraph=True) # raises if a graph break occurs
|
||||
x = torch.randn(B, S, D, device = "cuda", dtype = torch.bfloat16)
|
||||
cond = torch.randn(B, COND, device = "cuda", dtype = torch.bfloat16)
|
||||
compiled = torch.compile(m, fullgraph = True) # raises if a graph break occurs
|
||||
with torch.inference_mode():
|
||||
out = compiled(x, cond)
|
||||
assert out.shape == (B, S, D)
|
||||
|
|
|
|||
|
|
@ -18,7 +18,7 @@ gguf_utils = pytest.importorskip("diffusers.quantizers.gguf.utils")
|
|||
from core.inference import diffusion_gguf_compile as gc # noqa: E402
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
@pytest.fixture(autouse = True)
|
||||
def _clean():
|
||||
# Always start and end from a clean, unpatched state so tests do not leak the
|
||||
# process-wide patch into each other.
|
||||
|
|
@ -59,7 +59,7 @@ def test_compiled_dequant_kill_switch(monkeypatch):
|
|||
|
||||
def test_compiled_dequant_on_by_default(monkeypatch):
|
||||
# The compiled dequant is the real win, so it is ON without any env opt-in.
|
||||
monkeypatch.delenv("UNSLOTH_DIFFUSION_GGUF_COMPILE_DEQUANT", raising=False)
|
||||
monkeypatch.delenv("UNSLOTH_DIFFUSION_GGUF_COMPILE_DEQUANT", raising = False)
|
||||
assert gc.install_compiled_dequant() is True
|
||||
assert gc.is_compiled_dequant_installed() is True
|
||||
|
||||
|
|
|
|||
|
|
@ -233,7 +233,11 @@ def _stub_release(monkeypatch, *, zip_bytes: bytes, digest: str):
|
|||
release = {
|
||||
"tag_name": "master-1-deadbee",
|
||||
"assets": [
|
||||
{"name": name, "browser_download_url": f"https://example.invalid/{name}", "digest": digest}
|
||||
{
|
||||
"name": name,
|
||||
"browser_download_url": f"https://example.invalid/{name}",
|
||||
"digest": digest,
|
||||
}
|
||||
],
|
||||
}
|
||||
monkeypatch.setattr(sdmod, "_fetch_release", lambda *a, **k: release)
|
||||
|
|
@ -245,7 +249,9 @@ def _stub_release(monkeypatch, *, zip_bytes: bytes, digest: str):
|
|||
|
||||
def test_install_downloads_verifies_extracts(tmp_path, monkeypatch):
|
||||
zb = _zip_with_sd_cli()
|
||||
name = _stub_release(monkeypatch, zip_bytes = zb, digest = "sha256:" + hashlib.sha256(zb).hexdigest())
|
||||
name = _stub_release(
|
||||
monkeypatch, zip_bytes = zb, digest = "sha256:" + hashlib.sha256(zb).hexdigest()
|
||||
)
|
||||
sd_cli = install(install_dir = tmp_path)
|
||||
assert sd_cli.name == "sd-cli" and sd_cli.is_file()
|
||||
assert not (tmp_path / name).exists() # archive cleaned up after extract
|
||||
|
|
|
|||
|
|
@ -60,6 +60,7 @@ def _pinned_tag() -> Optional[str]:
|
|||
val = os.environ.get("UNSLOTH_SD_CPP_TAG", DEFAULT_TAG).strip()
|
||||
return val or None
|
||||
|
||||
|
||||
# accelerator -> the token that must appear in a Linux/Windows asset name.
|
||||
_LINUX_ACCEL_TOKEN = {"rocm": "rocm", "vulkan": "vulkan"}
|
||||
_WINDOWS_ACCEL_TOKEN = {
|
||||
|
|
@ -131,8 +132,11 @@ def resolve_release_asset(
|
|||
|
||||
|
||||
def _fetch_release(
|
||||
tag: Optional[str] = None, *, repo: Optional[str] = None,
|
||||
token: Optional[str] = None, timeout: float = 30.0,
|
||||
tag: Optional[str] = None,
|
||||
*,
|
||||
repo: Optional[str] = None,
|
||||
token: Optional[str] = None,
|
||||
timeout: float = 30.0,
|
||||
) -> dict:
|
||||
"""GET a release JSON from GitHub. With ``tag`` set, fetch that exact release (and fall
|
||||
back to latest if the tag is gone upstream); otherwise fetch latest. ``token`` is
|
||||
|
|
@ -154,7 +158,9 @@ def _fetch_release(
|
|||
except urllib.error.HTTPError as exc: # pinned tag removed upstream -> latest
|
||||
if exc.code != 404:
|
||||
raise
|
||||
print(f"sd-cli: pinned tag {tag} not found on {repo}; falling back to latest", flush = True)
|
||||
print(
|
||||
f"sd-cli: pinned tag {tag} not found on {repo}; falling back to latest", flush = True
|
||||
)
|
||||
return _get(f"{base}/latest")
|
||||
|
||||
|
||||
|
|
@ -172,7 +178,9 @@ def _verify_sha256(path: Path, expected_digest: Optional[str]) -> None:
|
|||
return
|
||||
algo, _, want = expected_digest.partition(":")
|
||||
if algo.lower() != "sha256" or not want:
|
||||
print(f"sd-cli: WARNING unrecognised digest {expected_digest!r}; skipping check", flush = True)
|
||||
print(
|
||||
f"sd-cli: WARNING unrecognised digest {expected_digest!r}; skipping check", flush = True
|
||||
)
|
||||
return
|
||||
h = hashlib.sha256()
|
||||
with open(path, "rb") as f:
|
||||
|
|
@ -180,9 +188,7 @@ def _verify_sha256(path: Path, expected_digest: Optional[str]) -> None:
|
|||
h.update(chunk)
|
||||
got = h.hexdigest()
|
||||
if got != want.lower():
|
||||
raise RuntimeError(
|
||||
f"sha256 mismatch for {path.name}: expected {want.lower()}, got {got}"
|
||||
)
|
||||
raise RuntimeError(f"sha256 mismatch for {path.name}: expected {want.lower()}, got {got}")
|
||||
|
||||
|
||||
def default_install_dir() -> Path:
|
||||
|
|
@ -199,7 +205,12 @@ def _make_executable(path: Path) -> None:
|
|||
path.chmod(mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH)
|
||||
|
||||
|
||||
def _download(url: str, dest: Path, *, timeout: float = 300.0) -> None:
|
||||
def _download(
|
||||
url: str,
|
||||
dest: Path,
|
||||
*,
|
||||
timeout: float = 300.0,
|
||||
) -> None:
|
||||
"""Stream a release asset to ``dest`` with a timeout. ``urlretrieve`` has no timeout,
|
||||
so a stalled connection would hang the lazy first-load (ensure_sd_cpp_binary) forever.
|
||||
Anonymous, matching the public release URL -- the API fetch carries any token."""
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue