Update models.py

models.py

@@ -1,698 +1,425 @@
 """
-Generation logic for Pixagram AI Pixel Art Generator
+Model loading and initialization for Pixagram AI Pixel Art Generator
+FIXED VERSION with proper IP-Adapter and BLIP-2 support
 """
 import torch
-import numpy as np
-import cv2
-from PIL import Image
-import torch.nn.functional as F
-from torchvision import transforms
+import time
+from diffusers import (
+    StableDiffusionXLControlNetImg2ImgPipeline,
+    ControlNetModel,
+    AutoencoderKL,
+    LCMScheduler
+)
+from diffusers.models.attention_processor import AttnProcessor2_0
+from transformers import CLIPVisionModelWithProjection
+from insightface.app import FaceAnalysis
+# --- MODIFIED: Import detectors (LeReSDetector removed) ---
+from controlnet_aux import ZoeDetector, OpenposeDetector, MidasDetector
+from huggingface_hub import hf_hub_download
+from compel import Compel, ReturnedEmbeddingsType
+
+# Use reference implementation's attention processor
+from attention_processor import IPAttnProcessor2_0, AttnProcessor
+from resampler import Resampler
+
+# --- ERROR WAS HERE ---
+# The "from models import (...)" block that was here has been removed
+# as it was causing a circular import.
+# --- END FIX ---
 
 from config import (
-    device, dtype, TRIGGER_WORD, MULTI_SCALE_FACTORS,
-    ADAPTIVE_THRESHOLDS, ADAPTIVE_PARAMS, CAPTION_CONFIG, IDENTITY_BOOST_MULTIPLIER
-)
-from utils import (
-    sanitize_text, enhanced_color_match, color_match, create_face_mask,
-    draw_kps, get_demographic_description, calculate_optimal_size, enhance_face_crop
-)
-from models import (
-    load_face_analysis, load_controlnets, load_image_encoder,
-    load_sdxl_pipeline, load_lora, setup_ip_adapter, setup_compel,
-    setup_scheduler, optimize_pipeline, load_caption_model, set_clip_skip,
-    load_openpose_detector, load_depth_models
+    device, dtype, MODEL_REPO, MODEL_FILES, HUGGINGFACE_TOKEN,
+    FACE_DETECTION_CONFIG, CLIP_SKIP, DOWNLOAD_CONFIG
 )
 
 
-class RetroArtConverter:
-    """Main class for retro art generation"""
+def download_model_with_retry(repo_id, filename, max_retries=None):
+    """Download model with retry logic and proper token handling."""
+    if max_retries is None:
+        max_retries = DOWNLOAD_CONFIG['max_retries']
     
-    def __init__(self):
-        self.device = device
-        self.dtype = dtype
-        self.models_loaded = {
-            'custom_checkpoint': False,
-            'lora': False,
-            'instantid': False,
-            'depth': False,
-            'ip_adapter': False,
-            'openpose': False
-        }
-        
-        # Initialize face analysis
-        self.face_app, self.face_detection_enabled = load_face_analysis()
-        
-        # Load Depth Detector Chain (Zoe -> MiDaS)
-        self.depth_detector, self.depth_detector_name, depth_success = load_depth_models()
-        self.models_loaded['depth'] = depth_success
-        
-        # Load OpenPose detector
-        self.openpose_detector, openpose_success = load_openpose_detector()
-        self.models_loaded['openpose'] = openpose_success
-        
-        # Load ControlNets
-        controlnet_depth, self.controlnet_instantid, self.controlnet_openpose, instantid_success = load_controlnets(self.depth_detector_name)
-        self.controlnet_depth = controlnet_depth
-        self.instantid_enabled = instantid_success
-        self.models_loaded['instantid'] = instantid_success
-        
-        # Load image encoder
-        if self.instantid_enabled:
-            self.image_encoder = load_image_encoder()
-        else:
-            self.image_encoder = None
-        
-        # Robust ControlNet Loading
-        self.instantid_active = self.instantid_enabled and self.controlnet_instantid is not None
-        self.depth_active = self.depth_detector is not None and self.controlnet_depth is not None
-        self.openpose_active = self.openpose_detector is not None and self.controlnet_openpose is not None
-        
-        controlnets = []
-        if self.instantid_active:
-            controlnets.append(self.controlnet_instantid)
-            print("  [CN] InstantID (Identity) active")
-        else:
-            print("  [CN] InstantID (Identity) DISABLED")
+    for attempt in range(max_retries):
+        try:
+            print(f"  Attempting to download {filename} (attempt {attempt + 1}/{max_retries})...")
             
-        if self.depth_active:
-            controlnets.append(self.controlnet_depth)
-            print(f"  [CN] Depth ({self.depth_detector_name}) active")
-        else:
-            print(f"  [CN] Depth ({self.depth_detector_name}) DISABLED (Detector or ControlNet missing)")
-
-        if self.openpose_active:
-            controlnets.append(self.controlnet_openpose)
-            print("  [CN] OpenPose (Expression) active")
-        else:
-            print("  [CN] OpenPose (Expression) DISABLED (Detector or ControlNet missing)")
+            kwargs = {"repo_type": "model"}
+            if HUGGINGFACE_TOKEN:
+                kwargs["token"] = HUGGINGFACE_TOKEN
+            
+            path = hf_hub_download(
+                repo_id=repo_id,
+                filename=filename,
+                **kwargs
+            )
+            print(f"  [OK] Downloaded: {filename}")
+            return path
             
