"""
Frozen CLIP text encoder for prompt conditioning. We reuse a pretrained encoder
rather than training our own — the "terrible" in T.I.T.S. lives entirely in the
image model, not the text understanding, otherwise prompts would just be noise
and you'd never be able to tell if a bad image was a bad *rendering* of your
prompt or the model just not understanding the prompt at all.
"""
from typing import List

import torch
import torch.nn as nn
from transformers import CLIPTokenizer, CLIPTextModelWithProjection


class FrozenCLIPTextEncoder(nn.Module):
    def __init__(self, model_name: str = "openai/clip-vit-base-patch32", device: str = "cuda"):
        super().__init__()
        self.tokenizer = CLIPTokenizer.from_pretrained(model_name)
        self.model = CLIPTextModelWithProjection.from_pretrained(model_name)
        self.model.to(device)
        self.model.eval()
        for p in self.model.parameters():
            p.requires_grad = False
        self.device = device
        self.embed_dim = self.model.config.projection_dim  # 512 for base-patch32

    @torch.no_grad()
    def encode(self, prompts: List[str]) -> torch.Tensor:
        tokens = self.tokenizer(
            prompts, padding=True, truncation=True, max_length=77, return_tensors="pt"
        ).to(self.device)
        out = self.model(**tokens)
        return out.text_embeds  # (batch, embed_dim), already pooled/projected