Update generator.py

generator.py

@@ -20,7 +20,10 @@ from utils import (
 )
 from models import (
     load_face_analysis, load_depth_detector, load_controlnets, load_image_encoder,
-    load_sdxl_pipeline, load_loras, setup_ip_adapter, setup_compel,
+    load_sdxl_pipeline, load_loras, setup_ip_adapter, 
+    # --- START FIX: Import setup_cappella ---
+    setup_cappella, 
+    # --- END FIX ---
     setup_scheduler, optimize_pipeline, load_caption_model, set_clip_skip,
     load_openpose_detector, load_mediapipe_face_detector
 )
@@ -127,8 +130,9 @@ class RetroArtConverter:
             self.models_loaded['ip_adapter'] = False
             self.image_proj_model = None
         
-        # Setup Compel
-        self.compel, self.use_compel = setup_compel(self.pipe)
+        # --- START FIX: Setup Cappella ---
+        self.cappella, self.use_cappella = setup_cappella(self.pipe)
+        # --- END FIX ---
         
         # Setup LCM scheduler
         setup_scheduler(self.pipe)
@@ -773,29 +777,28 @@ class RetroArtConverter:
         
         pipe_kwargs["generator"] = generator
         
-        # --- START FIX: Use CompelForSDXL as per documentation ---
-        # This fixes the "TypeError" and the "93 vs 77" error
-        if self.use_compel and self.compel is not None:
+        # --- START FIX: Use our new Cappella module ---
+        if self.use_cappella and self.cappella is not None:
             try:
-                print("Encoding prompts with Compel...")
+                print("Encoding prompts with Cappella...")
                 
-                # Call Compel with prompt and negative_prompt kwargs
-                # Compel will handle the padding internally
-                conditioning = self.compel(prompt, negative_prompt=negative_prompt)
+                # Call Cappella once. It handles truncation and padding.
+                conditioning = self.cappella(prompt, negative_prompt)
 
-                # Unpack the results from the returned object
+                # Unpack the results
                 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}")
+                print(f"[OK] Cappella 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}")
+                print(f"Cappella encoding failed, using standard prompts: {e}")
                 traceback.print_exc()
                 pipe_kwargs["prompt"] = prompt
                 pipe_kwargs["negative_prompt"] = negative_prompt
         else:
+            print("[WARNING] Cappella not found, using standard prompt encoding.")
             pipe_kwargs["prompt"] = prompt
             pipe_kwargs["negative_prompt"] = negative_prompt
         # --- END FIX ---
@@ -842,8 +845,8 @@ class RetroArtConverter:
                         
                         print(f"  - Face embedding: {face_proj_embeds.shape}, Scale: {boosted_scale:.2f}")
                         
-                        # --- START FIX 2: Your padding solution ---
-                        # This fixes the "109 vs 93" error
+                        # --- START FIX: Your padding solution ---
+                        # This fixes the "109 vs 77" error
                         if 'prompt_embeds' in pipe_kwargs:
                             original_embeds = pipe_kwargs['prompt_embeds']
                             
@@ -870,7 +873,7 @@ class RetroArtConverter:
                             
                             print(f"  [OK] Face embeddings concatenated successfully! Prompt: {combined_embeds.shape}")
                         else:
-                            print(f"  [WARNING] Can't concatenate - no prompt_embeds (use Compel)")
+                            print(f"  [WARNING] Can't concatenate - no prompt_embeds (use Cappella)")
                         # --- END FIX 2 ---
                 
                 elif has_detected_faces:
@@ -947,4 +950,4 @@ class RetroArtConverter:
         return generated_image
 
 
-print("[OK] Generator class ready")
+print("[OK] Generator class ready")