-        if not controlnets:
-            print("[WARNING] No ControlNets loaded!")
-        
-        print(f"Initializing with {len(controlnets)} active ControlNet(s)")
-        
-        # Load SDXL pipeline
-        self.pipe, checkpoint_success = load_sdxl_pipeline(controlnets if controlnets else None)
-        
-        self.models_loaded['custom_checkpoint'] = checkpoint_success
-        
-        # Load LORA
-        lora_success = load_lora(self.pipe)
-        self.models_loaded['lora'] = lora_success
-        
-        # Setup IP-Adapter
-        if self.instantid_active and self.image_encoder is not None:
-            self.image_proj_model, ip_adapter_success = setup_ip_adapter(self.pipe, self.image_encoder)
-            self.models_loaded['ip_adapter'] = ip_adapter_success
-        else:
-            print("[INFO] Face preservation: IP-Adapter disabled (InstantID model failed or encoder failed)")
-            self.models_loaded['ip_adapter'] = False
-            self.image_proj_model = None
-        
-        # Setup Compel
-        self.compel, self.use_compel = setup_compel(self.pipe)
-        
-        # Setup LCM scheduler
-        setup_scheduler(self.pipe)
-        
-        # Optimize pipeline
-        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")
-            else:
-                print("  [OK] Caption model loaded")
-        
-        
-        # Set CLIP skip
-        set_clip_skip(self.pipe)
-        
-        # Track controlnet configuration
-        self.using_multiple_controlnets = isinstance(controlnets, list)
-        print(f"Pipeline initialized with {'multiple' if self.using_multiple_controlnets else 'single'} ControlNet(s)")
-        
-        # Print model status
-        self._print_status()
-        
-        print("  [OK] Model initialization complete!")
-    
-    def _print_status(self):
-        """Print model loading status"""
-        print("\n=== MODEL STATUS ===")
-        for model, loaded in self.models_loaded.items():
-            status = "[OK] LOADED" if loaded else "[FALLBACK/DISABLED]"
-            print(f"{model}: {status}")
-        print("===================\n")
-        
-        print("=== UPGRADE VERIFICATION ===")
-        try:
-            # Check for enhanced classes if they exist
-            pass
         except Exception as e:
-            print(f"[INFO] Verification skipped: {e}")
-        print("============================\n")
-    
-    def get_depth_map(self, image):
-            """Generate depth map using the loaded detector (Zoe/MiDaS)"""
-            if self.depth_detector is not None:
-                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_resized = image.resize(size_for_depth, Image.LANCZOS)
-
-                    # --- FIX for numpy.int64 error ---
-                    # .copy() forces PIL to create a new image, stripping numpy-typed metadata
-                    image_for_depth = image_resized.copy()
-                    # --- END FIX ---
-                    
-                    if target_width != orig_width or target_height != orig_height:
-                        print(f"[DEPTH] Resized for {self.depth_detector_name}Detector: {orig_width}x{orig_height} -> {target_width}x{target_height}")
-                    
-                    with torch.no_grad():
-                        depth_image = self.depth_detector(image_for_depth)
-                    
-                    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)
-                    
-                    print(f"[DEPTH] {self.depth_detector_name} depth map generated: {orig_width}x{orig_height}")
-                    return depth_image
-                    
-                except Exception as e:
-                    print(f"[DEPTH] {self.depth_detector_name}Detector failed ({e}), falling back to grayscale depth")
-                    gray = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2GRAY)
-                    depth_colored = cv2.cvtColor(gray, cv2.COLOR_GRAY2RGB)
-                    return Image.fromarray(depth_colored)
+            print(f"  [WARNING] Download attempt {attempt + 1} failed: {e}")
+            
+            if attempt < max_retries - 1:
+                print(f"  Retrying in {DOWNLOAD_CONFIG['retry_delay']} seconds...")
+                time.sleep(DOWNLOAD_CONFIG['retry_delay'])
             else:
-                print("[DEPTH] No depth detector active, falling back to grayscale depth")
-                gray = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2GRAY)
-                depth_colored = cv2.cvtColor(gray, cv2.COLOR_GRAY2RGB)
-                return Image.fromarray(depth_colored)    
+                print(f"  [ERROR] Failed to download {filename} after {max_retries} attempts")
+                raise
+    
+    return None
 
