Update generator.py

generator.py

@@ -1,22 +1,20 @@
 """
-Generation logic for Pixagram AI Pixel Art Generator
-UPDATED VERSION with simplified InstantID face preservation
+Generation logic for Pixagram AI Pixel Art Generator  
+UPDATED VERSION with InstantID pipeline integration
 """
 import torch
 import numpy as np
 import cv2
 from PIL import Image
-import torch.nn.functional as F
-from torchvision import transforms
+import gc
 
 from config import (
-    device, dtype, TRIGGER_WORD, MULTI_SCALE_FACTORS,
-    ADAPTIVE_THRESHOLDS, ADAPTIVE_PARAMS, CAPTION_CONFIG, IDENTITY_BOOST_MULTIPLIER
+    device, dtype, TRIGGER_WORD,
+    ADAPTIVE_THRESHOLDS, ADAPTIVE_PARAMS, CAPTION_CONFIG
 )
 from utils import (
-    sanitize_text, enhanced_color_match, color_match, create_face_mask,
-    draw_kps, get_demographic_description, calculate_optimal_size, enhance_face_crop,
-    safe_image_size, ensure_int
+    sanitize_text, enhanced_color_match, color_match,
+    get_demographic_description, calculate_optimal_size, safe_image_size
 )
 from models import (
     load_face_analysis, load_depth_detector, load_controlnets,
@@ -26,7 +24,7 @@ from models import (
 
 
 class RetroArtConverter:
-    """Main class for retro art generation with InstantID face preservation"""
+    """Main class for retro art generation with InstantID"""
     
     def __init__(self):
         self.device = device
@@ -38,14 +36,14 @@ class RetroArtConverter:
             'zoe_depth': False
         }
         
-        # Initialize face analysis
+        # Load face analysis
         self.face_app, self.face_detection_enabled = load_face_analysis()
         
-        # Load Zoe Depth detector
+        # Load depth detector
         self.zoe_depth, zoe_success = load_depth_detector()
         self.models_loaded['zoe_depth'] = zoe_success
         
-        # Load ControlNets - ALWAYS as list for InstantID pipeline
+        # Load ControlNets AS LIST
         controlnet_instantid, controlnet_depth = load_controlnets()
         controlnets = [controlnet_instantid, controlnet_depth]
         self.models_loaded['instantid'] = True
@@ -63,29 +61,22 @@ class RetroArtConverter:
         # Setup Compel
         self.compel, self.use_compel = setup_compel(self.pipe)
         
-        # Setup LCM scheduler
+        # Setup scheduler
         setup_scheduler(self.pipe)
         
-        # Optimize pipeline
+        # Optimize
         optimize_pipeline(self.pipe)
         
         # Load caption model
         self.caption_processor, self.caption_model, self.caption_enabled, self.caption_model_type = load_caption_model()
         
-        # Report caption model status
-        if self.caption_enabled and self.caption_model is not None:
-            if self.caption_model_type == "git":
-                print("  [OK] Using GIT for detailed captions")
-            elif self.caption_model_type == "blip":
-                print("  [OK] Using BLIP for standard captions")
-        
         # Set CLIP skip
         set_clip_skip(self.pipe)
         
-        # Print model status
+        # Print status
         self._print_status()
         
-        print("  [OK] Model initialization complete with InstantID!")
+        print("  [OK] RetroArtConverter initialized with InstantID!")
     
     def _print_status(self):
         """Print model loading status"""
@@ -93,110 +84,110 @@ class RetroArtConverter:
         for model, loaded in self.models_loaded.items():
             status = "[OK] LOADED" if loaded else "[FALLBACK/DISABLED]"
             print(f"{model}: {status}")
-        print("InstantID Pipeline: [OK] Active with built-in IP-Adapter")
+        print("InstantID Pipeline: [OK] ACTIVE")
+        print("IP-Adapter: [OK] Built into pipeline")
         print("===================\n")
     
     def get_depth_map(self, image):
-            """
-            Generate depth map using available depth detector.
-            Supports: LeresDetector, ZoeDetector, or MidasDetector.
-            """
-            # --- FIX 1: Check for self.zoe_depth, not self.depth_detector ---
-            if self.zoe_depth is not None:
+        """Generate depth map using Zoe Depth"""
+        if self.zoe_depth is not None:
+            try:
+                if image.mode != 'RGB':
+                    image = image.convert('RGB')
+                
+                # Use safe size helper to avoid numpy.int64 issues
+                orig_width, orig_height = safe_image_size(image)
+                
+                # Use multiples of 64
+                target_width = int((orig_width // 64) * 64)
+                target_height = int((orig_height // 64) * 64)
+                
+                target_width = int(max(64, target_width))
+                target_height = int(max(64, target_height))
+                
+                size_for_depth = (int(target_width), int(target_height))
+                image_for_depth = image.resize(size_for_depth, Image.LANCZOS)
+                
+                depth_array = self.zoe_depth(image_for_depth, detect_resolution=512, image_resolution=1024)
+                depth_image = Image.fromarray(depth_array)
+                
+                if depth_image.size != image.size:
+                    depth_image = depth_image.resize(image.size, Image.LANCZOS)
+                
+                print(f"[DEPTH] Generated depth map: {depth_image.size}")
+                return depth_image, depth_array
+            except Exception as e:
+                print(f"[DEPTH] Generation failed: {e}, using grayscale")
+                return image.convert('L').convert('RGB'), None
+        else:
+            print("[DEPTH] Detector not available, using grayscale")
+            return image.convert('L').convert('RGB'), None
+    
+    def add_trigger_word(self, prompt):
+        """Add trigger word to prompt if not present"""
+        if TRIGGER_WORD.lower() not in prompt.lower():
+            if not prompt or not prompt.strip():
+                return TRIGGER_WORD
+            return f"{TRIGGER_WORD}, {prompt}"
+        return prompt
+    
+    def detect_face_quality(self, face):
+        """Detect face quality and adaptively adjust parameters"""
+        try:
+            bbox = face.bbox
+            face_size = (bbox[2] - bbox[0]) * (bbox[3] - bbox[1])
+            det_score = float(face.det_score) if hasattr(face, 'det_score') else 1.0
+            
+            # Small face -> boost preservation
+            if face_size < ADAPTIVE_THRESHOLDS['small_face_size']:
+                return ADAPTIVE_PARAMS['small_face'].copy()
+            
+            # Low confidence -> boost preservation
+            elif det_score < ADAPTIVE_THRESHOLDS['low_confidence']:
+                return ADAPTIVE_PARAMS['low_confidence'].copy()
+            
+            # Check for profile view
+            elif hasattr(face, 'pose') and len(face.pose) > 1:
                 try:
-                    if image.mode != 'RGB':
-                        image = image.convert('RGB')
-                    
-                    orig_width, orig_height = image.size
-                    orig_width = int(orig_width)
-                    orig_height = int(orig_height)
-                    
-                    target_width = int((orig_width // 64) * 64)
-                    target_height = int((orig_height // 64) * 64)
-                    
-                    target_width = int(max(64, target_width))
-                    target_height = int(max(64, target_height))
-                    
-                    size_for_depth = (int(target_width), int(target_height))
-                    image_for_depth = image.resize(size_for_depth, Image.LANCZOS)
-                    
-                    if target_width != orig_width or target_height != orig_height:
-                        # --- FIX 2: Use "ZOE" instead of undefined self.depth_type ---
-                        print(f"[DEPTH] Resized for ZOEDetector: {orig_width}x{orig_height} -> {target_width}x{target_height}")
-                    
-                    # Use torch.no_grad() and clear cache
-                    with torch.no_grad():
-                        # --- FIX 1: Use self.zoe_depth ---
-                        self.zoe_depth.to(self.device)
-                        depth_image = self.zoe_depth(image_for_depth)
-                        self.zoe_depth.to("cpu")
-                    
-                    # ADDED: Clear GPU cache after depth detection
-                    if torch.cuda.is_available():
-                        torch.cuda.empty_cache()
-                    
-                    depth_width, depth_height = depth_image.size
-                    if depth_width != orig_width or depth_height != orig_height:
-                        depth_image = depth_image.resize((int(orig_width), int(orig_height)), Image.LANCZOS)
-                    
-                    # --- FIX 2: Use "ZOE" instead of undefined self.depth_type ---
-                    print(f"[DEPTH] ZOE depth map generated: {orig_width}x{orig_height}")
-                    return depth_image
-                    
-                except Exception as e:
-                    # --- FIX 2: Use "ZOE" instead of undefined self.depth_type ---
-                    print(f"[DEPTH] ZOEDetector failed ({e}), falling back to grayscale depth")
-                    # ADDED: Clear cache on error
-                    if torch.cuda.is_available():
-                        torch.cuda.empty_cache()
-                    
-                    gray = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2GRAY)
-                    depth_colored = cv2.cvtColor(gray, cv2.COLOR_GRAY2RGB)
-                    return Image.fromarray(depth_colored)
-            else:
-                print("[DEPTH] No depth detector available, using grayscale fallback")
-                gray = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2GRAY)
-                depth_colored = cv2.cvtColor(gray, cv2.COLOR_GRAY2RGB)
-                return Image.fromarray(depth_colored)
-
-
+                    yaw = float(face.pose[1])
+                    if abs(yaw) > ADAPTIVE_THRESHOLDS['profile_angle']:
+                        return ADAPTIVE_PARAMS['profile_view'].copy()
+                except (ValueError, TypeError, IndexError):
+                    pass
+            
+            return None
+        
+        except Exception as e:
+            print(f"[ADAPTIVE] Quality detection failed: {e}")
+            return None
+    
     def generate_caption(self, image):
-        """Generate caption for image using loaded caption model"""
+        """Generate caption for image"""
         if not self.caption_enabled or self.caption_model is None:
             return None
         
         try:
             if self.caption_model_type == 'git':
-                # GIT model
-                pixel_values = self.caption_processor(images=image, return_tensors="pt").pixel_values
-                pixel_values = pixel_values.to(device=self.device, dtype=self.dtype)
-                
-                generated_ids = self.caption_model.generate(
-                    pixel_values=pixel_values,
-                    max_length=CAPTION_CONFIG['max_length']
-                )
+                inputs = self.caption_processor(images=image, return_tensors="pt").to(self.device)
+                generated_ids = self.caption_model.generate(**inputs, max_length=CAPTION_CONFIG['max_length'])
                 caption = self.caption_processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
-                
             elif self.caption_model_type == 'blip':
-                # BLIP model
-                inputs = self.caption_processor(image, return_tensors="pt").to(self.device, self.dtype)
-                out = self.caption_model.generate(**inputs, max_new_tokens=CAPTION_CONFIG['max_length'])
-                caption = self.caption_processor.decode(out[0], skip_special_tokens=True)
-            
+                inputs = self.caption_processor(image, return_tensors="pt").to(self.device)
+                generated_ids = self.caption_model.generate(**inputs, max_length=CAPTION_CONFIG['max_length'])
+                caption = self.caption_processor.decode(generated_ids[0], skip_special_tokens=True)
             else:
                 return None
             
             return sanitize_text(caption)
-            
         except Exception as e:
-            print(f"Caption generation failed: {e}")
+            print(f"[CAPTION] Generation failed: {e}")
             return None
     
     def generate_retro_art(
         self,
         input_image,
-        prompt,
-        negative_prompt="",
+        prompt=" ",
+        negative_prompt=" ",
         num_inference_steps=12,
         guidance_scale=1.3,
         depth_control_scale=0.75,
@@ -208,232 +199,190 @@ class RetroArtConverter:
         consistency_mode=True,
         seed=-1
     ):
-        """
-        Generate retro art with InstantID face preservation.
-        
-        UPDATED: Simplified face embedding handling using InstantID pipeline.
-        """
+        """Generate retro art with InstantID face preservation"""
         
-        # Validate and adjust parameters if consistency mode is enabled
-        if consistency_mode:
-            # Ensure guidance scale is in optimal range for LCM
-            if guidance_scale < 1.0:
-                guidance_scale = 1.0
-            elif guidance_scale > 1.8:
-                guidance_scale = 1.8
+        try:
+            # Add trigger word
+            prompt = self.add_trigger_word(prompt)
+            prompt = sanitize_text(prompt)
+            negative_prompt = sanitize_text(negative_prompt)
             
-            # Ensure identity preservation and lora scale balance
-            if identity_preservation > 1.5 and lora_scale > 1.2:
-                lora_scale = min(lora_scale, 1.0)
+            print(f"[PROMPT] {prompt}")
             
-            # Ensure strength is reasonable
-            if strength < 0.3:
-                strength = 0.3
-            elif strength > 0.8:
-                strength = 0.8
-        
-        # Calculate optimal size
-        orig_width, orig_height = safe_image_size(input_image)
-        optimal_width, optimal_height = calculate_optimal_size(orig_width, orig_height)
-        
-        # Resize image
-        resized_image = input_image.resize((optimal_width, optimal_height), Image.LANCZOS)
-        
-        # Generate depth map
-        print("Generating depth map...")
-        # --- FIX 3: get_depth_map only returns one value ---
-        depth_image = self.get_depth_map(resized_image)
-        
-        if depth_image is None:
-            raise RuntimeError("Failed to generate depth map")
-        
-        # Detect faces
-        print("Detecting faces...")
-        has_detected_faces = False
-        face_kps_image = None
-        face_embeddings = None
-        face_bbox_original = None
-        
-        if self.face_detection_enabled and self.face_app is not None:
-            try:
-                faces = self.face_app.get(np.array(resized_image))
-                
-                if len(faces) > 0:
-                    has_detected_faces = True
-                    face = faces[0]
-                    
-                    # Draw keypoints
-                    face_kps_image = draw_kps(resized_image, face.kps)
-                    
-                    # Get face embeddings (512D vector from InsightFace)
-                    face_embeddings = face.embedding
-                    
-                    # Get face bounding box for color matching
-                    face_bbox_original = face.bbox
-                    
-                    print(f"  [OK] Face detected")
-                    # --- FIX 4: Clarify this is the numpy shape ---
-                    print(f"  - Embedding shape (numpy): {face_embeddings.shape}")
-                    print(f"  - Keypoints: {face.kps.shape}")
-                    print(f"  - Bbox: {face_bbox_original}")
-                    
-                    # Check for adaptive parameter adjustment
-                    face_area = (face.bbox[2] - face.bbox[0]) * (face.bbox[3] - face.bbox[1])
-                    det_score = face.det_score if hasattr(face, 'det_score') else 1.0
-                    
-                    # Apply adaptive adjustments
-                    if face_area < ADAPTIVE_THRESHOLDS['small_face_size']:
-                        print("  [ADAPTIVE] Small face detected - boosting preservation")
-                        identity_preservation = max(identity_preservation, ADAPTIVE_PARAMS['small_face']['identity_preservation'])
-                        identity_control_scale = max(identity_control_scale, ADAPTIVE_PARAMS['small_face']['identity_control_scale'])
+            # Calculate optimal size
+            orig_width, orig_height = safe_image_size(input_image)
+            optimal_width, optimal_height = calculate_optimal_size(orig_width, orig_height)
+            
+            # Resize image
+            resized_image = input_image.resize((optimal_width, optimal_height), Image.LANCZOS)
+            print(f"[SIZE] Resized to {optimal_width}x{optimal_height}")
+            
+            # Generate depth map
+            depth_image, depth_array = self.get_depth_map(resized_image)
+            
+            # Detect faces
+            has_detected_faces = False
+            face_kps_image = None
+            face_embeddings = None
+            face_bbox_original = None
+            
+            if self.face_detection_enabled and self.face_app is not None:
+                try:
+                    image_array = cv2.cvtColor(np.array(resized_image), cv2.COLOR_RGB2BGR)
+                    faces = self.face_app.get(image_array)
                     
-                    elif det_score < ADAPTIVE_THRESHOLDS['low_confidence']:
-                        print("  [ADAPTIVE] Low confidence - increasing identity weight")
-                        identity_preservation = max(identity_preservation, ADAPTIVE_PARAMS['low_confidence']['identity_preservation'])
-                        identity_control_scale = max(identity_control_scale, ADAPTIVE_PARAMS['low_confidence']['identity_control_scale'])
+                    if len(faces) > 0:
+                        has_detected_faces = True
+                        face = faces[0]
+                        
+                        # Get face embeddings (512D array)
+                        face_embeddings = face.normed_embedding
+                        
+                        # Draw keypoints
+                        from pipeline_stable_diffusion_xl_instantid_img2img import draw_kps
+                        face_kps_image = draw_kps(resized_image, face.kps)
+                        
+                        # Get bbox for color matching
+                        face_bbox_original = face.bbox
+                        
+                        # Adaptive parameter adjustment
+                        adaptive_params = self.detect_face_quality(face)
+                        if adaptive_params:
+                            print(f"[ADAPTIVE] {adaptive_params['reason']}")
+                            identity_preservation = adaptive_params.get('identity_preservation', identity_preservation)
+                            identity_control_scale = adaptive_params.get('identity_control_scale', identity_control_scale)
+                            guidance_scale = adaptive_params.get('guidance_scale', guidance_scale)
+                            lora_scale = adaptive_params.get('lora_scale', lora_scale)
+                        
+                        print(f"[FACE] Detected face with {face.det_score:.2f} confidence")
+                        print(f"[FACE] Embeddings shape: {face_embeddings.shape}")
+                    else:
+                        print("[FACE] No faces detected")
                 
-                else:
-                    print("  No faces detected in image")
+                except Exception as e:
+                    print(f"[FACE] Detection failed: {e}")
+            
+            # Set LORA scale
+            if hasattr(self.pipe, 'set_adapters') and self.models_loaded['lora']:
+                try:
+                    self.pipe.set_adapters(["retroart"], adapter_weights=[lora_scale])
+                    print(f"[LORA] Scale: {lora_scale}")
+                except Exception as e:
+                    print(f"[LORA] Could not set scale: {e}")
+            
+            # Prepare generation kwargs
+            pipe_kwargs = {
+                "image": resized_image,
+                "strength": strength,
+                "num_inference_steps": num_inference_steps,
+                "guidance_scale": guidance_scale,
+            }
+            
+            # Setup generator with seed
+            if seed == -1:
+                generator = torch.Generator(device=self.device)
+                actual_seed = generator.seed()
+                print(f"[SEED] Random: {actual_seed}")
+            else:
+                generator = torch.Generator(device=self.device).manual_seed(seed)
+                actual_seed = seed
+                print(f"[SEED] Fixed: {actual_seed}")
+            
+            pipe_kwargs["generator"] = generator
+            
+            # Use Compel for prompt encoding
+            if self.use_compel and self.compel is not None:
+                try:
+                    conditioning = self.compel(prompt)
+                    negative_conditioning = self.compel(negative_prompt)
                     
-            except Exception as e:
-                print(f"Face detection error: {e}")
-                has_detected_faces = False
-        
-        # Enhance prompt with trigger word
-        if TRIGGER_WORD not in prompt.lower():
-            prompt = f"{TRIGGER_WORD}, {prompt}"
-        
-        # Set LORA scale
-        if hasattr(self.pipe, 'set_adapters') and self.models_loaded['lora']:
-            try:
-                self.pipe.set_adapters(["retroart"], adapter_weights=[lora_scale])
-                print(f"LORA scale: {lora_scale}")
-            except Exception as e:
-                print(f"Could not set LORA scale: {e}")
-        
-        # Prepare generation kwargs
-        pipe_kwargs = {
-            "image": resized_image,
-            "strength": strength,
-            "num_inference_steps": num_inference_steps,
-            "guidance_scale": guidance_scale,
-        }
-        
-        # Setup generator with seed control
-        if seed == -1:
-            generator = torch.Generator(device=self.device)
-            actual_seed = generator.seed()
-            print(f"[SEED] Using random seed: {actual_seed}")
-        else:
-            generator = torch.Generator(device=self.device).manual_seed(seed)
-            actual_seed = seed
-            print(f"[SEED] Using fixed seed: {actual_seed}")
-        
-        pipe_kwargs["generator"] = generator
-        
-        # Use Compel for prompt encoding if available
-        if self.use_compel and self.compel is not None:
-            try:
-                print("Encoding prompts with Compel...")
-                conditioning = self.compel(prompt)
-                negative_conditioning = self.compel(negative_prompt)
-                
-                pipe_kwargs["prompt_embeds"] = conditioning[0]
-                pipe_kwargs["pooled_prompt_embeds"] = conditioning[1]
-                pipe_kwargs["negative_prompt_embeds"] = negative_conditioning[0]
-                pipe_kwargs["negative_pooled_prompt_embeds"] = negative_conditioning[1]
-                
-                print("[OK] Using Compel-encoded prompts")
-            except Exception as e:
-                print(f"Compel encoding failed, using standard prompts: {e}")
+                    pipe_kwargs["prompt_embeds"] = conditioning[0]
+                    pipe_kwargs["pooled_prompt_embeds"] = conditioning[1]
+                    pipe_kwargs["negative_prompt_embeds"] = negative_conditioning[0]
+                    pipe_kwargs["negative_pooled_prompt_embeds"] = negative_conditioning[1]
+                    
+                    print("[OK] Using Compel-encoded prompts")
+                except Exception as e:
+                    print(f"[COMPEL] Failed, using standard prompts: {e}")
+                    pipe_kwargs["prompt"] = prompt
+                    pipe_kwargs["negative_prompt"] = negative_prompt
+            else:
                 pipe_kwargs["prompt"] = prompt
                 pipe_kwargs["negative_prompt"] = negative_prompt
-        else:
-            pipe_kwargs["prompt"] = prompt
-            pipe_kwargs["negative_prompt"] = negative_prompt
-        
-        # Add CLIP skip
-        if hasattr(self.pipe, 'text_encoder'):
-            pipe_kwargs["clip_skip"] = 2
-        
-        # ========================================
-        # SIMPLIFIED: Configure ControlNets + IP-Adapter
-        # ========================================
-        if has_detected_faces and face_kps_image is not None:
-            print("Using InstantID (keypoints + embeddings) + Depth ControlNets")
             
-            # Control images: [face keypoints, depth map]
-            pipe_kwargs["control_image"] = [face_kps_image, depth_image]
-            
-            # Conditioning scales: [identity, depth]
-            pipe_kwargs["controlnet_conditioning_scale"] = [
-                identity_control_scale,
-                depth_control_scale
-            ]
-            
-            # CRITICAL: Pass face embeddings for IP-Adapter
-            # The InstantID pipeline handles the Resampler internally!
-            if face_embeddings is not None:
-                print(f"Adding face embeddings for IP-Adapter...")
+            # Configure ControlNets + IP-Adapter (SIMPLIFIED!)
+            if has_detected_faces and face_kps_image is not None:
+                print("Using InstantID (keypoints + embeddings) + Depth ControlNets")
                 
-                # --- FIX 4: Convert numpy array to torch tensor, add batch dim, and move to device ---
-                face_embeds_tensor = torch.tensor(face_embeddings, dtype=self.dtype, device=self.device).unsqueeze(0)
-                pipe_kwargs["image_embeds"] = face_embeds_tensor
+                # Control images: [face keypoints, depth map]
+                pipe_kwargs["control_image"] = [face_kps_image, depth_image]
                 
-                # Control IP-Adapter strength
-                boosted_scale = identity_preservation * IDENTITY_BOOST_MULTIPLIER
-                pipe_kwargs["ip_adapter_scale"] = boosted_scale
+                # Conditioning scales: [identity, depth]
+                pipe_kwargs["controlnet_conditioning_scale"] = [
+                    identity_control_scale,
+                    depth_control_scale
+                ]
                 
-                # --- FIX 4: Update log to show tensor shape ---
-                print(f"  - Face embeddings tensor shape: {face_embeds_tensor.shape}")
-                print(f"  - IP-Adapter scale: {boosted_scale:.2f}")
-                print(f"  [OK] Face embeddings configured")
+                # IP-Adapter face embeddings (SIMPLE - pipeline handles everything!)
+                if face_embeddings is not None:
+                    print(f"Adding face embeddings for IP-Adapter...")
+                    
+                    # Just pass the embeddings - pipeline does the rest!
+                    pipe_kwargs["image_embeds"] = face_embeddings
+                    
+                    # Control IP-Adapter strength
+                    pipe_kwargs["ip_adapter_scale"] = identity_preservation
+                    
+                    print(f"  - Face embeddings shape: {face_embeddings.shape}")
+                    print(f"  - IP-Adapter scale: {identity_preservation}")
+                    print(f"  [OK] Face embeddings configured")
+                else:
+                    print("  [WARNING] No face embeddings - using keypoints only")
+            
             else:
-                print("  [WARNING] No face embeddings - using keypoints only")
-        
-        else:
-            print("No faces detected - using Depth ControlNet only")
+                print("No faces detected - using Depth ControlNet only")
+                
+                # Use depth for both ControlNet slots (identity scale = 0)
+                pipe_kwargs["control_image"] = [depth_image, depth_image]
+                pipe_kwargs["controlnet_conditioning_scale"] = [0.0, depth_control_scale]
             
-            # Use depth for both ControlNet slots (identity scale = 0)
-            pipe_kwargs["control_image"] = [depth_image, depth_image]
-            pipe_kwargs["controlnet_conditioning_scale"] = [0.0, depth_control_scale]
-        
-        # Generate
-        print(f"Generating with LCM: Steps={num_inference_steps}, CFG={guidance_scale}, Strength={strength}")
-        print(f"ControlNet scales - Identity: {identity_control_scale}, Depth: {depth_control_scale}")
-        result = self.pipe(**pipe_kwargs)
-        
-        generated_image = result.images[0]
-        
-        # Post-processing: Color matching
-        if enable_color_matching and has_detected_faces:
-            print("Applying enhanced face-aware color matching...")
-            try:
-                if face_bbox_original is not None:
-                    generated_image = enhanced_color_match(
-                        generated_image, 
-                        resized_image, 
-                        face_bbox=face_bbox_original
-                    )
-                    print("[OK] Enhanced color matching applied (face-aware)")
-                else:
+            # Generate
+            print(f"Generating: Steps={num_inference_steps}, CFG={guidance_scale}, Strength={strength}")
+            result = self.pipe(**pipe_kwargs)
+            
+            generated_image = result.images[0]
+            
+            # Post-processing: Color matching
+            if enable_color_matching and has_detected_faces:
+                print("Applying enhanced face-aware color matching...")
+                try:
+                    if face_bbox_original is not None:
+                        generated_image = enhanced_color_match(
+                            generated_image,
+                            resized_image,
+                            face_bbox=face_bbox_original
+                        )
+                        print("[OK] Enhanced color matching applied")
+                    else:
+                        generated_image = color_match(generated_image, resized_image, mode='mkl')
+                        print("[OK] Standard color matching applied")
+                except Exception as e:
+                    print(f"[COLOR] Matching failed: {e}")
+            elif enable_color_matching:
+                print("Applying standard color matching...")
+                try:
                     generated_image = color_match(generated_image, resized_image, mode='mkl')
                     print("[OK] Standard color matching applied")
-            except Exception as e:
-                print(f"Color matching failed: {e}")
-        elif enable_color_matching:
-            print("Applying standard color matching...")
-            try:
-                generated_image = color_match(generated_image, resized_image, mode='mkl')
-                print("[OK] Standard color matching applied")
-            except Exception as e:
-                print(f"Color matching failed: {e}")
-        
-        # Memory cleanup
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
+                except Exception as e:
+                    print(f"[COLOR] Matching failed: {e}")
+            
+            return generated_image
         
-        return generated_image
+        finally:
+            # Memory cleanup
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+            gc.collect()
 
 
-print("[OK] Generator class ready with InstantID pipeline")
+print("[OK] Generator class ready with InstantID support")