Update modeling_text_encoder.py

modeling_text_encoder.py

@@ -1,49 +1,66 @@
 import torch
 import torch.nn as nn
 import os
-from transformers import CLIPTextModelWithProjection, CLIPTokenizer, T5EncoderModel, T5TokenizerFast
+from transformers import (
+    CLIPTextModelWithProjection,
+    CLIPTokenizer,
+    T5EncoderModel,
+    T5TokenizerFast,
+)
 from typing import Union, List, Optional
-from functools import lru_cache
 
 class SD3TextEncoderWithMask(nn.Module):
     def __init__(self, model_path, torch_dtype):
         super().__init__()
-        
-        # Use lazy loading for models
-        self.model_path = model_path
-        self.torch_dtype = torch_dtype
-        
-        # CLIP-L
+
+        # Initialization of models and tokenizers, but delay moving to device
         self.tokenizer = CLIPTokenizer.from_pretrained(os.path.join(model_path, 'tokenizer'))
-        self.tokenizer_max_length = self.tokenizer.model_max_length
-        
-        # CLIP-G
         self.tokenizer_2 = CLIPTokenizer.from_pretrained(os.path.join(model_path, 'tokenizer_2'))
-        
-        # T5
         self.tokenizer_3 = T5TokenizerFast.from_pretrained(os.path.join(model_path, 'tokenizer_3'))
-        
+
+        # Lazy loading of models for memory efficiency
+        self.text_encoder = None
+        self.text_encoder_2 = None
+        self.text_encoder_3 = None
+        self.model_path = model_path
+        self.torch_dtype = torch_dtype
+        self.tokenizer_max_length = self.tokenizer.model_max_length
+
+        # Freeze parameters to avoid unnecessary gradient computation
         self._freeze()
 
     def _freeze(self):
+        """ Freeze all model parameters to avoid training overhead. """
         for param in self.parameters():
             param.requires_grad = False
 
-    @lru_cache(maxsize=128)
+    def _load_models_if_needed(self):
+        """ Load models only if they haven't been loaded already. """
+        if 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("cuda")
+
+        if 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("cuda")
+
+        if self.text_encoder_3 is None:
+            self.text_encoder_3 = T5EncoderModel.from_pretrained(
+                os.path.join(self.model_path, 'text_encoder_3'), torch_dtype=self.torch_dtype
+            ).to("cuda")
+
     def _get_t5_prompt_embeds(
         self,
-        prompt: Union[str, tuple],
+        prompt: Union[str, List[str]] = None,
         num_images_per_prompt: int = 1,
         device: Optional[torch.device] = None,
         max_sequence_length: int = 128,
     ):
-        if isinstance(prompt, tuple):
-            prompt = list(prompt)
-        else:
-            prompt = [prompt]
-        
-        batch_size = len(prompt)
-
+        """ Get embeddings from T5 model. """
+        self._load_models_if_needed()  # Lazy loading
+        prompt = [prompt] if isinstance(prompt, str) else prompt
         text_inputs = self.tokenizer_3(
             prompt,
             padding="max_length",
@@ -52,57 +69,35 @@ class SD3TextEncoderWithMask(nn.Module):
             add_special_tokens=True,
             return_tensors="pt",
         )
-        text_input_ids = text_inputs.input_ids
-        prompt_attention_mask = text_inputs.attention_mask.to(device)
-
-        if not hasattr(self, 'text_encoder_3'):
-            self.text_encoder_3 = T5EncoderModel.from_pretrained(
-                os.path.join(self.model_path, 'text_encoder_3'), 
-                torch_dtype=self.torch_dtype
-            ).to(device)
 
-        with torch.no_grad():
-            prompt_embeds = self.text_encoder_3(text_input_ids.to(device), attention_mask=prompt_attention_mask)[0]
-        
-        prompt_embeds = prompt_embeds.to(dtype=self.torch_dtype, device=device)
+        text_input_ids = text_inputs.input_ids.to(device)
+        prompt_attention_mask = text_inputs.attention_mask.to(device)
+        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)
 
+        # Duplicate embeddings for each image generation
+        batch_size = len(prompt)
         _, 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)
+        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1).view(batch_size * num_images_per_prompt, seq_len, -1)
+        prompt_attention_mask = prompt_attention_mask.view(batch_size, -1).repeat(num_images_per_prompt, 1)
 
         return prompt_embeds, prompt_attention_mask
 
-    @lru_cache(maxsize=128)
     def _get_clip_prompt_embeds(
         self,
-        prompt: Union[str, tuple],
+        prompt: Union[str, List[str]],
         num_images_per_prompt: int = 1,
         device: Optional[torch.device] = None,
         clip_model_index: int = 0,
     ):
-        if isinstance(prompt, tuple):
-            prompt = list(prompt)
-        else:
-            prompt = [prompt]
-        
-        batch_size = len(prompt)
+        """ Get embeddings from CLIP model. """
+        self._load_models_if_needed()  # Lazy loading
 
         clip_tokenizers = [self.tokenizer, self.tokenizer_2]
-        clip_text_encoders_attr = ['text_encoder', 'text_encoder_2']
+        clip_text_encoders = [self.text_encoder, self.text_encoder_2]
 
         tokenizer = clip_tokenizers[clip_model_index]
-        text_encoder_attr = clip_text_encoders_attr[clip_model_index]
-
-        if not hasattr(self, text_encoder_attr):
-            setattr(self, text_encoder_attr, CLIPTextModelWithProjection.from_pretrained(
-                os.path.join(self.model_path, text_encoder_attr), 
-                torch_dtype=self.torch_dtype
-            ).to(device))
-
-        text_encoder = getattr(self, text_encoder_attr)
+        text_encoder = clip_text_encoders[clip_model_index]
 
         text_inputs = tokenizer(
             prompt,
@@ -112,48 +107,36 @@ class SD3TextEncoderWithMask(nn.Module):
             return_tensors="pt",
         )
 
-        text_input_ids = text_inputs.input_ids
+        text_input_ids = text_inputs.input_ids.to(device)
+        prompt_embeds = text_encoder(text_input_ids, output_hidden_states=True)[0]
 
-        with torch.no_grad():
-            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)
+        # 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
 
-    @torch.no_grad()
     def encode_prompt(self, 
         prompt, 
         num_images_per_prompt=1, 
-        device=None,
+        device=None
     ):
-        if isinstance(prompt, str):
-            prompt = (prompt,)
-        elif isinstance(prompt, list):
-            prompt = tuple(prompt)
-
-        pooled_prompt_embed = self._get_clip_prompt_embeds(
-            prompt=prompt,
-            device=device,
-            num_images_per_prompt=num_images_per_prompt,
-            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_model_index=1,
-        )
+        """ 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)
 
-        prompt_embeds, prompt_attention_mask = self._get_t5_prompt_embeds(
-            prompt=prompt,
-            num_images_per_prompt=num_images_per_prompt,
-            device=device,
-        )
+        # 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):
-        return self.encode_prompt(input_prompts, 1, device=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