Update generator.py

generator.py

@@ -762,36 +762,31 @@ class RetroArtConverter:
         
         pipe_kwargs["generator"] = generator
         
-        # --- START FIX 1: Correct Compel batching and slicing ---
+        # --- START FIX: Use Compel as per documentation ---
         if self.use_compel and self.compel is not None:
             try:
                 print("Encoding prompts with Compel...")
                 
-                # Pass both prompts as a list to be batched
-                conditioning_batch, pooled_batch = self.compel([prompt, negative_prompt])
+                # Call Compel with prompt and negative_prompt kwargs
+                # Compel will handle the padding internally
+                conditioning = self.compel(prompt, negative_prompt=negative_prompt)
 
-                # Store positive and negative embeds separately for now
-                positive_prompt_embeds = conditioning_batch[0:1]
-                positive_pooled_embeds = pooled_batch[0:1]
-                negative_prompt_embeds = conditioning_batch[1:2]
-                negative_pooled_embeds = pooled_batch[1:2]
-                
-                print(f"[OK] Compel encoded - Pos: {positive_prompt_embeds.shape}, Neg: {negative_prompt_embeds.shape}")
-                
-                # Put the positive embeds in pipe_kwargs for the *next* step
-                pipe_kwargs["prompt_embeds"] = positive_prompt_embeds
-                pipe_kwargs["pooled_prompt_embeds"] = positive_pooled_embeds
+                # Unpack the results from the returned object
+                pipe_kwargs["prompt_embeds"] = conditioning.embeds
+                pipe_kwargs["pooled_prompt_embeds"] = conditioning.pooled_embeds
+                pipe_kwargs["negative_prompt_embeds"] = conditioning.negative_embeds
+                pipe_kwargs["negative_pooled_prompt_embeds"] = conditioning.negative_pooled_embeds
                 
+                print(f"[OK] Compel encoded - Prompt: {pipe_kwargs['prompt_embeds'].shape}, Negative: {pipe_kwargs['negative_prompt_embeds'].shape}")
             except Exception as e:
                 print(f"Compel encoding failed, using standard prompts: {e}")
                 traceback.print_exc()
                 pipe_kwargs["prompt"] = prompt
                 pipe_kwargs["negative_prompt"] = negative_prompt
-                self.use_compel = False # Fallback to standard
         else:
             pipe_kwargs["prompt"] = prompt
             pipe_kwargs["negative_prompt"] = negative_prompt
-        # --- END FIX 1 ---
+        # --- END FIX ---
         
         # Add CLIP skip
         if hasattr(self.pipe, 'text_encoder'):
@@ -835,38 +830,38 @@ class RetroArtConverter:
                         
                         print(f"  - Face embedding: {face_proj_embeds.shape}, Scale: {boosted_scale:.2f}")
                         
-                        # --- START FIX 2: Correct CFG and Negative Padding ---
-                        if self.use_compel and 'prompt_embeds' in pipe_kwargs:
-                            # 1. Get the Compel-generated embeds
-                            positive_embeds = pipe_kwargs['prompt_embeds']
-                            
-                            # 2. Concatenate face embeddings to POSITIVE prompt
-                            final_positive_embeds = torch.cat([positive_embeds, face_proj_embeds], dim=1)
-                            
-                            # 3. Create zero padding for NEGATIVE prompt (YOUR FIX)
-                            neg_padding = torch.zeros_like(face_proj_embeds)
-                            
-                            # 4. Concatenate zero padding to NEGATIVE prompt
-                            final_negative_embeds = torch.cat([negative_prompt_embeds, neg_padding], dim=1)
+                        # --- START FIX 2: Your padding solution ---
+                        if 'prompt_embeds' in pipe_kwargs:
+                            original_embeds = pipe_kwargs['prompt_embeds']
                             
-                            # 5. Create the final CFG batch (shape [2, 109, 2048])
-                            pipe_kwargs['prompt_embeds'] = torch.cat([final_negative_embeds, final_positive_embeds], dim=0)
+                            # Handle CFG by creating a [2, 16, 2048] tensor
+                            # [0] is zeros for negative, [1] is face embeds for positive
+                            face_proj_embeds_cfg = torch.cat([torch.zeros_like(face_proj_embeds), face_proj_embeds], dim=0)
                             
-                            # 6. Do the same for the pooled embeds (shape [2, 1280])
-                            pipe_kwargs['pooled_prompt_embeds'] = torch.cat([negative_pooled_embeds, positive_pooled_embeds], dim=0)
-                            
-                            # 7. CRITICAL: Remove the separate negative_prompt_embeds
+                            # Concatenate face embeddings to POSITIVE prompt
+                            combined_embeds = torch.cat([original_embeds, face_proj_embeds_cfg[1:2]], dim=1) # [1, 93, 2048] -> [1, 109, 2048]
+                            pipe_kwargs['prompt_embeds'] = combined_embeds
+
+                            # CRITICAL: Pad negative_prompt_embeds by the same amount
                             if 'negative_prompt_embeds' in pipe_kwargs:
-                                del pipe_kwargs['negative_prompt_embeds']
+                                negative_embeds = pipe_kwargs['negative_prompt_embeds']
+                                # Create zero padding [1, 16, 2048]
+                                neg_padding = torch.zeros(
+                                    (
+                                        negative_embeds.shape[0],      # 1
+                                        face_proj_embeds.shape[1], # 16
+                                        negative_embeds.shape[2],      # 2048
+                                    ),
+                                    device=negative_embeds.device,
+                                    dtype=negative_embeds.dtype
+                                )
+                                # Concatenate zero padding to NEGATIVE prompt
+                                pipe_kwargs['negative_prompt_embeds'] = torch.cat([negative_embeds, neg_padding], dim=1)
+                                print(f"  [OK] Negative prompt padded to match: {pipe_kwargs['negative_prompt_embeds'].shape}")
                             
-                            print(f"  [OK] CFG batch created! Embeds: {pipe_kwargs['prompt_embeds'].shape}, Pooled: {pipe_kwargs['pooled_prompt_embeds'].shape}")
-                        
+                            print(f"  [OK] Face embeddings concatenated successfully! Prompt: {combined_embeds.shape}")
                         else:
-                            # Fallback if Compel failed
-                            print(f"  [WARNING] Can't concatenate - Compel failed. Using standard prompt.")
-                            pipe_kwargs['prompt'] = prompt
-                            pipe_kwargs['negative_prompt'] = negative_prompt
-
+                            print(f"  [WARNING] Can't concatenate - no prompt_embeds (use Compel)")
                         # --- END FIX 2 ---
                 
                 elif has_detected_faces: