Update cappella.py

cappella.py

@@ -32,65 +32,155 @@ class Cappella:
         self.text_encoder_2: CLIPTextModelWithProjection = pipe.text_encoder_2
         self.device = device
 
+    
+    # In cappella.py
     @torch.no_grad()
     def __call__(self, prompt: str, negative_prompt: str) -> CappellaResult:
         """
         Encodes the positive and negative prompts.
+        Ensures both embedding tensors have the same sequence length.
         """
         # Encode the positive prompt
         pos_embeds, pos_pooled = self._encode_one(prompt)
         
         # Encode the negative prompt
         neg_embeds, neg_pooled = self._encode_one(negative_prompt)
-
+    
+        # --- START FIX: Pad shorter embeds ---
+        # Ensure embeds and negative_embeds have the same sequence length
+        seq_len_pos = pos_embeds.shape[1]
+        seq_len_neg = neg_embeds.shape[1]
+        
+        if seq_len_pos > seq_len_neg:
+            # Pad negative embeds
+            pad_len = seq_len_pos - seq_len_neg
+            padding = torch.zeros(
+                (neg_embeds.shape[0], pad_len, neg_embeds.shape[2]),
+                device=self.device, dtype=neg_embeds.dtype
+            )
+            neg_embeds = torch.cat([neg_embeds, padding], dim=1)
+            
+        elif seq_len_neg > seq_len_pos:
+            # Pad positive embeds
+            pad_len = seq_len_neg - seq_len_pos
+            padding = torch.zeros(
+                (pos_embeds.shape[0], pad_len, pos_embeds.shape[2]),
+                device=self.device, dtype=pos_embeds.dtype
+            )
+            pos_embeds = torch.cat([pos_embeds, padding], dim=1)
+        
+        # Now seq_len_pos and seq_len_neg are guaranteed to be equal
+        # --- END FIX ---
+    
         return CappellaResult(
             embeds=pos_embeds,
             pooled_embeds=pos_pooled,
             negative_embeds=neg_embeds,
             negative_pooled_embeds=neg_pooled
         )
-
+    
     def _encode_one(self, prompt: str) -> Tuple[torch.Tensor, torch.Tensor]:
         """
-        Runs a single prompt string through both text encoders,
-        ensuring truncation and padding to 77 tokens.
+        Runs a single prompt string through both text encoders.
+        Handles prompts longer than 77 tokens by chunking.
         """
-        # --- Tokenizer 1 (CLIP-L) ---
-        tok_1_inputs = self.tokenizer(
-            prompt,
-            padding="max_length",
-            max_length=self.tokenizer.model_max_length,
-            truncation=True,
-            return_tensors="pt"
-        )
         
