| """
|
| Unified Sparse-AST / BWM Architecture & Universal Subfolder Loader
|
| Domain: Blender 3D Mathematics & Blender Python Libraries (bpy, mathutils, bmesh, numpy, gpu)
|
| """
|
|
|
| import os
|
| import json
|
| import torch
|
| import torch.nn as nn
|
| import torch.nn.functional as F
|
| import safetensors.torch as st
|
|
|
| class RMSNorm(nn.Module):
|
| def __init__(self, d):
|
| super().__init__()
|
| self.w = nn.Parameter(torch.ones(d))
|
| def forward(self, x):
|
| return x * torch.rsqrt(x.float().square().mean(-1, keepdim=True) + 1e-6).to(x.dtype) * self.w
|
|
|
| class SparseASTBlock(nn.Module):
|
| def __init__(self, d=800, h=1600, al_hidden=24, n_heads=8):
|
| super().__init__()
|
| self.n1 = RMSNorm(d)
|
| self.n2 = RMSNorm(d)
|
| self.a = nn.MultiheadAttention(d, n_heads, batch_first=True)
|
| self.up = nn.Linear(d, h)
|
| self.r = nn.Linear(h, 3)
|
| self.down = nn.Linear(2*h, d)
|
| self.al = nn.Sequential(nn.Linear(3*h + 3, al_hidden), nn.SiLU(), nn.Linear(al_hidden, 1))
|
| self.m = nn.Linear(d, d, bias=False)
|
|
|
| def forward(self, x):
|
| t = x.shape[1]
|
| q = self.n1(x)
|
| mask = torch.ones(t, t, device=x.device, dtype=torch.bool).triu(1)
|
| x = torch.clamp(x + self.a(q, q, q, attn_mask=mask, need_weights=False)[0], -50, 50)
|
| z = torch.clamp(self.up(self.n2(x)), -12, 12)
|
| g = F.gumbel_softmax(self.r(z), tau=2.0, hard=True, dim=-1)
|
| e = z.sign() * torch.expm1(z.abs().clamp(max=2.0))
|
| u = (torch.stack((z, z.sign() * torch.log1p(z.abs()), e), -1) * g.unsqueeze(-2)).sum(-1).clamp(-8, 8)
|
| p, n = F.relu(u), F.relu(-u)
|
| d = torch.exp(-F.softplus(self.al(torch.cat((p, n, u, g), -1))).squeeze(-1)).clamp(1e-4, 0.9999)
|
| y = torch.clamp(x + self.down(torch.cat((p, n), -1)), -50, 50)
|
| state = torch.zeros_like(y[:, 0])
|
| o = []
|
| for j in range(t):
|
| state = torch.clamp(state * d[:, j:j+1] + y[:, j] * (1 - d[:, j:j+1]), -50, 50)
|
| o.append(torch.clamp(y[:, j] + self.m(state), -50, 50))
|
| return torch.stack(o, 1)
|
|
|
| class SparseAST(nn.Module):
|
| def __init__(self, d=800, h=1600, layers=28, seq_len=32, vocab=512, al_hidden=24, n_heads=8):
|
| super().__init__()
|
| self.d = d
|
| self.h = h
|
| self.layers = layers
|
| self.seq_len = seq_len
|
| self.vocab = vocab
|
|
|
| self.e = nn.Embedding(vocab, d)
|
| self.p = nn.Embedding(seq_len, d)
|
| self.b = nn.ModuleList([SparseASTBlock(d=d, h=h, al_hidden=al_hidden, n_heads=n_heads) for _ in range(layers)])
|
| self.n = RMSNorm(d)
|
| self.h_out = nn.Linear(d, vocab, bias=False)
|
| self.h_out.weight = self.e.weight
|
|
|
| def forward(self, i):
|
| if i.shape[1] > self.seq_len:
|
| i = i[:, -self.seq_len:]
|
| seq = i.shape[1]
|
| pos = torch.arange(seq, device=i.device)[None]
|
| x = self.e(i) + self.p(pos)
|
| for b in self.b:
|
| x = b(x)
|
| return self.h_out(self.n(x))
|
|
|
| @classmethod
|
| def from_pretrained(cls, repo_or_dir, subfolder="100M-32", device="cpu", target_context=None):
|
| target_dir = os.path.join(repo_or_dir, subfolder) if subfolder else repo_or_dir
|
| config_path = os.path.join(target_dir, "config.json")
|
| with open(config_path, "r", encoding="utf-8") as f:
|
| cfg = json.load(f)
|
|
|
| native_seq = cfg.get("native_seq_len", 32)
|
| effective_seq = max(native_seq, target_context) if target_context else native_seq
|
|
|
| model = cls(
|
| d=cfg["d_model"],
|
| h=cfg["d_hidden"],
|
| layers=cfg["num_layers"],
|
| seq_len=effective_seq,
|
| vocab=cfg["vocab_size"],
|
| al_hidden=cfg.get("al_hidden", 24),
|
| n_heads=cfg.get("num_attention_heads", 8)
|
| )
|
|
|
| weights_path = os.path.join(target_dir, "model.safetensors")
|
| state_dict = st.load_file(weights_path, device="cpu")
|
|
|
| if "h.weight" in state_dict and "h_out.weight" not in state_dict:
|
| state_dict["h_out.weight"] = state_dict["h.weight"]
|
| elif "h_out.weight" not in state_dict and "e.weight" in state_dict:
|
| state_dict["h_out.weight"] = state_dict["e.weight"]
|
|
|
| if effective_seq > native_seq and "p.weight" in state_dict:
|
| old_p = state_dict["p.weight"]
|
| new_p = F.interpolate(
|
| old_p.T.unsqueeze(0),
|
| size=effective_seq,
|
| mode="linear",
|
| align_corners=True
|
| ).squeeze(0).T
|
| state_dict["p.weight"] = new_p
|
|
|
| model.load_state_dict(state_dict)
|
| model.to(device)
|
| model.eval()
|
| return model
|
|
|
| class TopKRouter(nn.Module):
|
| def __init__(self, vocab_size=512, hidden_dim=64, num_experts=4):
|
| super().__init__()
|
| self.embed = nn.Embedding(vocab_size, hidden_dim)
|
| self.norm = nn.LayerNorm(hidden_dim)
|
| self.mlp = nn.Sequential(
|
| nn.Linear(hidden_dim, hidden_dim),
|
| nn.SiLU(),
|
| nn.Linear(hidden_dim, num_experts)
|
| )
|
| def forward(self, x):
|
| emb = self.norm(self.embed(x))
|
| return self.mlp(emb)
|
|
|
| class TopKSparseASTEnsemble(nn.Module):
|
| def __init__(self, router, experts=None, expert_names=None, k=2, tau=1.0, device="cpu"):
|
| super().__init__()
|
| self.router = router
|
| self.experts = nn.ModuleList(experts if experts else [])
|
| self.expert_names = expert_names or []
|
| self.k = k
|
| self.tau = tau
|
| self.device = device
|
|
|
| @classmethod
|
| def from_pretrained(cls, repo_or_dir, subfolder="TopK-MoE", device="cpu", load_experts=True):
|
| target_dir = os.path.join(repo_or_dir, subfolder) if subfolder else repo_or_dir
|
| cfg_path = os.path.join(target_dir, "config.json")
|
| with open(cfg_path, "r", encoding="utf-8") as f:
|
| cfg = json.load(f)
|
|
|
| router = TopKRouter(
|
| vocab_size=cfg["vocab_size"],
|
| hidden_dim=cfg["hidden_dim"],
|
| num_experts=cfg["num_experts"]
|
| )
|
| st_path = os.path.join(target_dir, "model.safetensors")
|
| router.load_state_dict(st.load_file(st_path, device=device))
|
| router.to(device)
|
| router.eval()
|
|
|
| experts = []
|
| names = cfg.get("expert_names", ["3M-32", "10M-32", "100M-32", "200M-32"])
|
| if load_experts:
|
| for exp_name in names:
|
| exp_sub = exp_name if exp_name.endswith("-32") else f"{exp_name}-32"
|
| exp_path = os.path.join(repo_or_dir, exp_sub)
|
| if os.path.exists(exp_path):
|
| print(f"Loading MoE expert [{exp_name}] from {exp_sub}...")
|
| m = SparseAST.from_pretrained(repo_or_dir, subfolder=exp_sub, device=device)
|
| experts.append(m)
|
|
|
| return cls(router=router, experts=experts, expert_names=names, k=cfg.get("k", 2), tau=cfg.get("tau", 1.0), device=device)
|
|
|
| def forward(self, x, k=None):
|
| if k is None:
|
| k = self.k
|
| router_logits = self.router(x)
|
| topk_scores, topk_indices = torch.topk(router_logits, k=k, dim=-1)
|
| gate_weights = F.softmax(topk_scores / self.tau, dim=-1)
|
| return router_logits, topk_indices, gate_weights
|
|
|