+
+def load_face_analysis():
+    """Load face analysis model with proper error handling."""
+    print("Loading face analysis model...")
+    try:
+        face_app = FaceAnalysis(
+            name=FACE_DETECTION_CONFIG['model_name'],
+            root='./models/insightface',
+            providers=['CUDAExecutionProvider', 'CPUExecutionProvider']
+        )
+        face_app.prepare(
+            ctx_id=FACE_DETECTION_CONFIG['ctx_id'], 
+            det_size=FACE_DETECTION_CONFIG['det_size']
+        )
+        print("  [OK] Face analysis model loaded successfully")
+        return face_app, True
+    except Exception as e:
+        print(f"  [WARNING] Face detection not available: {e}")
+        return None, False
+
+# --- MODIFIED FUNCTION: Depth Detector Fallback Chain (Zoe -> MiDaS) ---
+def load_depth_models():
+    """
+    Load depth detector with fallback: Zoe -> MiDaS.
+    """
+    print("Loading depth detector...")
+
+    # 1. Try Zoe
+    try:
+        detector = ZoeDetector.from_pretrained("lllyasviel/Annotators")
+        detector.to(device)
+        print("  [OK] Using Zoe Depth detector")
+        return detector, "zoe", True
+    except Exception as e_zoe:
+        print(f"  [INFO] Zoe failed ({e_zoe}), falling back to MiDaS...")
+        
+    # 2. Try MiDaS
+    try:
+        detector = MidasDetector.from_pretrained("lllyasviel/Annotators")
+        detector.to(device)
+        print("  [OK] Using MiDaS Depth detector")
+        return detector, "midas", True
+    except Exception as e_midas:
+        print(f"  [WARNING] All depth detectors failed ({e_midas})")
+        return None, "none", False
+# --- END MODIFICATION ---
+
+
+def load_openpose_detector():
+    """Load OpenPose detector."""
+    print("Loading OpenPose detector...")
+    try:
+        openpose = OpenposeDetector.from_pretrained("lllyasviel/Annotators")
+        openpose.to(device)
+        print("  [OK] OpenPose loaded successfully")
+        return openpose, True
+    except Exception as e:
+        print(f"  [WARNING] OpenPose not available: {e}")
+        return None, False
+
+
+def load_controlnets(depth_detector_name="zoe"):
+    """Load ControlNet models."""
     
-    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
+    # This logic correctly handles the "zoe" or "midas" name
+    depth_model_repo = {
+        "zoe": "diffusers/controlnet-zoe-depth-sdxl-1.0",
+        "midas": "diffusers/controlnet-midas-sdxl-1.0",
+        "none": None
+    }
     
-    def extract_multi_scale_face(self, face_crop, face):
-        """
-        Extract face features at multiple scales for better detail.
-        +1-2% improvement in face preservation.
-        """
-        try:
-            multi_scale_embeds = []
-            
-            for scale in MULTI_SCALE_FACTORS:
-                # Resize
-                w, h = face_crop.size
-                scaled_size = (int(w * scale), int(h * scale))
-                scaled_crop = face_crop.resize(scaled_size, Image.LANCZOS)
-                
-                # Pad/crop back to original
-                scaled_crop = scaled_crop.resize((w, h), Image.LANCZOS)
-                
-                # Extract features
-                scaled_array = cv2.cvtColor(np.array(scaled_crop), cv2.COLOR_RGB2BGR)
-                scaled_faces = self.face_app.get(scaled_array)
-                
-                if len(scaled_faces) > 0:
-                    multi_scale_embeds.append(scaled_faces[0].normed_embedding)
-            
-            # Average embeddings
-            if len(multi_scale_embeds) > 0:
-                averaged = np.mean(multi_scale_embeds, axis=0)
-                # Renormalize
-                averaged = averaged / np.linalg.norm(averaged)
-                print(f"[MULTI-SCALE] Combined {len(multi_scale_embeds)} scales")
-                return averaged
-            
-            return face.normed_embedding
-        
-        except Exception as e:
-            print(f"[MULTI-SCALE] Failed: {e}, using single scale")
-            return face.normed_embedding
+    repo_id = depth_model_repo.get(depth_detector_name)
+    controlnet_depth = None
     
-    def detect_face_quality(self, face):
-        """
-        Detect face quality and adaptively adjust parameters.
-        +2-3% consistency improvement.
-        """
+    if repo_id:
+        print(f"Loading ControlNet Depth model for {depth_detector_name} ({repo_id})...")
         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 identity 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/side view (if pose available)
-            elif hasattr(face, 'pose') and len(face.pose) > 1:
-                try:
-                    yaw = float(face.pose[1])
-                    if abs(yaw) > ADAPTIVE_THRESHOLDS['profile_angle']:
-                        return ADAPTIVE_PARAMS['profile_view'].copy()
-                except (ValueError, TypeError, IndexError):
-                    pass
-            
-            # Good quality face - use provided parameters
-            return None
-        
+            controlnet_depth = ControlNetModel.from_pretrained(
+                repo_id,
+                torch_dtype=dtype
+            ).to(device)
+            print(f"  [OK] ControlNet {depth_detector_name} Depth loaded")
         except Exception as e:
-            print(f"[ADAPTIVE] Quality detection failed: {e}")
-            return None
+            print(f"  [WARNING] Could not load {depth_detector_name} ControlNet: {e}")
+    else:
+        print("  [INFO] No depth detector loaded, skipping depth ControlNet.")
     
