pixeldit/text_encoder_qwen.py

"""
Qwen3-2B text encoder for PixelDiT.
Requires a trained projection (train_qwen_proj.py) to map 2048→2304.

Usage:
    from pixeldit.text_encoder_qwen import QwenEncoder
    enc  = QwenEncoder(proj_path="pixeldit/qwen_proj.pt")
    cond = enc.encode(["a dragon at sunset"])  # [1, 300, 2304]
    null = enc.encode_null(1)                  # [1, 300, 2304]
"""

import torch
import torch.nn as nn
from transformers import AutoTokenizer, AutoModel

_QWEN_ID   = "Qwen/Qwen3-2B"
_QWEN_DIM  = 2048
_GEMMA_DIM = 2304
_TXT_MAX   = 300


class QwenEncoder:
    def __init__(
        self,
        model_id=_QWEN_ID,
        proj_path=None,           # path to trained qwen_proj.pt
        output_device="cuda",
        output_dtype=torch.bfloat16,
    ):
        self.output_device = torch.device(output_device)
        self.output_dtype  = output_dtype

        print(f"[QwenEncoder] loading {model_id} (CPU)")
        self.tokenizer = AutoTokenizer.from_pretrained(model_id)
        self.tokenizer.padding_side = "right"
        self._model = AutoModel.from_pretrained(model_id, torch_dtype=torch.float32).eval()

        self.proj = nn.Linear(_QWEN_DIM, _GEMMA_DIM, bias=False)
        if proj_path:
            sd = torch.load(proj_path, map_location="cpu", weights_only=True)
            self.proj.load_state_dict(sd)
            print(f"[QwenEncoder] loaded projection: {proj_path}")
        else:
            with torch.no_grad():
                w = torch.zeros(_GEMMA_DIM, _QWEN_DIM)
                w[:_QWEN_DIM] = torch.eye(_QWEN_DIM)
                self.proj.weight.copy_(w)
            print("[QwenEncoder] projection: identity init — run train_qwen_proj.py for real quality")
        self.proj = self.proj.to(self.output_device).to(output_dtype)
        print("[QwenEncoder] ready")

    @torch.no_grad()
    def encode(self, texts: list[str]) -> torch.Tensor:
        """Returns [B, 300, 2304]."""
        tok = self.tokenizer(
            texts, max_length=_TXT_MAX,
            padding="max_length", truncation=True, return_tensors="pt",
        )
        emb = self._model(**tok).last_hidden_state
        emb = emb.to(self.output_device).to(self.output_dtype)
        return self.proj(emb)

    @torch.no_grad()
    def encode_null(self, batch_size: int) -> torch.Tensor:
        """Returns [B, 300, 2304] for empty string (CFG unconditional)."""
        tok = self.tokenizer(
            [""] * batch_size, max_length=_TXT_MAX,
            padding="max_length", truncation=True, return_tensors="pt",
        )
        emb = self._model(**tok).last_hidden_state
        emb = emb.to(self.output_device).to(self.output_dtype)
        return self.proj(emb)