Update cappella.py

cappella.py

@@ -1,6 +1,7 @@
 import torch
 from dataclasses import dataclass
 from transformers import CLIPTokenizer, CLIPTextModel, CLIPTextModelWithProjection
+from typing import Tuple
 
 @dataclass
 class CappellaResult:
@@ -21,7 +22,7 @@ class Cappella:
     It correctly:
     1. Uses both SDXL tokenizers and text encoders.
     2. Truncates prompts that are too long (fixes "78 vs 77" error).
-    3. Pads prompts that are too short (fixes "93 vs 77" error).
+    3. Pads prompts (by using max_length) to ensure they are all 77 tokens.
     4. Returns all 4 required embedding tensors.
     """
     def __init__(self, pipe, device):
@@ -49,9 +50,10 @@ class Cappella:
             negative_pooled_embeds=neg_pooled
         )
 
-    def _encode_one(self, prompt: str) -> (torch.Tensor, torch.Tensor):
+    def _encode_one(self, prompt: str) -> Tuple[torch.Tensor, torch.Tensor]:
         """
-        Runs a single prompt string through both text encoders.
+        Runs a single prompt string through both text encoders,
+        ensuring truncation and padding to 77 tokens.
         """
         # --- Tokenizer 1 (CLIP-L) ---
         tok_1_inputs = self.tokenizer(