| """BiRefNet (ZhengPeng7/BiRefNet, BiRefNet_lite; MIT) in MLX: a port of the Hub repo's `birefnet.py` for the |
| released configuration (Swin v1 backbone run at full and half resolution and concatenated, 3 context features, |
| ASPPDeformable decoder blocks, image-patch inputs, gradient attention). |
| |
| import mlx_model |
| model = mlx_model.load("bf16") # a folder with config.json + model.safetensors |
| alpha = mlx_model.matte(model, rgb) # rgb: HxWx3 uint8 -> HxW uint8 alpha (255 = foreground) |
| |
| Runs on its own next to its weights (mlx and numpy). NHWC throughout. Each BatchNorm is folded into the conv before |
| it at conversion; torchvision's deform_conv2d (DCNv2, with modulation) is written out as bilinear gathers. |
| Quantized builds (8bit / 6bit / 4bit) quantize the Swin backbone's linears whose input size is a multiple of 64; |
| the convolutions stay bf16. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import json |
| from pathlib import Path |
|
|
| import mlx.core as mx |
| import mlx.nn as nn |
| import numpy as np |
|
|
| SWIN = {"swin_v1_t": dict(embed_dim=96, depths=[2, 2, 6, 2], num_heads=[3, 6, 12, 24], window_size=7), |
| "swin_v1_l": dict(embed_dim=192, depths=[2, 2, 18, 2], num_heads=[6, 12, 24, 48], window_size=12)} |
| CONFIG = {"model_type": "birefnet", "backbone": "swin_v1_t", "input_size": 1024} |
| PRECISIONS = ["bf16", "fp16", "8bit", "6bit", "4bit"] |
| MEAN, STD = np.array([0.485, 0.456, 0.406], np.float32), np.array([0.229, 0.224, 0.225], np.float32) |
| GROUP = 64 |
|
|
|
|
| def quantizable(path: str, module) -> bool: |
| return isinstance(module, nn.Linear) and module.weight.shape[-1] % GROUP == 0 |
|
|
|
|
| |
|
|
| def resize(x, size, align_corners: bool): |
| """torch F.interpolate(mode="bilinear") for NHWC (no antialiasing).""" |
| b, h, w, c = x.shape |
| oh, ow = int(size[0]), int(size[1]) |
| if (oh, ow) == (h, w): |
| return x |
|
|
| def axis(n_in, n_out): |
| if align_corners: |
| s = np.arange(n_out, dtype=np.float64) * ((n_in - 1) / (n_out - 1) if n_out > 1 else 0.0) |
| else: |
| s = np.clip((np.arange(n_out, dtype=np.float64) + 0.5) * (n_in / n_out) - 0.5, 0, None) |
| i0 = np.minimum(np.floor(s).astype(np.int32), n_in - 1) |
| i1 = np.minimum(i0 + 1, n_in - 1) |
| return mx.array(i0), mx.array(i1), mx.array((s - i0).astype(np.float32)) |
|
|
| y0, y1, fy = axis(h, oh) |
| x0, x1, fx = axis(w, ow) |
| fy = fy.astype(x.dtype)[None, :, None, None] |
| fx = fx.astype(x.dtype)[None, None, :, None] |
| top, bot = x[:, y0], x[:, y1] |
| top = top[:, :, x0] * (1 - fx) + top[:, :, x1] * fx |
| bot = bot[:, :, x0] * (1 - fx) + bot[:, :, x1] * fx |
| return top * (1 - fy) + bot * fy |
|
|
|
|
| def to_patches(x, ref_hw): |
| """einops 'b c (hg h) (wg w) -> b (c hg wg) h w' in NHWC.""" |
| b, H, W, c = x.shape |
| h, w = ref_hw |
| hg, wg = H // h, W // w |
| x = x.reshape(b, hg, h, wg, w, c).transpose(0, 2, 4, 5, 1, 3) |
| return x.reshape(b, h, w, c * hg * wg) |
|
|
|
|
| def deform_conv(x, offset, mask, weight, bias, k: int): |
| """torchvision.ops.deform_conv2d (stride 1, padding k // 2, dilation 1, one offset group) with a modulation mask. |
| x (1, H, W, C), offset (1, H, W, 2 k k) as (dy, dx) per kernel position, mask (1, H, W, k k), weight (O, k, k, C).""" |
| b, h, w, c = x.shape |
| assert b == 1, "batch 1" |
| pad = k // 2 |
| flat = x.reshape(h * w, c) |
| gy = mx.arange(h, dtype=mx.float32)[:, None] |
| gx = mx.arange(w, dtype=mx.float32)[None, :] |
| off = offset[0].astype(mx.float32) |
| out = None |
| for i in range(k * k): |
| ky, kx = divmod(i, k) |
| py = gy + (ky - pad) + off[:, :, 2 * i] |
| px = gx + (kx - pad) + off[:, :, 2 * i + 1] |
| y0, x0 = mx.floor(py), mx.floor(px) |
| ly, lx = py - y0, px - x0 |
| acc = None |
| for yy, xx, wt in ((y0, x0, (1 - ly) * (1 - lx)), (y0, x0 + 1, (1 - ly) * lx), |
| (y0 + 1, x0, ly * (1 - lx)), (y0 + 1, x0 + 1, ly * lx)): |
| ok = (yy >= 0) & (yy <= h - 1) & (xx >= 0) & (xx <= w - 1) |
| idx = (mx.clip(yy, 0, h - 1) * w + mx.clip(xx, 0, w - 1)).astype(mx.int32).reshape(-1) |
| v = mx.take(flat, idx, axis=0) * mx.where(ok, wt, 0.0).reshape(-1, 1).astype(x.dtype) |
| acc = v if acc is None else acc + v |
| acc = acc * mask[0, :, :, i].reshape(-1, 1) |
| y = acc @ weight[:, ky, kx, :].T |
| out = y if out is None else out + y |
| return (out + bias).reshape(1, h, w, -1) |
|
|
|
|
| |
|
|
| def window_partition(x, ws): |
| b, h, w, c = x.shape |
| x = x.reshape(b, h // ws, ws, w // ws, ws, c).transpose(0, 1, 3, 2, 4, 5) |
| return x.reshape(-1, ws * ws, c) |
|
|
|
|
| def window_reverse(win, ws, h, w): |
| c = win.shape[-1] |
| b = win.shape[0] // ((h // ws) * (w // ws)) |
| x = win.reshape(b, h // ws, w // ws, ws, ws, c).transpose(0, 1, 3, 2, 4, 5) |
| return x.reshape(b, h, w, c) |
|
|
|
|
| def relative_index(ws: int) -> np.ndarray: |
| coords = np.stack(np.meshgrid(np.arange(ws), np.arange(ws), indexing="ij")).reshape(2, -1) |
| rel = (coords[:, :, None] - coords[:, None, :]).transpose(1, 2, 0) + (ws - 1) |
| return rel[:, :, 0] * (2 * ws - 1) + rel[:, :, 1] |
|
|
|
|
| def shift_mask(hp: int, wp: int, ws: int, shift: int) -> np.ndarray: |
| img = np.zeros((1, hp, wp, 1), np.float32) |
| cnt = 0 |
| for hs in (slice(0, -ws), slice(-ws, -shift), slice(-shift, None)): |
| for wsl in (slice(0, -ws), slice(-ws, -shift), slice(-shift, None)): |
| img[:, hs, wsl, :] = cnt |
| cnt += 1 |
| m = img.reshape(1, hp // ws, ws, wp // ws, ws, 1).transpose(0, 1, 3, 2, 4, 5).reshape(-1, ws * ws) |
| m = m[:, None, :] - m[:, :, None] |
| return np.where(m != 0, -100.0, 0.0).astype(np.float32) |
|
|
|
|
| class WindowAttention(nn.Module): |
| def __init__(self, dim, ws, heads): |
| super().__init__() |
| self.heads, self.scale = heads, (dim // heads) ** -0.5 |
| self.relative_position_bias_table = mx.zeros(((2 * ws - 1) ** 2, heads)) |
| self._index = mx.array(relative_index(ws).reshape(-1)) |
| self.qkv = nn.Linear(dim, dim * 3) |
| self.proj = nn.Linear(dim, dim) |
|
|
| def __call__(self, x, mask=None): |
| b, n, c = x.shape |
| qkv = self.qkv(x).reshape(b, n, 3, self.heads, c // self.heads).transpose(2, 0, 3, 1, 4) |
| q, k, v = qkv[0] * self.scale, qkv[1], qkv[2] |
| bias = self.relative_position_bias_table[self._index].reshape(n, n, -1).transpose(2, 0, 1) |
| attn = (q @ k.transpose(0, 1, 3, 2)) + bias[None].astype(q.dtype) |
| if mask is not None: |
| nw = mask.shape[0] |
| attn = attn.reshape(b // nw, nw, self.heads, n, n) + mask[None, :, None].astype(q.dtype) |
| attn = attn.reshape(b, self.heads, n, n) |
| attn = mx.softmax(attn, axis=-1, precise=True) |
| return self.proj((attn @ v).transpose(0, 2, 1, 3).reshape(b, n, c)) |
|
|
|
|
| class Mlp(nn.Module): |
| def __init__(self, dim): |
| super().__init__() |
| self.fc1, self.fc2 = nn.Linear(dim, dim * 4), nn.Linear(dim * 4, dim) |
|
|
| def __call__(self, x): |
| return self.fc2(nn.gelu(self.fc1(x))) |
|
|
|
|
| class SwinBlock(nn.Module): |
| def __init__(self, dim, heads, ws, shift): |
| super().__init__() |
| self.ws, self.shift = ws, shift |
| self.norm1, self.norm2 = nn.LayerNorm(dim), nn.LayerNorm(dim) |
| self.attn = WindowAttention(dim, ws, heads) |
| self.mlp = Mlp(dim) |
|
|
| def __call__(self, x, h, w, mask): |
| b, L, c = x.shape |
| ws, s = self.ws, self.shift |
| y = self.norm1(x).reshape(b, h, w, c) |
| pr, pb = (ws - w % ws) % ws, (ws - h % ws) % ws |
| if pr or pb: |
| y = mx.pad(y, [(0, 0), (0, pb), (0, pr), (0, 0)]) |
| hp, wp = h + pb, w + pr |
| if s: |
| y = mx.roll(y, (-s, -s), axis=(1, 2)) |
| y = window_reverse(self.attn(window_partition(y, ws), mask if s else None), ws, hp, wp) |
| if s: |
| y = mx.roll(y, (s, s), axis=(1, 2)) |
| y = y[:, :h, :w].reshape(b, h * w, c) |
| x = x + y |
| return x + self.mlp(self.norm2(x)) |
|
|
|
|
| class PatchMerging(nn.Module): |
| def __init__(self, dim): |
| super().__init__() |
| self.norm = nn.LayerNorm(4 * dim) |
| self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False) |
|
|
| def __call__(self, x, h, w): |
| b, L, c = x.shape |
| x = x.reshape(b, h, w, c) |
| if h % 2 or w % 2: |
| x = mx.pad(x, [(0, 0), (0, h % 2), (0, w % 2), (0, 0)]) |
| x = mx.concatenate([x[:, 0::2, 0::2], x[:, 1::2, 0::2], x[:, 0::2, 1::2], x[:, 1::2, 1::2]], -1) |
| return self.reduction(self.norm(x.reshape(b, -1, 4 * c))) |
|
|
|
|
| class SwinLayer(nn.Module): |
| def __init__(self, dim, depth, heads, ws, down): |
| super().__init__() |
| self.ws = ws |
| self.blocks = [SwinBlock(dim, heads, ws, 0 if i % 2 == 0 else ws // 2) for i in range(depth)] |
| if down: |
| self.downsample = PatchMerging(dim) |
| self._masks = {} |
|
|
| def __call__(self, x, h, w): |
| ws = self.ws |
| hp, wp = -(-h // ws) * ws, -(-w // ws) * ws |
| if (hp, wp) not in self._masks: |
| self._masks[(hp, wp)] = mx.array(shift_mask(hp, wp, ws, ws // 2)) |
| mask = self._masks[(hp, wp)] |
| for blk in self.blocks: |
| x = blk(x, h, w, mask) |
| if "downsample" in self: |
| return x, self.downsample(x, h, w), (h + 1) // 2, (w + 1) // 2 |
| return x, x, h, w |
|
|
|
|
| class PatchEmbed(nn.Module): |
| def __init__(self, dim): |
| super().__init__() |
| self.proj = nn.Conv2d(3, dim, 4, stride=4) |
| self.norm = nn.LayerNorm(dim) |
|
|
| def __call__(self, x): |
| b, h, w, _ = x.shape |
| if w % 4 or h % 4: |
| x = mx.pad(x, [(0, 0), (0, (4 - h % 4) % 4), (0, (4 - w % 4) % 4), (0, 0)]) |
| return self.norm(self.proj(x)) |
|
|
|
|
| class Swin(nn.Module): |
| def __init__(self, embed_dim, depths, num_heads, window_size): |
| super().__init__() |
| self.patch_embed = PatchEmbed(embed_dim) |
| self.dims = [embed_dim * 2 ** i for i in range(4)] |
| self.layers = [SwinLayer(self.dims[i], depths[i], num_heads[i], window_size, i < 3) for i in range(4)] |
| for i, d in enumerate(self.dims): |
| setattr(self, f"norm{i}", nn.LayerNorm(d)) |
|
|
| def __call__(self, x): |
| x = self.patch_embed(x) |
| b, h, w, c = x.shape |
| x = x.reshape(b, h * w, c) |
| outs = [] |
| for i, layer in enumerate(self.layers): |
| xo, x, nh, nw = layer(x, h, w) |
| outs.append(getattr(self, f"norm{i}")(xo).reshape(b, h, w, self.dims[i])) |
| h, w = nh, nw |
| return outs |
|
|
|
|
| |
|
|
| class DeformConv(nn.Module): |
| def __init__(self, cin, cout, k): |
| super().__init__() |
| self.k = k |
| self.offset_conv = nn.Conv2d(cin, 2 * k * k, k, padding=k // 2) |
| self.modulator_conv = nn.Conv2d(cin, k * k, k, padding=k // 2) |
| self.regular_conv = nn.Conv2d(cin, cout, k, padding=k // 2) |
|
|
| def __call__(self, x): |
| off = self.offset_conv(x) |
| mask = 2 * mx.sigmoid(self.modulator_conv(x)) |
| return deform_conv(x, off, mask, self.regular_conv.weight, self.regular_conv.bias, self.k) |
|
|
|
|
| class ASPPDeformable(nn.Module): |
| def __init__(self, cin, mid=256, sizes=(1, 3, 7)): |
| super().__init__() |
| self.aspp1 = DeformConv(cin, mid, 1) |
| self.aspp_deforms = [DeformConv(cin, mid, k) for k in sizes] |
| self.gap_conv = nn.Conv2d(cin, mid, 1) |
| self.conv1 = nn.Conv2d(mid * (2 + len(sizes)), cin, 1) |
|
|
| def __call__(self, x): |
| b, h, w, _ = x.shape |
| parts = [nn.relu(self.aspp1(x))] + [nn.relu(d(x)) for d in self.aspp_deforms] |
| g = x.astype(mx.float32).mean(axis=(1, 2), keepdims=True).astype(x.dtype) |
| g = nn.relu(self.gap_conv(g)) |
| parts.append(mx.broadcast_to(g, (b, h, w, g.shape[-1]))) |
| return nn.relu(self.conv1(mx.concatenate(parts, -1))) |
|
|
|
|
| class DecBlk(nn.Module): |
| def __init__(self, cin, cout, inter=64): |
| super().__init__() |
| self.conv_in = nn.Conv2d(cin, inter, 3, padding=1) |
| self.dec_att = ASPPDeformable(inter) |
| self.conv_out = nn.Conv2d(inter, cout, 3, padding=1) |
|
|
| def __call__(self, x): |
| return self.conv_out(self.dec_att(nn.relu(self.conv_in(x)))) |
|
|
|
|
| class SimpleConvs(nn.Module): |
| def __init__(self, cin, cout, inter=64): |
| super().__init__() |
| self.conv1 = nn.Conv2d(cin, inter, 3, padding=1) |
| self.conv_out = nn.Conv2d(inter, cout, 3, padding=1) |
|
|
| def __call__(self, x): |
| return self.conv_out(self.conv1(x)) |
|
|
|
|
| class Decoder(nn.Module): |
| def __init__(self, ch): |
| super().__init__() |
| self.ipt_blk5 = SimpleConvs(2 ** 10 * 3, ch[0] // 8) |
| self.ipt_blk4 = SimpleConvs(2 ** 8 * 3, ch[0] // 8) |
| self.ipt_blk3 = SimpleConvs(2 ** 6 * 3, ch[1] // 8) |
| self.ipt_blk2 = SimpleConvs(2 ** 4 * 3, ch[2] // 8) |
| self.ipt_blk1 = SimpleConvs(3, ch[3] // 8) |
| self.decoder_block4 = DecBlk(ch[0] + ch[0] // 8, ch[1]) |
| self.decoder_block3 = DecBlk(ch[1] + ch[0] // 8, ch[2]) |
| self.decoder_block2 = DecBlk(ch[2] + ch[1] // 8, ch[3]) |
| self.decoder_block1 = DecBlk(ch[3] + ch[2] // 8, ch[3] // 2) |
| self.conv_out1 = nn.Conv2d(ch[3] // 2 + ch[3] // 8, 1, 1) |
| self.lateral_block4 = nn.Conv2d(ch[1], ch[1], 1) |
| self.lateral_block3 = nn.Conv2d(ch[2], ch[2], 1) |
| self.lateral_block2 = nn.Conv2d(ch[3], ch[3], 1) |
| for i, c in ((4, ch[1]), (3, ch[2]), (2, ch[3])): |
| setattr(self, f"gdt_convs_{i}", nn.Conv2d(c, 16, 3, padding=1)) |
| setattr(self, f"gdt_attn_{i}", nn.Conv2d(16, 1, 1)) |
|
|
| def gdt(self, i, p): |
| return p * mx.sigmoid(getattr(self, f"gdt_attn_{i}")(nn.relu(getattr(self, f"gdt_convs_{i}")(p)))) |
|
|
| def __call__(self, x, x1, x2, x3, x4): |
| up = lambda a, ref: resize(a, ref.shape[1:3], True) |
| x4 = mx.concatenate([x4, self.ipt_blk5(to_patches(x, x4.shape[1:3]))], -1) |
| p4 = self.gdt(4, self.decoder_block4(x4)) |
| p3 = up(p4, x3) + self.lateral_block4(x3) |
| p3 = mx.concatenate([p3, self.ipt_blk4(to_patches(x, p3.shape[1:3]))], -1) |
| p3 = self.gdt(3, self.decoder_block3(p3)) |
| p2 = up(p3, x2) + self.lateral_block3(x2) |
| p2 = mx.concatenate([p2, self.ipt_blk3(to_patches(x, p2.shape[1:3]))], -1) |
| p2 = self.gdt(2, self.decoder_block2(p2)) |
| p1 = up(p2, x1) + self.lateral_block2(x1) |
| p1 = mx.concatenate([p1, self.ipt_blk2(to_patches(x, p1.shape[1:3]))], -1) |
| p1 = up(self.decoder_block1(p1), x) |
| p1 = mx.concatenate([p1, self.ipt_blk1(x)], -1) |
| return self.conv_out1(p1) |
|
|
|
|
| class Model(nn.Module): |
| def __init__(self, config: dict = CONFIG): |
| super().__init__() |
| sw = SWIN[config["backbone"]] |
| self.size = config["input_size"] |
| self.bb = Swin(**sw) |
| ch = [d * 2 for d in self.bb.dims[::-1]] |
| cxt = ch[1:][::-1] |
| self.squeeze_module = DecBlk(ch[0] + sum(cxt), ch[0]) |
| self.decoder = Decoder(ch) |
|
|
| def __call__(self, x): |
| """x: (1, S, S, 3) ImageNet-normalized -> (1, S, S, 1) logits.""" |
| b, h, w, _ = x.shape |
| f = self.bb(x) |
| g = self.bb(resize(x, (h // 2, w // 2), True)) |
| f = [mx.concatenate([a, resize(c, a.shape[1:3], True)], -1) for a, c in zip(f, g)] |
| x1, x2, x3, x4 = f |
| x4 = mx.concatenate([resize(x1, x4.shape[1:3], True), resize(x2, x4.shape[1:3], True), |
| resize(x3, x4.shape[1:3], True), x4], -1) |
| x4 = self.squeeze_module(x4) |
| return self.decoder(x, x1, x2, x3, x4) |
|
|
|
|
| def load(folder) -> Model: |
| folder = Path(folder) |
| config = json.loads((folder / "config.json").read_text()) |
| model = Model(config) |
| q = config.get("quantization") |
| if q: |
| nn.quantize(model, group_size=q["group_size"], bits=q["bits"], class_predicate=quantizable) |
| model.load_weights(str(folder / "model.safetensors"), strict=True) |
| mx.eval(model.parameters()) |
| return model |
|
|
|
|
| def dtype(model: Model): |
| return model.decoder.conv_out1.weight.dtype |
|
|
|
|
| def predict(model: Model, rgb: np.ndarray) -> np.ndarray: |
| """HxWx3 uint8 -> HxW float32 foreground probability: resized to the model size (bilinear, as PIL in the |
| original handler) and back.""" |
| from PIL import Image |
| h, w, _ = rgb.shape |
| s = model.size |
| x = np.asarray(Image.fromarray(rgb).resize((s, s), Image.BILINEAR), np.float32) / 255.0 |
| x = mx.array(((x - MEAN) / STD)[None]).astype(dtype(model)) |
| y = mx.sigmoid(model(x).astype(mx.float32)) |
| return np.array(resize(y, (h, w), True))[0, :, :, 0] |
|
|
|
|
| def matte(model: Model, rgb: np.ndarray) -> np.ndarray: |
| return (np.clip(predict(model, rgb), 0, 1) * 255).round().astype(np.uint8) |
|
|
|
|
| |
|
|
| def _fold(sd, conv, bn): |
| w = sd[conv + ".weight"].astype(np.float64) |
| b = sd.get(conv + ".bias", np.zeros(w.shape[0])).astype(np.float64) |
| if bn is not None: |
| g = sd[bn + ".weight"] / np.sqrt(sd[bn + ".running_var"].astype(np.float64) + 1e-5) |
| w = w * g[:, None, None, None] |
| b = (b - sd[bn + ".running_mean"]) * g + sd[bn + ".bias"] |
| return w.astype(np.float32).transpose(0, 2, 3, 1), b.astype(np.float32) |
|
|
|
|
| def convert(source: Path) -> dict: |
| """runs/source/model.safetensors -> MLX names: BatchNorms folded, Sequential indices renamed, buffers and the |
| training-only heads (conv_ms_spvn_*, gdt_convs_pred_*) dropped.""" |
| from safetensors.numpy import load_file |
| sd = load_file(str(Path(source) / "model.safetensors")) |
| out = {} |
| for k, v in sd.items(): |
| if k.startswith("bb."): |
| if k.endswith("relative_position_index"): |
| continue |
| out[k] = v.transpose(0, 2, 3, 1) if v.ndim == 4 else v |
| blocks = sorted({k.split(".dec_att.")[0] for k in sd if ".dec_att." in k}) |
| for blk in blocks: |
| dst = "squeeze_module" if blk == "squeeze_module.0" else blk |
| out[dst + ".conv_in.weight"], out[dst + ".conv_in.bias"] = _fold(sd, blk + ".conv_in", blk + ".bn_in") |
| out[dst + ".conv_out.weight"], out[dst + ".conv_out.bias"] = _fold(sd, blk + ".conv_out", blk + ".bn_out") |
| a = blk + ".dec_att" |
| for src_m, dst_m in [(a + ".aspp1", dst + ".dec_att.aspp1")] + [ |
| (f"{a}.aspp_deforms.{i}", f"{dst}.dec_att.aspp_deforms.{i}") for i in range(3)]: |
| for c in ("offset_conv", "modulator_conv"): |
| out[f"{dst_m}.{c}.weight"], out[f"{dst_m}.{c}.bias"] = _fold(sd, f"{src_m}.atrous_conv.{c}", None) |
| out[f"{dst_m}.regular_conv.weight"], out[f"{dst_m}.regular_conv.bias"] = _fold( |
| sd, f"{src_m}.atrous_conv.regular_conv", f"{src_m}.bn") |
| out[dst + ".dec_att.gap_conv.weight"], out[dst + ".dec_att.gap_conv.bias"] = _fold( |
| sd, a + ".global_avg_pool.1", a + ".global_avg_pool.2") |
| out[dst + ".dec_att.conv1.weight"], out[dst + ".dec_att.conv1.bias"] = _fold(sd, a + ".conv1", a + ".bn1") |
| d = "decoder." |
| for i in range(1, 6): |
| for c in ("conv1", "conv_out"): |
| out[f"{d}ipt_blk{i}.{c}.weight"], out[f"{d}ipt_blk{i}.{c}.bias"] = _fold(sd, f"{d}ipt_blk{i}.{c}", None) |
| for i in (2, 3, 4): |
| out[f"{d}lateral_block{i}.weight"], out[f"{d}lateral_block{i}.bias"] = _fold(sd, f"{d}lateral_block{i}.conv", None) |
| out[f"{d}gdt_convs_{i}.weight"], out[f"{d}gdt_convs_{i}.bias"] = _fold(sd, f"{d}gdt_convs_{i}.0", f"{d}gdt_convs_{i}.1") |
| out[f"{d}gdt_attn_{i}.weight"], out[f"{d}gdt_attn_{i}.bias"] = _fold(sd, f"{d}gdt_convs_attn_{i}.0", None) |
| out[d + "conv_out1.weight"], out[d + "conv_out1.bias"] = _fold(sd, d + "conv_out1.0", None) |
| return out |
|
|