-    def validate_and_adjust_parameters(self, strength, guidance_scale, lora_scale, 
-                                       identity_preservation, identity_control_scale,
-                                       depth_control_scale, consistency_mode=True,
-                                       expression_control_scale=0.6):
-        """
-        Enhanced parameter validation with stricter rules for consistency.
-        """
-        if consistency_mode:
-            print("[CONSISTENCY] Applying strict parameter validation...")
-            adjustments = []
-            
-            # Rule 1: Strong inverse relationship between identity and LORA
-            if identity_preservation > 1.2:
-                original_lora = lora_scale
-                lora_scale = min(lora_scale, 1.0)
-                if abs(lora_scale - original_lora) > 0.01:
-                    adjustments.append(f"LORA: {original_lora:.2f}->{lora_scale:.2f} (high identity)")
-            
-            # Rule 2: Strength-based profile activation
-            if strength < 0.5:
-                # Maximum preservation mode
-                if identity_preservation < 1.3:
-                    original_identity = identity_preservation
-                    identity_preservation = 1.3
-                    adjustments.append(f"Identity: {original_identity:.2f}->{identity_preservation:.2f} (max preservation)")
-                if lora_scale > 0.9:
-                    original_lora = lora_scale
-                    lora_scale = 0.9
-                    adjustments.append(f"LORA: {original_lora:.2f}->{lora_scale:.2f} (max preservation)")
-                if guidance_scale > 1.3:
-                    original_cfg = guidance_scale
-                    guidance_scale = 1.3
-                    adjustments.append(f"CFG: {original_cfg:.2f}->{guidance_scale:.2f} (max preservation)")
-                    
-            elif strength > 0.7:
-                # Artistic transformation mode
-                if identity_preservation > 1.0:
-                    original_identity = identity_preservation
-                    identity_preservation = 1.0
-                    adjustments.append(f"Identity: {original_identity:.2f}->{identity_preservation:.2f} (artistic mode)")
-                if lora_scale < 1.2:
-                    original_lora = lora_scale
-                    lora_scale = 1.2
-                    adjustments.append(f"LORA: {original_lora:.2f}->{lora_scale:.2f} (artistic mode)")
-            
-            # Rule 3: CFG-LORA relationship
-            if guidance_scale > 1.4 and lora_scale > 1.2:
-                original_lora = lora_scale
-                lora_scale = 1.1
-                adjustments.append(f"LORA: {original_lora:.2f}->{lora_scale:.2f} (high CFG detected)")
-            
-            # Rule 4: LCM sweet spot enforcement
-            original_cfg = guidance_scale
-            guidance_scale = max(1.0, min(guidance_scale, 1.5))
-            if abs(guidance_scale - original_cfg) > 0.01:
-                adjustments.append(f"CFG: {original_cfg:.2f}->{guidance_scale:.2f} (LCM optimal)")
-            
-            # Rule 5: ControlNet balance
-            total_control = 0
-            if self.instantid_active:
-                total_control += identity_control_scale
-            if self.depth_active:
-                total_control += depth_control_scale
-            if self.openpose_active:
-                total_control += expression_control_scale
-            
-            if total_control > 2.0:
-                scale_factor = 2.0 / total_control
-                original_id_ctrl = identity_control_scale
-                original_depth_ctrl = depth_control_scale
-                original_expr_ctrl = expression_control_scale
-                
-                if self.instantid_active:
-                    identity_control_scale *= scale_factor
-                if self.depth_active:
-                    depth_control_scale *= scale_factor
-                if self.openpose_active:
-                    expression_control_scale *= scale_factor
-                
-                adjustments.append(f"ControlNets balanced: ID {original_id_ctrl:.2f}->{identity_control_scale:.2f}, Depth {original_depth_ctrl:.2f}->{depth_control_scale:.2f}, Expr {original_expr_ctrl:.2f}->{expression_control_scale:.2f}")
-            
-            if adjustments:
-                print("  [OK] Applied adjustments:")
-                for adj in adjustments:
-                    print(f"    - {adj}")
-            else:
-                print("  [OK] Parameters already optimal")
-        
-        return strength, guidance_scale, lora_scale, identity_preservation, identity_control_scale, depth_control_scale, expression_control_scale
+    print("Loading ControlNet OpenPose model...")
+    try:
+        controlnet_openpose = ControlNetModel.from_pretrained(
+            "diffusers/controlnet-openpose-sdxl-1.0",
+            torch_dtype=dtype
+        ).to(device)
+        print("  [OK] ControlNet OpenPose loaded")
+    except Exception as e:
+        print(f"  [WARNING] ControlNet OpenPose not available: {e}")
+        controlnet_openpose = None
     
