import re def parse_layers(spec): out = set() for part in spec.replace(" ", "").split(","): if not part: continue if "-" in part: a, b = part.split("-") out.update(range(int(a), int(b) + 1)) else: out.add(int(part)) return out class MOSSBlockSwap: """Copy exact weights from `donor` into `model` for chosen transformer layers and weight types. layers: "21-41,3,4"; empty = none. parts: comma list matched against the key after `layers.N.` (qkv, o, ffn_up_gate, ffn_down, norm) or "all". include_other also swaps every non-layer key (embedders, time embedding, shared modulation, final head, text embedder). """ @classmethod def INPUT_TYPES(s): return {"required": {"model": ("MODEL",), "donor": ("MODEL",), "layers": ("STRING", {"default": "0-41"}), "parts": ("STRING", {"default": "all"}), "include_other": ("BOOLEAN", {"default": False}), "ratio": ("FLOAT", {"default": 1.0, "min": -2.0, "max": 2.0, "step": 0.05}), }} RETURN_TYPES = ("MODEL",) FUNCTION = "swap" CATEGORY = "model/merging" def swap(self, model, donor, layers, parts, include_other, ratio): want = parse_layers(layers) parts = [p.strip() for p in parts.split(",") if p.strip()] m = model.clone() kp = donor.get_key_patches("diffusion_model.") n = 0 for k in kp: mm = re.search(r"\.layers\.(\d+)\.(.+)$", k) if mm: if int(mm.group(1)) not in want: continue if "all" not in parts and not any(p in mm.group(2) for p in parts): continue elif not include_other: continue m.add_patches({k: kp[k]}, ratio, 1.0 - ratio) n += 1 print(f"MOSSBlockSwap: swapped {n} tensors (layers={layers}, parts={parts}, other={include_other}, ratio={ratio})") return (m,) NODE_CLASS_MAPPINGS = {"MOSSBlockSwap": MOSSBlockSwap} NODE_DISPLAY_NAME_MAPPINGS = {"MOSSBlockSwap": "MOSS Block Swap"} import os import numpy as np import torch import folder_paths import comfy.utils PROBES = {} def column_spectrum(t): """t: (B, C, H, W). Remove each row's mean along W, average down the rows, FFT along W, sum power over channels -> (B, W//2+1).""" t = t - t.mean(dim=3, keepdim=True) col = t.mean(dim=2) return (torch.fft.rfft(col, dim=-1).abs() ** 2).sum(dim=1) class MOSSStripeProbe: """Records, for every transformer layer and model call, the column-stripe spectrum of (a) the logit lens: the residual stream passed through the model's own output head, (b) the lens of what this layer added, and (c) the raw residual stream. Writes output/moss/probe/.npz after every model call.""" @classmethod def INPUT_TYPES(s): return {"required": {"model": ("MODEL",), "head_from": (folder_paths.get_filename_list("diffusion_models"),), "tag": ("STRING", {"default": "probe"}), "width": ("INT", {"default": 1024, "min": 16, "max": 8192, "step": 16}), "height": ("INT", {"default": 1024, "min": 16, "max": 8192, "step": 16}), }} RETURN_TYPES = ("MODEL",) FUNCTION = "probe" CATEGORY = "model/debug" def probe(self, model, head_from, tag, width, height): sd = comfy.utils.load_torch_file(folder_paths.get_full_path_or_raise("diffusion_models", head_from)) head = {k: sd[k].float() for k in ("final_norm_video.weight", "final_linear_video.weight", "video_head_modulation.weight", "video_head_modulation.bias")} del sd H, W = height // 16, width // 16 n_layers = model.model.diffusion_model.layers.__len__() out_path = os.path.join(folder_paths.get_output_directory(), "moss", "probe", tag + ".npz") os.makedirs(os.path.dirname(out_path), exist_ok=True) state = PROBES[tag] = {"calls": [], "cur": None} def lens(x, vec, dev): hw = {k: v.to(dev) for k, v in head.items()} mod = torch.nn.functional.linear(vec.float(), hw["video_head_modulation.weight"], hw["video_head_modulation.bias"]) shift, scale = mod.unflatten(-1, (2, -1)).unsqueeze(2).unbind(1) y = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + 1e-6) * hw["final_norm_video.weight"] * (1 + scale) + shift return torch.nn.functional.linear(y, hw["final_linear_video.weight"]) # (B, HW, 128) def grid(t): # (B, HW, C) -> (B, C, H, W) return t.transpose(1, 2).reshape(t.shape[0], t.shape[2], H, W) def make(i): def wrap(args, extra): to = args["transformer_options"] x_in = args["img"][:, :H * W].float().clone() # blocks update x in place out = extra["original_block"](args) x_out = out["img"][:, :H * W].float() l_in, l_out = lens(x_in, args["vec"], x_out.device), lens(x_out, args["vec"], x_out.device) row = torch.stack([column_spectrum(grid(l_out)), column_spectrum(grid(l_out - l_in)), column_spectrum(grid(x_out))], dim=1) # (B, 3, nb) if i == 0 or state["cur"] is None: state["cur"] = {"sigma": float(to["sigmas"].flatten()[0]), "cou": list(to.get("cond_or_uncond", [])), "spec": np.full((x_out.shape[0], n_layers, 3, W // 2 + 1), np.nan, np.float32)} state["cur"]["spec"][:, i] = row.cpu().numpy() if i == n_layers - 1: state["calls"].append(state["cur"]) state["cur"] = None np.savez(out_path, sigma=np.array([c["sigma"] for c in state["calls"]]), cou=np.array([str(c["cou"]) for c in state["calls"]]), spec=np.array([c["spec"] for c in state["calls"]], dtype=object), allow_pickle=True) return out return wrap m = model.clone() for i in range(n_layers): m.set_model_patch_replace(make(i), "dit", "single_block", i) return (m,) NODE_CLASS_MAPPINGS["MOSSStripeProbe"] = MOSSStripeProbe NODE_DISPLAY_NAME_MAPPINGS["MOSSStripeProbe"] = "MOSS Stripe Probe"