import torch import torch.nn.functional as F import numpy as np import torch.nn as nn def tokenize_prompt(tokenizer, prompt, max_sequence_length): text_inputs = tokenizer( prompt, padding="max_length", max_length=max_sequence_length, truncation=True, return_length=False, return_overflowing_tokens=False, return_tensors="pt", ) text_input_ids = text_inputs.input_ids return text_input_ids def _encode_prompt_with_t5( text_encoder, tokenizer, max_sequence_length=512, prompt=None, num_images_per_prompt=1, device=None, text_input_ids=None, ): prompt = [prompt] if isinstance(prompt, str) else prompt batch_size = len(prompt) if tokenizer is not None: text_inputs = tokenizer( prompt, padding="max_length", max_length=max_sequence_length, truncation=True, return_length=False, return_overflowing_tokens=False, return_tensors="pt", ) text_input_ids = text_inputs.input_ids else: if text_input_ids is None: raise ValueError("text_input_ids must be provided when the tokenizer is not specified") prompt_embeds = text_encoder(text_input_ids.to(device))[0] if hasattr(text_encoder, "module"): dtype = text_encoder.module.dtype else: dtype = text_encoder.dtype prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) _, seq_len, _ = prompt_embeds.shape # duplicate text embeddings and attention mask for each generation per prompt, using mps friendly method prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1) prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1) return prompt_embeds def _encode_prompt_with_clip( text_encoder, tokenizer, prompt: str, device=None, text_input_ids=None, num_images_per_prompt: int = 1, ): prompt = [prompt] if isinstance(prompt, str) else prompt batch_size = len(prompt) if tokenizer is not None: text_inputs = tokenizer( prompt, padding="max_length", max_length=77, truncation=True, return_overflowing_tokens=False, return_length=False, return_tensors="pt", ) text_input_ids = text_inputs.input_ids else: if text_input_ids is None: raise ValueError("text_input_ids must be provided when the tokenizer is not specified") prompt_embeds = text_encoder(text_input_ids.to(device), output_hidden_states=False) if hasattr(text_encoder, "module"): dtype = text_encoder.module.dtype else: dtype = text_encoder.dtype # Use pooled output of CLIPTextModel prompt_embeds = prompt_embeds.pooler_output prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) # duplicate text embeddings for each generation per prompt, using mps friendly method prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1) prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, -1) return prompt_embeds def encode_prompt( text_encoders, tokenizers, prompt: str, max_sequence_length, device=None, num_images_per_prompt: int = 1, text_input_ids_list=None, ): prompt = [prompt] if isinstance(prompt, str) else prompt if hasattr(text_encoders[0], "module"): dtype = text_encoders[0].module.dtype else: dtype = text_encoders[0].dtype pooled_prompt_embeds = _encode_prompt_with_clip( text_encoder=text_encoders[0], tokenizer=tokenizers[0], prompt=prompt, device=device if device is not None else text_encoders[0].device, num_images_per_prompt=num_images_per_prompt, text_input_ids=text_input_ids_list[0] if text_input_ids_list else None, ) prompt_embeds = _encode_prompt_with_t5( text_encoder=text_encoders[1], tokenizer=tokenizers[1], max_sequence_length=max_sequence_length, prompt=prompt, num_images_per_prompt=num_images_per_prompt, device=device if device is not None else text_encoders[1].device, text_input_ids=text_input_ids_list[1] if text_input_ids_list else None, ) text_ids = torch.zeros(prompt_embeds.shape[1], 3).to(device=device, dtype=dtype) return prompt_embeds, pooled_prompt_embeds, text_ids from transformers import T5EncoderModel, T5Tokenizer, CLIPTokenizer, CLIPTextModel class T5Embedder(torch.nn.Module): def __init__(self, device, max_length=300): super().__init__() self.device = device self.max_length = max_length dtype = torch.bfloat16 self.dtype = dtype t5_version = './t5-v1_1-xxl' self.t5_tokenizer = T5Tokenizer.from_pretrained(t5_version, max_length=max_length) self.t5_encoder = T5EncoderModel.from_pretrained(t5_version, dtype=dtype).to(device=device).to(dtype) self.t5_encoder = self.t5_encoder.eval().requires_grad_(False) self.num_shared = max_length @torch.no_grad() def forward(self, text): if isinstance(text, str): text = [text] batch_encoding = self.t5_tokenizer( text, truncation=True, max_length=self.max_length, return_length=False, return_overflowing_tokens=False, padding="max_length", return_tensors="pt", ) prompt_embeds = self.t5_encoder( input_ids=batch_encoding["input_ids"].to(self.device), attention_mask=None, output_hidden_states=False, )['last_hidden_state'] prompt_attention_mask = batch_encoding['attention_mask'].to(self.device) new_text = [x.split('.')[0] for x in text] batch_encoding = self.t5_tokenizer( new_text, truncation=True, max_length=self.num_shared, return_length=False, return_overflowing_tokens=False, padding="max_length", return_tensors="pt", ) shared_prompt_embeds = self.t5_encoder( input_ids=batch_encoding["input_ids"].to(self.device), attention_mask=None, output_hidden_states=False, )['last_hidden_state'] return prompt_embeds, shared_prompt_embeds, prompt_attention_mask import random from torch.utils.checkpoint import checkpoint from peft import LoraConfig, set_peft_model_state_dict import torch from torch import nn, func from typing import Optional from transformers import T5Tokenizer, CLIPVisionModelWithProjection from transformers.models.t5.modeling_t5 import T5PreTrainedModel, T5Stack, T5Config class JacobianT5Encoder(T5PreTrainedModel): """ T5 Encoder with Jacobian loss on last block only. No attention masking applied during forward pass. """ _tied_weights_keys = ["encoder.embed_tokens.weight"] _keys_to_ignore_on_load_unexpected = [r"decoder"] def __init__(self, config: T5Config, num_jacobian_samples=1, max_length=512): super().__init__(config) self.num_jacobian_samples = num_jacobian_samples self.max_length = max_length # Create shared embeddings self.shared = nn.Embedding(config.vocab_size, config.d_model) # Create encoder encoder_config = config encoder_config.use_cache = False encoder_config.is_encoder_decoder = False self.encoder = T5Stack(encoder_config, self.shared) self.post_init() # Freeze everything except last block self._freeze_all_except_last_block() # Print stats trainable_params = sum(p.numel() for p in self.parameters() if p.requires_grad) total_params = sum(p.numel() for p in self.parameters()) print( f"\nTrainable params: {trainable_params:,} / {total_params:,} ({100 * trainable_params / total_params:.2f}%)") def _freeze_all_except_last_block(self): """Freeze all parameters except the last transformer block.""" # Freeze everything first for param in self.parameters(): param.requires_grad = False num_blocks = len(self.encoder.block) # Unfreeze last block for param in self.encoder.block[-1].parameters(): param.requires_grad = True print(f"Unfroze last block (block {num_blocks - 1})") # Unfreeze final layer norm for param in self.encoder.final_layer_norm.parameters(): param.requires_grad = True # Unfreeze embeddings (optional - comment out if you don't want this) #self.shared.weight.requires_grad = True def get_input_embeddings(self): return self.shared def set_input_embeddings(self, new_embeddings): self.shared = new_embeddings self.encoder.set_input_embeddings(new_embeddings) def get_encoder(self): return self.encoder def _process_through_blocks( self, input_ids: torch.LongTensor, stop_at_block: int = None, ): """ Process input through encoder blocks up to stop_at_block. NO ATTENTION MASKING APPLIED - all tokens can attend to all tokens. """ inputs_embeds = self.encoder.embed_tokens(input_ids) batch_size, seq_length = input_ids.shape # No attention mask - all tokens attend to all tokens extended_attention_mask = None cache_position = torch.arange(seq_length, device=input_ids.device) hidden_states = self.encoder.dropout(inputs_embeds) position_bias = None num_blocks = len(self.encoder.block) end_block = stop_at_block if stop_at_block is not None else num_blocks for i in range(end_block): layer_outputs = self.encoder.block[i]( hidden_states, extended_attention_mask, # None - no masking position_bias, None, None, None, past_key_values=None, use_cache=False, output_attentions=False, return_dict=True, cache_position=cache_position, ) hidden_states = layer_outputs[0] if len(layer_outputs) > 1 and layer_outputs[1] is not None: position_bias = layer_outputs[1] return hidden_states, position_bias, cache_position def forward( self, input_ids: torch.LongTensor, attention_mask: Optional[torch.FloatTensor] = None, compute_jacobian: bool = False, ): """ Forward pass with optional Jacobian computation. NO ATTENTION MASKING is applied during the forward pass. attention_mask is only used to identify valid (non-padding) tokens for Jacobian sampling. Args: input_ids: Input token IDs attention_mask: Only used for sampling valid tokens in Jacobian (not for masking) compute_jacobian: Whether to compute Jacobian loss Returns: last_hidden_state: Final encoder output jacobian_loss: Dict with same_token and cross_token losses (if compute_jacobian=True) """ num_blocks = len(self.encoder.block) # Create attention_mask for sampling purposes if not provided if attention_mask is None: attention_mask = torch.ones_like(input_ids) # Process through all blocks except last (no gradients needed, NO MASKING) with torch.no_grad(): second_last_output, position_bias, cache_position = self._process_through_blocks( input_ids, stop_at_block=num_blocks - 1 ) # Detach and optionally require gradients second_last_output = second_last_output.detach() # Need gradients if training or computing Jacobian needs_grad = True if needs_grad: second_last_output.requires_grad_(True) # Process through last block (NO MASKING) if needs_grad: layer_outputs = self.encoder.block[-1]( second_last_output, None, # No attention mask position_bias, None, None, None, past_key_values=None, use_cache=False, output_attentions=False, return_dict=True, cache_position=cache_position, ) hidden_states = layer_outputs[0] hidden_states = self.encoder.final_layer_norm(hidden_states) hidden_states = self.encoder.dropout(hidden_states) else: with torch.no_grad(): layer_outputs = self.encoder.block[-1]( second_last_output, None, # No attention mask position_bias, None, None, None, past_key_values=None, use_cache=False, output_attentions=False, return_dict=True, cache_position=cache_position, ) hidden_states = layer_outputs[0] hidden_states = self.encoder.final_layer_norm(hidden_states) hidden_states = self.encoder.dropout(hidden_states) # Compute Jacobian loss if requested jacobian_loss = { "same_token": torch.tensor(0.0, device=input_ids.device), "cross_token": torch.tensor(0.0, device=input_ids.device) } return hidden_states, jacobian_loss # Wrapper class for compatibility class JacobianLoraT5Embedder(nn.Module): """Wrapper to match your original interface""" def __init__(self, device, max_length=512, use_gradient_checkpointing=True, num_jacobian_samples=1): super().__init__() self.device = device self.max_length = max_length # Load T5 config from transformers import T5Config config = T5Config.from_pretrained('google/t5-v1_1-xxl') # Create encoder model self.t5_encoder = JacobianT5Encoder.from_pretrained( 'google/t5-v1_1-xxl', config=config, num_jacobian_samples=num_jacobian_samples, max_length=max_length ).to(torch.bfloat16).to(device) self.dtype = torch.bfloat16 # Tokenizer self.t5_tokenizer = T5Tokenizer.from_pretrained('google/t5-v1_1-xxl', max_length=max_length) # Image encoder image_encoder_path = 'openai/clip-vit-large-patch14' self.image_encoder = CLIPVisionModelWithProjection.from_pretrained( image_encoder_path ).to(device=device).to(torch.bfloat16) self.image_encoder = self.image_encoder.eval().requires_grad_(False) def forward(self, text, image=None, compute_jacobian=False): """ Forward pass with optional Jacobian regularization. Args: text: Input text (string or list of strings) image: Optional image input compute_jacobian: Whether to compute Jacobian loss Returns: prompt_embeds: T5 encoder output clip_image_embeds: CLIP image embeddings jacobian_loss: Jacobian sparsity loss dict attn_mask: Attention mask (for reference, not used in forward) """ if isinstance(text, str): text = [text] batch_encoding = self.t5_tokenizer( text, truncation=True, max_length=self.max_length, return_length=False, return_overflowing_tokens=False, padding="max_length", return_tensors="pt", ) input_ids = batch_encoding["input_ids"].to(self.device) attn_mask = batch_encoding["attention_mask"].to(self.device) # Forward through T5 encoder (no masking applied) prompt_embeds, jacobian_loss = self.t5_encoder( input_ids=input_ids, attention_mask=attn_mask, # Only used for Jacobian sampling compute_jacobian=compute_jacobian ) # Encode image with torch.no_grad(): if image is not None: clip_image_embeds = self.image_encoder(image.to(self.device)).image_embeds else: clip_image_embeds = None return prompt_embeds, clip_image_embeds, jacobian_loss, attn_mask