-    def generate_caption(self, image, max_length=None, num_beams=None):
-        """Generate a descriptive caption for the image (supports BLIP-2, GIT, BLIP)."""
-        if not self.caption_enabled or self.caption_model is None:
-            return None
-        
-        if max_length is None:
-            if self.caption_model_type == "blip2":
-                max_length = 50
-            elif self.caption_model_type == "git":
-                max_length = 40
-            else:
-                max_length = CAPTION_CONFIG['max_length']
-        
-        if num_beams is None:
-            num_beams = CAPTION_CONFIG['num_beams']
-        
-        try:
-            if self.caption_model_type == "blip2":
-                inputs = self.caption_processor(image, return_tensors="pt").to(self.device, self.dtype)
-                with torch.no_grad():
-                    output = self.caption_model.generate(
-                        **inputs, max_length=max_length, num_beams=num_beams, min_length=10,
-                        length_penalty=1.0, repetition_penalty=1.5, early_stopping=True
-                    )
-                caption = self.caption_processor.decode(output[0], skip_special_tokens=True)
-                
-            elif self.caption_model_type == "git":
-                inputs = self.caption_processor(images=image, return_tensors="pt").to(self.device, self.dtype)
-                with torch.no_grad():
-                    output = self.caption_model.generate(
-                        pixel_values=inputs.pixel_values, max_length=max_length, num_beams=num_beams, min_length=10,
-                        length_penalty=1.0, repetition_penalty=1.5, early_stopping=True
-                    )
-                caption = self.caption_processor.batch_decode(output, skip_special_tokens=True)[0]
-                
-            else:
-                inputs = self.caption_processor(image, return_tensors="pt").to(self.device, self.dtype)
-                with torch.no_grad():
-                    output = self.caption_model.generate(
-                        **inputs, max_length=max_length, num_beams=num_beams, early_stopping=True
-                    )
-                caption = self.caption_processor.decode(output[0], skip_special_tokens=True)
-            
-            return caption.strip()
-        
-        except Exception as e:
-            print(f"Caption generation failed: {e}")
-            return None
+    print("Loading InstantID ControlNet...")
+    try:
+        controlnet_instantid = ControlNetModel.from_pretrained(
+            "InstantX/InstantID",
+            subfolder="ControlNetModel",
+            torch_dtype=dtype
+        ).to(device)
+        print("  [OK] InstantID ControlNet loaded successfully")
+        return controlnet_depth, controlnet_instantid, controlnet_openpose, True
+    except Exception as e:
+        print(f"  [WARNING] InstantID ControlNet not available: {e}")
+        return controlnet_depth, None, controlnet_openpose, False
+
+
+def load_image_encoder():
+    """Load CLIP Image Encoder for IP-Adapter."""
+    print("Loading CLIP Image Encoder for IP-Adapter...")
+    try:
+        image_encoder = CLIPVisionModelWithProjection.from_pretrained(
+            "h94/IP-Adapter",
+            subfolder="models/image_encoder",
+            torch_dtype=dtype
+        ).to(device)
+        print("  [OK] CLIP Image Encoder loaded successfully")
+        return image_encoder
+    except Exception as e:
+        print(f"  [ERROR] Could not load image encoder: {e}")
+        return None
+
+
+def load_sdxl_pipeline(controlnets):
+    """Load SDXL checkpoint from HuggingFace Hub."""
+    print("Loading SDXL checkpoint (horizon) with bundled VAE from HuggingFace Hub...")
+    try:
+        model_path = download_model_with_retry(MODEL_REPO, MODEL_FILES['checkpoint'])
+        
+        pipe = StableDiffusionXLControlNetImg2ImgPipeline.from_single_file(
+            model_path,
+            controlnet=controlnets,
+            torch_dtype=dtype,
+            use_safetensors=True
+        ).to(device)
+        print("  [OK] Custom checkpoint loaded successfully (VAE bundled)")
+        return pipe, True
+    except Exception as e:
+        print(f"  [WARNING] Could not load custom checkpoint: {e}")
+        print("  Using default SDXL base model")
+        pipe = StableDiffusionXLControlNetImg2ImgPipeline.from_pretrained(
+            "stabilityai/stable-diffusion-xl-base-1.0",
+            controlnet=controlnets,
+            torch_dtype=dtype,
+            use_safetensors=True
+        ).to(device)
+        return pipe, False
+
+
+def load_lora(pipe):
+    """Load LORA from HuggingFace Hub."""
+    print("Loading LORA (retroart) from HuggingFace Hub...")
+    try:
+        lora_path = download_model_with_retry(MODEL_REPO, MODEL_FILES['lora'])
+        pipe.load_lora_weights(lora_path, adapter_name="retroart")
+        print(f"  [OK] LORA loaded successfully")
+        return True
+    except Exception as e:
+        print(f"  [WARNING] Could not load LORA: {e}")
+        return False
+
+
+def setup_ip_adapter(pipe, image_encoder):
+    """
+    Setup IP-Adapter for InstantID face embeddings - PROPER IMPLEMENTATION.
+    Based on the reference InstantID pipeline.
+    """
+    if image_encoder is None:
+        return None, False
     
