Upload models/clip_vae.py with huggingface_hub

models/clip_vae.py

@@ -0,0 +1,116 @@
+"""
+CLIP-conditioned VAE.
+
+The encoder concatenates a projected CLIP text embedding with the CNN
+image features before the latent bottleneck.  The decoder is identical
+to the baseline VAE.
+
+CLIP is loaded once and kept frozen; only the linear projector and the
+rest of the VAE are trained.
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from models.vae import ConvDecoder
+
+CLIP_DIM = 512      # openai/clip-vit-base-patch32 text embedding size
+PROJ_DIM = 64       # projected text feature size
+
+
+class TextProjector(nn.Module):
+    def __init__(self, in_dim: int = CLIP_DIM, out_dim: int = PROJ_DIM):
+        super().__init__()
+        self.fc = nn.Linear(in_dim, out_dim)
+
+    def forward(self, text_emb: torch.Tensor) -> torch.Tensor:
+        return F.relu(self.fc(text_emb), inplace=True)
+
+
+class ClipCondEncoder(nn.Module):
+    def __init__(
+        self,
+        z_dim: int = 4,
+        img_channels: int = 3,
+        proj_dim: int = PROJ_DIM,
+    ):
+        super().__init__()
+        self.conv = nn.Sequential(
+            nn.Conv2d(img_channels, 32, 4, stride=2, padding=1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(32, 64, 4, stride=2, padding=1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(64, 128, 4, stride=2, padding=1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(128, 256, 4, stride=2, padding=1),
+            nn.ReLU(inplace=True),
+        )
+        self.text_proj = TextProjector(CLIP_DIM, proj_dim)
+        feat_dim = 256 * 4 * 4 + proj_dim
+        self.fc_mu = nn.Linear(feat_dim, z_dim)
+        self.fc_logvar = nn.Linear(feat_dim, z_dim)
+
+    def forward(self, x: torch.Tensor, text_emb: torch.Tensor):
+        img_feat = self.conv(x).flatten(1)
+        txt_feat = self.text_proj(text_emb)
+        h = torch.cat([img_feat, txt_feat], dim=1)
+        return self.fc_mu(h), self.fc_logvar(h)
+
+
+class ClipVAE(nn.Module):
+    def __init__(self, z_dim: int = 4, img_channels: int = 3):
+        super().__init__()
+        self.encoder = ClipCondEncoder(z_dim, img_channels)
+        self.decoder = ConvDecoder(z_dim, img_channels)
+        self.z_dim = z_dim
+
+    def reparameterize(self, mu: torch.Tensor, logvar: torch.Tensor) -> torch.Tensor:
+        if self.training:
+            std = (0.5 * logvar).exp()
+            return mu + std * torch.randn_like(std)
+        return mu
+
+    def encode(self, x: torch.Tensor, text_emb: torch.Tensor):
+        mu, logvar = self.encoder(x, text_emb)
+        z = self.reparameterize(mu, logvar)
+        return z, mu, logvar
+
+    def decode(self, z: torch.Tensor) -> torch.Tensor:
+        return self.decoder(z)
+
+    def forward(self, x: torch.Tensor, text_emb: torch.Tensor):
+        z, mu, logvar = self.encode(x, text_emb)
+        recon = self.decode(z)
+        return recon, mu, logvar
+
+
+# ------------------------------------------------------------------
+# CLIP helper — load once, freeze, cache text embeddings
+# ------------------------------------------------------------------
+
+_clip_model = None
+_clip_tokenizer = None
+
+
+def get_clip(device: torch.device):
+    global _clip_model, _clip_tokenizer
+    if _clip_model is None:
+        from transformers import CLIPModel, CLIPTokenizer
+        _clip_tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-base-patch32")
+        _clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
+        _clip_model.eval()
+        for p in _clip_model.parameters():
+            p.requires_grad_(False)
+    return _clip_model.to(device), _clip_tokenizer
+
+
+@torch.no_grad()
+def encode_text(texts: list, device: torch.device) -> torch.Tensor:
+    model, tokenizer = get_clip(device)
+    tokens = tokenizer(texts, padding=True, return_tensors="pt").to(device)
+    out = model.get_text_features(**tokens)
+    # get_text_features returns a plain tensor in newer transformers
+    emb = out if isinstance(out, torch.Tensor) else out.pooler_output
+    emb = emb / emb.norm(dim=-1, keepdim=True)     # L2-normalise
+    return emb.float()