"""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) # read by conv(); set from config.json on load 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]