-    def generate_retro_art(
-        self,
-        input_image,
-        prompt="retro game character, vibrant colors, detailed",
-        negative_prompt="blurry, low quality, ugly, distorted",
-        num_inference_steps=12,
-        guidance_scale=1.0,
-        depth_control_scale=0.8,
-        identity_control_scale=0.85,
-        expression_control_scale=0.6,
-        lora_scale=1.0,
-        identity_preservation=0.8,
-        strength=0.75,
-        enable_color_matching=False,
-        consistency_mode=True,
-        seed=-1
-    ):
-        """Generate retro art with img2img pipeline and enhanced InstantID"""
-        
-        # --- FIX for Compel tensor mismatch error ---
-        prompt = sanitize_text(prompt)
-        if not prompt or not prompt.strip():
-            prompt = "" # Ensure prompt is not None or just whitespace
-            
-        negative_prompt = sanitize_text(negative_prompt)
-        if not negative_prompt or not negative_prompt.strip():
-            negative_prompt = "" # Ensure negative_prompt is "" if blank
-        # --- END FIX ---
-        
-        if consistency_mode:
-            print("\n[CONSISTENCY] Validating and adjusting parameters...")
-            strength, guidance_scale, lora_scale, identity_preservation, identity_control_scale, depth_control_scale, expression_control_scale = \
-                self.validate_and_adjust_parameters(
-                    strength, guidance_scale, lora_scale, identity_preservation, 
-                    identity_control_scale, depth_control_scale, consistency_mode,
-                    expression_control_scale
-                )
-        
-        prompt = self.add_trigger_word(prompt)
-        
-        original_width, original_height = input_image.size
-        target_width, target_height = calculate_optimal_size(original_width, original_height)
-        
-        print(f"Resizing from {original_width}x{original_height} to {target_width}x{target_height}")
-        print(f"Prompt: {prompt}")
-        print(f"Img2Img Strength: {strength}")
-        
-        resized_image = input_image.resize((int(target_width), int(target_height)), Image.LANCZOS)
-        
-        depth_image = None
-        if self.depth_active:
-            depth_image = self.get_depth_map(resized_image)
-            if depth_image.size != (target_width, target_height):
-                depth_image = depth_image.resize((int(target_width), int(target_height)), Image.LANCZOS)
-        
-        openpose_image = None
-        if self.openpose_active:
-            print("Generating OpenPose map...")
+    print("Setting up IP-Adapter for InstantID face embeddings (proper implementation)...")
+    try:
+        # Download InstantID weights
+        ip_adapter_path = download_model_with_retry(
+            "InstantX/InstantID",
+            "ip-adapter.bin"
+        )
+        
+        # Load full state dict
+        state_dict = torch.load(ip_adapter_path, map_location="cpu")
+        
+        # Extract image_proj and ip_adapter weights
+        image_proj_state_dict = {}
+        ip_adapter_state_dict = {}
+        
+        for key, value in state_dict.items():
+            if key.startswith("image_proj."):
+                image_proj_state_dict[key.replace("image_proj.", "")] = value
+            elif key.startswith("ip_adapter."):
+                ip_adapter_state_dict[key.replace("ip_adapter.", "")] = value
+        
+        print("Creating Resampler (Perceiver architecture) with custom settings...")
+        image_proj_model = Resampler(
+            dim=1280,
+            depth=8,
+            dim_head=64,
+            heads=20,
+            num_queries=32,
+            embedding_dim=512,
+            output_dim=pipe.unet.config.cross_attention_dim,
+            ff_mult=4
+        )
+        
+        image_proj_model.eval()
+        image_proj_model = image_proj_model.to(device, dtype=dtype)
+        
+        # Load image_proj weights
+        if image_proj_state_dict:
             try:
-                openpose_image = self.openpose_detector(resized_image, face_only=True)
+                image_proj_model.load_state_dict(image_proj_state_dict, strict=True)
+                print("  [OK] Resampler loaded with pretrained weights")
             except Exception as e:
-                print(f"OpenPose failed, using blank map: {e}")
-                openpose_image = Image.new("RGB", (target_width, target_height), (0,0,0))
-        
+                print(f"  [WARNING] Could not load Resampler weights: {e}")
+                print("  Using randomly initialized Resampler")
+        else:
+            print("  [WARNING] No image_proj weights found, using random initialization")
         
-        face_kps_image = None
-        face_embeddings = None
-        face_crop_enhanced = None
-        has_detected_faces = False
-        face_bbox_original = None
+        # Setup IP-Adapter attention processors
+        print("Setting up IP-Adapter attention processors...")
+        attn_procs = {}
+        num_tokens = 32
         
-        if self.instantid_active and self.face_app is not None:
-            print("Detecting faces and extracting keypoints...")
-            img_array = cv2.cvtColor(np.array(resized_image), cv2.COLOR_RGB2BGR)
-            faces = self.face_app.get(img_array)
+        for name in pipe.unet.attn_processors.keys():
+            cross_attention_dim = None if name.endswith("attn1.processor") else pipe.unet.config.cross_attention_dim
             
