Update modeling_text_encoder.py

modeling_text_encoder.py

@@ -1,168 +1,172 @@
+now cam you optimise this file
+
 import torch
 import torch.nn as nn
 import os
+
+
 from transformers import (
-    CLIPTextModelWithProjection,
-    CLIPTokenizer,
-    T5EncoderModel,
-    T5TokenizerFast,
-    BitsAndBytesConfig
+    CLIPTextModelWithProjection,
+    CLIPTokenizer,
+    T5EncoderModel,
+    T5TokenizerFast,
 )
+
+
 from typing import Any, Callable, Dict, List, Optional, Union
 
+
+
 class SD3TextEncoderWithMask(nn.Module):
-    def __init__(self, model_path, torch_dtype, use_8bit_quantization=True):
-        super().__init__()
-        self.model_path = model_path
-        self.torch_dtype = torch_dtype
-        self.device = "cpu"
-        self.use_8bit_quantization = use_8bit_quantization
-
-        # Initialize tokenizers
-        self.tokenizer = CLIPTokenizer.from_pretrained(os.path.join(model_path, 'tokenizer'))
-        self.tokenizer_max_length = self.tokenizer.model_max_length
-        self.tokenizer_2 = CLIPTokenizer.from_pretrained(os.path.join(model_path, 'tokenizer_2'))
-        self.tokenizer_3 = T5TokenizerFast.from_pretrained(os.path.join(model_path, 'tokenizer_3'))
-
-        # Initialize encoders (kept on CPU initially)
-        self.text_encoder = None
-        self.text_encoder_2 = None
-        self.text_encoder_3 = None
-
-        self._freeze()
-
-    def _freeze(self):
-        for param in self.parameters():
-            param.requires_grad = False
-
-    def to(self, device):
-        self.device = device
-        return self
-
-    def _load_and_move_encoder(self, encoder_name):
-        if encoder_name == 'text_encoder' and self.text_encoder is None:
-            self.text_encoder = CLIPTextModelWithProjection.from_pretrained(
-                os.path.join(self.model_path, 'text_encoder'),
-                torch_dtype=self.torch_dtype
-            ).to(self.device)
-        elif encoder_name == 'text_encoder_2' and self.text_encoder_2 is None:
-            self.text_encoder_2 = CLIPTextModelWithProjection.from_pretrained(
-                os.path.join(self.model_path, 'text_encoder_2'),
-                torch_dtype=self.torch_dtype
-            ).to(self.device)
-        elif encoder_name == 'text_encoder_3' and self.text_encoder_3 is None:
-            if self.use_8bit_quantization:
-                quantization_config = BitsAndBytesConfig(load_in_8bit=True)
-                self.text_encoder_3 = T5EncoderModel.from_pretrained(
-                    os.path.join(self.model_path, 'text_encoder_3'),
-                    quantization_config=quantization_config,
-                    device_map="auto"  # This will automatically decide the best device mapping
-                )
-            else:
-                self.text_encoder_3 = T5EncoderModel.from_pretrained(
-                    os.path.join(self.model_path, 'text_encoder_3'),
-                    torch_dtype=self.torch_dtype
-                ).to(self.device)
-
-    def _unload_encoder(self, encoder_name):
-        if encoder_name == 'text_encoder':
-            self.text_encoder = None
-        elif encoder_name == 'text_encoder_2':
-            self.text_encoder_2 = None
-        elif encoder_name == 'text_encoder_3':
-            self.text_encoder_3 = None
-        torch.cuda.empty_cache()
-
-    def _get_t5_prompt_embeds(
-        self,
-        prompt: Union[str, List[str]] = None,
-        num_images_per_prompt: int = 1,
-        max_sequence_length: int = 128,
-    ):
-        self._load_and_move_encoder('text_encoder_3')
-        
-        prompt = [prompt] if isinstance(prompt, str) else prompt
-        batch_size = len(prompt)
-
-        text_inputs = self.tokenizer_3(
-            prompt,
-            padding="max_length",
-            max_length=max_sequence_length,
-            truncation=True,
-            add_special_tokens=True,
-            return_tensors="pt",
-        )
-        text_input_ids = text_inputs.input_ids.to(self.device)
-        prompt_attention_mask = text_inputs.attention_mask.to(self.device)
-
-        prompt_embeds = self.text_encoder_3(text_input_ids, attention_mask=prompt_attention_mask)[0]
-        if not self.use_8bit_quantization:
-            prompt_embeds = prompt_embeds.to(dtype=self.torch_dtype)
-
-        _, seq_len, _ = prompt_embeds.shape
-
-        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)
-        prompt_attention_mask = prompt_attention_mask.view(batch_size, -1)
-        prompt_attention_mask = prompt_attention_mask.repeat(num_images_per_prompt, 1)
-
-        self._unload_encoder('text_encoder_3')
-        return prompt_embeds, prompt_attention_mask
-
-    def _get_clip_prompt_embeds(
-        self,
-        prompt: Union[str, List[str]],
-        num_images_per_prompt: int = 1,
-        device: Optional[torch.device] = None,
-        clip_model_index: int = 0,
-    ):
-        """ Get embeddings from CLIP model. """
-        self._load_models_if_needed()  # Lazy loading
-
-        clip_tokenizers = [self.tokenizer, self.tokenizer_2]
-        clip_text_encoders = [self.text_encoder, self.text_encoder_2]
-
-        tokenizer = clip_tokenizers[clip_model_index]
-        text_encoder = clip_text_encoders[clip_model_index]
-
-        text_inputs = tokenizer(
-            prompt,
-            padding="max_length",
-            max_length=self.tokenizer_max_length,
-            truncation=True,
-            return_tensors="pt",
-        )
-
-        text_input_ids = text_inputs.input_ids.to(device)
-        prompt_embeds = text_encoder(text_input_ids, output_hidden_states=True)[0]
-
-        # Duplicate embeddings for each image generation
-        batch_size = len(prompt)
-        pooled_prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1).view(batch_size * num_images_per_prompt, -1)
-
-        return pooled_prompt_embeds
-
-    def encode_prompt(self, 
-        prompt, 
-        num_images_per_prompt=1, 
-        device=None
-    ):
-        """ Encode the prompt using both CLIP and T5 models. """
-        prompt = [prompt] if isinstance(prompt, str) else prompt
-
-        # Get embeddings from both CLIP models
-        pooled_prompt_embed = self._get_clip_prompt_embeds(prompt, num_images_per_prompt=num_images_per_prompt, device=device, clip_model_index=0)
-        pooled_prompt_2_embed = self._get_clip_prompt_embeds(prompt, num_images_per_prompt=num_images_per_prompt, device=device, clip_model_index=1)
-        pooled_prompt_embeds = torch.cat([pooled_prompt_embed, pooled_prompt_2_embed], dim=-1)
-
-        # Get T5 embeddings
-        prompt_embeds, prompt_attention_mask = self._get_t5_prompt_embeds(prompt, num_images_per_prompt=num_images_per_prompt, device=device)
-
-        return prompt_embeds, prompt_attention_mask, pooled_prompt_embeds
-
-    def forward(self, input_prompts, device):
-        """ Forward pass for encoding prompts. """
-        with torch.no_grad():
-            prompt_embeds, prompt_attention_mask, pooled_prompt_embeds = self.encode_prompt(input_prompts, num_images_per_prompt=1, device=device)
-
-        return prompt_embeds, prompt_attention_mask, pooled_prompt_embeds
+    def __init__(self, model_path, torch_dtype):
+        super().__init__()
+
+        self.device_0 = torch.device('cuda:0')  # GPU 0 for text encoder
+        self.device_1 = torch.device('cuda:1')  # GPU 1 for other tasks
+
+        # CLIP-L
+        self.tokenizer = CLIPTokenizer.from_pretrained(os.path.join(model_path, 'tokenizer'))
+        self.tokenizer_max_length = self.tokenizer.model_max_length
+        self.text_encoder = CLIPTextModelWithProjection.from_pretrained(os.path.join(model_path, 'text_encoder'), torch_dtype=torch_dtype)
+
+
+        # CLIP-G
+        self.tokenizer_2 = CLIPTokenizer.from_pretrained(os.path.join(model_path, 'tokenizer_2'))
+        self.text_encoder_2 = CLIPTextModelWithProjection.from_pretrained(os.path.join(model_path, 'text_encoder_2'), torch_dtype=torch_dtype)
+
+
+        # T5
+        self.tokenizer_3 = T5TokenizerFast.from_pretrained(os.path.join(model_path, 'tokenizer_3'))
+        self.text_encoder_3 = T5EncoderModel.from_pretrained(os.path.join(model_path, 'text_encoder_3'), torch_dtype=torch_dtype).to(self.device_1)
+    
+        self._freeze()
+
+
+    def _freeze(self):
+        for param in self.parameters():
+            param.requires_grad = False
+
+
+    def _get_t5_prompt_embeds(
+        self,
+        prompt: Union[str, List[str]] = None,
+        num_images_per_prompt: int = 1,
+        device: Optional[torch.device] = None,
+        max_sequence_length: int = 128,
+    ):
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+        batch_size = len(prompt)
+
+
+        text_inputs = self.tokenizer_3(
+            prompt,
+            padding="max_length",
+            max_length=max_sequence_length,
+            truncation=True,
+            add_special_tokens=True,
+            return_tensors="pt",
+        )
+        text_input_ids = text_inputs.input_ids
+        prompt_attention_mask = text_inputs.attention_mask
+        prompt_attention_mask = prompt_attention_mask.to(device)
+        prompt_embeds = self.text_encoder_3(text_input_ids.to(device), attention_mask=prompt_attention_mask)[0]
+        dtype = self.text_encoder_3.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)
+        prompt_attention_mask = prompt_attention_mask.view(batch_size, -1)
+        prompt_attention_mask = prompt_attention_mask.repeat(num_images_per_prompt, 1)
+
+
+        return prompt_embeds, prompt_attention_mask
+
+
+    def _get_clip_prompt_embeds(
+        self,
+        prompt: Union[str, List[str]],
+        num_images_per_prompt: int = 1,
+        device: Optional[torch.device] = None,
+        clip_skip: Optional[int] = None,
+        clip_model_index: int = 0,
+    ):
+
+
+        clip_tokenizers = [self.tokenizer, self.tokenizer_2]
+        clip_text_encoders = [self.text_encoder, self.text_encoder_2]
+
+
+        tokenizer = clip_tokenizers[clip_model_index]
+        text_encoder = clip_text_encoders[clip_model_index]
+
+
+        batch_size = len(prompt)
+
+
+        text_inputs = tokenizer(
+            prompt,
+            padding="max_length",
+            max_length=self.tokenizer_max_length,
+            truncation=True,
+            return_tensors="pt",
+        )
+
+
+        text_input_ids = text_inputs.input_ids
+        prompt_embeds = text_encoder(text_input_ids.to(device), output_hidden_states=True)
+        pooled_prompt_embeds = prompt_embeds[0]
+        pooled_prompt_embeds = pooled_prompt_embeds.repeat(1, num_images_per_prompt, 1)
+        pooled_prompt_embeds = pooled_prompt_embeds.view(batch_size * num_images_per_prompt, -1)
+
+
+        return pooled_prompt_embeds
+
+
+    def encode_prompt(self, 
+        prompt, 
+        num_images_per_prompt=1, 
+        clip_skip: Optional[int] = None,
+        device=None,
+    ):
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+
+
+        pooled_prompt_embed = self._get_clip_prompt_embeds(
+            prompt=prompt,
+            device=device,
+            num_images_per_prompt=num_images_per_prompt,
+            clip_skip=clip_skip,
+            clip_model_index=0,
+        )
+        pooled_prompt_2_embed = self._get_clip_prompt_embeds(
+            prompt=prompt,
+            device=device,
+            num_images_per_prompt=num_images_per_prompt,
+            clip_skip=clip_skip,
+            clip_model_index=1,
+        )
+        pooled_prompt_embeds = torch.cat([pooled_prompt_embed, pooled_prompt_2_embed], dim=-1)
+
+
+        prompt_embeds, prompt_attention_mask = self._get_t5_prompt_embeds(
+            prompt=prompt,
+            num_images_per_prompt=num_images_per_prompt,
+            device=device,
+        )
+        return prompt_embeds, prompt_attention_mask, pooled_prompt_embeds
+
+
+    def forward(self, input_prompts, device):
+        with torch.no_grad():
+            prompt_embeds, prompt_attention_mask, pooled_prompt_embeds = self.encode_prompt(input_prompts, 1, clip_skip=None, device=device)
+
+
+        return prompt_embeds, prompt_attention_mask, pooled_prompt_embeds
+
+
+implement cpu offload