| """ |
| ArtiGen V1.0 — Main Model |
| CARTEL backbone with PHI-SCAN, AdaLN conditioning, ASDL heads. |
| """ |
| import math |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| try: |
| from .cartel_block import CARTELBlock |
| from .asdl_head import StyleHead, ContentHead, ConceptHead, MoodHead, CompositionHead |
| from .phi_scan import build_scan_permutations, apply_scan, unscan, get_scan_pattern |
| except ImportError: |
| from cartel_block import CARTELBlock |
| from asdl_head import StyleHead, ContentHead, ConceptHead, MoodHead, CompositionHead |
| from phi_scan import build_scan_permutations, apply_scan, unscan, get_scan_pattern |
|
|
| class PatchEmbed(nn.Module): |
| def __init__(self, in_ch, embed_dim, patch_size=2): |
| super().__init__() |
| self.proj = nn.Conv2d(in_ch, embed_dim, kernel_size=patch_size, stride=patch_size) |
| self.norm = nn.LayerNorm(embed_dim) |
| def forward(self, x): |
| x = self.proj(x) |
| B, C, H, W = x.shape |
| x = x.permute(0, 2, 3, 1).reshape(B, H * W, C) |
| return self.norm(x), H, W |
|
|
| class AdaLN(nn.Module): |
| def __init__(self, dim, cond_dim=512): |
| super().__init__() |
| self.modulation = nn.Sequential( |
| nn.SiLU(), |
| nn.Linear(cond_dim, dim * 2), |
| ) |
| def forward(self, x, cond): |
| scale, shift = self.modulation(cond).chunk(2, dim=-1) |
| return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1) |
|
|
| class ArtiGen(nn.Module): |
| def __init__( |
| self, |
| latent_ch=4, |
| latent_h=32, |
| latent_w=32, |
| embed_dim=256, |
| num_layers=12, |
| d_state=16, |
| expand=2, |
| text_dim=768, |
| style_classes=128, |
| content_objects=1024, |
| mood_classes=64, |
| ): |
| super().__init__() |
| self.embed_dim = embed_dim |
| self.num_layers = num_layers |
| self.latent_h = latent_h |
| self.latent_w = latent_w |
| self.patch_embed = PatchEmbed(latent_ch, embed_dim, patch_size=1) |
| self.t_embed = nn.Sequential( |
| nn.Linear(1, text_dim), |
| nn.SiLU(), |
| nn.Linear(text_dim, text_dim), |
| ) |
| self.cond_proj = nn.Linear(text_dim, text_dim) |
| self.cond_transform = nn.Sequential( |
| nn.SiLU(), |
| nn.Linear(text_dim, text_dim), |
| ) |
| self.token_pos = nn.Parameter(torch.randn(1, latent_h * latent_w, embed_dim) * 0.02) |
| self.scans = build_scan_permutations(latent_h, latent_w) |
| self.blocks = nn.ModuleList([ |
| CARTELBlock(embed_dim, d_state=d_state, expand=expand) |
| for _ in range(num_layers) |
| ]) |
| self.adalns = nn.ModuleList([ |
| AdaLN(embed_dim, cond_dim=text_dim) |
| for _ in range(num_layers) |
| ]) |
| self.skip_connect = nn.Sequential( |
| nn.Linear(embed_dim, embed_dim), |
| nn.SiLU(), |
| nn.Linear(embed_dim, embed_dim), |
| ) |
| self.final_proj = nn.Sequential( |
| nn.LayerNorm(embed_dim), |
| nn.Linear(embed_dim, embed_dim * 4), |
| nn.SiLU(), |
| nn.Linear(embed_dim * 4, embed_dim), |
| nn.Linear(embed_dim, latent_ch), |
| ) |
| self.style_head = StyleHead(embed_dim, num_style_classes=style_classes) |
| self.content_head = ContentHead(embed_dim, num_objects=content_objects) |
| self.concept_head = ConceptHead(embed_dim) |
| self.mood_head = MoodHead(embed_dim, num_moods=mood_classes) |
| self.comp_head = CompositionHead(embed_dim) |
| self.apply(self._init_weights) |
|
|
| def _init_weights(self, m): |
| if isinstance(m, nn.Linear): |
| nn.init.xavier_uniform_(m.weight) |
| if m.bias is not None: |
| nn.init.zeros_(m.bias) |
|
|
| def forward(self, z_t, t, text_embed, return_asdl=False): |
| B = z_t.shape[0] |
| x, H, W = self.patch_embed(z_t) |
| x = x + self.token_pos[:, :x.shape[1], :] |
| t_emb = self.t_embed(t.view(B, 1).float()) |
| cond = self.cond_proj(text_embed) + t_emb |
| cond = self.cond_transform(cond) |
| x_shallow = x |
| for i, (block, adaln) in enumerate(zip(self.blocks, self.adalns)): |
| x = adaln(x, cond) |
| scan_name = get_scan_pattern(i) |
| perm, inv = self.scans[scan_name] |
| x_scanned = apply_scan(x, perm) |
| x_scanned = block(x_scanned) |
| x = unscan(x_scanned, inv) |
| if i == self.num_layers // 4: |
| x_shallow = x |
| x = x + self.skip_connect(x_shallow) |
| v = self.final_proj(x).transpose(1, 2).reshape(B, -1, H, W) |
| asdl = {} |
| s, s_logits = self.style_head(x) |
| c, c_logits = self.content_head(x) |
| n = self.concept_head(x) |
| m, m_logits = self.mood_head(x) |
| p = self.comp_head(x) |
| asdl = { |
| "style_vec": s, "style_logits": s_logits, |
| "content_vec": c, "content_logits": c_logits, |
| "concept_vec": n, |
| "mood_vec": m, "mood_logits": m_logits, |
| "comp_vec": p, |
| } |
| if return_asdl: |
| return v, asdl |
| return v, None |
|
|