-            if len(faces) > 0:
-                has_detected_faces = True
-                print(f"Detected {len(faces)} face(s)")
-                
-                face = sorted(faces, key=lambda x: (x.bbox[2] - x.bbox[0]) * (x.bbox[3] - x.bbox[1]))[-1]
-                
-                adaptive_params = self.detect_face_quality(face)
-                if adaptive_params is not None:
-                    print(f"[ADAPTIVE] {adaptive_params['reason']}")
-                    identity_preservation = adaptive_params['identity_preservation']
-                    identity_control_scale = adaptive_params['identity_control_scale']
-                    guidance_scale = adaptive_params['guidance_scale']
-                    lora_scale = adaptive_params['lora_scale']
-                
-                face_embeddings_base = face.normed_embedding
-                
-                bbox = face.bbox.astype(int)
-                x1, y1, x2, y2 = bbox[0], bbox[1], bbox[2], bbox[3]
-                face_bbox_original = [x1, y1, x2, y2]
-                
-                face_width = x2 - x1
-                face_height = y2 - y1
-                padding_x = int(face_width * 0.3)
-                padding_y = int(face_height * 0.3)
-                x1 = max(0, x1 - padding_x)
-                y1 = max(0, y1 - padding_y)
-                x2 = min(resized_image.width, x2 + padding_x)
-                y2 = min(resized_image.height, y2 + padding_y)
-                
-                face_crop = resized_image.crop((x1, y1, x2, y2))
-                
-                face_embeddings = self.extract_multi_scale_face(face_crop, face)
-                face_crop_enhanced = enhance_face_crop(face_crop)
-                face_kps = face.kps
-                face_kps_image = draw_kps(resized_image, face_kps)
-                
-                from utils import get_facial_attributes, build_enhanced_prompt
-                facial_attrs = get_facial_attributes(face)
-                prompt = build_enhanced_prompt(prompt, facial_attrs, TRIGGER_WORD)
-                
-                age = facial_attrs['age']
-                gender_code = facial_attrs['gender']
-                det_score = facial_attrs['quality']
-                
-                gender_str = 'M' if gender_code == 1 else ('F' if gender_code == 0 else 'N/A')
-                print(f"Face info: bbox={face.bbox}, age={age if age else 'N/A'}, gender={gender_str}")
-                print(f"Face crop size: {face_crop.size}, enhanced: {face_crop_enhanced.size if face_crop_enhanced else 'N/A'}")
-        
-        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}")
-        
-        pipe_kwargs = {
-            "image": resized_image,
-            "strength": strength,
-            "num_inference_steps": num_inference_steps,
-            "guidance_scale": guidance_scale,
-        }
-        
-        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}")
+            if name.startswith("mid_block"):
+                hidden_size = pipe.unet.config.block_out_channels[-1]
+            elif name.startswith("up_blocks"):
+                block_id = int(name[len("up_blocks.")])
+                hidden_size = list(reversed(pipe.unet.config.block_out_channels))[block_id]
+            elif name.startswith("down_blocks"):
+                block_id = int(name[len("down_blocks.")])
+                hidden_size = pipe.unet.config.block_out_channels[block_id]
+            else:
+                hidden_size = pipe.unet.config.block_out_channels[-1]
+            
+            if cross_attention_dim is None:
+                attn_procs[name] = AttnProcessor2_0()
+            else:
+                attn_procs[name] = IPAttnProcessor2_0(
+                    hidden_size=hidden_size,
+                    cross_attention_dim=cross_attention_dim,
+                    scale=1.0,
+                    num_tokens=num_tokens
+                ).to(device, dtype=dtype)
         
-        pipe_kwargs["generator"] = generator
+        pipe.unet.set_attn_processor(attn_procs)
         
-        if self.use_compel and self.compel is not None:
+        if ip_adapter_state_dict:
             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")
+                ip_layers = torch.nn.ModuleList(pipe.unet.attn_processors.values())
+                ip_layers.load_state_dict(ip_adapter_state_dict, strict=False)
+                print("  [OK] IP-Adapter attention weights loaded")
             except Exception as e:
