| """Minecraft texture diffusion model: standalone inference code. |
| |
| Usage: |
| from model import TextureGenerator |
| gen = TextureGenerator("ArshPathan/TexelDiff-16-62M") # or a local folder |
| imgs = gen("minecraft item texture: gold ingot", res=16, n=8, guidance=4.0) |
| imgs[0].resize((256, 256), 0).save("ingot.png") # 0 = nearest neighbour |
| """ |
| import json, math |
| from pathlib import Path |
| from types import SimpleNamespace |
| import numpy as np, torch, torch.nn as nn, torch.nn.functional as F |
| from PIL import Image |
|
|
| CFG = SimpleNamespace(tileable=False) |
|
|
| def sinusoidal(t, dim): |
| half = dim // 2 |
| freqs = torch.exp(-math.log(10000) * torch.arange(half, device=t.device) / half) |
| a = t[:, None].float() * freqs[None] |
| return torch.cat([a.sin(), a.cos()], -1) |
|
|
| def conv(i, o, k=3): |
| return nn.Conv2d(i, o, k, padding=k // 2, padding_mode="circular" if CFG.tileable else "zeros") |
|
|
| class ResBlock(nn.Module): |
| def __init__(self, i, o, emb): |
| super().__init__() |
| self.n1 = nn.GroupNorm(8, i); self.c1 = conv(i, o) |
| self.emb = nn.Linear(emb, o * 2) |
| self.n2 = nn.GroupNorm(8, o); self.c2 = conv(o, o) |
| self.skip = nn.Conv2d(i, o, 1) if i != o else nn.Identity() |
| def forward(self, x, e): |
| h = self.c1(F.silu(self.n1(x))) |
| sc, sh = self.emb(F.silu(e))[:, :, None, None].chunk(2, 1) |
| h = self.n2(h) * (1 + sc) + sh |
| h = self.c2(F.silu(h)) |
| return h + self.skip(x) |
|
|
| class AttnBlock(nn.Module): |
| def __init__(self, ch, ctx_dim, head_dim): |
| super().__init__() |
| heads = max(1, ch // head_dim) |
| self.norm = nn.GroupNorm(8, ch) |
| self.ln1, self.ln2, self.ln3 = nn.LayerNorm(ch), nn.LayerNorm(ch), nn.LayerNorm(ch) |
| self.sa = nn.MultiheadAttention(ch, heads, batch_first=True) |
| self.ca = nn.MultiheadAttention(ch, heads, kdim=ctx_dim, vdim=ctx_dim, batch_first=True) |
| self.ff = nn.Sequential(nn.Linear(ch, ch * 4), nn.GELU(), nn.Linear(ch * 4, ch)) |
| def forward(self, x, ctx): |
| B, C, H, W = x.shape |
| h = self.norm(x).flatten(2).transpose(1, 2) |
| q = self.ln1(h); h = h + self.sa(q, q, q, need_weights=False)[0] |
| q = self.ln2(h); h = h + self.ca(q, ctx, ctx, need_weights=False)[0] |
| h = h + self.ff(self.ln3(h)) |
| return x + h.transpose(1, 2).reshape(B, C, H, W) |
|
|
| class Down(nn.Module): |
| def __init__(self, c): super().__init__(); self.c = conv(c, c) |
| def forward(self, x): return self.c(F.avg_pool2d(x, 2)) |
|
|
| class Up(nn.Module): |
| def __init__(self, c): super().__init__(); self.c = conv(c, c) |
| def forward(self, x): return self.c(F.interpolate(x, scale_factor=2, mode="nearest")) |
|
|
| class UNet(nn.Module): |
| def __init__(self, cfg): |
| super().__init__() |
| ch, emb, mults = cfg.base_ch, cfg.base_ch * 4, cfg.ch_mult |
| self.ch = ch |
| mlp = lambda i: nn.Sequential(nn.Linear(i, emb), nn.SiLU(), nn.Linear(emb, emb)) |
| self.t_mlp, self.r_mlp, self.p_mlp = mlp(ch), mlp(ch), mlp(cfg.text_dim) |
| self.inc = conv(4, ch) |
|
|
| self.down = nn.ModuleList(); chans = [ch]; cur = ch |
| for li, m in enumerate(mults): |
| out = ch * m |
| for _ in range(cfg.n_res_blocks): |
| blk = nn.ModuleList([ResBlock(cur, out, emb)]) |
| if li in cfg.attn_levels: blk.append(AttnBlock(out, cfg.text_dim, cfg.head_dim)) |
| self.down.append(blk); cur = out; chans.append(cur) |
| if li != len(mults) - 1: |
| self.down.append(nn.ModuleList([Down(cur)])); chans.append(cur) |
|
|
| self.mid = nn.ModuleList([ResBlock(cur, cur, emb), AttnBlock(cur, cfg.text_dim, cfg.head_dim), ResBlock(cur, cur, emb)]) |
|
|
| self.up = nn.ModuleList() |
| for li, m in reversed(list(enumerate(mults))): |
| out = ch * m |
| for j in range(cfg.n_res_blocks + 1): |
| skip = chans.pop() |
| blk = nn.ModuleList([ResBlock(cur + skip, out, emb)]) |
| if li in cfg.attn_levels: blk.append(AttnBlock(out, cfg.text_dim, cfg.head_dim)) |
| cur = out |
| if li != 0 and j == cfg.n_res_blocks: blk.append(Up(cur)) |
| self.up.append(blk) |
|
|
| self.out_norm = nn.GroupNorm(8, cur) |
| self.out = conv(cur, 4) |
| nn.init.zeros_(self.out.weight); nn.init.zeros_(self.out.bias) |
|
|
| def run(self, blk, x, e, ctx): |
| for m in blk: |
| if isinstance(m, ResBlock): x = m(x, e) |
| elif isinstance(m, AttnBlock): x = m(x, ctx) |
| else: x = m(x) |
| return x |
|
|
| def forward(self, x, t, res, pooled, ctx): |
| e = (self.t_mlp(sinusoidal(t, self.ch)) |
| + self.r_mlp(sinusoidal(torch.log2(res.float()) * 100, self.ch)) |
| + self.p_mlp(pooled)) |
| h = self.inc(x); hs = [h] |
| for blk in self.down: |
| h = self.run(blk, h, e, ctx); hs.append(h) |
| h = self.run(self.mid, h, e, ctx) |
| for blk in self.up: |
| h = self.run(blk, torch.cat([h, hs.pop()], 1), e, ctx) |
| return self.out(F.silu(self.out_norm(h))) |
|
|
|
|
| def _tuplify(c): |
| for k in ("ch_mult", "attn_levels", "resolutions"): |
| if k in c: c[k] = tuple(c[k]) |
| return c |
|
|
| def _cosine_alpha_bar(T, s=0.008): |
| t = torch.linspace(0, T, T + 1, dtype=torch.float64) / T |
| f = torch.cos((t + s) / (1 + s) * math.pi / 2) ** 2 |
| ac = f / f[0] |
| betas = (1 - ac[1:] / ac[:-1]).clamp(0, 0.999) |
| return torch.cumprod(1 - betas, 0).float() |
|
|
|
|
| class TextureGenerator: |
| def __init__(self, repo_or_dir, device=None): |
| from safetensors.torch import load_file |
| from transformers import CLIPTokenizer, CLIPTextModel |
| self.device = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu")) |
| p = Path(repo_or_dir) |
| if p.is_dir(): |
| cfg_path, w_path = p / "config.json", p / "model.safetensors" |
| else: |
| from huggingface_hub import hf_hub_download |
| cfg_path = hf_hub_download(repo_or_dir, "config.json") |
| w_path = hf_hub_download(repo_or_dir, "model.safetensors") |
| self.cfg = SimpleNamespace(**_tuplify(json.load(open(cfg_path)))) |
| CFG.tileable = self.cfg.tileable |
| self.model = UNet(self.cfg) |
| self.model.load_state_dict(load_file(str(w_path))) |
| self.model.to(self.device).eval() |
| self.tok = CLIPTokenizer.from_pretrained(self.cfg.text_model) |
| self.tenc = CLIPTextModel.from_pretrained(self.cfg.text_model).to(self.device).eval().requires_grad_(False) |
| self.ac = _cosine_alpha_bar(self.cfg.timesteps).to(self.device) |
|
|
| @torch.no_grad() |
| def _encode(self, texts): |
| t = self.tok(list(texts), padding="max_length", truncation=True, |
| max_length=self.cfg.text_len, return_tensors="pt").to(self.device) |
| o = self.tenc(**t) |
| return o.last_hidden_state.float(), o.pooler_output.float() |
|
|
| @torch.no_grad() |
| def sample(self, prompt, res=16, n=8, steps=50, guidance=4.0, seed=None): |
| """Returns a (n, 4, res, res) tensor in [-1, 1] (RGBA).""" |
| assert res in self.cfg.resolutions, f"res must be one of {self.cfg.resolutions}" |
| dev, T = self.device, self.cfg.timesteps |
| g = torch.Generator(device=dev) |
| if seed is not None: g.manual_seed(int(seed)) |
| ctx, pooled = self._encode([prompt] * n) |
| nctx, npool = self._encode([""]) |
| ctx_all = torch.cat([ctx, nctx.expand(n, -1, -1)]); pool_all = torch.cat([pooled, npool.expand(n, -1)]) |
| resb = torch.full((2 * n,), res, device=dev) |
| x = torch.randn(n, 4, res, res, device=dev, generator=g) |
| ts = torch.linspace(T - 1, 0, steps, device=dev).long() |
| for i, t in enumerate(ts): |
| a_t = self.ac[t] |
| a_prev = self.ac[ts[i + 1]] if i + 1 < len(ts) else torch.tensor(1.0, device=dev) |
| out = self.model(torch.cat([x, x]), torch.full((2 * n,), int(t), device=dev), resb, pool_all, ctx_all) |
| e_c, e_u = out[:n], out[n:] |
| eps = e_u + guidance * (e_c - e_u) |
| x0 = ((x - (1 - a_t).sqrt() * eps) / a_t.sqrt()).clamp(-1, 1) |
| eps = (x - a_t.sqrt() * x0) / (1 - a_t).sqrt() |
| x = a_prev.sqrt() * x0 + (1 - a_prev).sqrt() * eps |
| return x |
|
|
| def __call__(self, prompt, res=16, n=8, steps=50, guidance=4.0, seed=None, opaque=False, snap_alpha=True): |
| x = self.sample(prompt, res, n, steps, guidance, seed) |
| if opaque: |
| x[:, 3] = 1.0 |
| elif snap_alpha: |
| x[:, 3] = torch.where(x[:, 3] > 0, 1.0, -1.0) |
| x[:, :3] = x[:, :3] * (x[:, 3:4] > 0) |
| a = ((x.clamp(-1, 1) + 1) * 127.5).round().byte().permute(0, 2, 3, 1).cpu().numpy() |
| return [Image.fromarray(i, "RGBA") for i in a] |
|
|
