Update modeling_text_encoder.py

modeling_text_encoder.py

@@ -10,7 +10,7 @@ from transformers import (
 from typing import Union, List, Optional
 
 class SD3TextEncoderWithMask(nn.Module):
-    def __init__(self, model_path, torch_dtype=torch.float16):
+    def __init__(self, model_path, torch_dtype=torch.float16):  # Default to FP16
         super().__init__()
 
         # Tokenizers for CLIP and T5
@@ -47,10 +47,10 @@ class SD3TextEncoderWithMask(nn.Module):
             ).to("cuda")
 
         if self.text_encoder_3 is None:
-            # Load the FP8 T5 model (adjust path if needed for specific FP8 model)
+            # Load the T5 model in FP16
             self.text_encoder_3 = T5EncoderModel.from_pretrained(
-                os.path.join(self.model_path, 'text_encoder_3_fp8'),  # Use the FP8 version
-                torch_dtype=torch.float8  # FP8 precision
+                os.path.join(self.model_path, 'text_encoder_3'),  # Ensure you're using the correct path
+                torch_dtype=self.torch_dtype  # Use FP16 precision
             ).to("cuda")
 
     def _get_t5_prompt_embeds(
@@ -75,9 +75,9 @@ 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 in FP8
+        # Use the T5 model to generate embeddings in FP16
         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)  # Ensure correct dtype
+        prompt_embeds = prompt_embeds.to(dtype=self.torch_dtype)  # Ensure correct dtype
 
         # Duplicate embeddings for each image generation
         batch_size = len(prompt)
@@ -111,7 +111,7 @@ class SD3TextEncoderWithMask(nn.Module):
             return_tensors="pt",
         )
 
-        text_input_ids = text_inputs.input.ids.to(device)
+        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
@@ -133,7 +133,7 @@ class SD3TextEncoderWithMask(nn.Module):
         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 in FP8
+        # Get T5 embeddings in FP16
         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