Delete pixeldit/text_encoder_gemma.py with huggingface_hub

pixeldit/text_encoder_gemma.py

@@ -1,79 +0,0 @@
-"""
-Gemma-2-2B text encoder for PixelDiT.
-Handles chi_prompt prefix + select_index to match training exactly.
-
-Usage:
-    from pixeldit.text_encoder_gemma import GemmaEncoder
-    enc  = GemmaEncoder()
-    cond = enc.encode(["a dragon at sunset"])  # [1, 300, 2304]
-    null = enc.encode_null(1)                  # [1, 300, 2304]
-"""
-
-import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
-
-_GEMMA_ID  = "Efficient-Large-Model/gemma-2-2b-it"
-_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 GemmaEncoder:
-    def __init__(
-        self,
-        model_id=_GEMMA_ID,
-        output_device="cuda",
-        output_dtype=torch.bfloat16,
-    ):
-        self.output_device = torch.device(output_device)
-        self.output_dtype  = output_dtype
-
-        print(f"[GemmaEncoder] loading {model_id} (CPU)")
-        self.tokenizer = AutoTokenizer.from_pretrained(model_id)
-        self.tokenizer.padding_side = "right"
-        self._model = (
-            AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float32)
-            .get_decoder().eval()
-        )
-        self._num_chi_tokens = len(self.tokenizer.encode(_CHI_PROMPT))
-        print("[GemmaEncoder] 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(
-            input_ids=tok.input_ids,
-            attention_mask=tok.attention_mask,
-        ).last_hidden_state
-        emb = emb[:, _SELECT_IDX, :]
-        return emb.to(self.output_device).to(self.output_dtype)
-
-    @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(
-            input_ids=tok.input_ids,
-            attention_mask=tok.attention_mask,
-        ).last_hidden_state
-        return emb.to(self.output_device).to(self.output_dtype)