Spaces:
Sleeping
Sleeping
| """ | |
| 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") | |
| 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) | |
| 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) | |