Spaces:
Running on Zero
Running on Zero
| 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 | |
| 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 | |