from __future__ import annotations import os, shutil import json import re import numpy as np import random import torch import safetensors import filelock import hashlib from tqdm.auto import tqdm from typing import List, Tuple, NamedTuple from pathlib import Path import torch.nn.functional as F from collections import defaultdict from scipy.optimize import linear_sum_assignment from functools import lru_cache try: FP8_E4M3 = getattr(torch, "float8_e4m3fn", None) FP8_E5M2 = getattr(torch, "float8_e5m2", None) FP8_DTYPES = tuple(d for d in (FP8_E4M3, FP8_E5M2) if d is not None) except Exception: FP8_E4M3 = FP8_E5M2 = None FP8_DTYPES = () FP_SET = {torch.float32, torch.float16, torch.float64, torch.bfloat16} NUM_INPUT_BLOCKS = 12 NUM_MID_BLOCK = 1 NUM_OUTPUT_BLOCKS = 12 NUM_TOTAL_BLOCKS = NUM_INPUT_BLOCKS + NUM_MID_BLOCK + NUM_OUTPUT_BLOCKS BLOCKID = ["BASE"] + [f"IN{i:02}" for i in range(12)] + ["M00"] + [f"OUT{i:02}" for i in range(12)] BLOCKIDXLL = ["BASE"] + [f"IN{i:02}" for i in range(9)] + ["M00"] + [f"OUT{i:02}" for i in range(9)] + ["VAE"] BLOCKIDFLUX = ["CLIP", "T5", "IN"] + ["D{:002}".format(x) for x in range(19)] + ["S{:002}".format(x) for x in range(38)] + ["OUT"] # Len: 61 BLOCKIDZI = ["BASE","CONT","NOISE"] + [f"L{i:02}" for i in range(30)] + ["VAE"] BLOCKIDAM = ["BASE"] + [f"L{i:02}" for i in range(28)] + ["VAE"] # Anima Model has 28 blocks _re_inp = re.compile(r'\.input_blocks\.(\d+)\.') _re_mid = re.compile(r'\.middle_block\.(\d+)\.') _re_out = re.compile(r'\.output_blocks\.(\d+)\.') _re_flux_any_num = re.compile(r'\.(\d+)\.') _re_xl_tf = re.compile(r'transformer_blocks\.(\d+)\.') COLS = [[-1, 1/3, 2/3], [1, 1, 0], [0, -1, -1], [1, 0, 1]] COLSXL = [[0, 0, 1], [1, 0, 0], [-1, -1, 0], [-1, 1, 0]] BNB = ".quant_state.bitsandbytes__" QTYPES = ["fp4", "nf4"] FINETUNES = [ "model.diffusion_model.input_blocks.0.0.weight", "model.diffusion_model.input_blocks.0.0.bias", "model.diffusion_model.out.0.weight", "model.diffusion_model.out.0.bias", "model.diffusion_model.out.2.weight", "model.diffusion_model.out.2.bias", ] # --------------------------------------------------------------------- # LBLOCKS (per-arch) # - each entry corresponds 1:1 with that arch's BLOCKID* order # - each entry is a list of substrings (aliases) to match against "full" or "msd" # --------------------------------------------------------------------- def make_lblocks_sdxl(): """ Order aligns with BLOCKIDXLL: BASE, IN00..IN08, M00, OUT00..OUT08, VAE """ lblocks = [] # BASE (text/conditioner) lblocks.append([ "conditioner.embedders.", # SDXL conditioner "text_encoders.", # diffusers packs "clip_l.", "clip_g.", # alt roots "0_transformer_text_model_", # your convert() output for te1 "1_model_transformer_resblocks_", # your convert() output for te2 ]) # IN00..IN08 for i in range(9): lblocks.append([ f"diffusion_model_input_blocks_{i}_", # your convert() output f"model.diffusion_model.input_blocks.{i}.", # raw checkpoint ]) # M00 lblocks.append([ "diffusion_model_middle_block_", "model.diffusion_model.middle_block.", ]) # OUT00..OUT08 for i in range(9): lblocks.append([ f"diffusion_model_output_blocks_{i}_", f"model.diffusion_model.output_blocks.{i}.", ]) # SDXL has extra ".out." heads; fold them into OUT08 bucket lblocks[-1].extend([ "diffusion_model_out_", "model.diffusion_model.out.", ]) # VAE lblocks.append([ "first_stage_model.", "vae.", ]) return lblocks def make_lblocks_flux(): """ Order aligns with BLOCKIDFLUX: CLIP, T5, IN, D00..D18, S00..S37, OUT """ lblocks = [] # CLIP lblocks.append([ "text_encoders.clip", "clip.", "clip_l.", "clip_g.", "text_encoder.", "conditioner.embedders.", ]) # T5 lblocks.append([ "t5xxl", "t5.", "text_encoders.t5", "text_encoder_2.", ]) # IN (input projections / embeddings) lblocks.append([ "img_in", "txt_in", "time_in", "vector_in", "x_embedder", "t_embedder", ]) # D00..D18 (double blocks) for i in range(19): lblocks.append([f"double_blocks.{i}.", f"double_block.{i}."]) # S00..S37 (single blocks) for i in range(38): lblocks.append([f"single_blocks.{i}.", f"single_block.{i}."]) # OUT (final layer / heads) lblocks.append([ "final_layer", "out.", "vector_out", ]) return lblocks def make_lblocks_zi(): """ Order aligns with BLOCKIDZI: BASE, CONT, NOISE, L00..L29, VAE """ lblocks = [] # BASE (text enc / caption embedder) lblocks.append([ "text_encoders.qwen3_4b.", "qwen3_4b.", "model.diffusion_model.cap_embedder.", "cap_embedder.", ]) # CONT / NOISE lblocks.append(["context_refiner", "diffusion_model.context_refiner"]) lblocks.append(["noise_refiner", "diffusion_model.noise_refiner"]) # L00..L29 for i in range(30): lblocks.append([ f"diffusion_model.layers.{i}.", # LoRA example (no 'model.' prefix) f"model.diffusion_model.layers.{i}.", # model keys f"diffusion_model_layers_{i}_", # if you ever convert to underscore form ]) # VAE lblocks.append([ "vae.", "first_stage_model.", ]) return lblocks def make_lblocks_am(): """ Order aligns with BLOCKIDAM: BASE, L00..L27, VAE Supports checkpoint keys: - official: net.* (then normalized) - unofficial: no net. - AIO: cond_stage_model.qwen3_06b.* + model.diffusion_model.* + first_stage_model.* Supports LoRA keys (kohya-style): - TE: lora_te_layers_{n}_... - UNet: lora_unet_blocks_{n}_... - VAE: lora_vae_* or lora_first_stage_model_* """ lblocks = [] # ------------------------- # BASE (text + adapters/embeds) # ------------------------- lblocks.append([ # ---- AIO text encoder (checkpoint) ---- "cond_stage_model.qwen3_06b.", # ---- canonicalized text encoder (if you normalize AIO -> text_encoders.*) ---- "text_encoders.qwen3_06b.", "text_encoders.qwen3_06b_base.", "qwen3_06b.", "qwen3_06b_base.", # ---- diffusion-side adapters / time/pos ---- "model.diffusion_model.llm_adapter.", "diffusion_model.llm_adapter.", "llm_adapter.", "model_diffusion_model_llm_adapter_", "diffusion_model_llm_adapter_", "net_llm_adapter_", "model.diffusion_model.t_embedder.", "diffusion_model.t_embedder.", "t_embedder.", "model_diffusion_model_t_embedder_", "diffusion_model_t_embedder_", "net_t_embedder_", "model.diffusion_model.t_embedding_norm.", "diffusion_model.t_embedding_norm.", "t_embedding_norm.", "model.diffusion_model.pos_embedder.", "diffusion_model.pos_embedder.", "pos_embedder.", # ---- official may come as net.* ---- "net.llm_adapter.", "net.t_embedder.", "net.t_embedding_norm.", "net.pos_embedder.", # ---- LoRA: TE side (all goes to BASE, like ZI) ---- "lora_te_", # generic "lora_te_layers_", # matches: lora_te_layers_7_... "lora_text_encoder_", # some trainers use this # ---- LoRA: sometimes text is named as cond_stage_model in lora packs ---- "lora_cond_stage_model_", # just in case ]) # ------------------------- # L00..L27 (diffusion blocks) # ------------------------- for i in range(28): lblocks.append([ # checkpoint keys f"model.diffusion_model.blocks.{i}.", f"blocks.{i}.", f"net.blocks.{i}.", # LoRA keys (kohya-style + direct module-path style) f"lora_unet_blocks_{i}_", # matches: lora_unet_blocks_0_self_attn_q_proj... f"lora_unet_blocks.{i}.", # safety (rare) f"diffusion_model.blocks.{i}.", f"model.diffusion_model.blocks.{i}.", f"net.blocks.{i}.", f"diffusion_model_blocks_{i}_", # convert_diffusers_name_to_compvis() direct-key form f"model_diffusion_model_blocks_{i}_", f"net_blocks_{i}_", ]) # Add x_embedder into L00 bucket (input-like) lblocks[1].extend([ "model.diffusion_model.x_embedder.", "diffusion_model.x_embedder.", "x_embedder.", "net.x_embedder.", "model_diffusion_model_x_embedder_", "diffusion_model_x_embedder_", "net_x_embedder_", "lora_unet_x_embedder", "lora_unet_x_embedder_", # just in case ]) # Add final_layer into L27 bucket (output-like) lblocks[28].extend([ "model.diffusion_model.final_layer.", "diffusion_model.final_layer.", "final_layer.", "net.final_layer.", "model_diffusion_model_final_layer_", "diffusion_model_final_layer_", "net_final_layer_", "lora_unet_final_layer", "lora_unet_final_layer_", # just in case "lora_unet_out", "lora_unet_out_", # some packs use out naming ]) # ------------------------- # VAE # ------------------------- lblocks.append([ "vae.", "first_stage_model.", # LoRA (if present) "lora_vae_", "lora_first_stage_model_", "lora_fs_", "lora_vae", ]) return lblocks # Convenient ready-to-use constants LBLOCKS_SDXL = make_lblocks_sdxl() LBLOCKS_FLUX = make_lblocks_flux() LBLOCKS_ZI = make_lblocks_zi() LBLOCKS_AM = make_lblocks_am() LBLOCKS26 = [ "encoder", "diffusion_model_input_blocks_0_","diffusion_model_input_blocks_1_","diffusion_model_input_blocks_2_", "diffusion_model_input_blocks_3_","diffusion_model_input_blocks_4_","diffusion_model_input_blocks_5_", "diffusion_model_input_blocks_6_","diffusion_model_input_blocks_7_","diffusion_model_input_blocks_8_", "diffusion_model_input_blocks_9_","diffusion_model_input_blocks_10_","diffusion_model_input_blocks_11_", "diffusion_model_middle_block_", "diffusion_model_output_blocks_0_","diffusion_model_output_blocks_1_","diffusion_model_output_blocks_2_", "diffusion_model_output_blocks_3_","diffusion_model_output_blocks_4_","diffusion_model_output_blocks_5_", "diffusion_model_output_blocks_6_","diffusion_model_output_blocks_7_","diffusion_model_output_blocks_8_", "diffusion_model_output_blocks_9_","diffusion_model_output_blocks_10_","diffusion_model_output_blocks_11_", "embedders", ] checkpoint_dict_replacements = { 'cond_stage_model.transformer.embeddings.': 'cond_stage_model.transformer.text_model.embeddings.', 'cond_stage_model.transformer.encoder.': 'cond_stage_model.transformer.text_model.encoder.', 'cond_stage_model.transformer.final_layer_norm.': 'cond_stage_model.transformer.text_model.final_layer_norm.', } checkpoint_dict_skip_on_merge = set(["cond_stage_model.transformer.text_model.embeddings.position_ids"]) def normalize_path(path: str) -> str: path = os.path.abspath(path) if os.name == "nt": if not path.startswith("\\\\?\\"): path = "\\\\?\\" + path return str(Path(path).expanduser().resolve()) def tagdict(presets: str) -> dict: """Parse presets text into a dict if value part has exactly 26 items.""" wdict = {} for line in presets.splitlines(): parts = re.split(r'[:\t]', line, maxsplit=1) if len(parts) == 2: key, w = parts if len(w.split(",")) == 26: wdict[key.strip()] = w.strip() return wdict def base_path(path: str=None) -> str: if path is None: return os.path.dirname(os.path.realpath(__file__)) return os.path.join(os.path.dirname(os.path.realpath(__file__)), path) file_path = base_path("mbwpresets.txt") if not os.path.isfile(file_path): shutil.copyfile(base_path("mbwpresets_master.txt"), file_path) weights_presets_list = tagdict(open(file_path).read()) _SPLIT = re.compile(r"[,\n]+") def _split(s:str): return [t.strip() for t in _SPLIT.split(s) if t.strip()] def _get_cast(xs, i, cast, default): try: return cast(xs[i]) except (IndexError, ValueError): return default def wgt(x, dp): useblocks = False if isinstance(x, (int, float)):return float(x), dp, useblocks useblocks = True nums, rest = deepblock(x if isinstance(x, list) else [x]) return (nums[0] if len(nums) == 1 else nums), rest, useblocks def deepblock(items:List[str])->Tuple[List[float],List[str]]: nums:List[float]=[];rest:List[str]=[];stack=list(items) while stack: s = stack.pop() src = weights_presets_list.get(s, s) for t in _split(src): if t in weights_presets_list: stack.append(t); continue try: nums.append(float(t)) except ValueError: rest.append(t) return nums, rest def rinfo(s:str,seed:int)->str: core, _, rest=s.replace(" ", "").partition("[") fe = rest[:-1] if rest.endswith("]") else None toks = _split(core) rmin = _get_cast(toks, 0, float, 0.0) rmax = _get_cast(toks, 1, float, 1.0) get = _get_cast(toks, 2, int, seed) return f"({rmin},{rmax},{get},[{fe}])" def roundeep(term): if not term: return None out=[] for d in term: try: a, b, c = d.split(":", 2) out.append(f"{a}:{b}:{round(float(c), 3)}") except ValueError: out.append(d) return out def rand_ratio(s:str): core, _, rest = s.partition("[") deep = _split(rest[:-1]) if rest.endswith("]") else [] toks = _split(core.replace(" ","")) rmin = _get_cast(toks, 0, float, 0.0) rmax = _get_cast(toks, 1, float, 1.0) seed = _get_cast(toks, 2, int, random.randint(1, 4294967295)) np.random.seed(seed) ratios = np.random.uniform(rmin, rmax, 26).tolist() deep_res = [] for d in deep: if "PRESET" in d: try: _, pack = d.split(":",1) name, drat_s = pack.split("(") base_vals = [float(x) for x in _split(weights_presets_list[name])] drat = float(drat_s.rstrip(")")) ratios = [r * (1 - drat) + b * drat for r, b in zip(ratios, base_vals)] except Exception: pass continue if d.count(":") != 2: continue dbs_s, dws, dr_s = d.split(":", 2) dbs = dbs_s.split() if "(" in dr_s: v, drat_s = dr_s.split("(") v = float(v) drat = float(drat_s.rstrip(")")) if dws == "ALL": for db in dbs: i = BLOCKID.index(db) ratios[i] = ratios[i] * (1 - drat) + v * drat else: for db in dbs: cur = ratios[BLOCKID.index(db)] deep_res.append(f"{db}:{dws}:{cur * (1 - drat) + v * drat}") else: v = float(dr_s) if dws == "ALL": for db in dbs: ratios[BLOCKID.index(db)] = v else: for db in dbs: deep_res.append(f"{db}:{dws}:{v}") info = rinfo(core, seed) ratios, deep_res = wgt(ratios, deep_res) return ratios, seed, deep_res, info def colorcalc(cols, isxl): M = COLSXL if isxl else COLS return [0.02 * sum(v * cols[i] for i, v in enumerate(col)) for col in zip(*M)] def fineman(fine, arch): if arch.get("FLUX", False): mul = { "double_block": 1.0 + (fine[0] * 0.01) if len(fine) > 0 else 1.0, "img_in": 1.0 + (fine[1] * 0.01) if len(fine) > 1 else 1.0, "txt_in": 1.0 + (fine[2] * 0.01) if len(fine) > 2 else 1.0, "time": 1.0 + (fine[3] * 0.01) if len(fine) > 3 else 1.0, "out": 1.0 + (fine[4] * 0.01) if len(fine) > 4 else 1.0, } add = (fine[5] * 0.02) if len(fine) > 5 else 0.0 return {"mul": mul, "add": add} if isinstance(fine, (list, tuple, np.ndarray)) and len(fine) >= 8: return { "mode": "image_tone_v2", "noise1": float(fine[0]), "noise2": float(fine[1]), "noise3": float(fine[2]), "contrast": float(fine[3]), "brightness": float(fine[4]), "r": float(fine[5]), "g": float(fine[6]), "b": float(fine[7]), # Optional 9th value: direct saturation factor. # 1.0 = off, 0.75 = 25% desaturate, 1.15 = boost. "saturation": float(fine[8]) if len(fine) > 8 else 1.0, } r = [ 1 - fine[0] * 0.01, 1 + fine[0] * 0.02, 1 - fine[1] * 0.01, 1 + fine[1] * 0.02, 1 - fine[2] * 0.01, [fine[3] * 0.02] + colorcalc(fine[4:8], arch.get("XL", False)) ] return r def weighttoxl(weight): return weight[:9] + weight[12:22] +[0] def parse_ratio(ratios, info, dp): if isinstance(ratios, list): ratio, *weights = ratios rounded = [round(a, 3) for a in weights] round_deep = roundeep(dp) prefix = f"preset:[{info}]," if info else "" info = f"{prefix}{round(ratio,3)},[{rounded},[{round_deep}]]" else: ratio, weights, info = ratios, [ratios]*34, f"{round(ratios,3)}" return weights, ratio, info DTYPES = {torch.float32, torch.float64, torch.bfloat16} def to_half(tensor, enable): return tensor.half() if enable and getattr(tensor, "dtype", None) in DTYPES else tensor def upcast_fp8_state_dict(theta: dict, target_dtype: torch.dtype = torch.float16): if not FP8_DTYPES: return theta fp8_keys = [] for k, v in theta.items(): if isinstance(v, torch.Tensor) and v.dtype in FP8_DTYPES: fp8_keys.append(k) if not fp8_keys: return theta for k in tqdm(fp8_keys, total=len(fp8_keys), desc=f"Upcasting {len(fp8_keys)} fp8 tensors..."): theta[k] = theta[k].to(target_dtype) return theta @torch.inference_mode() def prepare_state_dict_for_save( theta: dict, args, *, arch: dict, vae_prefix: str | None, prune: bool = True, make_cpu: bool = True, make_contiguous: bool = True, ): # roots if arch.get("FLUX", False): cond_prefixes = ("clip.cond_stage_model.",) elif arch.get("XL", False): cond_prefixes = ("conditioner.",) elif arch.get("ZI", False): cond_prefixes = ("text_encoders.qwen3_4b.",) elif arch.get("AM", False): cond_prefixes = ( "text_encoders.qwen3_06b.", "text_encoders.qwen3_06b_base.", "cond_stage_model.qwen3_06b.", "qwen3_06b.", "qwen3_06b_base.", ) else: cond_prefixes = ("cond_stage_model.",) roots = ( "model.diffusion_model.", "depth_model.", "first_stage_model.", "vae.", *cond_prefixes, ) want_fp8 = bool(getattr(args, "save_quarter", False)) and bool(FP8_DTYPES) want_fp16 = bool(getattr(args, "save_half", False)) and not want_fp8 want_bf16 = bool(getattr(args, "save_bhalf", False)) and not want_fp8 fp8_dtype = FP8_E4M3 if FP8_E4M3 is not None else (FP8_E5M2 if FP8_E5M2 is not None else None) if want_fp8 and fp8_dtype is None: want_fp8 = False cpu = torch.device("cpu") key_list = list(theta.keys()) for k in tqdm(key_list, desc="Preparing merged model for save...", unit="param"): v = theta.get(k, None) # prune if prune and ((".scale" in k and arch.get("ZI", False)) or ("pos_embedder" in k and arch.get("AM", False)) or (not any(k.startswith(r) for r in roots))): theta.pop(k, None) continue if not isinstance(v, torch.Tensor): continue # keep_ema if getattr(args, "keep_ema", False): k_ema = "model_ema." + k[6:].replace(".", "") v_ema = theta.get(k_ema, None) if isinstance(v_ema, torch.Tensor): v = v_ema # dtype is_vae = bool(vae_prefix) and k.startswith(vae_prefix) if (not is_vae) and v.is_floating_point(): dt = v.dtype if want_fp8: if dt in FP_SET: v = v.to(fp8_dtype) elif want_fp16: if dt in {torch.float32, torch.float64, torch.bfloat16, *FP8_DTYPES}: v = v.to(torch.float16) elif want_bf16: if dt in {torch.float32, torch.float64, torch.float16, *FP8_DTYPES}: v = v.to(torch.bfloat16) else: if dt in {torch.float16, torch.float64, torch.bfloat16, *FP8_DTYPES}: v = v.to(torch.float32) v = v.detach() if make_cpu and v.device.type != "cpu": v = v.to(cpu, non_blocking=False) if make_contiguous and not v.is_contiguous(): v = v.contiguous() theta[k] = v return theta def _safe_load_json(path: str): try: with open(path, "r", encoding="utf-8") as f: return json.load(f) except FileNotFoundError: return {} except json.JSONDecodeError: return {} def merge_cache_json(cache_path, model_path): model_cache_path = os.path.join(model_path, "cache.json") base = _safe_load_json(cache_path if os.path.exists(cache_path) else model_cache_path) update = _safe_load_json(model_cache_path) if os.path.exists(model_cache_path) else {} if isinstance(base, dict) and isinstance(update, dict): base.update(update) elif isinstance(base, list) and isinstance(update, list): base.extend(update) with filelock.FileLock(f"{cache_path}.lock"): with open(cache_path, "w", encoding="utf-8") as f: json.dump(base, f, ensure_ascii=False, indent=2) def dump_cache(cache_data, cache_path=None): with filelock.FileLock(f"{cache_path}.lock"): with open(cache_path, "w", encoding="utf8") as f: json.dump(cache_data or {}, f, indent=4) def cache(subsection, cache_data, cache_path=None): if cache_data is None: with filelock.FileLock(f"{cache_path}.lock"): cache_data = _safe_load_json(cache_path) if subsection not in cache_data or not isinstance(cache_data.get(subsection), dict): cache_data[subsection] = {} dump_cache(cache_data, cache_path) return cache_data def model_hash(filename: str) -> str: try: with open(filename, "rb") as f: f.seek(0x100000) return hashlib.sha256(f.read(0x10000)).hexdigest()[:8] except FileNotFoundError: return "NOFILE" def sha256_from_cache(filename: str, title: str, cache_data, cache_path=None): cache_data = cache("hashes", cache_data, cache_path) hsect = cache_data.get("hashes", {}) h = hsect.get(title) if not h: return None, None, cache_data cur_mtime = os.path.getmtime(filename) if os.path.exists(filename) else None if h.get("mtime") != cur_mtime: return None, None, cache_data return h.get("sha256"), h.get("model_hash"), cache_data def calculate_sha256(filename: str, chunk_size: int = 8 * 1024 * 1024) -> str: hasher = hashlib.sha256() with open(filename, "rb") as f: for chunk in tqdm(iter(lambda: f.read(chunk_size), b""), desc="Hashing", unit="MB", unit_scale=chunk_size / (1024 * 1024)): hasher.update(chunk) return hasher.hexdigest() def sha256(filename: str, title: str, cache_path: str | None) -> str: cache_data = cache("hashes", None, cache_path) s256_cached, _, cache_data = sha256_from_cache(filename, title, cache_data) if s256_cached: return s256_cached print(f"Calculating sha256 for {filename}: ", end="") sha_val = calculate_sha256(filename) mhash = model_hash(filename) print(sha_val) cache_data["hashes"][title] = { "mtime": os.path.getmtime(filename) if os.path.exists(filename) else None, "sha256": sha_val, "model_hash": mhash, } dump_cache(cache_data, cache_path) return sha_val def read_metadata_from_safetensors(filename): with open(filename, "rb") as f: size = int.from_bytes(f.read(8), "little") start = f.read(2) assert size > 2 and start in (b'{"', b"{'"), f"{filename} is not a safetensors file" meta = json.loads(start + f.read(size - 2)) res = {} for k, v in meta.get("__metadata__", {}).items(): if isinstance(v, str) and v.startswith("{"): try: v = json.loads(v) except: pass res[k] = v return res def qdtyper(sd): if any("fp4" in k for k in sd): return "fp4" if any("nf4" in k for k in sd): return "nf4" for v in sd.values(): dt = getattr(v, "dtype", None) if dt is not None: return dt def to_qdtype(sd1, sd2, qd1, qd2, device): t1 = t2 = torch.float16 if (qd1 in QTYPES and qd2 in QTYPES) else None if qd1 in QTYPES: sd1, _ = q_dequantize(sd1, qd1, device, t1) if qd2 in QTYPES: sd2, _ = q_dequantize(sd2, qd2, device, t2) return sd1, sd2 def maybe_to_qdtype(a, b, qa, qb, device, isflux: bool = True): return to_qdtype(a, b, qa, qb, device) if isflux and qa != qb else (a, b) def detect_arch(theta): keys = list(theta.keys()) # ------------------------- # Anima (AM) detection # ------------------------- def _is_am_key(k: str) -> bool: # diffusion signature (DiT blocks + adaln + proj names) if (".blocks." in k or k.startswith(("blocks.", "net.blocks.", "model.diffusion_model.blocks."))): if ("adaln_modulation_" in k) or (".q_proj.weight" in k) or (".k_proj.weight" in k) or (".v_proj.weight" in k): return True if (".self_attn." in k) or (".cross_attn." in k): return True # diffusion head / adapters if ( "final_layer.linear.weight" in k or "final_layer.adaln_modulation" in k or "llm_adapter." in k or "t_embedding_norm" in k or (k.startswith(("diffusion_model.final_layer.", "model.diffusion_model.final_layer.", "net.final_layer.", "final_layer.")) and ".lora_" in k) or (k.startswith(("diffusion_model.llm_adapter.", "model.diffusion_model.llm_adapter.", "net.llm_adapter.", "llm_adapter.")) and ".lora_" in k) ): return True # AIO text encoder signature if k.startswith("cond_stage_model.qwen3_06b."): return True return False isam = any(_is_am_key(k) for k in keys) # ------------------------- # Z-Image (ZI) detection (avoid x_embedder/t_embedder false positives) # ------------------------- iszi = (not isam) and ( ("model.diffusion_model.cap_embedder.0.weight" in theta) or ("cap_embedder.0.weight" in theta) or any(k.startswith(("model.diffusion_model.layers.", "diffusion_model.layers.", "layers.")) for k in keys) or any("context_refiner" in k for k in keys) or any("noise_refiner" in k for k in keys) ) isxl = ("conditioner.embedders.1.model.transformer.resblocks.9.mlp.c_proj.weight" in theta) # Flux detection (keep both keys for compatibility) isflux = any(("double_blocks" in k) or ("double_block" in k) or ("single_blocks" in k) or ("single_block" in k) for k in keys) arch = { "XL": bool(isxl), "FLUX": bool(isflux), "ZI": bool(iszi), "AM": bool(isam), } # ------------------------- # normalize keys for ZI # ------------------------- if iszi: def _zi_key(k: str) -> str: if k.startswith(("vae.", "text_encoders.", "model.diffusion_model.")): return k if k.startswith("diffusion_model."): return "model." + k if k.startswith(("layers.", "context_refiner.", "noise_refiner.", "final_layer.", "cap_embedder.")): return "model.diffusion_model." + k return k theta = {_zi_key(k): v for k, v in theta.items()} # ------------------------- # normalize keys for AM (format conversion) # official: net.* -> model.diffusion_model.* # unofficial: blocks.* -> model.diffusion_model.blocks.* # AIO text: cond_stage_model.qwen3_06b.* -> text_encoders.qwen3_06b.* # ------------------------- if isam: def _am_key(k: str) -> str: # VAE keep if k.startswith(("vae.", "first_stage_model.")): return k # AIO text encoder -> canonical text_encoders.* if k.startswith("text_encoders.qwen3_06b."): return "cond_stage_model.qwen3_06b." + k[len("text_encoders.qwen3_06b."):] # already canonical diffusion if k.startswith("model.diffusion_model."): return k # strip net. if present k2 = k[4:] if k.startswith("net.") else k # diffusion modules: map to model.diffusion_model.* if k2.startswith(( "blocks.", "x_embedder.", "t_embedder.", "t_embedding_norm.", "pos_embedder.", "final_layer.", "llm_adapter.", )): return "model.diffusion_model." + k2 # some packs may use diffusion_model.* without model. if k2.startswith("diffusion_model."): return "model." + k2 # split-file style text encoders already ok if k2.startswith(("cond_stage_model.", "qwen3_")): return k2 return k theta = {_am_key(k): v for k, v in theta.items()} return arch, theta def _common_dtype(a: torch.Tensor, b: torch.Tensor, c: torch.Tensor): dt = torch.promote_types(torch.promote_types(a.dtype, b.dtype), c.dtype) if dt in FP8_DTYPES: dt = torch.float16 return dt def q_dequantize(sd, qtype, device, dtype, setbnb=True): from bitsandbytes.functional import dequantize_4bit dels = [] calc = "cuda:0" if torch.cuda.is_available() else ("mps:0" if torch.backends.mps.is_available() else "cpu") for k, v in list(sd.items()): qk = k + BNB + qtype if ("weight" in k) and ("weight." not in k) and (qk in sd): qs = q_tensor_to_dict(sd[qk]) out = torch.empty(qs["shape"], device = calc) deq = dequantize_4bit(v.to(calc), out = out, absmax = sd[k + ".absmax"].to(calc), blocksize = qs["blocksize"], quant_type = qs["quant_type"]) sd[k] = deq.to(device, dtype) if dtype else deq.to(device) dels += [k + ".absmax", k + ".quant_map"] + ([qk] if setbnb else []) elif isinstance(v, torch.Tensor) and dtype: sd[k] = v.to(dtype) for k in dels: sd.pop(k, None) return sd, dtype def q_tensor_to_dict(t): return json.loads(bytes(t.tolist()).decode("utf-8")) _blocker_cache = {} def blocker(blocks: str, blockids: list[str]) -> str: ck = (id(blockids), blocks) hit = _blocker_cache.get(ck) if hit is not None: return hit out = [] idx = {b:i for i,b in enumerate(blockids)} for w in blocks.split(): if "-" in w: a, b = (t.strip() for t in w.split("-", 1)) i, j = idx[a], idx[b] lo, hi = (i, j) if i <= j else (j, i) out.extend(blockids[lo:hi + 1]) else: out.append(w) res = " ".join(out) _blocker_cache[ck] = res return res def _parse_int_after(s: str, marker: str): i = s.find(marker) if i < 0: return None j = i + len(marker) k = j while k < len(s) and s[k].isdigit(): k += 1 if k == j: return None try: return int(s[j:k]) except ValueError: return None def _digits_concat(s: str) -> str: return "".join(ch for ch in s if ch.isdigit()) @lru_cache(maxsize=250_000) def _blockfromkey_cached_flags(key: str, xl: bool, flux: bool, zi: bool, am: bool) -> Tuple[str, str]: arch = {"XL": bool(xl), "FLUX": bool(flux), "ZI": bool(zi), "AM": bool(am)} # ------------------------- # SD1.5 / SD2.x (non-XL/non-Flux/non-ZI/non-AM) # ------------------------- if not (arch["XL"] or arch["FLUX"] or arch["ZI"] or arch["AM"]): if "time_embed" in key: idx = -2 elif ".out." in key: idx = NUM_TOTAL_BLOCKS - 1 else: m = _re_inp.search(key) if m: idx = int(m.group(1)) else: if _re_mid.search(key): idx = NUM_INPUT_BLOCKS else: m2 = _re_out.search(key) if m2: idx = NUM_INPUT_BLOCKS + NUM_MID_BLOCK + int(m2.group(1)) else: return "Not Merge", "Not Merge" b = BLOCKID[idx + 1] return b, b # ------------------------- # Flux # ------------------------- if arch.get("FLUX", False): if "vae" in key: return "VAE", "Not Merge" if "t5xxl" in key: return "T5", "T5" if "text_encoders.clip" in key: return "CLIP", "CLIP" di = _parse_int_after(key, "double_blocks.") if di is None: di = _parse_int_after(key, "double_block.") if di is not None: tag = f"D{di:02d}" return tag, tag si = _parse_int_after(key, "single_blocks.") if si is None: si = _parse_int_after(key, "single_block.") if si is not None: tag = f"S{si:02d}" return tag, tag if "_in" in key: return "IN", "IN" if "final_layer" in key: return "OUT", "OUT" m = _re_flux_any_num.search(key) if m and "double_blocks" in key: tag = f"D{int(m.group(1)):02d}" return tag, tag if m and "single_blocks" in key: tag = f"S{int(m.group(1)):02d}" return tag, tag return "Not Merge", "Not Merge" # ------------------------- # SDXL # ------------------------- if arch.get("XL", False): if not ("weight" in key or "bias" in key): return "Not Merge", "Not Merge" if "label_emb" in key or "time_embed" in key: return "Not Merge", "Not Merge" if "conditioner.embedders" in key: return "BASE", "BASE" if "first_stage_model" in key: return "VAE", "BASE" if "model.diffusion_model.out." in key: return "OUT8", "OUT08" if "model.diffusion_model" in key: if ".input_blocks." in key: blk = "IN" elif ".middle_block." in key: blk = "MID" elif ".output_blocks." in key: blk = "OUT" else: return "Not Merge", "Not Merge" nums = _digits_concat(key) if not nums: return "Not Merge", "Not Merge" tag = (nums[:1] + "0") if blk == "MID" else nums[:2] add = "" mi = _parse_int_after(key, "transformer_blocks.") if mi is None: m = _re_xl_tf.search(key) if m: mi = int(m.group(1)) if mi is not None: add = str(mi) left = blk + tag + add right = ("M00" if blk == "MID" else f"{blk}0{tag[0]}") return left, right return "Not Merge", "Not Merge" # ------------------------- # Anima (AM) # ------------------------- if arch.get("AM", False): if ("qwen3_06b" in key) or key.startswith("text_encoders.") or key.startswith("cond_stage_model.qwen3_06b."): return "BASE", "BASE" if "vae" in key or key.startswith("first_stage_model."): return "VAE", "VAE" li = _parse_int_after(key, "blocks.") if li is not None and 0 <= li < 28: tag = f"L{li:02d}" return tag, tag if "x_embedder" in key: return "L00", "L00" if "final_layer" in key: return "L27", "L27" if any(s in key for s in ("t_embedder", "t_embedding_norm", "pos_embedder", "llm_adapter")): return "BASE", "BASE" return "Not Merge", "Not Merge" # ------------------------- # Z-Image # ------------------------- if arch.get("ZI", False): if ("qwen3_4b" in key) or ("cap_embedder" in key): return "BASE", "BASE" if not ("weight" in key or "bias" in key): return "Not Merge", "Not Merge" if "t_embedder" in key or "x_embedder" in key: return "Not Merge", "Not Merge" if "vae" in key: return "VAE", "VAE" if "norm_final" in key: return "L29", "L29" if "model.diffusion_model" in key: if "model.diffusion_model.final_layer" in key: return "L29", "L29" if "model.diffusion_model.context_refiner" in key: return "CONT", "CONT" if "model.diffusion_model.noise_refiner" in key: return "NOISE", "NOISE" li = _parse_int_after(key, "layers.") if li is not None: tag = f"L{li:02d}" return tag, tag return "Not Merge", "Not Merge" return "Not Merge", "Not Merge" def blockfromkey(key: str, arch: dict) -> Tuple[str, str]: xl = bool(arch.get("XL", False)) flux = bool(arch.get("FLUX", False) or arch.get("Flux", False)) zi = bool(arch.get("ZI", False)) am = bool(arch.get("AM", False)) return _blockfromkey_cached_flags(key, xl, flux, zi, am) EXTRA_ELEM_TAGS = ("LABEL", "TIME", "OUT", "CLIP", "CLIP-L", "CLIP-G", "T5") def extra_tag_for_key(key: str, *, arch: dict) -> str | None: kl = key.lower() # --- explicit pseudo blocks --- if "label_emb" in kl: return "LABEL" if "time_embed" in kl: return "TIME" if ".out." in kl or "final_layer" in kl: return "OUT" # --- CLIP grouping (for elementals targeting) --- if arch.get("XL", False): # SDXL: CLIP-L / CLIP-G if ("conditioner.embedders.0." in kl) or ("text_encoders.encoder_l." in kl) or ("clip_l." in kl): return "CLIP-L" if ("conditioner.embedders.1." in kl) or ("text_encoders.encoder_g." in kl) or ("clip_g." in kl): return "CLIP-G" if ("conditioner.embedders." in kl) or ("text_encoders." in kl) or ("clip_l." in kl) or ("clip_g." in kl): return "CLIP" elif arch.get("FLUX", False): # Flux: CLIP / T5 if ("t5" in kl) or ("t5xxl" in kl) or ("text_encoders.t5" in kl): return "T5" if ("clip" in kl) or ("text_encoder" in kl) or ("conditioner.embedders" in kl): return "CLIP" elif arch.get("ZI", False): # ZI: Qwen + cap_embedder (treat as CLIP tag for elementals convenience) if ("qwen3_4b" in kl) or ("cap_embedder" in kl) or ("text_encoders.qwen3_4b" in kl): return "CLIP" elif arch.get("AM", False): if ("qwen3_06b" in kl) or ("text_encoders.qwen3_06b" in kl) or ("cond_stage_model.qwen3_06b" in kl) or ("llm_adapter" in kl): return "CLIP" else: # SD1.x/2.x: cond_stage_model / clip if kl.startswith("cond_stage_model.") or kl.startswith("clip."): return "CLIP" return None def _extend_blockids_for_elementals(blockids: list[str]) -> list[str]: out = list(blockids) for t in EXTRA_ELEM_TAGS: if t not in out: out.append(t) return out def elementals2( key: str, weight_index: int, deep: list[str], current_alpha: float, *, blockids=BLOCKID, arch: dict, ) -> float: """ elementals() compatible, but: - works even when weight_index < 0 ("Not Merge") if extra_tag_for_key() returns a pseudo tag. - appends BOTH (block tag) and (extra tag) to the match key, so deep rules can target: CLIP / CLIP-L / CLIP-G / LABEL / TIME / OUT """ if not deep: return current_alpha block_tag = blockids[weight_index] if (0 <= weight_index < len(blockids)) else "" extra_tag = extra_tag_for_key(key, arch=arch) # matching string (case-sensitive tags are uppercase; keys are usually lowercase) skey = f"{key}|{block_tag}|{extra_tag or ''}" blockids_ex = _extend_blockids_for_elementals(list(blockids)) def _neg(tokens: list[str]): return (True, tokens[1:]) if tokens and tokens[0] == "NOT" else (False, tokens) for d in deep: if d.count(":") != 2: continue dbs_s, dws_s, dr_s = d.split(":", 2) # allow ranges for normal blocks; LABEL/TIME/OUT/CLIP are also accepted dbs = blocker(dbs_s, blockids_ex).split() dws = dws_s.split() dbn, dbs = _neg(dbs) dwn, dws = _neg(dws) ok = (any(db in skey for db in dbs) ^ dbn) if ok: ok = (any(dw in skey for dw in dws) ^ dwn) if ok: current_alpha = float(dr_s) return current_alpha def diff_inplace(dst, src, func, desc): for k in tqdm(dst.keys(), desc=desc, total=len(dst)): if ("model" not in k) and ("text_encoders" not in k): continue v2 = src.get(k) if v2 is None: dst[k] = torch.zeros_like(dst[k]) else: dst[k] = func(dst[k], v2) def _normalize_components_list(alpha_text: str): # "UNet, CLIP-L, VAE" -> {'unet','clip-l','vae'} if not alpha_text: return set() tokens = [t.strip().lower() for t in alpha_text.replace(";", ",").split(",") if t.strip()] syn = { "u": "unet", "unet": "unet", "v": "vae", "vae": "vae", "clip": "clip", "text": "clip", "te": "clip", "clip-l": "clip-l", "clipl": "clip-l", "clip_l": "clip-l", "l": "clip-l", "text-l": "clip-l", "clip-g": "clip-g", "clipg": "clip-g", "clip_g": "clip-g", "g": "clip-g", "text-g": "clip-g", "denoiser": "transformer", "denoise": "transformer", "transformer": "transformer", "mmdit": "transformer", "t5": "text", "t5-xxl": "text", "text1": "text", "text2": "text2", "all": "all", } mapped = [] for t in tokens: tt = syn.get(t, None) if tt is None: try: float(t) continue except Exception: continue mapped.append(tt) if (not mapped) or ("all" in mapped): return {"unet", "vae", "clip-l", "clip-g", "clip", "transformer", "text", "text2"} return set(mapped) def _parse_components_with_only(alpha_text: str): if not alpha_text: return set(), set() tokens = [t.strip().lower() for t in alpha_text.replace(";", ",").split(",") if t.strip()] syn = { "u": "unet", "unet": "unet", "v": "vae", "vae": "vae", "clip": "clip", "text": "clip", "te": "clip", "clip-l": "clip-l", "clipl": "clip-l", "clip_l": "clip-l", "l": "clip-l", "text-l": "clip-l", "clip-g": "clip-g", "clipg": "clip-g", "clip_g": "clip-g", "g": "clip-g", "text-g": "clip-g", "denoiser":"transformer", "denoise":"transformer", "transformer":"transformer", "mmdit":"transformer", "t5":"text", "t5-xxl":"text", "text1":"text", "text2":"text2", "all": "all", } all_set = {"unet", "vae", "clip-l", "clip-g", "clip", "transformer", "text", "text2"} comps = set() only = set() for t in tokens: is_only = False if t.endswith("_only"): is_only = True t = t[:-5] try: float(t) continue except Exception: pass t = syn.get(t, None) if t is None: continue if t == "all": return set(all_set), set() if t in all_set: comps.add(t) if is_only: only.add(t) return comps, only def _component_prefix_map(arch: dict) -> dict[str, list[str]]: if arch.get("FLUX", False): return { "transformer": ["transformer."], "vae": ["vae.", "first_stage_model."], "text": ["text_encoder.", "conditioner.embedders.", "clip.", "t5."], "text2": ["text_encoder_2."], "unet": ["transformer."], "clip": ["text_encoder.", "text_encoder_2.", "conditioner.embedders.", "clip.", "t5."], "clip-l": ["text_encoder."], "clip-g": ["text_encoder_2."], } if arch.get("XL", False): return { "unet": ["model.diffusion_model."], "vae": ["first_stage_model."], "clip-l": [ "conditioner.embedders.0.", "text_encoders.encoder_l.", "clip_l.", "clip_l/" ], "clip-g": [ "conditioner.embedders.1.", "text_encoders.encoder_g.", "clip_g.", "clip_g/" ], "clip": [ "conditioner.embedders.", "text_encoders.", "clip_l.", "clip_g." ], "text": [ "conditioner.embedders.", "text_encoders." ], "text2": [], "transformer": ["model.diffusion_model."], } if arch.get("ZI", False): return { "unet": ["model.diffusion_model."], "transformer": ["model.diffusion_model."], "vae": ["first_stage_model.", "vae."], "text": ["text_encoders.qwen3_4b.transformer.", "qwen3_4b.", "cap_embedder.", "text_encoders.qwen3_4b."], "text2": [], "clip": ["text_encoders.qwen3_4b.transformer.", "qwen3_4b.", "cap_embedder.", "model.diffusion_model.cap_embedder.", "text_encoders.qwen3_4b."], "clip-l": ["qwen3_4b.", "text_encoders.qwen3_4b."], "clip-g": ["cap_embedder.", "model.diffusion_model.cap_embedder."], } if arch.get("AM", False): return { "unet": ["model.diffusion_model."], "transformer": ["model.diffusion_model."], "vae": ["first_stage_model.", "vae."], "text": ["text_encoders.qwen3_06b.", "text_encoders.qwen3_06b_base.", "cond_stage_model.qwen3_06b."], "text2": [], "clip": ["text_encoders.qwen3_06b.", "text_encoders.qwen3_06b_base.", "cond_stage_model.qwen3_06b."], "clip-l": ["text_encoders.qwen3_06b.", "text_encoders.qwen3_06b_base.", "cond_stage_model.qwen3_06b."], "clip-g": [], } # SD1.x / SD2.x (non-XL, non-Flux, non-ZI) return { "unet": ["model.diffusion_model."], "vae": ["first_stage_model."], "clip": ["cond_stage_model.", "clip."], "text": ["cond_stage_model.", "clip."], "text2": [], "transformer": ["model.diffusion_model."], } def _key_belongs_to_component(key: str, prefixes: list[str]) -> bool: for p in prefixes: if key.startswith(p): return True if p.endswith(".") and key.startswith(p[:-1] + "/"): return True return False def _best_prefix_match(k: str, prefixes: list[str]) -> str | None: best = None best_len = -1 for p in prefixes: if k.startswith(p) and len(p) > best_len: best = p best_len = len(p) return best def _suffix_candidates(k: str): cands = [k] if k.startswith("transformer."): cands.append(k[len("transformer."):]) if k.startswith("transformer.model."): cands.append(k[len("transformer."):]) if not k.startswith("model."): cands.append("model." + k) seen = set() out = [] for x in cands: if x not in seen: seen.add(x) out.append(x) return out def normalize_external_text_encoder(theta_src: dict, arch: dict) -> dict: if not isinstance(theta_src, dict) or not theta_src: return theta_src if arch.get("AM", False): te_root = "text_encoders.qwen3_06b." raw_model_prefix = te_root + "transformer." ok_prefixes = ( "text_encoders.qwen3_06b.", "text_encoders.qwen3_06b_base.", "cond_stage_model.qwen3_06b.", "qwen3_06b.", "qwen3_06b_base.", ) elif arch.get("ZI", False): te_root = "text_encoders.qwen3_4b." raw_model_prefix = te_root + "transformer." ok_prefixes = ( "text_encoders.qwen3_4b.", "cond_stage_model.qwen3_4b.", "qwen3_4b.", ) else: return theta_src def is_raw_qwen_key(k: str) -> bool: return k.startswith(( "model.layers.", "model.embed_tokens.", "model.norm.", "lm_head.", "transformer.model.layers.", "transformer.model.embed_tokens.", "transformer.model.norm.", "transformer.lm_head.", )) or (k in ("logit_scale",)) def map_key(k: str) -> str: if k.startswith("cond_stage_model.qwen3_06b.") and arch.get("AM", False): return te_root + k[len("cond_stage_model.qwen3_06b."):] if k.startswith("cond_stage_model.qwen3_4b.") and arch.get("ZI", False): return te_root + k[len("cond_stage_model.qwen3_4b."):] if k.startswith("qwen3_06b.") and arch.get("AM", False): return te_root + k[len("qwen3_06b."):] if k.startswith("qwen3_4b.") and arch.get("ZI", False): return te_root + k[len("qwen3_4b."):] if k.startswith(ok_prefixes): return k if k == "logit_scale": return te_root + "logit_scale" if k.startswith("transformer."): return te_root + k if k.startswith("model.") or k.startswith("lm_head."): return raw_model_prefix + k return k has_raw = any(is_raw_qwen_key(k) for k in theta_src.keys()) has_any_ok = any(k.startswith(ok_prefixes) for k in theta_src.keys()) if not has_raw and has_any_ok: return theta_src out = {} for k, v in theta_src.items(): nk = map_key(k) out[nk] = v return out @torch.inference_mode() def _swap_components_inplace( theta_dst: dict, theta_src: dict, components: set[str], arch: dict, *, src_only: set[str] | None = None, ): if all(not v for v in arch.values()): arch, theta_src = detect_arch(theta_src) pref = _component_prefix_map(arch) selected = {c for c in (components or set()) if c in pref and pref.get(c)} prefixes = [p for c in selected for p in pref.get(c, [])] if not prefixes: return 0, 0, 0, theta_dst only = set(src_only or set()) do_only = bool(only & selected) suffix_map = None if do_only: p_sorted = sorted(prefixes, key=len, reverse=True) suffix_map = {} for kd in theta_dst.keys(): mp = _best_prefix_match(kd, p_sorted) if mp is None: continue suf = kd[len(mp):] if suf and (suf not in suffix_map): suffix_map[suf] = kd moved, created, skipped_shape = 0, 0, 0 for k, v in tqdm(theta_src.items(), desc="Swapping components...", unit="param"): if not do_only: if not _key_belongs_to_component(k, prefixes): continue k_use = k else: k_use = None if k in theta_dst: k_use = k else: for cand in _suffix_candidates(k): kk = suffix_map.get(cand) if suffix_map is not None else None if kk is not None: k_use = kk break if k_use is None: k_use = prefixes[0] + k if k_use in theta_dst: dv = theta_dst[k_use] if hasattr(dv, "shape") and hasattr(v, "shape") and tuple(dv.shape) != tuple(v.shape): skipped_shape += 1 continue if isinstance(dv, torch.Tensor) and isinstance(v, torch.Tensor): vv = v if vv.dtype != dv.dtype or vv.device != dv.device: vv = vv.to(device=dv.device, dtype=dv.dtype) dv.copy_(vv) theta_dst[k_use] = dv else: theta_dst[k_use] = v moved += 1 else: theta_dst[k_use] = v created += 1 return moved, created, skipped_shape, theta_dst def _is_clip_key(key: str, arch: dict) -> bool: if arch.get("FLUX", False): prefixes = ["text_encoder.", "text_encoder_2.", "conditioner.embedders.", "clip.", "t5."] elif arch.get("XL", False): prefixes = ["conditioner.embedders.", "text_encoders.", "clip_l.", "clip_g."] elif arch.get("ZI", False): prefixes = ["qwen3_4b.", "cap_embedder.","text_encoders.qwen3_4b.","model.diffusion_model.cap_embedder."] elif arch.get("AM", False): prefixes = [ "text_encoders.qwen3_06b.", "text_encoders.qwen3_06b_base.", "cond_stage_model.qwen3_06b.", "qwen3_06b.", "qwen3_06b_base.", ] else: prefixes = ["cond_stage_model.", "clip."] return any(key.startswith(p) for p in prefixes) def _tensor_cosine(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: a32 = a.detach().float().view(-1); b32 = b.detach().float().view(-1) if a32.numel() == 0: return torch.tensor(1.0, device=a.device) num = torch.dot(a32, b32) den = (a32.norm() * b32.norm()).clamp_min(1e-12) return (num / den).clamp(-1.0, 1.0) def _maybe_skip_small_norm_bias_for_clipxor(key: str, tens: torch.Tensor) -> bool: if tens.numel() < 128: return True lk = key.lower() if lk.endswith(".bias") or ".bias" in lk: return True if ("norm" in lk) or (".ln" in lk) or (".bn" in lk): return True if ("emb" in lk) or ("pos" in lk): return True if ("text_projection" in lk) or ("logit_scale" in lk): return True return False def _is_small_or_norm_or_bias(key, tens): n = tens.numel() if n < 128: return True k = key.lower() if (k.endswith(".bias") or ".bias" in k or "norm" in k or "ln" in k or "bn" in k): return True if "emb" in k or "pos" in k: return True return False def _collect_clipxor_targets(theta_base: dict, theta_other: dict, arch: dict) -> list[str]: targets = [] for k, A in theta_base.items(): if not _is_clip_key(k, arch): continue B = theta_other.get(k) if getattr(A, "shape", None) != getattr(B, "shape", None): continue if _maybe_skip_small_norm_bias_for_clipxor(k, A): continue targets.append(k) return targets def _clip_roots_for_arch(arch: dict) -> list[str]: if arch.get("FLUX", False): # Flux: T5 / text encoders return [ "text_encoder.", "text_encoder_2.", "conditioner.embedders.", "clip.", "t5." ] if arch.get("XL", False): # SDXL: CLIP-L / CLIP-G return [ "conditioner.embedders.0.", "conditioner.embedders.1.", "text_encoders.encoder_l.", "text_encoders.encoder_g.", "clip_l.", "clip_g." ] if arch.get("ZI", False): # Z-Image: Qwen + caption embedder return [ "qwen3_4b.", "cap_embedder.", ] if arch.get("AM", False): return [ "text_encoders.qwen3_06b.", "text_encoders.qwen3_06b_base.", "cond_stage_model.qwen3_06b.", "qwen3_06b.", "qwen3_06b_base.", ] # SD1.x / SD2.x (non-XL) return ["cond_stage_model.", "clip."] def _iter_clip_items(sd: dict, arch: dict): roots = _clip_roots_for_arch(arch) for k, v in sd.items(): for r in roots: if k.startswith(r): # canonical suffix after the first root occurrence suffix = k[len(r):] yield (k, r, suffix, v) break def _collect_clip_pairs_by_suffix(sd_a: dict, sd_b: dict, arch_a: dict, arch_b: dict): # map suffix -> (orig_key, tensor) for A and B separately map_a = {} for k, _, suf, v in _iter_clip_items(sd_a, arch_a): map_a[suf] = (k, v) map_b = {} for k, _, suf, v in _iter_clip_items(sd_b, arch_b): map_b[suf] = (k, v) # intersect by suffix and by matching shape pairs = [] for suf, (ka, va) in map_a.items(): kb_v = map_b.get(suf) if kb_v is None: continue kb, vb = kb_v if getattr(va, "shape", None) == getattr(vb, "shape", None): pairs.append((suf, ka, kb)) # use A's key for writing back return pairs # list of (suffix, key_in_A, key_in_B) def _clip_tier_for_xl(key: str) -> str: k = key.lower() if ("clip_l" in k) or ("text_model" in k and "clip" in k) or ("conditioner.embedders.0" in k): return "clip-l" if ("clip_g" in k) or ("text2_model" in k and "clip" in k) or ("open_clip" in k) or ("conditioner.embedders.1" in k): return "clip-g" return "other" def _clip_tier_for_flux(key: str) -> str: k = key.lower() if ("t5" in k) or ("text_encoder" in k) or ("textencoder" in k) or ("conditioner.embedders" in k): return "t5" if ("clip" in k) and ("text" in k or "emb" in k or "proj" in k): return "clip" return "other" def _clip_tier_for_zi(key: str) -> str: k = key.lower() if ("qwen3_4b" in k): return "qwen3_4b" if ("cap_embedder" in k): return "cap_embedder" return "other" def _norm_stats(t: torch.Tensor, eps: float = 1e-6): v = t.detach().float().reshape(-1) if v.numel() == 0: return v.new_tensor(0.0), v.new_tensor(1.0) std, mean = torch.std_mean(v, correction=0) # PyTorch 2.x return mean, std.clamp_min(eps) def _apply_stat_alignment(out: torch.Tensor, ref: torch.Tensor, strength: float = 0.5): m_ref, s_ref = _norm_stats(ref) m_out, s_out = _norm_stats(out) aligned = (out - m_out) / s_out * s_ref + m_ref return out * (1 - strength) + aligned * strength def _topk_mask(delta: torch.Tensor, k_frac: float) -> torch.Tensor: if delta.numel() == 0: return torch.zeros_like(delta, dtype=torch.float32) k = max(int(delta.numel() * float(k_frac)), 1) flat = delta.detach().float().abs().view(-1) if k >= flat.numel(): mask = torch.ones_like(flat) else: thresh = torch.kthvalue(flat, flat.numel() - k + 1).values mask = (flat >= thresh).float() return mask.view_as(delta) def _clipxor_semi_hard_blend( A: torch.Tensor, B: torch.Tensor, *, hardness: float = 0.7, use_cosine_gate: bool = True, keep_stats: bool = True ) -> torch.Tensor: same_sign = torch.sign(A) == torch.sign(B) overlap = torch.where(same_sign, torch.sign(A) * torch.minimum(A.abs(), B.abs()), torch.zeros_like(A)) U = A + B - overlap delta = (U - A) if delta.numel() == 0: return A k_frac = 0.10 + 0.45 * float(hardness) alpha = 0.12 + 0.38 * float(hardness) tau = 0.35 - 0.20 * float(hardness) sharp = 0.12 + 0.08 * (1.0 - float(hardness)) tr_budget = 0.12 + 0.38 * float(hardness) stat_strength = 0.65 - 0.30 * float(hardness) m = _topk_mask(delta, k_frac=k_frac) delta_sel = (delta.detach().float() * m).to(A.dtype) if use_cosine_gate: cos = _tensor_cosine(A, U) g = torch.sigmoid((cos - tau) / max(sharp, 1e-3)).to(A.dtype) else: g = A.new_tensor(1.0) out = A + (alpha * g) * delta_sel a_norm = A.detach().float().norm() d_norm = (out - A).detach().float().norm().clamp_min(1e-12) budget = tr_budget * a_norm if d_norm > budget: scale = (budget / d_norm).to(out.dtype) out = A + (out - A) * scale if keep_stats and stat_strength > 0.0: out = _apply_stat_alignment(out, A, strength=float(stat_strength)) return out def _hp_kernel_inplace(w: torch.Tensor, strength: float) -> torch.Tensor: if w.ndim != 4: return w kh, kw = w.shape[-2], w.shape[-1] if kh < 3 or kw < 3: return w x = w.float() C = x.shape[0] * x.shape[1] x2 = x.reshape(1, C, kh, kw) blur = F.avg_pool2d(x2, kernel_size=3, stride=1, padding=1) hp = x2 - blur x3 = x2 + (strength * hp) return x3.reshape_as(x).to(dtype=w.dtype) def _detail_kernel_inplace(w: torch.Tensor, strength: float) -> torch.Tensor: if w.ndim != 4: return w kh, kw = w.shape[-2], w.shape[-1] if kh < 3 or kw < 3: return w x = w.float() C = x.shape[0] * x.shape[1] x2 = x.reshape(1, C, kh, kw) blur = F.avg_pool2d(x2, kernel_size=3, stride=1, padding=1) hp = x2 - blur if strength >= 0: y = x2 + strength * hp else: amt = min(abs(float(strength)), 1.0) * 0.22 y = torch.lerp(x2, blur, amt) return y.reshape_as(x).to(dtype=w.dtype) def _is_conv_weight(key: str, tens: torch.Tensor) -> bool: return tens.ndim == 4 and (key.endswith(".weight") or ".weight" in key) def _is_bias(key: str, tens: torch.Tensor) -> bool: return tens.ndim == 1 and (key.endswith(".bias") or ".bias" in key) def _is_weight_like(key: str, tens: torch.Tensor) -> bool: return tens.ndim >= 2 and (key.endswith(".weight") or ".weight" in key) def _is_norm_key(key: str) -> bool: kl = key.lower() if not (key.endswith(".weight") or key.endswith(".bias") or ".weight" in key or ".bias" in key): return False return ("norm" in kl) or (".ln" in kl) or (".bn" in kl) def _is_attn1_vout(key: str) -> bool: kl = key.lower() return ( "attn1.to_v" in kl or "attn1.to_out" in kl or "self_attn.v_proj" in kl or "self_attn.out_proj" in kl ) def _is_resnetish(key: str) -> bool: return ("resnets" in key) or ("in_layers" in key) or ("out_layers" in key) or (".conv1." in key) or (".conv2." in key) def _is_proj_ff(key: str) -> bool: return ("proj_in" in key) or ("proj_out" in key) or ("ff.net" in key) def _is_unet_out(key: str) -> bool: kl = key.lower() return ( ("model.diffusion_model.out.0." in kl) or ("model.diffusion_model.out.2." in kl) or # Anima / DiT-like final RGB head ("model.diffusion_model.final_layer." in kl) or kl.startswith("final_layer.") or kl.startswith("net.final_layer.") ) def _is_vae_rgb(key: str) -> bool: kl = key.lower() return ( ("first_stage_model.decoder.conv_out." in kl) or ("vae.decoder.conv_out." in kl) or ("model.vae.decoder.conv_out." in kl) or ("model.first_stage_model.decoder.conv_out." in kl) ) def _am_block_band(blk: str) -> float: """Broad tone weighting for Anima L00-L27 blocks.""" if not isinstance(blk, str) or not blk.startswith("L"): return 0.0 try: n = int(blk[1:]) except Exception: return 0.0 if 20 <= n <= 27: return 1.20 if 12 <= n <= 19: return 1.00 if 8 <= n <= 11: return 0.75 if 4 <= n <= 7: return 0.55 return 0.35 def _am_saturation_finetune_tensor(key: str, tens: torch.Tensor, blk: str, sat: float): """ Gentle diffusion-side tone correction for Anima. Real RGB saturation is handled by VAE conv_out via apply_vae_saturation_inplace(). This only damp/boosts AM color/tone sensitive transformer tensors so NoIn/merge fine passes have AM coverage. """ if not isinstance(tens, torch.Tensor) or (not tens.is_floating_point()): return tens band = _am_block_band(blk) if band <= 0: return tens kl = key.lower() down = max(0.0, 1.0 - float(sat)) up = max(0.0, float(sat) - 1.0) if down == 0.0 and up == 0.0: return tens coeff = 0.0 if ("mlp.layer1" in kl) or ("mlp.layer2" in kl) or ("adaln_modulation_mlp" in kl): coeff = 0.10 elif (("v_proj" in kl) or ("output_proj" in kl)) and (("self_attn" in kl) or ("cross_attn" in kl)): coeff = 0.07 elif ("q_norm" in kl) or ("k_norm" in kl) or ("t_embedding_norm" in kl): coeff = 0.04 elif ("adaln_modulation_self_attn" in kl) or ("adaln_modulation_cross_attn" in kl): coeff = 0.035 if coeff <= 0.0: return tens scale = 1.0 - (down * coeff * band) + (up * coeff * band) scale = max(0.70, min(1.30, scale)) return (tens.float() * scale).to(dtype=tens.dtype) def _finetune_inplace(key, tens, fine, arch: dict): if fine == "" or fine is None: return tens if isinstance(fine, (list, tuple)) and len(fine) >= 8: fine = { "mode": "image_tone_v2", "noise1": float(fine[0]), "noise2": float(fine[1]), "noise3": float(fine[2]), "contrast": float(fine[3]), "brightness": float(fine[4]), "r": float(fine[5]), "g": float(fine[6]), "b": float(fine[7]), "saturation": float(fine[8]) if len(fine) > 8 else 1.0, } if isinstance(fine, dict): if fine.get("mode") == "image_tone_v2": dn1 = float(fine.get("noise1", 0.0)) dn2 = float(fine.get("noise2", 0.0)) dn3 = float(fine.get("noise3", 0.0)) ct = float(fine.get("contrast", 0.0)) br = float(fine.get("brightness", 0.0)) rr = float(fine.get("r", 0.0)) gg = float(fine.get("g", 0.0)) bb = float(fine.get("b", 0.0)) sat = float(fine.get("saturation", 1.0)) # detail / clarity band controls dn1_k = dn1 * 0.06 dn2_k = dn2 * 0.05 dn3_k = dn3 * 0.04 # signed tone controls soft = max(-ct, 0.0) # contrast down / black lift up hard = max(ct, 0.0) # contrast up / black deepen ct_scale_out = 1.0 - soft * 0.018 + hard * 0.03 ct_scale_out08 = 1.0 - soft * 0.015 + hard * 0.025 bright_add = br * 0.008 black_lift = soft * (4.0 / 255.0) - hard * (3.0 / 255.0) gamma_soft = 1.0 + soft * 0.018 - hard * 0.02 r_k = 1.0 + rr * 0.05 g_k = 1.0 + gg * 0.05 b_k = 1.0 + bb * 0.05 left, right = blockfromkey(key, arch) blk = right # --- VAE RGB / small brightness carry / direct saturation --- if _is_vae_rgb(key): if _is_conv_weight(key, tens) and tens.shape[0] >= 3: w = tens.float() if abs(sat - 1.0) >= 1e-6: A = _rgb_sat_matrix(sat, w.device, torch.float32) w[:3] = torch.einsum("ij,jchw->ichw", A, w[:3]) w[0] *= r_k w[1] *= g_k w[2] *= b_k # w *= (1.0 + ct * 0.01) return w.to(dtype=tens.dtype) if _is_bias(key, tens) and tens.shape[0] >= 3: b = tens.float() if abs(sat - 1.0) >= 1e-6: A = _rgb_sat_matrix(sat, b.device, torch.float32) b[:3] = torch.einsum("ij,j->i", A, b[:3]) base_add = bright_add + black_lift b[0] = b[0] * r_k + base_add b[1] = b[1] * g_k + base_add b[2] = b[2] * b_k + base_add return b.to(dtype=tens.dtype) return tens # --- Anima diffusion-side saturation/tone support --- if arch.get("AM", False) and abs(sat - 1.0) >= 1e-6: tuned = _am_saturation_finetune_tensor(key, tens, blk, sat) if tuned is not tens: return tuned # --- UNet final out: primary global contrast / black point --- if _is_unet_out(key): if _is_weight_like(key, tens): scale = ct_scale_out * gamma_soft return (tens.float() * scale).to(dtype=tens.dtype) if _is_bias(key, tens): # add = bright_add + black_lift + shadow_lift * 0.5 # return (tens.float() + add).to(dtype=tens.dtype) return tens # --- tone shaping on norm blocks --- if _is_norm_key(key) and blk in ("IN00", "OUT00", "OUT01", "OUT02", "OUT08"): if key.endswith(".weight") or ".weight" in key: scale = 1.0 - soft * 0.015 + hard * 0.02 if blk == "OUT08": scale += (-soft * 0.01 + hard * 0.01) return (tens.float() * scale).to(dtype=tens.dtype) if key.endswith(".bias") or ".bias" in key: # add = bright_add * 0.35 # add += shadow_lift * (0.35 if blk in ("OUT00", "OUT01", "OUT02") else 0.2) # if blk == "OUT08": # add += black_lift * 0.5 # elif blk == "IN00": # add += black_lift * 0.2 # return (tens.float() + add).to(dtype=tens.dtype) return tens # --- OUT08 proj/ff: gamma / black rolloff --- if blk == "OUT08" and _is_proj_ff(key) and _is_weight_like(key, tens): scale = ct_scale_out08 return (tens.float() * scale).to(dtype=tens.dtype) # --- OUT07/OUT08 attn1 v/out: lighting / perceived contrast --- if blk == "OUT07" and hard > 0 and _is_attn1_vout(key) and _is_weight_like(key, tens): scale = 1.0 + hard * 0.01 return (tens.float() * scale).to(dtype=tens.dtype) # --- detail / local contrast / clarity --- if not _is_conv_weight(key, tens): return tens if not _is_resnetish(key): return tens strength = 0.0 # NOISE1: OUT07-OUT08 if blk == "OUT08": strength = dn1_k # NOISE2: OUT05-OUT07 elif blk in ("OUT05", "OUT06"): strength = dn2_k # NOISE3: OUT03-OUT05 + IN00 elif blk in ("OUT03", "OUT04", "IN00"): strength = dn3_k if strength != 0.0: return _detail_kernel_inplace(tens, strength) return tens mul = fine.get("mul", {}) or {} m = 1.0 if any(s in key for s in ("double_block", "double_blocks", "db.")): m *= float(mul.get("double_block", 1.0)) if any(s in key for s in (".img_in", "image_in", "image_proj")): m *= float(mul.get("img_in", 1.0)) if any(s in key for s in (".txt_in", "context_in", "text_in", "clip_proj")): m *= float(mul.get("txt_in", 1.0)) if any(s in key for s in ("time_in", "time_embed", "timestep", "vector_in")): m *= float(mul.get("time", 1.0)) if any(s in key for s in (".out.", "final_layer", "vector_out")) or key.endswith(".out"): m *= float(mul.get("out", 1.0)) if m != 1.0: tens = tens * torch.as_tensor(m, device=tens.device, dtype=tens.dtype) add = float(fine.get("add", 0.0) or 0.0) if add and (key.endswith(".bias") or ".bias" in key): tens = tens + torch.as_tensor(add, device=tens.device, dtype=tens.dtype) return tens # ---- list mode ---- if isinstance(fine, list): # ===== NEW: 8 sliders for SDXL/SD15 ===== if len(fine) >= 8: dn1, dn2, dn3, ct, br, rr, gg, bb = [float(x) for x in fine[:8]] sat = float(fine[8]) if len(fine) > 8 else 1.0 dn1_k = dn1 * 0.12 dn2_k = dn2 * 0.10 dn3_k = dn3 * 0.08 ct_k = 1.0 + ct * 0.06 br_k = br * 0.02 r_k, g_k, b_k = (1.0 + rr * 0.05, 1.0 + gg * 0.05, 1.0 + bb * 0.05) # --- VAE: Brightness / RGB / direct saturation --- if _is_vae_rgb(key): if _is_conv_weight(key, tens) and tens.shape[0] >= 3: w = tens.float() if abs(sat - 1.0) >= 1e-6: A = _rgb_sat_matrix(sat, w.device, torch.float32) w[:3] = torch.einsum("ij,jchw->ichw", A, w[:3]) w[0] *= r_k w[1] *= g_k w[2] *= b_k return w.to(dtype=tens.dtype) if _is_bias(key, tens) and tens.shape[0] >= 3: b = tens.float() if abs(sat - 1.0) >= 1e-6: A = _rgb_sat_matrix(sat, b.device, torch.float32) b[:3] = torch.einsum("ij,j->i", A, b[:3]) b[0] = b[0] * r_k + br_k b[1] = b[1] * g_k + br_k b[2] = b[2] * b_k + br_k return b.to(dtype=tens.dtype) return tens # --- Anima diffusion-side saturation/tone support --- left, right = blockfromkey(key, arch) blk = right if arch.get("AM", False) and abs(sat - 1.0) >= 1e-6: tuned = _am_saturation_finetune_tensor(key, tens, blk, sat) if tuned is not tens: return tuned # --- UNet Contrast / Brightness --- if _is_unet_out(key): if _is_conv_weight(key, tens): return (tens.float() * ct_k).to(dtype=tens.dtype) if _is_bias(key, tens): return (tens.float() + br_k).to(dtype=tens.dtype) # --- Detail/Noise --- left, right = blockfromkey(key, arch) blk = right if not _is_conv_weight(key, tens): return tens if not (_is_resnetish(key) or _is_proj_ff(key)): return tens strength = 0.0 # DN1: OUT07-OUT08 if blk in ("OUT07", "OUT08"): strength = dn1_k # DN2: OUT05-OUT07 elif blk in ("OUT05", "OUT06", "OUT07"): strength = dn2_k # DN3: OUT03-OUT05 + IN00 elif blk in ("OUT03", "OUT04", "OUT05", "IN00"): strength = dn3_k if strength != 0.0: return _hp_kernel_inplace(tens, strength) return tens # ===== OLD: keep existing behavior (FINETUNES 6 keys) ===== idx = next((i for i, pat in enumerate(FINETUNES) if pat in key), -1) if idx == -1: return tens if idx < 5: return tens * torch.as_tensor(fine[idx], device=tens.device, dtype=tens.dtype) else: try: add = torch.as_tensor(fine[5], device=tens.device, dtype=tens.dtype) return tens + add except Exception: if isinstance(fine[5], (list, tuple)) and len(fine[5]) > 0: add = torch.as_tensor(fine[5][0], device=tens.device, dtype=tens.dtype) return tens + add else: add = torch.as_tensor(float(fine[5]) if fine[5] is not None else 0.0, device=tens.device, dtype=tens.dtype) return tens + add return tens def trim_delta(delta: torch.Tensor, percentile: float = 0.5) -> torch.Tensor: if delta.dim() == 4 and min(delta.shape[-2:]) > 2: blurred = F.avg_pool2d(delta, kernel_size=3, stride=1, padding=1) delta = delta * (1 - percentile) + blurred * percentile return delta def prune_extras_vs_model1(theta_base, theta_ref): to_delete = [] for k in theta_base.keys(): if ("model" not in k) and ("text_encoders" not in k): continue if k not in theta_ref: to_delete.append(k) for k in to_delete: del theta_base[k] if to_delete: print(f"[TF] Pruned {len(to_delete)} keys that were not present in model_1") return theta_base class PermutationSpec(NamedTuple): perm_to_axes: dict # {perm_name: [(weight_key, axis), ...]} axes_to_perm: dict # {weight_key: (perm_for_axis0, perm_for_axis1, ...)} def permutation_spec_from_axes_to_perm(axes_to_perm: dict) -> PermutationSpec: perm_to_axes = defaultdict(list) for wk, axis_perms in axes_to_perm.items(): for axis, perm in enumerate(axis_perms): if perm is not None: perm_to_axes[perm].append((wk, axis)) return PermutationSpec(perm_to_axes=dict(perm_to_axes), axes_to_perm=axes_to_perm) def unet_permutation_spec(isxl: bool) -> PermutationSpec: def conv(name: str, p_in, p_out): # weight: (out, in), bias: (out,) return { f"{name}.weight": (p_out, p_in), f"{name}.bias": (p_out,), } def norm(name: str, p): # weight/bias: (p,) return { f"{name}.weight": (p,), f"{name}.bias": (p,), } def dense(name: str, p_in, p_out, bias: bool = True): d = {f"{name}.weight": (p_out, p_in)} if bias: d[f"{name}.bias"] = (p_out,) return d def easyblock(name: str, p_in, p_out): p_inner = f"P_{name}_inner" p_inner2 = f"P_{name}_inner2" p_inner3 = f"P_{name}_inner3" p_inner4 = f"P_{name}_inner4" return { **norm(f"{name}.in_layers.0", p_in), **conv(f"{name}.in_layers.2", p_in, p_inner), **dense(f"{name}.emb_layers.1", p_inner2, p_inner3, bias=True), **norm(f"{name}.out_layers.0", p_inner4), **conv(f"{name}.out_layers.3", p_inner4, p_out), } filename = "sdxl_perm.json" if isxl else "sd_perm.json" with open(base_path(filename), "r", encoding="utf-8") as f: spec = json.load(f) axes_to_perm = {} def bg(idx: int): return None if idx < 0 else f"P_bg{idx}" for key, value in spec.items(): if "skip" in value: axes_to_perm[key] = (None, None, None, None) continue if "conv" in value: i = int(value["conv"]) axes_to_perm.update(conv(key, bg(i), f"P_bg{i + 1}")) elif "norm" in value: i = int(value["norm"]) axes_to_perm.update(norm(key, bg(i))) elif "dense" in value: i = int(value["dense"]) axes_to_perm.update( dense(key, bg(i), f"P_bg{i + 1}", bool(value.get("bias", True))) ) elif "eb" in value: i = int(value["eb"]) axes_to_perm.update(easyblock(key, bg(i), f"P_bg{i + 1}")) return permutation_spec_from_axes_to_perm(axes_to_perm) def get_permuted_param(ps: PermutationSpec, perm, k: str, params, except_axis=None): w = params[k] for axis, p in enumerate(ps.axes_to_perm.get(k, [])): if axis == except_axis: continue if not p: continue if p not in perm: idx = torch.arange(w.shape[axis], device=w.device) perm[p] = idx else: idx = perm[p].to(w.device).long() axis_dim = w.shape[axis] if ( idx.numel() != axis_dim or idx.min().item() < 0 or idx.max().item() >= axis_dim ): idx = torch.arange(axis_dim, device=w.device) perm[p] = idx w = torch.index_select(w, axis, idx) return w def apply_permutation(ps: PermutationSpec, perm: dict, params: dict) -> dict: return {k: get_permuted_param(ps, perm, k, params) for k in params} def inner_matching( n: int, ps: PermutationSpec, p: str, params_a: dict, params_b: dict, usefp16: bool, progress: bool, number: int, linear_sum: float, perm: dict, device, ): dtype = torch.float16 if usefp16 else torch.float32 A = torch.zeros((n, n), dtype=dtype, device=device) for wk, axis in ps.perm_to_axes.get(p, []): if wk not in params_a or wk not in params_b: continue w_a = params_a[wk] w_b = get_permuted_param(ps, perm, wk, params_b, except_axis=axis) w_a = torch.moveaxis(w_a, axis, 0).reshape(n, -1).to(device) w_b = torch.moveaxis(w_b, axis, 0).reshape(n, -1).T.to(device) if usefp16: w_a = w_a.half() w_b = w_b.half() try: A = A + (w_a @ w_b) except RuntimeError: A = A + (torch.dequantize(w_a) @ torch.dequantize(w_b)) # Hungarian A_cpu = A.detach().cpu() ri, ci = linear_sum_assignment(A_cpu.numpy(), maximize=True) ri = torch.as_tensor(ri) ci = torch.as_tensor(ci) assert torch.equal(ri, torch.arange(len(ri))), "Unexpected row indices" eye = torch.eye(n, device=device) A_flat = A.flatten().float() oldL = torch.vdot(A_flat, eye[perm[p].long()].flatten()) newL = torch.vdot(A_flat, eye[ci.long(), :].flatten()) if usefp16: oldL = oldL.half() newL = newL.half() if (newL - oldL) != 0: linear_sum += float(abs(newL - oldL)) number += 1 improved = bool(newL > oldL + 1e-12) progress = progress or improved perm[p] = ci.to(device).float() return linear_sum, number, perm, progress def weight_matching( ps: PermutationSpec, params_a: dict, params_b: dict, max_iter: int = 1, init_perm: dict | None = None, usefp16: bool = False, device: str | torch.device = "cpu", groups: list[str] | None = None, ): perm_sizes = { p: params_a[axes[0][0]].shape[axes[0][1]] for p, axes in ps.perm_to_axes.items() if axes and axes[0][0] in params_a } if init_perm is None: perm = {p: torch.arange(n, device=device) for p, n in perm_sizes.items()} else: perm = {p: v.to(device).long() for p, v in init_perm.items()} for p, axes in ps.perm_to_axes.items(): if not axes: continue size = None for tensor_name, axis in axes: t = params_a.get(tensor_name) if t is None: t = params_b.get(tensor_name) if t is not None: size = t.shape[axis] break if size is not None: perm_sizes[p] = size if init_perm is None: perm = { p: torch.arange(n, device=device).float() for p, n in perm_sizes.items() } else: perm = {p: v.to(device).float() for p, v in init_perm.items()} special_layers = ["P_bg324"] target_groups = groups if groups is not None else special_layers target_groups = [g for g in target_groups if g in perm_sizes] linear_sum: float = 0.0 number: int = 0 if not target_groups or max_iter <= 0: return perm, 0.0 for _ in tqdm(range(max_iter), desc="Weight matching"): random.shuffle(target_groups) progress = False for p in target_groups: n = perm_sizes[p] linear_sum, number, perm, progress = inner_matching( n, ps, p, params_a, params_b, usefp16, progress, number, linear_sum, perm, device, ) if not progress: break average = float(linear_sum) / float(number) if number > 0 else 0.0 return perm, average def _filter_state_dict_by_components(theta: dict, components: set[str], arch: dict): pref = _component_prefix_map(arch) prefixes = [p for c in components for p in pref.get(c, []) if p] total = len(theta) if not prefixes: return {}, 0, total kept = {k: v for k, v in theta.items() if _key_belongs_to_component(k, prefixes)} return kept, len(kept), total def _turbo_force_copy_key(k: str) -> bool: kl = k.lower() return ( "alphas_cumprod" in kl or "betas" in kl or "posterior_" in kl or "sqrt_" in kl or "log_" in kl or "sigmas" in kl or "model_ema." in kl or "num_batches_tracked" in kl ) @torch.inference_mode() def _scalar_ls_scale(tb: torch.Tensor, tc: torch.Tensor, sample: int = 2048) -> float: b = tb.detach().to(torch.float32).reshape(-1) c = tc.detach().to(torch.float32).reshape(-1) n = b.numel() if n == 0: return 1.0 if n > sample: if sample <= 1: idx = torch.zeros((1,), device=b.device, dtype=torch.int64) else: idx = torch.arange(sample, device=b.device, dtype=torch.int64) idx = (idx * (n - 1)) // (sample - 1) idx.clamp_(0, n - 1) b = b.index_select(0, idx) c = c.index_select(0, idx) denom = float(torch.dot(c, c).clamp_min(1e-12)) num = float(torch.dot(b, c)) s = num / denom if not np.isfinite(s): s = 1.0 return s def _is_cond_sensitive_key(k: str) -> bool: kl = k.lower() return ( "attn2" in kl or "cross_attn" in kl or "crossattn" in kl or ".to_q" in kl or ".to_k" in kl or ".to_v" in kl or ".to_out" in kl or "q_proj" in kl or "k_proj" in kl or "v_proj" in kl or "out_proj" in kl or "encoder_hid_proj" in kl or "conditioner." in kl or "cond_stage_model." in kl or "text_encoders" in kl or "context_refiner" in kl or "cap_embedder" in kl ) @torch.inference_mode() def turbo_convert_inplace( A: dict, B: dict, C: dict, resolver, *, deturbo: bool, vae_key: str, bake_vae_enabled: bool, fine=None, ): def allow_vae(k: str) -> bool: return bake_vae_enabled or (vae_key not in k) keys = set(A.keys()) | (set(B.keys()) & set(C.keys())) for k in tqdm(keys, desc=("DeTurbo" if deturbo else "Turbo") + " converting..."): tb = B.get(k, None) tc = C.get(k, None) if not (isinstance(tb, torch.Tensor) and isinstance(tc, torch.Tensor)): continue if tb.shape != tc.shape: continue if not allow_vae(k): continue ent = resolver(k) if ent is None: continue _, cur_a, _ = ent a = float(cur_a) if a == 0.0: continue ta = A.get(k, None) if _turbo_force_copy_key(k) or (not tb.is_floating_point()) or (not tc.is_floating_point()): tgt = tc if deturbo else tb if isinstance(ta, torch.Tensor) and ta.shape == tgt.shape and ta.is_floating_point() and tgt.is_floating_point(): out = torch.lerp(ta.to(torch.float32), tgt.to(torch.float32), a).to(ta.dtype) else: out = tgt A[k] = _finetune_inplace(k, out, fine) if fine else out continue s = _scalar_ls_scale(tb, tc) hi = 6.0 if deturbo and _is_cond_sensitive_key(k) else 4.0 s = float(np.clip(s, 0.25, hi)) invs = 1.0 / max(s, 1e-6) if isinstance(ta, torch.Tensor) and ta.is_floating_point() and ta.shape == tb.shape: A32 = ta.detach().to(torch.float32) else: A32 = (tc.detach().to(torch.float32) if deturbo else tb.detach().to(torch.float32)) B32 = tb.detach().to(torch.float32) C32 = tc.detach().to(torch.float32) if deturbo: # A' = A*(1-a + a*(1/s)) + a*(C - B/s) out32 = A32.mul(1.0 - a + a * invs).add_(C32 - B32 * invs, alpha=a) else: # A' = A*(1-a + a*s) + a*(B - s*C) out32 = A32.mul(1.0 - a + a * s).add_(B32 - C32 * s, alpha=a) def _soft_match_mean_std(out32: torch.Tensor, ref32: torch.Tensor, strength: float = 0.5, eps: float = 1e-6): varO, meanO = torch.var_mean(out32, unbiased=False) varR, meanR = torch.var_mean(ref32, unbiased=False) stdO = varO.sqrt().clamp_min(eps) stdR = varR.sqrt().clamp_min(eps) normed = (out32 - meanO) / stdO matched = normed * stdR + meanR return torch.lerp(out32, matched, float(strength)) if deturbo and _is_cond_sensitive_key(k) and (not _is_small_or_norm_or_bias(k, tb)): COND_DAMP = 0.35 out32 = C32 + (out32 - C32) * COND_DAMP out32 = _soft_match_mean_std(out32, C32, strength=0.5) out = out32.to(ta.dtype if isinstance(ta, torch.Tensor) and ta.is_floating_point() else tb.dtype) A[k] = _finetune_inplace(k, out, fine) if fine else out fill = C if deturbo else B for k, v in fill.items(): if k in A: continue if not allow_vae(k): continue A[k] = v return A def _rgb_sat_matrix(s: float, device, dtype): s = float(s) m = (1.0 - s) / 3.0 A = torch.tensor([ [s + m, m, m], [ m, s + m, m], [ m, m, s + m], ], device=device, dtype=dtype) return A @torch.inference_mode() def apply_vae_saturation_inplace(sd: dict, vae_key: str, sat: float): if abs(float(sat) - 1.0) < 1e-12: return sd weight_suffixes = ( "decoder.conv_out.weight", "decoder.conv_out.0.weight", ) bias_suffixes = ( "decoder.conv_out.bias", "decoder.conv_out.0.bias", ) w_keys = [k for k in sd.keys() if any(k.endswith(suf) for suf in weight_suffixes)] b_keys = [k for k in sd.keys() if any(k.endswith(suf) for suf in bias_suffixes)] nW = nB = 0 for k in w_keys: W = sd.get(k, None) if not isinstance(W, torch.Tensor) or (not W.is_floating_point()): continue if W.dim() != 4 or W.shape[0] != 3: # out_ch must be RGB=3 continue A = _rgb_sat_matrix(sat, W.device, torch.float32) W32 = W.detach().to(torch.float32) # W' = A @ W (mix output RGB channels) W2 = torch.einsum("ij,jchw->ichw", A, W32) sd[k] = W2.to(W.dtype) nW += 1 for k in b_keys: b = sd.get(k, None) if not isinstance(b, torch.Tensor) or (not b.is_floating_point()): continue if b.dim() != 1 or b.shape[0] != 3: continue A = _rgb_sat_matrix(sat, b.device, torch.float32) b32 = b.detach().to(torch.float32) b2 = torch.einsum("ij,j->i", A, b32) sd[k] = b2.to(b.dtype) nB += 1 print(f"[vae_sat] applied sat={sat} to VAE conv_out: weights={nW}, bias={nB}") return sd def _model_path(root: str, name: str | None) -> str | None: return None if name is None else normalize_path(os.path.join(root, name)) def _model_stem(path: str) -> str: return os.path.splitext(os.path.basename(path))[0] def _load_umodel(path: str, *, name: str | None = None, device: str = "cpu", model_type: str = "checkpoint", verify_hash: bool = True, cache_path: str | None = None) -> UnifiedModel: from model import UnifiedModel return UnifiedModel.from_file(path, name=name, device=device, model_type=model_type, verify_hash=verify_hash, cache_path=cache_path) def _clone_info(src: UnifiedModel | None, *, name: str | None = None, path: str | None = None) -> ModelInfo: from model import ModelInfo info = src.info.clone() if src is not None else ModelInfo(model_type="checkpoint") if name is not None: info.name = name if path is not None: info.path = path return info def _build_output_path(args) -> tuple[str, str, str]: fmt = "safetensors" if args.save_safetensors else "ckpt" output_name = args.output output_file = f"{output_name}.{fmt}" output_path = normalize_path(os.path.join(args.model_path, output_file)) if os.path.isfile(output_path): if args.force: print(f"[force] Overwriting existing file: {output_path}") try: os.remove(output_path) except Exception as e: print(f"[force] Failed to remove existing file: {e}") else: i = 0 while os.path.isfile(output_path): output_name = f"{args.output}_{i:02}" output_file = f"{output_name}.{fmt}" output_path = normalize_path(os.path.join(args.model_path, output_file)) i += 1 print(f"Assigned result checkpoint name as {output_file}\n") return output_name, output_file, output_path def _save_umodel(model: UnifiedModel, path: str, *, args, metadata: dict | None = None) -> None: if metadata: model.info.metadata.update(metadata) model.save( path, no_metadata=bool(args.no_metadata), save_half=bool(args.save_half), save_quarter=bool(args.save_quarter), save_bhalf=bool(args.save_bhalf), prune=bool(args.prune), args=args, ) def _register_merge_parents(dst: UnifiedModel, *parents: UnifiedModel | None) -> None: for i, p in enumerate(parents): if p is not None: dst.register_parent(p, role=f"model_{i}")