-                print(f"Compel encoding 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
-        
-        if hasattr(self.pipe, 'text_encoder'):
-            pipe_kwargs["clip_skip"] = 2
-        
-        control_images = []
-        conditioning_scales = []
-        scale_debug_str = []
-        
-        if self.instantid_active:
-            if has_detected_faces and face_kps_image is not None:
-                control_images.append(face_kps_image)
-                conditioning_scales.append(identity_control_scale)
-                scale_debug_str.append(f"Identity: {identity_control_scale:.2f}")
-
-                if face_embeddings is not None and self.models_loaded.get('ip_adapter', False) and face_crop_enhanced is not None:
-                    print(f"Processing InstantID face embeddings with Resampler...")
-                    
-                    with torch.no_grad():
-                        face_emb_tensor = torch.from_numpy(face_embeddings).to(device=self.device, dtype=self.dtype)
-                        face_emb_tensor = face_emb_tensor.reshape(1, -1, 512)
-                        face_proj_embeds = self.image_proj_model(face_emb_tensor)
-                        
-                        boosted_scale = identity_preservation * IDENTITY_BOOST_MULTIPLIER
-                        face_proj_embeds = face_proj_embeds * boosted_scale
-                        
-                        print(f"  - Face embedding: {face_emb_tensor.shape} -> {face_proj_embeds.shape}, Scale: {boosted_scale:.2f}")
-                        
-                        if 'prompt_embeds' in pipe_kwargs:
-                            original_embeds = pipe_kwargs['prompt_embeds']
-                            
-                            if original_embeds.shape[0] > 1: # Handle CFG
-                                face_proj_embeds = torch.cat([torch.zeros_like(face_proj_embeds), face_proj_embeds], dim=0)
-                            
-                            combined_embeds = torch.cat([original_embeds, face_proj_embeds], dim=1)
-                            pipe_kwargs['prompt_embeds'] = combined_embeds
-                            print(f"  [OK] Face embeddings concatenated successfully! New shape: {combined_embeds.shape}")
-                        else:
-                            print(f"  [WARNING] Can't concatenate - no prompt_embeds (use Compel)")
-                
-                elif has_detected_faces:
-                    print("  Face detected but IP-Adapter/embeddings unavailable, using keypoints only")
-
-            else:
-                print("Using blank map for InstantID (no face/disabled)")
-                control_images.append(Image.new("RGB", (target_width, target_height), (0,0,0)))
-                conditioning_scales.append(0.0)
-                scale_debug_str.append("Identity: 0.00")
-
-        if self.depth_active:
-            control_images.append(depth_image)
-            conditioning_scales.append(depth_control_scale)
-            scale_debug_str.append(f"Depth ({self.depth_detector_name}): {depth_control_scale:.2f}")
-
-        if self.openpose_active:
-            control_images.append(openpose_image)
-            conditioning_scales.append(expression_control_scale)
-            scale_debug_str.append(f"Expression: {expression_control_scale:.2f}")
-
-        if control_images:
-            pipe_kwargs["control_image"] = control_images
-            pipe_kwargs["controlnet_conditioning_scale"] = conditioning_scales
-            print(f"Active ControlNets: {len(control_images)}")
+                print(f"  [WARNING] Could not load IP-Adapter weights: {e}")
         else:
-            print("No active ControlNets, running standard Img2Img")
+            print("  [WARNING] No ip_adapter weights found")
         
+        pipe.image_encoder = image_encoder
         
-        print(f"Generating with LCM: Steps={num_inference_steps}, CFG={guidance_scale}, Strength={strength}")
-        print(f"Controlnet scales - {' | '.join(scale_debug_str)}")
-        result = self.pipe(**pipe_kwargs)
+        print("  [OK] IP-Adapter fully loaded with InstantID architecture")
+        print(f"  - Resampler: 8 layers, 20 heads, 32 output tokens")
+        print(f"  - Face embeddings: 512D -> 32x2048D")
         
-        generated_image = result.images[0]
+        return image_proj_model, True
         
-        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:
-                    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}")
+    except Exception as e:
+        print(f"  [ERROR] Could not setup IP-Adapter: {e}")
+        import traceback
+        traceback.print_exc()
+        return None, False
+
+
+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]
+        )
+        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):
+    """Setup LCM scheduler."""
+    print("Setting up LCM scheduler...")
+    pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
+    print("  [OK] LCM scheduler configured")
+
+
+def optimize_pipeline(pipe):
+    """Apply optimizations to pipeline."""
+    if device == "cuda":
+        try:
+            pipe.enable_xformers_memory_efficient_attention()
+            print("  [OK] xformers enabled")
+        except Exception as e:
+            print(f"  [INFO] xformers not available: {e}")
+
+
+def load_caption_model():
+    """
+    Load caption model with proper error handling.
+    Tries multiple models in order of quality.
+    """
+    print("Loading caption model...")
+    
+    try:
+        from transformers import AutoProcessor, AutoModelForCausalLM
+        
+        print("  Attempting GIT-Large (recommended)...")
+        caption_processor = AutoProcessor.from_pretrained("microsoft/git-large-coco")
+        caption_model = AutoModelForCausalLM.from_pretrained(
+            "microsoft/git-large-coco",
+            torch_dtype=dtype
+        ).to(device)
+        print("  [OK] GIT-Large model loaded (produces detailed captions)")
+        return caption_processor, caption_model, True, 'git'
+    except Exception as e1:
+        print(f"  [INFO] GIT-Large not available: {e1}")
         
-        return generated_image
+        try:
+            from transformers import BlipProcessor, BlipForConditionalGeneration
+            
+            print("  Attempting BLIP base (fallback)...")
+            caption_processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
+            caption_model = BlipForConditionalGeneration.from_pretrained(
+                "Salesforce/blip-image-captioning-base",
+                torch_dtype=dtype
+            ).to(device)
+            print("  [OK] BLIP base model loaded (standard captions)")
+            return caption_processor, caption_model, True, 'blip'
+        except Exception as e2:
+            print(f"  [WARNING] Caption models not available: {e2}")
+            print("  Caption generation will be disabled")
+            return None, None, False, 'none'
+
+
+def set_clip_skip(pipe):
+    """Set CLIP skip value."""
+    if hasattr(pipe, 'text_encoder'):
+        print(f"  [OK] CLIP skip set to {CLIP_SKIP}")
 
 
-print("[OK] Generator class ready")
+print("[OK] Model loading functions ready")