Update models.py

models.py

@@ -18,8 +18,8 @@ from insightface.app import FaceAnalysis
 from controlnet_aux import ZoeDetector, OpenposeDetector, LeresDetector, MidasDetector, MediapipeFaceDetector
 from huggingface_hub import hf_hub_download, snapshot_download
 
-# --- START FIX: Import CompelForSDXL ---
-from compel import CompelForSDXL 
+# --- START FIX: Import our new Cappella module ---
+from cappella import Cappella
 # --- END FIX ---
 
 # Use reference implementation's attention processor
@@ -385,17 +385,16 @@ def setup_ip_adapter(pipe, image_encoder):
         return None, False
 
 
-# --- START FIX: Use CompelForSDXL and ENABLE TRUNCATION ---
-def setup_compel(pipe):
-    """Setup CompelForSDXL for better SDXL prompt handling."""
-    print("Setting up Compel for enhanced prompt processing (using CompelForSDXL)...")
+# --- START FIX: Use our new Cappella module ---
+def setup_cappella(pipe):
+    """Setup Cappella for our custom prompt encoding."""
+    print("Setting up Cappella (custom prompt encoder)...")
     try:
-        # Pass the whole pipeline and enable truncation
-        compel = CompelForSDXL(pipe, truncate_long_prompts=True) 
-        print("  [OK] CompelForSDXL loaded successfully (with truncation)")
-        return compel, True
+        cappella = Cappella(pipe, device) 
+        print("  [OK] Cappella loaded successfully.")
+        return cappella, True
     except Exception as e:
-        print(f"  [WARNING] Compel not available: {e}")
+        print(f"  [WARNING] Cappella not available: {e}")
         return None, False
 # --- END FIX ---
 
@@ -469,4 +468,4 @@ def set_clip_skip(pipe):
         print(f"  [OK] CLIP skip set to {CLIP_SKIP}")
 
 
-print("[OK] Model loading functions ready")
+print("[OK] Model loading functions ready")