| """Chitra Tiny: a 212M-parameter text-to-image model (256x256) by UdgamLabs. Standalone inference code (Apache-2.0). |
| |
| from chitra_tiny import ChitraTiny |
| pipe = ChitraTiny.from_pretrained("UdgamLabs/Chitra-Tiny") # or a local folder |
| images = pipe("a lighthouse on a rocky coast at sunset", n=4) # list of PIL images |
| |
| Model: diffusion transformer (patch 2 on 32x32x4 latents) trained from scratch with rectified flow. Text encoder: |
| google/flan-t5-base (Apache-2.0). Image decoder: stabilityai/sd-vae-ft-mse (MIT). Both are downloaded on first use. |
| """ |
| import json |
| import math |
| import os |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
|
|
| |
| def _sincos_2d(d, n): |
| y, x = torch.meshgrid(torch.arange(n, dtype=torch.float32), torch.arange(n, dtype=torch.float32), indexing="ij") |
| q = d // 4 |
| f = torch.exp(-math.log(10000.0) * torch.arange(q, dtype=torch.float32) / q) |
| return torch.cat([torch.sin(x.flatten()[:, None] * f), torch.cos(x.flatten()[:, None] * f), |
| torch.sin(y.flatten()[:, None] * f), torch.cos(y.flatten()[:, None] * f)], dim=1) |
|
|
|
|
| def _timestep_embedding(t, dim=256): |
| half = dim // 2 |
| f = torch.exp(-math.log(10000.0) * torch.arange(half, device=t.device, dtype=torch.float32) / half) |
| a = t.float()[:, None] * 1000.0 * f[None] |
| return torch.cat([a.cos(), a.sin()], dim=-1) |
|
|
|
|
| def _modulate(x, shift, scale): |
| return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1) |
|
|
|
|
| class _Attention(nn.Module): |
| def __init__(self, d, h, cross=False): |
| super().__init__() |
| self.h, self.cross = h, cross |
| if cross: |
| self.q = nn.Linear(d, d); self.kv = nn.Linear(d, 2 * d) |
| else: |
| self.qkv = nn.Linear(d, 3 * d) |
| self.q_norm = nn.RMSNorm(d // h, eps=1e-6); self.k_norm = nn.RMSNorm(d // h, eps=1e-6) |
| self.proj = nn.Linear(d, d) |
|
|
| def forward(self, x, ctx=None, mask=None): |
| B, N, D = x.shape |
| if self.cross: |
| q = self.q(x); k, v = self.kv(ctx).chunk(2, dim=-1) |
| else: |
| q, k, v = self.qkv(x).chunk(3, dim=-1) |
| sh = lambda t: t.view(B, -1, self.h, D // self.h).transpose(1, 2) |
| q, k, v = self.q_norm(sh(q)), self.k_norm(sh(k)), sh(v) |
| o = F.scaled_dot_product_attention(q, k, v, attn_mask=None if mask is None else mask[:, None, None, :]) |
| return self.proj(o.transpose(1, 2).reshape(B, N, D)) |
|
|
|
|
| class _Block(nn.Module): |
| def __init__(self, d, h, mlp_ratio): |
| super().__init__() |
| self.n1 = nn.LayerNorm(d, elementwise_affine=False, eps=1e-6); self.attn = _Attention(d, h) |
| self.n2 = nn.LayerNorm(d, elementwise_affine=True, eps=1e-6); self.xattn = _Attention(d, h, cross=True) |
| self.n3 = nn.LayerNorm(d, elementwise_affine=False, eps=1e-6) |
| hid = int(d * mlp_ratio) |
| self.mlp = nn.Sequential(nn.Linear(d, hid), nn.GELU(approximate="tanh"), nn.Linear(hid, d)) |
| self.ada = nn.Linear(d, 6 * d) |
|
|
| def forward(self, x, c, cmask, temb): |
| sh1, sc1, g1, sh2, sc2, g2 = self.ada(F.silu(temb)).chunk(6, dim=-1) |
| x = x + g1.unsqueeze(1) * self.attn(_modulate(self.n1(x), sh1, sc1)) |
| x = x + self.xattn(self.n2(x), c, cmask) |
| return x + g2.unsqueeze(1) * self.mlp(_modulate(self.n3(x), sh2, sc2)) |
|
|
|
|
| class ChitraTinyDiT(nn.Module): |
| """Predicts the flow velocity v = noise - x0 for latents x_t = (1 - t) x0 + t noise (t=1 is pure noise).""" |
|
|
| def __init__(self, width=768, depth=16, heads=12, patch=2, in_ch=4, grid=32, txt_dim=768, mlp_ratio=4.0, **_): |
| super().__init__() |
| self.patch, self.in_ch, self.n_side = patch, in_ch, grid // patch |
| d = width |
| self.patch_in = nn.Linear(in_ch * patch * patch, d) |
| self.register_buffer("pos", _sincos_2d(d, self.n_side)[None], persistent=False) |
| self.txt_in = nn.Sequential(nn.Linear(txt_dim, d), nn.GELU(approximate="tanh"), nn.Linear(d, d)) |
| self.t_in = nn.Sequential(nn.Linear(256, d), nn.SiLU(), nn.Linear(d, d)) |
| self.pool_in = nn.Linear(d, d) |
| self.blocks = nn.ModuleList(_Block(d, heads, mlp_ratio) for _ in range(depth)) |
| self.norm_out = nn.LayerNorm(d, elementwise_affine=False, eps=1e-6) |
| self.ada_out = nn.Linear(d, 2 * d) |
| self.out = nn.Linear(d, in_ch * patch * patch) |
|
|
| def forward(self, x, t, txt, txt_mask): |
| B, C, H, W = x.shape; p = self.patch |
| h = x.view(B, C, H // p, p, W // p, p).permute(0, 2, 4, 1, 3, 5).reshape(B, (H // p) * (W // p), C * p * p) |
| h = self.patch_in(h) + self.pos |
| c = self.txt_in(txt) |
| m = txt_mask.to(c.dtype).unsqueeze(-1) |
| temb = self.t_in(_timestep_embedding(t)) + self.pool_in((c * m).sum(1) / m.sum(1).clamp(min=1.0)) |
| for blk in self.blocks: |
| h = blk(h, c, txt_mask, temb) |
| sh, sc = self.ada_out(F.silu(temb)).chunk(2, dim=-1) |
| v = self.out(_modulate(self.norm_out(h), sh, sc)) |
| n = self.n_side |
| return v.view(B, n, n, C, p, p).permute(0, 3, 1, 4, 2, 5).reshape(B, C, n * p, n * p) |
|
|
|
|
| |
| class ChitraTiny: |
| def __init__(self, model, config, device): |
| from diffusers import AutoencoderKL |
| from transformers import AutoTokenizer, T5EncoderModel |
| self.model, self.cfg, self.device = model, config, device |
| self.tok = AutoTokenizer.from_pretrained(config["text_encoder"]) |
| self.te = T5EncoderModel.from_pretrained(config["text_encoder"]).to(device).eval() |
| self.vae = AutoencoderKL.from_pretrained(config["vae"]).to(device).eval() |
|
|
| @classmethod |
| def from_pretrained(cls, path_or_repo="UdgamLabs/Chitra-Tiny", device=None): |
| from safetensors.torch import load_file |
| if not os.path.isdir(path_or_repo): |
| from huggingface_hub import snapshot_download |
| path_or_repo = snapshot_download(path_or_repo, allow_patterns=["config.json", "model.safetensors"]) |
| device = device or ("cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu") |
| cfg = json.load(open(os.path.join(path_or_repo, "config.json"))) |
| model = ChitraTinyDiT(**cfg["model"]) |
| model.load_state_dict(load_file(os.path.join(path_or_repo, "model.safetensors"))) |
| return cls(model.to(device).eval(), cfg, device) |
|
|
| @torch.no_grad() |
| def _encode(self, texts): |
| e = self.tok(texts, max_length=self.cfg["max_text_tokens"], truncation=True, padding="max_length", return_tensors="pt") |
| ids, mask = e.input_ids.to(self.device), e.attention_mask.to(self.device) |
| return self.te(input_ids=ids, attention_mask=mask).last_hidden_state.float(), mask.bool() |
|
|
| @torch.no_grad() |
| def __call__(self, prompt, n=1, guidance=None, steps=None, seed=0, negative_prompt=""): |
| """Returns a list of n PIL images (256x256). guidance: classifier-free guidance scale (default from config).""" |
| from PIL import Image |
| guidance = self.cfg["default_guidance"] if guidance is None else guidance |
| steps = self.cfg["default_steps"] if steps is None else steps |
| txt, m = self._encode([prompt] * n); nt, nm = self._encode([negative_prompt] * n) |
| g = self.cfg["latent_grid"] |
| x = torch.randn((n, self.cfg["latent_channels"], g, g), generator=torch.Generator().manual_seed(seed)).to(self.device) |
| ts = torch.linspace(1.0, 0.0, steps + 1, device=self.device) |
| for i in range(steps): |
| v = self.model(torch.cat([x, x]), ts[i].expand(2 * n), torch.cat([txt, nt]), torch.cat([m, nm])) |
| vc, vu = v.chunk(2) |
| x = x + (ts[i + 1] - ts[i]) * (vu + guidance * (vc - vu)) |
| img = self.vae.decode(x / self.cfg["vae_scale"]).sample.clamp(-1, 1) |
| img = (img * 127.5 + 127.5).round().byte().permute(0, 2, 3, 1).cpu().numpy() |
| return [Image.fromarray(a) for a in img] |
|
|
|
|
| if __name__ == "__main__": |
| import argparse |
| ap = argparse.ArgumentParser(description="Generate images with Chitra Tiny") |
| ap.add_argument("prompt"); ap.add_argument("--n", type=int, default=4); ap.add_argument("--guidance", type=float) |
| ap.add_argument("--steps", type=int); ap.add_argument("--seed", type=int, default=0) |
| ap.add_argument("--model", default="UdgamLabs/Chitra-Tiny"); ap.add_argument("--out", default="chitra") |
| a = ap.parse_args() |
| pipe = ChitraTiny.from_pretrained(a.model) |
| for i, im in enumerate(pipe(a.prompt, n=a.n, guidance=a.guidance, steps=a.steps, seed=a.seed)): |
| im.save(f"{a.out}_{i}.png"); print(f"{a.out}_{i}.png") |
|
|
