Upload 2 files

generator.py

@@ -18,7 +18,7 @@ from utils import (
 )
 from models import (
     load_face_analysis, load_depth_detector, load_controlnets, load_image_encoder,
-    load_sdxl_pipeline, load_lora, setup_ip_adapter,
+    load_sdxl_pipeline, load_lora, setup_ip_adapter, setup_compel,
     setup_scheduler, optimize_pipeline, load_caption_model, set_clip_skip
 )
 
@@ -83,10 +83,8 @@ class RetroArtConverter:
         
         # Setup Compel
         # TEMPORARILY DISABLED - SDXL token mismatch issue
-        # Skip Compel - use native SDXL encoding instead
-        self.compel = None
-        self.use_compel = False
-        print("  [INFO] Using native SDXL prompt encoding (more reliable than Compel)")
+        # Setup Compel
+        self.compel, self.use_compel = setup_compel(self.pipe)
         print("  [INFO] Compel temporarily disabled - using standard prompts")
         
         # Setup LCM scheduler
@@ -577,10 +575,37 @@ class RetroArtConverter:
         
         pipe_kwargs["generator"] = generator
         
-        # Use native SDXL prompt encoding (more reliable than Compel)
-        print("Using native SDXL prompt encoding...")
-        pipe_kwargs["prompt"] = prompt
-        pipe_kwargs["negative_prompt"] = negative_prompt if negative_prompt and negative_prompt.strip() else None
+        # Use Compel for prompt encoding (critical for quality)
+        negative_conditioning = None  # Initialize for later use
+        if self.use_compel and self.compel is not None:
+            try:
+                print("Encoding prompts with Compel...")
+                
+                # Direct tuple unpacking as in working example
+                conditioning, pooled = self.compel(prompt)
+                
+                # Handle negative prompt conditionally
+                if negative_prompt and negative_prompt.strip():
+                    negative_conditioning, negative_pooled = self.compel(negative_prompt)
+                else:
+                    negative_conditioning, negative_pooled = None, None
+                
+                # Set embeddings for pipeline
+                pipe_kwargs["prompt_embeds"] = conditioning
+                pipe_kwargs["pooled_prompt_embeds"] = pooled
+                pipe_kwargs["negative_prompt_embeds"] = negative_conditioning
+                pipe_kwargs["negative_pooled_prompt_embeds"] = negative_pooled
+                
+                print("[OK] Using Compel-encoded prompts")
+            except Exception as e:
+                print(f"[FALLBACK] Compel failed ({e}), using standard encoding")
+                pipe_kwargs["prompt"] = prompt
+                pipe_kwargs["negative_prompt"] = negative_prompt if negative_prompt and negative_prompt.strip() else None
+        else:
+            # Fallback to native SDXL encoding
+            print("Using standard SDXL prompt encoding...")
+            pipe_kwargs["prompt"] = prompt
+            pipe_kwargs["negative_prompt"] = negative_prompt if negative_prompt and negative_prompt.strip() else None
         
         # Add CLIP skip
         if hasattr(self.pipe, 'text_encoder'):
@@ -620,9 +645,30 @@ class RetroArtConverter:
                     print(f"  - Resampler output: {face_proj_embeds.shape}")
                     print(f"  - Scale: {boosted_scale:.2f}")
                     
-                    # Pass face embeddings through image_embeds parameter (works without Compel)
-                    pipe_kwargs['image_embeds'] = face_proj_embeds
-                    print(f"  [OK] Face embeddings set for IP-Adapter!")
+                    # Handle face embeddings with or without Compel
+                    if 'prompt_embeds' in pipe_kwargs:
+                        # Compel is being used - concatenate embeddings
+                        original_embeds = pipe_kwargs['prompt_embeds']
+                        
+                        # Handle CFG (classifier-free guidance)
+                        if negative_conditioning is not None:
+                            # Duplicate for negative + positive
+                            face_proj_embeds = torch.cat([
+                                torch.zeros_like(face_proj_embeds),  # Negative
+                                face_proj_embeds                      # Positive
+                            ], dim=0)
+                        
+                        # Concatenate: [batch, text_tokens, 2048] + [batch, 16, 2048]
+                        combined_embeds = torch.cat([original_embeds, face_proj_embeds], dim=1)
+                        pipe_kwargs['prompt_embeds'] = combined_embeds
+                        
+                        print(f"  - Text embeds: {original_embeds.shape}")
+                        print(f"  - Combined embeds: {combined_embeds.shape}")
+                        print(f"  [OK] Face embeddings concatenated with text embeddings!")
+                    else:
+                        # Native encoding - use image_embeds parameter
+                        pipe_kwargs['image_embeds'] = face_proj_embeds
+                        print(f"  [OK] Face embeddings set via image_embeds!")
             
             elif has_detected_faces and self.models_loaded.get('ip_adapter', False):
                 # Face detected but embeddings unavailable

models.py

@@ -15,7 +15,7 @@ from transformers import CLIPVisionModelWithProjection
 from insightface.app import FaceAnalysis
 from controlnet_aux import LeresDetector
 from huggingface_hub import hf_hub_download
-# removed compel - using native SDXL encoding
+from compel import Compel, ReturnedEmbeddingsType
 
 # Use reference implementation's attention processor
 from attention_processor import IPAttnProcessor2_0, AttnProcessor
@@ -326,23 +326,21 @@ def setup_ip_adapter(pipe, image_encoder):
         return None, False
 
 
-# Removed setup_compel - using native SDXL encoding instead
-# def setup_compel(pipe):
-#     """Setup Compel for better SDXL prompt handling."""
-#     print("Setting up Compel for enhanced prompt processing...")
-#     try:
-#         compel = Compel(
-#             tokenizer=[pipe.tokenizer, pipe.tokenizer_2],
-#             text_encoder=[pipe.text_encoder, pipe.text_encoder_2],
-#             returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED,
-#             requires_pooled=[False, True],
-#             truncate_long_prompts=False  # Important for SDXL compatibility
-#         )
-#         print("  [OK] Compel loaded successfully")
-#         return compel, True
-#     except Exception as e:
-#         print(f"  [WARNING] Compel not available: {e}")
-#         return None, False
+def setup_compel(pipe):
+    """Setup Compel for SDXL prompt handling - based on working example."""
+    print("Setting up Compel for enhanced prompt processing...")
+    try:
+        compel = Compel(
+            tokenizer=[pipe.tokenizer, pipe.tokenizer_2],
+            text_encoder=[pipe.text_encoder, pipe.text_encoder_2],
+            returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED,
+            requires_pooled=[False, True]
+        )
+        print("  [OK] Compel loaded successfully")
+        return compel, True
+    except Exception as e:
+        print(f"  [WARNING] Compel not available: {e}")
+        return None, False
 
 
 def setup_scheduler(pipe):