| 1 | import functools |
| 2 | from pathlib import Path |
| 3 | |
| 4 | import torch |
| 5 | from transformers import AutoImageProcessor, Gemma3ForConditionalGeneration, Gemma3Processor |
| 6 | |
| 7 | from ltx_core.loader.module_ops import ModuleOps |
| 8 | from ltx_core.text_encoders.gemma.tokenizer import LTXVGemmaTokenizer |
| 9 | from ltx_core.utils import find_matching_file |
| 10 | |
| 11 | |
| 12 | class GemmaTextEncoder(torch.nn.Module): |
| 13 | """Pure Gemma text encoder — runs the LLM and returns raw hidden states. |
| 14 | Prompt enhancement (generate) is also supported since the full |
| 15 | Gemma3ForConditionalGeneration model (including lm_head) is loaded. |
| 16 | """ |
| 17 | |
| 18 | def __init__( |
| 19 | self, |
| 20 | model: Gemma3ForConditionalGeneration | None = None, |
| 21 | tokenizer: LTXVGemmaTokenizer | None = None, |
| 22 | processor: Gemma3Processor | None = None, |
| 23 | dtype: torch.dtype = torch.bfloat16, |
| 24 | ): |
| 25 | super().__init__() |
| 26 | self.model = model |
| 27 | self.tokenizer = tokenizer |
| 28 | self.processor = processor |
| 29 | self._dtype = dtype |
| 30 | |
| 31 | def encode( |
| 32 | self, |
| 33 | text: str, |
| 34 | padding_side: str = "left", # noqa: ARG002 |
| 35 | ) -> tuple[tuple[torch.Tensor, ...], torch.Tensor]: |
| 36 | """Run Gemma LLM and return raw hidden states + attention mask. |
| 37 | Calls the inner model (self.model.model) to skip lm_head logits computation (~500 MiB saving). |
| 38 | Returns: |
| 39 | (hidden_states, attention_mask) where hidden_states is a tuple of per-layer tensors. |
| 40 | """ |
| 41 | token_pairs = self.tokenizer.tokenize_with_weights(text)["gemma"] |
| 42 | input_ids = torch.tensor([[t[0] for t in token_pairs]], device=self.model.device) |
| 43 | attention_mask = torch.tensor([[w[1] for w in token_pairs]], device=self.model.device) |
| 44 | outputs = self.model.model(input_ids=input_ids, attention_mask=attention_mask, output_hidden_states=True) |
| 45 | hidden_states = outputs.hidden_states |
| 46 | del outputs |
| 47 | return hidden_states, attention_mask |
| 48 | |
| 49 | # --- Prompt enhancement methods --- |
| 50 | |
| 51 | def _enhance( |
| 52 | self, |
| 53 | messages: list[dict[str, str]], |
| 54 | image: torch.Tensor | None = None, |
| 55 | max_new_tokens: int = 512, |
| 56 | seed: int = 10, |
| 57 | ) -> str: |
| 58 | text = self.processor.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) |
| 59 | |
| 60 | model_inputs = self.processor( |
| 61 | text=text, |
| 62 | images=image, |
| 63 | return_tensors="pt", |
| 64 | ).to(self.model.device) |
| 65 | pad_token_id = self.processor.tokenizer.pad_token_id if self.processor.tokenizer.pad_token_id is not None else 0 |
| 66 | model_inputs = _pad_inputs_for_attention_alignment(model_inputs, pad_token_id=pad_token_id) |
| 67 | |
| 68 | with torch.inference_mode(), torch.random.fork_rng(devices=[self.model.device]): |
| 69 | torch.manual_seed(seed) |
| 70 | outputs = self.model.generate( |
| 71 | **model_inputs, |
| 72 | max_new_tokens=max_new_tokens, |
| 73 | do_sample=True, |
| 74 | temperature=0.7, |
| 75 | ) |
| 76 | generated_ids = outputs[0][len(model_inputs.input_ids[0]) :] |
| 77 | enhanced_prompt = self.processor.tokenizer.decode(generated_ids, skip_special_tokens=True) |
| 78 | |
| 79 | return enhanced_prompt |
| 80 | |
| 81 | def enhance_t2v( |
| 82 | self, |
| 83 | prompt: str, |
| 84 | max_new_tokens: int = 512, |
| 85 | system_prompt: str | None = None, |
| 86 | seed: int = 10, |
| 87 | ) -> str: |
| 88 | """Enhance a text prompt for T2V generation.""" |
| 89 | system_prompt = system_prompt or self.default_gemma_t2v_system_prompt |
| 90 | |
| 91 | messages = [ |
| 92 | {"role": "system", "content": system_prompt}, |
| 93 | {"role": "user", "content": f"user prompt: {prompt}"}, |
| 94 | ] |
| 95 | |
| 96 | return self._enhance(messages, max_new_tokens=max_new_tokens, seed=seed) |
| 97 | |
| 98 | def enhance_i2v( |
| 99 | self, |
| 100 | prompt: str, |
| 101 | image: torch.Tensor, |
| 102 | max_new_tokens: int = 512, |
| 103 | system_prompt: str | None = None, |
| 104 | seed: int = 10, |
| 105 | ) -> str: |
| 106 | """Enhance a text prompt for I2V generation using a reference image.""" |
| 107 | system_prompt = system_prompt or self.default_gemma_i2v_system_prompt |
| 108 | messages = [ |
| 109 | {"role": "system", "content": system_prompt}, |
| 110 | { |
| 111 | "role": "user", |
| 112 | "content": [ |
| 113 | {"type": "image"}, |
| 114 | {"type": "text", "text": f"User Raw Input Prompt: {prompt}."}, |
| 115 | ], |
| 116 | }, |
| 117 | ] |
| 118 | return self._enhance(messages, image=image, max_new_tokens=max_new_tokens, seed=seed) |
| 119 | |
| 120 | @functools.cached_property |
| 121 | def default_gemma_i2v_system_prompt(self) -> str: |
| 122 | return _load_system_prompt("gemma_i2v_system_prompt.txt") |
| 123 | |
| 124 | @functools.cached_property |
| 125 | def default_gemma_t2v_system_prompt(self) -> str: |
| 126 | return _load_system_prompt("gemma_t2v_system_prompt.txt") |
| 127 | |
| 128 | |
| 129 | # --- Standalone utility functions --- |
| 130 | |
| 131 | |
| 132 | @functools.lru_cache(maxsize=2) |
| 133 | def _load_system_prompt(prompt_name: str) -> str: |
| 134 | with open(Path(__file__).parent / "prompts" / f"{prompt_name}", "r") as f: |
| 135 | return f.read() |
| 136 | |
| 137 | |
| 138 | def _cat_with_padding( |
| 139 | tensor: torch.Tensor, |
| 140 | padding_length: int, |
| 141 | value: int | float, |
| 142 | ) -> torch.Tensor: |
| 143 | """Concatenate a tensor with a padding tensor of the given value.""" |
| 144 | return torch.cat( |
| 145 | [ |
| 146 | tensor, |
| 147 | torch.full( |
| 148 | (1, padding_length), |
| 149 | value, |
| 150 | dtype=tensor.dtype, |
| 151 | device=tensor.device, |
| 152 | ), |
| 153 | ], |
| 154 | dim=1, |
| 155 | ) |
| 156 | |
| 157 | |
| 158 | def _pad_inputs_for_attention_alignment( |
| 159 | model_inputs: dict[str, torch.Tensor], |
| 160 | pad_token_id: int = 0, |
| 161 | alignment: int = 8, |
| 162 | ) -> dict[str, torch.Tensor]: |
| 163 | """Pad sequence length to multiple of alignment for Flash Attention compatibility.""" |
| 164 | seq_len = model_inputs.input_ids.shape[1] |
| 165 | padded_len = ((seq_len + alignment - 1) // alignment) * alignment |
| 166 | padding_length = padded_len - seq_len |
| 167 | |
| 168 | if padding_length > 0: |
| 169 | model_inputs["input_ids"] = _cat_with_padding(model_inputs.input_ids, padding_length, pad_token_id) |
| 170 | model_inputs["attention_mask"] = _cat_with_padding(model_inputs.attention_mask, padding_length, 0) |
| 171 | if "token_type_ids" in model_inputs and model_inputs["token_type_ids"] is not None: |
| 172 | model_inputs["token_type_ids"] = _cat_with_padding(model_inputs["token_type_ids"], padding_length, 0) |
| 173 | |
| 174 | return model_inputs |
| 175 | |
| 176 | |
| 177 | def module_ops_from_gemma_root(gemma_root: str) -> tuple[ModuleOps, ...]: |
| 178 | tokenizer_root = str(find_matching_file(gemma_root, "tokenizer.model").parent) |
| 179 | processor_root = str(find_matching_file(gemma_root, "preprocessor_config.json").parent) |
| 180 | |
| 181 | def load_tokenizer(module: GemmaTextEncoder) -> GemmaTextEncoder: |
| 182 | module.tokenizer = LTXVGemmaTokenizer(tokenizer_root, 1024) |
| 183 | return module |
| 184 | |
| 185 | def load_processor(module: GemmaTextEncoder) -> GemmaTextEncoder: |
| 186 | image_processor = AutoImageProcessor.from_pretrained(processor_root, local_files_only=True) |
| 187 | if not module.tokenizer: |
| 188 | raise ValueError("Tokenizer model operation must be performed before processor model operation") |
| 189 | module.processor = Gemma3Processor(image_processor=image_processor, tokenizer=module.tokenizer.tokenizer) |
| 190 | return module |
| 191 | |
| 192 | tokenizer_load_ops = ModuleOps( |
| 193 | "TokenizerLoad", |
| 194 | matcher=lambda module: isinstance(module, GemmaTextEncoder) and module.tokenizer is None, |
| 195 | mutator=load_tokenizer, |
| 196 | ) |
| 197 | processor_load_ops = ModuleOps( |
| 198 | "ProcessorLoad", |
| 199 | matcher=lambda module: isinstance(module, GemmaTextEncoder) and module.processor is None, |
| 200 | mutator=load_processor, |
| 201 | ) |
| 202 | return (tokenizer_load_ops, processor_load_ops) |
| 203 |