Update modeling_text_encoder.py

modeling_text_encoder.py

@@ -6,6 +6,7 @@ from transformers import (
     CLIPTokenizer,
     T5EncoderModel,
     T5TokenizerFast,
+    BitsAndBytesConfig  # Import for 8-bit quantization
 )
 from typing import Union, List, Optional
 
@@ -13,12 +14,12 @@ class SD3TextEncoderWithMask(nn.Module):
     def __init__(self, model_path, torch_dtype):
         super().__init__()
 
-        # Initialization of models and tokenizers, but delay moving to device
+        # Tokenizers for CLIP and T5
         self.tokenizer = CLIPTokenizer.from_pretrained(os.path.join(model_path, 'tokenizer'))
         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'))
 
-        # Lazy loading of models for memory efficiency
+        # Lazy loading for memory efficiency
         self.text_encoder = None
         self.text_encoder_2 = None
         self.text_encoder_3 = None
@@ -26,6 +27,9 @@ class SD3TextEncoderWithMask(nn.Module):
         self.torch_dtype = torch_dtype
         self.tokenizer_max_length = self.tokenizer.model_max_length
 
+        # Quantization config for T5 model
+        self.quantization_config = BitsAndBytesConfig(load_in_8bit=True)  # Quantize T5 to 8-bit
+
         # Freeze parameters to avoid unnecessary gradient computation
         self._freeze()
 
@@ -47,9 +51,12 @@ class SD3TextEncoderWithMask(nn.Module):
             ).to("cuda")
 
         if self.text_encoder_3 is None:
+            # Load the T5 model with 8-bit quantization
             self.text_encoder_3 = T5EncoderModel.from_pretrained(
-                os.path.join(self.model_path, 'text_encoder_3'), torch_dtype=self.torch_dtype
-            ).to("cuda")
+                os.path.join(self.model_path, 'text_encoder_3'),
+                torch_dtype=self.torch_dtype,
+                quantization_config=self.quantization_config  # Apply quantization
+            )  # Do NOT use .to("cuda") for 8-bit quantized models, as they handle device placement automatically
 
     def _get_t5_prompt_embeds(
         self,
@@ -58,7 +65,7 @@ class SD3TextEncoderWithMask(nn.Module):
         device: Optional[torch.device] = None,
         max_sequence_length: int = 128,
     ):
-        """ Get embeddings from T5 model. """
+        """ Get embeddings from the T5 model. """
         self._load_models_if_needed()  # Lazy loading
         prompt = [prompt] if isinstance(prompt, str) else prompt
         text_inputs = self.tokenizer_3(
@@ -72,8 +79,10 @@ class SD3TextEncoderWithMask(nn.Module):
 
         text_input_ids = text_inputs.input_ids.to(device)
         prompt_attention_mask = text_inputs.attention_mask.to(device)
+
+        # Use the T5 model to generate embeddings (quantized)
         prompt_embeds = self.text_encoder_3(text_input_ids, attention_mask=prompt_attention_mask)[0]
-        prompt_embeds = prompt_embeds.to(dtype=self.text_encoder_3.dtype, device=device)
+        prompt_embeds = prompt_embeds.to(dtype=self.text_encoder_3.dtype)  # Ensure correct dtype
 
         # Duplicate embeddings for each image generation
         batch_size = len(prompt)