Spaces:
Sleeping
Sleeping
File size: 4,090 Bytes
c85ad6e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 | """
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
_CHI_PROMPT = "\n".join([
'Given a user prompt, generate an "Enhanced prompt" that provides detailed visual descriptions suitable for image generation. Evaluate the level of detail in the user prompt:',
'- If the prompt is simple, focus on adding specifics about colors, shapes, sizes, textures, and spatial relationships to create vivid and concrete scenes.',
'- If the prompt is already detailed, refine and enhance the existing details slightly without overcomplicating.',
'Here are examples of how to transform or refine prompts:',
'- User Prompt: A cat sleeping -> Enhanced: A small, fluffy white cat curled up in a round shape, sleeping peacefully on a warm sunny windowsill, surrounded by pots of blooming red flowers.',
'- User Prompt: A busy city street -> Enhanced: A bustling city street scene at dusk, featuring glowing street lamps, a diverse crowd of people in colorful clothing, and a double-decker bus passing by towering glass skyscrapers.',
'Please generate only the enhanced description for the prompt below and avoid including any additional commentary or evaluations:',
'User Prompt: ',
])
_SELECT_IDX = [0] + list(range(-(_TXT_MAX - 1), 0))
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._num_chi_tokens = len(self.tokenizer.encode(_CHI_PROMPT))
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]."""
texts_full = [_CHI_PROMPT + t for t in texts]
max_len = self._num_chi_tokens + _TXT_MAX - 2
tok = self.tokenizer(
texts_full, max_length=max_len,
padding="max_length", truncation=True, return_tensors="pt",
)
emb = self._model(**tok).last_hidden_state
emb = emb[:, _SELECT_IDX, :]
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)
|