-        # --- Tokenizer 2 (OpenCLIP-G) ---
-        tok_2_inputs = self.tokenizer_2(
-            prompt,
-            padding="max_length",
-            max_length=self.tokenizer_2.model_max_length,
-            truncation=True,
-            return_tensors="pt"
-        )
-
-        # --- Text Encoder 1 (CLIP-L) ---
-        # Gets last_hidden_state. Pooled output is not used.
-        embeds_1 = self.text_encoder(
-            tok_1_inputs.input_ids.to(self.device)
-        ).last_hidden_state
-
-        # --- Text Encoder 2 (OpenCLIP-G) ---
-        # Gets hidden_states[-2] and the pooled output.
-        output_2 = self.text_encoder_2(
-            tok_2_inputs.input_ids.to(self.device),
-            output_hidden_states=True
-        )
-        embeds_2 = output_2.hidden_states[-2]
-        pooled_embeds = output_2.text_embeds
-
-        # --- Concatenate ---
-        # The final embeddings are a concatenation of both.
+        # --- Get Tokenizers and Encoders ---
+        tokenizers = [self.tokenizer, self.tokenizer_2]
+        text_encoders = [self.text_encoder, self.text_encoder_2]
+        
+        prompt_embeds_list = []
+        pooled_prompt_embeds = None
+    
+        for tokenizer, text_encoder in zip(tokenizers, text_encoders):
+            # --- Tokenize ---
+            # Tokenize without padding or truncation first
+            text_inputs = tokenizer(
+                prompt,
+                padding=False,
+                truncation=False,
+                return_tensors="pt"
+            )
+            input_ids = text_inputs.input_ids.to(self.device)
+            
+            # --- Chunking ---
+            # Manually chunk the input_ids
+            max_length = tokenizer.model_max_length
+            bos = tokenizer.bos_token_id
+            eos = tokenizer.eos_token_id
+            
+            # We subtract 2 for BOS and EOS
+            chunk_length = max_length - 2
+            
+            # Get all token IDs *except* BOS and EOS
+            clean_input_ids = input_ids[0, 1:-1]
+            
+            # Split into chunks
+            chunks = [clean_input_ids[i:i + chunk_length] for i in range(0, len(clean_input_ids), chunk_length)]
+            
+            # --- Prepare Batches ---
+            batch_input_ids = []
+            for chunk in chunks:
+                # Add BOS and EOS
+                chunk_with_bos_eos = torch.cat([
+                    torch.tensor([bos], dtype=torch.long, device=self.device),
+                    chunk.to(torch.long),
+                    torch.tensor([eos], dtype=torch.long, device=self.device)
+                ])
+                
+                # Pad to max_length
+                pad_len = max_length - len(chunk_with_bos_eos)
+                if pad_len > 0:
+                    padding = torch.full((pad_len,), tokenizer.pad_token_id, dtype=torch.long, device=self.device)
+                    chunk_with_bos_eos = torch.cat([chunk_with_bos_eos, padding])
+                
+                batch_input_ids.append(chunk_with_bos_eos)
+            
+            if not batch_input_ids:
+                # Handle empty prompt
+                batch_input_ids.append(
+                    torch.full((max_length,), tokenizer.pad_token_id, dtype=torch.long, device=self.device)
+                )
+    
+            batch_input_ids = torch.stack(batch_input_ids)
+            
+            # --- Encode ---
+            if text_encoder == self.text_encoder:
+                # Text Encoder 1 (CLIP-L)
+                # We only need the last_hidden_state
+                encoder_output = text_encoder(
+                    batch_input_ids,
+                    output_hidden_states=False
+                )
+                # [num_chunks, 77, 768]
+                prompt_embeds = encoder_output.last_hidden_state
+                prompt_embeds_list.append(prompt_embeds)
+                
+            elif text_encoder == self.text_encoder_2:
+                # Text Encoder 2 (OpenCLIP-G)
+                # We need hidden_states[-2] and the pooled output from the FIRST chunk
+                encoder_output = text_encoder(
+                    batch_input_ids,
+                    output_hidden_states=True
+                )
+                # [num_chunks, 77, 1280]
+                prompt_embeds = encoder_output.hidden_states[-2]
+                prompt_embeds_list.append(prompt_embeds)
+                
+                # Pooled output comes from the FIRST chunk
+                # We use .text_embeds which is the pooled output
+                # [num_chunks, 1280]
+                all_pooled = encoder_output.text_embeds
+                pooled_prompt_embeds = all_pooled[0:1] # Keep as [1, 1280]
+    
+        # --- Concatenate Chunks ---
+        # Reshape from [num_chunks, 77, dim] to [1, num_chunks*77, dim]
+        # and then concatenate along the dim=-1
+        
+        embeds_1 = prompt_embeds_list[0].reshape(1, -1, prompt_embeds_list[0].shape[-1])
+        embeds_2 = prompt_embeds_list[1].reshape(1, -1, prompt_embeds_list[1].shape[-1])
+    
         prompt_embeds = torch.cat([embeds_1, embeds_2], dim=-1)
-
-        return prompt_embeds, pooled_embeds
-
+    
+        # pooled_prompt_embeds is already [1, 1280] from Encoder 2's first chunk
+        return prompt_embeds, pooled_prompt_embeds