| """Block-diffusion inference for this checkpoint. |
| Usage: |
| from blockdiff_inference import load, generate |
| model, tok, mask_id = load("REPO_ID") |
| print(generate(model, tok, mask_id, "The capital of France is")) |
| """ |
| import sys, torch, torch.nn as nn, torch.nn.functional as F |
| from transformers import AutoModelForCausalLM, AutoTokenizer |
|
|
| BLOCK_LEN = 32 |
| CENTERED_CONV = True |
| SHIFT = True |
| MASK_TOKEN = "<|mdm_mask|>" |
|
|
|
|
| def _mask(attention_mask, B, T, device, window=None): |
| i = torch.arange(T, device=device) |
| m = (i[None, :] // BLOCK_LEN) <= (i[:, None] // BLOCK_LEN) |
| if window is not None: |
| m &= (i[:, None] - i[None, :]).abs() < window |
| m = m[None, None].expand(B, 1, T, T).clone() |
| if attention_mask is not None and attention_mask.dim() == 2: |
| m &= attention_mask[:, -T:].bool()[:, None, None, :] |
| return m | torch.eye(T, dtype=torch.bool, device=device)[None, None] |
|
|
|
|
| def _mask_fn(sliding): |
| def fn(*args, **kw): |
| emb = kw.get("input_embeds", kw.get("inputs_embeds")) |
| if emb is None: |
| emb = next(a for a in args if torch.is_tensor(a) and a.dim() == 3) |
| win = (getattr(kw.get("config"), "sliding_window", None) or 256) if sliding else None |
| return _mask(kw.get("attention_mask"), emb.shape[0], emb.shape[1], emb.device, win) |
| return fn |
|
|
|
|
| class CenteredConv1d(nn.Module): |
| def __init__(self, conv): |
| super().__init__(); self.conv = conv |
| L = conv.kernel_size[0]; self.lp, self.rp = (L - 1) // 2, L - 1 - (L - 1) // 2 |
| conv.padding = (0,) |
| @property |
| def weight(self): return self.conv.weight |
| @property |
| def bias(self): return self.conv.bias |
| def forward(self, x): return self.conv(F.pad(x, (self.lp, self.rp))) |
|
|
|
|
| def load(repo_id, device="cuda", dtype=torch.bfloat16): |
| tok = AutoTokenizer.from_pretrained(repo_id, trust_remote_code=True) |
| model = AutoModelForCausalLM.from_pretrained(repo_id, dtype=dtype, trust_remote_code=True, |
| attn_implementation="sdpa").to(device).eval() |
| for t in {sys.modules[type(m).__module__] for m in model.modules()}: |
| for name in dir(t): |
| lo = name.lower() |
| if "mask" in lo and "causal" in lo and callable(getattr(t, name)): |
| setattr(t, name, _mask_fn("sliding" in lo)) |
| for m in model.modules(): |
| if hasattr(m, "is_causal"): m.is_causal = False |
| if CENTERED_CONV: |
| for parent in list(model.modules()): |
| for cn, ch in list(parent.named_children()): |
| if isinstance(ch, nn.Conv1d) and ch.groups == ch.in_channels and ch.kernel_size[0] > 1: |
| setattr(parent, cn, CenteredConv1d(ch)) |
| model.config.use_cache = False |
| return model, tok, tok.convert_tokens_to_ids(MASK_TOKEN) |
|
|
|
|
| def _logits(model, x): |
| lg = model(input_ids=x).logits |
| return torch.cat([lg[:, :1], lg[:, :-1]], 1) if SHIFT else lg |
|
|
|
|
| @torch.inference_mode() |
| def generate(model, tok, mask_id, prompt, gen_len=256, tau=0.9, temperature=0.0, repetition_penalty=1.1): |
| x = tok(prompt, return_tensors="pt").input_ids.to(model.device) |
| P, L, eos = x.shape[1], BLOCK_LEN, tok.eos_token_id |
| produced, n = 0, (L - P % L if P % L else L) |
| while produced < gen_len: |
| x = torch.cat([x, torch.full((1, n), mask_id, device=x.device)], 1) |
| s, e = x.shape[1] - n, x.shape[1] |
| for _ in range(L): |
| is_mask = x[0, s:e] == mask_id |
| if not is_mask.any(): break |
| lg = _logits(model, x)[0, s:e].float(); lg[:, mask_id] = -float("inf") |
| if repetition_penalty != 1.0 and s > P: |
| seen = x[0, P:s].unique(); v = lg[:, seen] |
| lg[:, seen] = torch.where(v > 0, v / repetition_penalty, v * repetition_penalty) |
| pr = F.softmax(lg / temperature if temperature > 0 else lg, -1) |
| pred = torch.multinomial(pr, 1)[:, 0] if temperature > 0 else pr.argmax(-1) |
| conf = pr.gather(-1, pred[:, None])[:, 0].masked_fill(~is_mask, -1.0) |
| sel = (conf >= tau) & is_mask |
| if not sel.any(): sel = torch.zeros_like(is_mask); sel[conf.argmax()] = True |
| x[0, s:e] = torch.where(sel, pred, x[0, s:e]) |
| produced += n; n = L |
| if (x[0, s:e] == eos).any(): break |
| out = x[0, P:] |
| hit = (out == eos).nonzero() |
| if len(hit): out = out[: hit[0, 0]] |
| return tok.decode(out, skip_special_tokens=True) |
|
|
