| 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/<tag>.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"]) |
|
|
| def grid(t): |
| 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() |
| 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) |
| 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" |
|
|
