src/fixed_realistic_vision_pipeline.py

"""
FIXED REALISTIC VISION + FACE SWAP PIPELINE
===========================================

Fixes:
1. Proper checkpoint loading using from_single_file() method
2. Integrated face swapping from your proven system
3. RealisticVision-optimized parameters
4. Complete pipeline with all working components
"""

import torch
import numpy as np
import cv2
from PIL import Image, ImageFilter, ImageEnhance
from typing import Optional, Union, Tuple, Dict
import os

from appearance_enhancer import ImprovedUnifiedGenderAppearanceEnhancer
from generation_validator import ImprovedGenerationValidator

# HTTP-safe setup (from your working system)
os.environ["HF_HUB_DISABLE_EXPERIMENTAL_HTTP_BACKEND"] = "1"
os.environ["HF_HUB_DISABLE_XET"] = "1"
os.environ["HF_HUB_DISABLE_HF_XET"] = "1"
os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "0"
os.environ["HF_HUB_DOWNLOAD_BACKEND"] = "requests"

class FixedRealisticVisionPipeline:
    """
    Fixed pipeline that properly loads RealisticVision and includes face swapping
    """
    
    def __init__(self, checkpoint_path: str, device: str = 'cuda'):
        self.checkpoint_path = checkpoint_path
        self.device = device
        
        print(f"🎯 Initializing FIXED RealisticVision Pipeline")
        print(f"   Checkpoint: {os.path.basename(checkpoint_path)}")
        print(f"   Fixes: Proper loading + Face swap integration")
        
        # Load pipeline with proper method
        self._load_pipeline_properly()
        
        # Initialize pose and face systems
        self._init_pose_system()
        self._init_face_system()
        
        print(f"✅ FIXED RealisticVision Pipeline ready!")
    
    def _load_pipeline_properly(self):
        """Load pipeline using proper from_single_file() method"""
        try:
            from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
            
            # Load ControlNet
            print("🔥 Loading ControlNet...")
            self.controlnet = ControlNetModel.from_pretrained(
                "lllyasviel/sd-controlnet-openpose",
                torch_dtype=torch.float16,
                cache_dir="./models",
                use_safetensors=True
            ).to(self.device)
            print("   ✅ ControlNet loaded")
            
            # CRITICAL FIX: Use from_single_file() for RealisticVision
            print("🔥 Loading RealisticVision using from_single_file()...")
            self.pipeline = StableDiffusionControlNetPipeline.from_single_file(
                self.checkpoint_path,
                controlnet=self.controlnet,
                torch_dtype=torch.float16,
                safety_checker=None,
                requires_safety_checker=False,
                use_safetensors=True,
                cache_dir="./models",
                original_config_file=None  # Let diffusers infer
            ).to(self.device)
            
            print("   ✅ RealisticVision loaded properly!")
            print("   Expected: Photorealistic style, single person bias")
            
            # Apply optimizations
            self.pipeline.enable_model_cpu_offload()
            try:
                self.pipeline.enable_xformers_memory_efficient_attention()
                print("   ✅ xformers enabled")
            except:
                print("   ⚠️ xformers not available")
            
        except Exception as e:
            print(f"❌ Pipeline loading failed: {e}")
            raise
    
    def _init_pose_system(self):
        """Initialize pose detection system"""
        print("🎯 Initializing pose system...")
        
        # Try controlnet_aux first (best quality)
        try:
            from controlnet_aux import OpenposeDetector
            self.openpose_detector = OpenposeDetector.from_pretrained('lllyasviel/ControlNet')
            self.pose_method = 'controlnet_aux'
            print("   ✅ controlnet_aux OpenPose loaded")
            return
        except Exception as e:
            print(f"   ⚠️ controlnet_aux failed: {e}")
        
        # Fallback to MediaPipe
        try:
            import mediapipe as mp
            self.mp_pose = mp.solutions.pose
            self.pose_detector = self.mp_pose.Pose(
                static_image_mode=True,
                model_complexity=2,
                enable_segmentation=False,
                min_detection_confidence=0.7
            )
            self.pose_method = 'mediapipe'
            print("   ✅ MediaPipe pose loaded")
            return
        except Exception as e:
            print(f"   ⚠️ MediaPipe failed: {e}")
        
        # Ultimate fallback
        self.pose_method = 'fallback'
        print("   ⚠️ Using fallback pose system")
    
    def _init_face_system(self):
        """Initialize face detection and swapping system"""
        print("👤 Initializing face system...")
        
        try:
            # Initialize face detection
            self.face_cascade = cv2.CascadeClassifier(
                cv2.data.haarcascades + 'haarcascade_frontalface_default.xml'
            )
            self.eye_cascade = cv2.CascadeClassifier(
                cv2.data.haarcascades + 'haarcascade_eye.xml'
            )
            print("   ✅ Face detection ready")
        except Exception as e:
            print(f"   ⚠️ Face detection failed: {e}")
            self.face_cascade = None
            self.eye_cascade = None
    
    def extract_pose(self, source_image: Union[str, Image.Image], 
                    target_size: Tuple[int, int] = (512, 512)) -> Image.Image:
        """Extract pose using best available method"""
        print("🎯 Extracting pose...")
        
        # Load and prepare image
        if isinstance(source_image, str):
            image = Image.open(source_image).convert('RGB')
        else:
            image = source_image.convert('RGB')
        
        image = image.resize(target_size, Image.Resampling.LANCZOS)
        
        # Method 1: controlnet_aux (best quality)
        if self.pose_method == 'controlnet_aux':
            try:
                pose_image = self.openpose_detector(image, hand_and_face=True)
                print("   ✅ High-quality pose extracted (controlnet_aux)")
                return pose_image
            except Exception as e:
                print(f"   ⚠️ controlnet_aux failed: {e}")
        
        # Method 2: MediaPipe fallback
        if self.pose_method == 'mediapipe':
            try:
                pose_image = self._extract_mediapipe_pose(image, target_size)
                print("   ✅ Pose extracted (MediaPipe)")
                return pose_image
            except Exception as e:
                print(f"   ⚠️ MediaPipe failed: {e}")
        
        # Method 3: Fallback
        print("   ⚠️ Using fallback pose extraction")
        return self._create_fallback_pose(image, target_size)
    
    def _extract_mediapipe_pose(self, image: Image.Image, target_size: Tuple[int, int]) -> Image.Image:
        """MediaPipe pose extraction with enhanced quality"""
        image_cv = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
        results = self.pose_detector.process(image_cv)
        
        h, w = target_size
        pose_image = np.zeros((h, w, 3), dtype=np.uint8)
        
        if results.pose_landmarks:
            # Enhanced keypoint drawing
            for landmark in results.pose_landmarks.landmark:
                x, y = int(landmark.x * w), int(landmark.y * h)
                confidence = landmark.visibility
                
                if 0 <= x < w and 0 <= y < h and confidence > 0.5:
                    radius = int(6 + 6 * confidence)
                    cv2.circle(pose_image, (x, y), radius, (255, 255, 255), -1)
            
            # Enhanced connection drawing
            connections = self.mp_pose.POSE_CONNECTIONS
            for connection in connections:
                start_idx, end_idx = connection
                start = results.pose_landmarks.landmark[start_idx]
                end = results.pose_landmarks.landmark[end_idx]
                
                start_x, start_y = int(start.x * w), int(start.y * h)
                end_x, end_y = int(end.x * w), int(end.y * h)
                
                if (0 <= start_x < w and 0 <= start_y < h and 
                    0 <= end_x < w and 0 <= end_y < h):
                    avg_confidence = (start.visibility + end.visibility) / 2
                    thickness = int(3 + 3 * avg_confidence)
                    cv2.line(pose_image, (start_x, start_y), (end_x, end_y), 
                            (255, 255, 255), thickness)
        
        return Image.fromarray(pose_image)
    
    def _create_fallback_pose(self, image: Image.Image, target_size: Tuple[int, int]) -> Image.Image:
        """Enhanced fallback pose using edge detection"""
        image_np = np.array(image)
        gray = cv2.cvtColor(image_np, cv2.COLOR_RGB2GRAY)
        
        # Multi-scale edge detection
        edges1 = cv2.Canny(gray, 50, 150)
        edges2 = cv2.Canny(gray, 100, 200)
        edges = cv2.addWeighted(edges1, 0.7, edges2, 0.3, 0)
        
        # Morphological operations for better structure
        kernel = np.ones((3, 3), np.uint8)
        edges = cv2.morphologyEx(edges, cv2.MORPH_CLOSE, kernel)
        edges = cv2.dilate(edges, kernel, iterations=2)
        
        # Convert to RGB
        pose_rgb = cv2.cvtColor(edges, cv2.COLOR_GRAY2RGB)
        pose_pil = Image.fromarray(pose_rgb)
        return pose_pil.resize(target_size, Image.Resampling.LANCZOS)
    
    def generate_outfit(self,
                       source_image_path: str,
                       outfit_prompt: str,
                       output_path: str = "realistic_outfit.jpg",
                       num_inference_steps: int = 50,
                       guidance_scale: float = 7.5,  # FIXED: Lower for RealisticVision
                       controlnet_conditioning_scale: float = 1.0,
                       seed: Optional[int] = None) -> Tuple[Image.Image, Dict]:
        """
        Generate outfit using properly loaded RealisticVision checkpoint
        """
        print(f"🎭 GENERATING WITH FIXED REALISTICVISION")
        print(f"   Source: {source_image_path}")
        print(f"   Target: {outfit_prompt}")
        print(f"   Expected: Photorealistic (not painting-like)")
        
        # Set seed
        if seed is None:
            seed = np.random.randint(0, 2**31 - 1)
        torch.manual_seed(seed)
        if torch.cuda.is_available():
            torch.cuda.manual_seed(seed)
        
        # Extract pose
        print("🎯 Extracting pose...")
        pose_image = self.extract_pose(source_image_path, target_size=(512, 512))
        
        # Save pose for debugging
        pose_debug_path = output_path.replace('.jpg', '_pose_debug.jpg')
        pose_image.save(pose_debug_path)
        print(f"   Pose saved: {pose_debug_path}")
        
        # Create RealisticVision-optimized prompts
        #enhanced_prompt = self._create_realistic_vision_prompt(outfit_prompt)
        enhanced_prompt = self._create_realistic_vision_prompt(outfit_prompt, source_image_path)
        negative_prompt = self._create_realistic_vision_negative()
        
        print(f"   Enhanced prompt: {enhanced_prompt[:70]}...")
        print(f"   Guidance scale: {guidance_scale} (RealisticVision optimized)")
        
        # Generate with properly loaded checkpoint
        try:
            with torch.no_grad():
                result = self.pipeline(
                    prompt=enhanced_prompt,
                    negative_prompt=negative_prompt,
                    image=pose_image,
                    num_inference_steps=num_inference_steps,
                    guidance_scale=guidance_scale,  # Lower for photorealistic
                    controlnet_conditioning_scale=controlnet_conditioning_scale,
                    height=512,
                    width=512
                )
            
            generated_image = result.images[0]
            generated_image.save(output_path)
            
            # Validate results
            validation = self._validate_generation_quality(generated_image)
            
            metadata = {
                'seed': seed,
                'checkpoint_loaded_properly': True,
                'validation': validation,
                'pose_debug_path': pose_debug_path,
                'enhanced_prompt': enhanced_prompt,
                'guidance_scale': guidance_scale,
                'method': 'from_single_file_fixed'
            }
            
            print(f"✅ OUTFIT GENERATION COMPLETED!")
            print(f"   Photorealistic: {validation['looks_photorealistic']}")
            print(f"   Single person: {validation['single_person']}")
            print(f"   Face quality: {validation['face_quality']:.2f}")
            print(f"   Output: {output_path}")
            
            return generated_image, metadata
            
        except Exception as e:
            print(f"❌ Generation failed: {e}")
            raise
    
    #def _create_realistic_vision_prompt(self, base_prompt: str) -> str:
    #    """Create prompt optimized for RealisticVision photorealistic output"""
    #    # Ensure single person
    #    if not any(word in base_prompt.lower() for word in ["woman", "person", "model", "lady"]):
    #        enhanced = f"a beautiful woman wearing {base_prompt}"
    #    else:
    #        enhanced = base_prompt
    #    
    #    # CRITICAL: RealisticVision-specific terms for photorealism
    #    enhanced += ", RAW photo, 8k uhd, dslr, soft lighting, high quality"
    #    enhanced += ", film grain, Fujifilm XT3, photorealistic, realistic"
    #    enhanced += ", professional photography, studio lighting"
    #    enhanced += ", detailed face, natural skin, sharp focus"
    #   
    #    return enhanced

    def _create_realistic_vision_prompt(self, base_prompt: str, source_image_path: str) -> str:
        """FIXED: Gender-aware prompt with appearance matching"""
    
        if not hasattr(self, 'appearance_enhancer'):
            self.appearance_enhancer = ImprovedUnifiedGenderAppearanceEnhancer()
    
        result = self.appearance_enhancer.create_unified_enhanced_prompt(
            base_prompt, source_image_path
        )
    
        return result['enhanced_prompt'] if result['success'] else base_prompt
    
    def _create_realistic_vision_negative(self) -> str:
        """Create negative prompt to prevent painting-like results"""
        return (
            # Prevent multiple people
            "multiple people, group photo, crowd, extra person, "
            # Prevent painting/artistic styles
            "painting, drawing, illustration, artistic, sketch, cartoon, "
            "anime, rendered, digital art, cgi, 3d render, "
            # Prevent low quality
            "low quality, worst quality, blurry, out of focus, "
            "bad anatomy, extra limbs, malformed hands, deformed, "
            "poorly drawn hands, distorted, ugly, disfigured"
        )
    
    def perform_face_swap(self, source_image_path: str, target_image: Image.Image, 
                         balance_mode: str = "natural") -> Image.Image:
        """
        Perform balanced face swap using PROVEN techniques from balanced_clear_color_face_swap.py
        """
        print("👤 Performing PROVEN balanced face swap...")
        print(f"   Balance mode: {balance_mode} (preserves source colors)")
        
        try:
            # Convert PIL to CV2 format (matching proven system)
            source_img = cv2.imread(source_image_path)
            target_img = cv2.cvtColor(np.array(target_image), cv2.COLOR_RGB2BGR)
            
            if source_img is None:
                raise ValueError(f"Could not load source image: {source_image_path}")
            
            # Use proven face detection method
            source_faces = self._detect_faces_quality_proven(source_img, "source")
            target_faces = self._detect_faces_quality_proven(target_img, "target")
            
            if not source_faces or not target_faces:
                print("   ⚠️ Face detection failed - returning target image")
                return target_image
            
            print(f"   Found {len(source_faces)} source faces, {len(target_faces)} target faces")
            
            # Select best faces using proven quality scoring
            source_face = max(source_faces, key=lambda f: f['quality_score'])
            target_face = max(target_faces, key=lambda f: f['quality_score'])
            
            # Balance mode parameters (from proven system)
            balance_params = {
                'natural': {
                    'color_preservation': 0.85,
                    'clarity_enhancement': 0.4,
                    'color_saturation': 1.0,
                    'skin_tone_protection': 0.9,
                },
                'optimal': {
                    'color_preservation': 0.75,
                    'clarity_enhancement': 0.6,
                    'color_saturation': 1.1,
                    'skin_tone_protection': 0.8,
                },
                'vivid': {
                    'color_preservation': 0.65,
                    'clarity_enhancement': 0.8,
                    'color_saturation': 1.2,
                    'skin_tone_protection': 0.7,
                }
            }
            
            if balance_mode not in balance_params:
                balance_mode = 'natural'
            
            params = balance_params[balance_mode]
            print(f"   Using proven parameters: {balance_mode}")
            
            # Perform proven balanced swap
            result = self._perform_balanced_swap_proven(
                source_img, target_img, source_face, target_face, params
            )
            
            # Apply final optimization (from proven system)
            result = self._optimize_color_clarity_balance_proven(result, target_face, params)
            
            # Convert back to PIL
            result_pil = Image.fromarray(cv2.cvtColor(result, cv2.COLOR_BGR2RGB))
            
            print("   ✅ PROVEN face swap completed successfully")
            return result_pil
            
        except Exception as e:
            print(f"   ⚠️ Face swap failed: {e}")
            return target_image
    
    def _detect_faces_with_quality(self, image: Image.Image) -> list:
        """Detect faces with quality scoring"""
        if self.face_cascade is None:
            return []
        
        image_np = np.array(image)
        gray = cv2.cvtColor(image_np, cv2.COLOR_RGB2GRAY)
        
        faces = self.face_cascade.detectMultiScale(
            gray, scaleFactor=1.05, minNeighbors=4, minSize=(60, 60)
        )
        
        face_data = []
        for (x, y, w, h) in faces:
            # Quality scoring
            face_area = w * h
            image_area = gray.shape[0] * gray.shape[1]
            size_ratio = face_area / image_area
            
            # Position quality (prefer centered, upper portion)
            center_x = x + w // 2
            center_y = y + h // 2
            position_score = 1.0 - abs(center_x - gray.shape[1] // 2) / (gray.shape[1] // 2)
            position_score *= 1.0 if center_y < gray.shape[0] * 0.6 else 0.5
            
            quality = size_ratio * position_score
            
            face_data.append({
                'bbox': (x, y, w, h),
                'quality': quality,
                'size_ratio': size_ratio,
                'center': (center_x, center_y)
            })
        
        return face_data
    
    def _detect_faces_quality_proven(self, image: np.ndarray, image_type: str) -> list:
        """PROVEN quality face detection from balanced_clear_color_face_swap.py"""
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        faces = self.face_cascade.detectMultiScale(
            gray, scaleFactor=1.05, minNeighbors=4, minSize=(60, 60)
        )
        
        face_data = []
        for (x, y, w, h) in faces:
            # Eye detection (proven method)
            face_roi = gray[y:y+h, x:x+w]
            eyes = self.eye_cascade.detectMultiScale(face_roi, scaleFactor=1.1, minNeighbors=3)
            
            # Quality scoring (proven method)
            quality_score = self._calculate_balanced_quality_proven(gray, (x, y, w, h), eyes)
            
            face_info = {
                'bbox': (x, y, w, h),
                'area': w * h,
                'eyes_count': len(eyes),
                'quality_score': quality_score,
                'center': (x + w//2, y + h//2)
            }
            
            face_data.append(face_info)
        
        print(f"   👤 {image_type} faces: {len(face_data)}")
        return face_data
    
    def _calculate_balanced_quality_proven(self, gray_image: np.ndarray, bbox: tuple, eyes: list) -> float:
        """PROVEN quality calculation from balanced_clear_color_face_swap.py"""
        x, y, w, h = bbox
        
        # Size score
        size_score = min(w * h / 8000, 1.0)
        
        # Eye detection score
        eye_score = min(len(eyes) / 2.0, 1.0)
        
        # Position score
        h_img, w_img = gray_image.shape
        center_x, center_y = x + w//2, y + h//2
        img_center_x, img_center_y = w_img // 2, h_img // 2
        
        distance = np.sqrt((center_x - img_center_x)**2 + (center_y - img_center_y)**2)
        max_distance = np.sqrt((w_img//2)**2 + (h_img//2)**2)
        position_score = 1.0 - (distance / max_distance)
        
        # Combine scores (proven formula)
        total_score = size_score * 0.4 + eye_score * 0.4 + position_score * 0.2
        
        return total_score
    
    def _perform_balanced_swap_proven(self, source_img: np.ndarray, target_img: np.ndarray,
                                     source_face: dict, target_face: dict, params: dict) -> np.ndarray:
        """PROVEN balanced face swap from balanced_clear_color_face_swap.py"""
        result = target_img.copy()
        
        sx, sy, sw, sh = source_face['bbox']
        tx, ty, tw, th = target_face['bbox']
        
        # Moderate padding for balance (proven method)
        padding_ratio = 0.12
        px = int(sw * padding_ratio)
        py = int(sh * padding_ratio)
        
        # Extract regions (proven method)
        sx1 = max(0, sx - px)
        sy1 = max(0, sy - py)
        sx2 = min(source_img.shape[1], sx + sw + px)
        sy2 = min(source_img.shape[0], sy + sh + py)
        
        source_face_region = source_img[sy1:sy2, sx1:sx2]
        
        tx1 = max(0, tx - px)
        ty1 = max(0, ty - py)
        tx2 = min(target_img.shape[1], tx + tw + px)
        ty2 = min(target_img.shape[0], ty + th + py)
        
        target_w = tx2 - tx1
        target_h = ty2 - ty1
        
        # High-quality resize (proven method)
        source_resized = cv2.resize(
            source_face_region, 
            (target_w, target_h), 
            interpolation=cv2.INTER_LANCZOS4
        )
        
        # PROVEN STEP 1: Preserve original colors first
        source_color_preserved = self._preserve_source_colors_proven(
            source_resized, target_img, target_face, params
        )
        
        # PROVEN STEP 2: Apply gentle color harmony (not replacement)
        source_harmonized = self._apply_color_harmony_proven(
            source_color_preserved, target_img, target_face, params
        )
        
        # PROVEN STEP 3: Enhance clarity without destroying colors
        source_clear = self._enhance_clarity_preserve_color_proven(source_harmonized, params)
        
        # PROVEN STEP 4: Create balanced blending mask
        mask = self._create_balanced_mask_proven(target_w, target_h, params)
        
        # PROVEN STEP 5: Apply balanced blend
        target_region = result[ty1:ty2, tx1:tx2]
        blended = self._color_preserving_blend_proven(source_clear, target_region, mask, params)
        
        # Apply result
        result[ty1:ty2, tx1:tx2] = blended
        
        return result
    
    def _preserve_source_colors_proven(self, source_face: np.ndarray, target_img: np.ndarray,
                                      target_face: dict, params: dict) -> np.ndarray:
        """PROVEN color preservation from balanced_clear_color_face_swap.py"""
        color_preservation = params['color_preservation']
        
        if color_preservation >= 0.8:  # High color preservation
            print(f"      🎨 High color preservation mode ({color_preservation})")
            # Return source with minimal changes
            return source_face
        
        # For lower preservation, apply very gentle color adjustment
        try:
            tx, ty, tw, th = target_face['bbox']
            target_face_region = target_img[ty:ty+th, tx:tx+tw]
            target_face_resized = cv2.resize(target_face_region, (source_face.shape[1], source_face.shape[0]))
            
            # Convert to LAB for gentle color adjustment (proven method)
            source_lab = cv2.cvtColor(source_face, cv2.COLOR_BGR2LAB).astype(np.float32)
            target_lab = cv2.cvtColor(target_face_resized, cv2.COLOR_BGR2LAB).astype(np.float32)
            
            # Very gentle L channel adjustment only (proven method)
            source_l_mean = np.mean(source_lab[:, :, 0])
            target_l_mean = np.mean(target_lab[:, :, 0])
            
            adjustment_strength = (1 - color_preservation) * 0.3  # Max 30% adjustment
            l_adjustment = (target_l_mean - source_l_mean) * adjustment_strength
            
            source_lab[:, :, 0] = source_lab[:, :, 0] + l_adjustment
            
            # Convert back
            result = cv2.cvtColor(source_lab.astype(np.uint8), cv2.COLOR_LAB2BGR)
            
            print(f"      🎨 Gentle color preservation applied")
            return result
            
        except Exception as e:
            print(f"      ⚠️ Color preservation failed: {e}")
            return source_face
    
    def _apply_color_harmony_proven(self, source_face: np.ndarray, target_img: np.ndarray,
                                   target_face: dict, params: dict) -> np.ndarray:
        """PROVEN color harmony from balanced_clear_color_face_swap.py"""
        try:
            # Extract target face for harmony reference
            tx, ty, tw, th = target_face['bbox']
            target_face_region = target_img[ty:ty+th, tx:tx+tw]
            target_face_resized = cv2.resize(target_face_region, (source_face.shape[1], source_face.shape[0]))
            
            # Convert to HSV for better color harmony control (proven method)
            source_hsv = cv2.cvtColor(source_face, cv2.COLOR_BGR2HSV).astype(np.float32)
            target_hsv = cv2.cvtColor(target_face_resized, cv2.COLOR_BGR2HSV).astype(np.float32)
            
            # Very subtle hue harmony (only if very different) - proven method
            source_hue_mean = np.mean(source_hsv[:, :, 0])
            target_hue_mean = np.mean(target_hsv[:, :, 0])
            
            hue_diff = abs(source_hue_mean - target_hue_mean)
            if hue_diff > 30:  # Only adjust if very different hues
                harmony_strength = 0.1  # Very subtle
                hue_adjustment = (target_hue_mean - source_hue_mean) * harmony_strength
                source_hsv[:, :, 0] = source_hsv[:, :, 0] + hue_adjustment
            
            # Convert back
            result = cv2.cvtColor(source_hsv.astype(np.uint8), cv2.COLOR_HSV2BGR)
            
            print(f"      🎨 Subtle color harmony applied")
            return result
            
        except Exception as e:
            print(f"      ⚠️ Color harmony failed: {e}")
            return source_face
    
    def _enhance_clarity_preserve_color_proven(self, source_face: np.ndarray, params: dict) -> np.ndarray:
        """PROVEN clarity enhancement from balanced_clear_color_face_swap.py"""
        clarity_enhancement = params['clarity_enhancement']
        
        if clarity_enhancement <= 0:
            return source_face
        
        # Method 1: Luminance-only sharpening (preserves color) - PROVEN
        # Convert to LAB to work on lightness only
        lab = cv2.cvtColor(source_face, cv2.COLOR_BGR2LAB).astype(np.float32)
        l_channel = lab[:, :, 0]
        
        # Apply unsharp mask to L channel only (proven method)
        blurred_l = cv2.GaussianBlur(l_channel, (0, 0), 1.0)
        sharpened_l = cv2.addWeighted(l_channel, 1.0 + clarity_enhancement, blurred_l, -clarity_enhancement, 0)
        
        # Clamp values
        sharpened_l = np.clip(sharpened_l, 0, 255)
        lab[:, :, 0] = sharpened_l
        
        # Convert back to BGR
        result = cv2.cvtColor(lab.astype(np.uint8), cv2.COLOR_LAB2BGR)
        
        # Method 2: Edge enhancement (very subtle) - PROVEN
        if clarity_enhancement > 0.5:
            gray = cv2.cvtColor(result, cv2.COLOR_BGR2GRAY)
            edges = cv2.Canny(gray, 50, 150)
            edges_bgr = cv2.cvtColor(edges, cv2.COLOR_GRAY2BGR)
            
            # Very subtle edge enhancement
            edge_strength = (clarity_enhancement - 0.5) * 0.02  # Max 1% edge enhancement
            result = cv2.addWeighted(result, 1.0, edges_bgr, edge_strength, 0)
        
        print(f"      🔍 Color-preserving clarity enhancement applied")
        return result
    
    def _create_balanced_mask_proven(self, width: int, height: int, params: dict) -> np.ndarray:
        """PROVEN mask creation from balanced_clear_color_face_swap.py"""
        mask = np.zeros((height, width), dtype=np.float32)
        
        # Create elliptical mask (proven method)
        center_x, center_y = width // 2, height // 2
        ellipse_w = int(width * 0.37)
        ellipse_h = int(height * 0.45)
        
        Y, X = np.ogrid[:height, :width]
        ellipse_mask = ((X - center_x) / ellipse_w) ** 2 + ((Y - center_y) / ellipse_h) ** 2 <= 1
        mask[ellipse_mask] = 1.0
        
        # Moderate blur for natural blending (proven method)
        blur_size = 19
        mask = cv2.GaussianBlur(mask, (blur_size, blur_size), 0)
        
        # Normalize
        if mask.max() > 0:
            mask = mask / mask.max()
        
        return mask
    
    def _color_preserving_blend_proven(self, source: np.ndarray, target: np.ndarray, 
                                      mask: np.ndarray, params: dict) -> np.ndarray:
        """PROVEN blending from balanced_clear_color_face_swap.py"""
        # Strong blend to preserve source colors (proven method)
        blend_strength = 0.9  # High to preserve source color
        
        mask_3d = np.stack([mask] * 3, axis=-1)
        blended = (source.astype(np.float32) * mask_3d * blend_strength + 
                  target.astype(np.float32) * (1 - mask_3d * blend_strength))
        
        return blended.astype(np.uint8)
    
    def _optimize_color_clarity_balance_proven(self, result: np.ndarray, target_face: dict, 
                                              params: dict) -> np.ndarray:
        """PROVEN final optimization from balanced_clear_color_face_swap.py"""
        tx, ty, tw, th = target_face['bbox']
        face_region = result[ty:ty+th, tx:tx+tw].copy()
        
        # Enhance saturation if specified (proven method)
        saturation_boost = params['color_saturation']
        if saturation_boost != 1.0:
            hsv = cv2.cvtColor(face_region, cv2.COLOR_BGR2HSV).astype(np.float32)
            hsv[:, :, 1] = hsv[:, :, 1] * saturation_boost  # Boost saturation
            hsv[:, :, 1] = np.clip(hsv[:, :, 1], 0, 255)  # Clamp
            face_region = cv2.cvtColor(hsv.astype(np.uint8), cv2.COLOR_HSV2BGR)
            
            print(f"      🎨 Saturation optimized ({saturation_boost})")
        
        # Skin tone protection (proven method)
        skin_protection = params['skin_tone_protection']
        if skin_protection > 0:
            # Apply bilateral filter for skin smoothing while preserving edges
            smooth_strength = int(9 * skin_protection)
            if smooth_strength > 0:
                bilateral_filtered = cv2.bilateralFilter(face_region, smooth_strength, 40, 40)
                
                # Blend with original for subtle effect
                alpha = 0.3 * skin_protection
                face_region = cv2.addWeighted(face_region, 1-alpha, bilateral_filtered, alpha, 0)
                
                print(f"      🎨 Skin tone protection applied")
        
        # Apply optimized face back
        result[ty:ty+th, tx:tx+tw] = face_region
        
        return result
    
    def _validate_generation_quality(self, generated_image):
        """Use improved lenient validation"""
        if not hasattr(self, 'improved_validator'):
            self.improved_validator = ImprovedGenerationValidator()
    
        return self.improved_validator.validate_generation_quality_improved(generated_image)
    
    def complete_fashion_transformation(self,
                                     source_image_path: str,
                                     outfit_prompt: str,
                                     output_path: str = "complete_transformation.jpg",
                                     **kwargs) -> Tuple[Image.Image, Dict]:
        """
        Complete fashion transformation: Generate outfit + Face swap
        """
        print(f"🎭 COMPLETE FASHION TRANSFORMATION")
        print(f"   Source: {source_image_path}")
        print(f"   Target: {outfit_prompt}")
        print(f"   Method: Fixed RealisticVision + Balanced face swap")
        
        # Step 1: Generate outfit with proper RealisticVision
        outfit_path = output_path.replace('.jpg', '_outfit_only.jpg')
        generated_image, generation_metadata = self.generate_outfit(
            source_image_path=source_image_path,
            outfit_prompt=outfit_prompt,
            output_path=outfit_path,
            **kwargs
        )
        
        print(f"   Step 1 completed: {outfit_path}")
        
        # Step 2: Perform face swap if generation quality is good
        if generation_metadata['validation']['single_person']:
            print("   ✅ Good generation quality - proceeding with PROVEN face swap")
            
            final_image = self.perform_face_swap(source_image_path, generated_image, balance_mode="natural")
            final_image.save(output_path)
            
            final_metadata = generation_metadata.copy()
            final_metadata['face_swap_applied'] = True
            final_metadata['face_swap_method'] = 'proven_balanced_clear_color'
            final_metadata['balance_mode'] = 'natural'
            final_metadata['final_output'] = output_path
            final_metadata['outfit_only_output'] = outfit_path
            
            print(f"✅ COMPLETE TRANSFORMATION FINISHED!")
            print(f"   Final result: {output_path}")
            print(f"   Face swap: PROVEN method with natural skin tones")
            
            return final_image, final_metadata
            
        else:
            print("   ⚠️ Generation quality insufficient for face swap")
            generated_image.save(output_path)
            
            final_metadata = generation_metadata.copy()
            final_metadata['face_swap_applied'] = False
            final_metadata['final_output'] = output_path
            
            return generated_image, final_metadata


# Easy usage functions
def fix_realistic_vision_issues(source_image_path: str,
                               checkpoint_path: str,
                               outfit_prompt: str = "red evening dress",
                               output_path: str = "fixed_result.jpg"):
    """
    Fix both RealisticVision loading and integrate face swapping
    """
    print(f"🔧 FIXING REALISTIC VISION ISSUES")
    print(f"   Issue 1: Painting-like results (checkpoint not loading)")
    print(f"   Issue 2: No face swapping integration")
    print(f"   Solution: Proper from_single_file() + integrated face swap")
    
    pipeline = FixedRealisticVisionPipeline(checkpoint_path)
    
    result_image, metadata = pipeline.complete_fashion_transformation(
        source_image_path=source_image_path,
        outfit_prompt=outfit_prompt,
        output_path=output_path
    )
    
    return result_image, metadata


if __name__ == "__main__":
    print("🔧 FIXED REALISTIC VISION + FACE SWAP PIPELINE")
    print("=" * 50)
    
    # Your specific files
    source_path = "woman_jeans_t-shirt.png"
    checkpoint_path = "realisticVisionV60B1_v51HyperVAE.safetensors"
    
    print(f"\n❌ CURRENT ISSUES:")
    print(f"   • Generated image looks like painting (not photorealistic)")
    print(f"   • RealisticVision checkpoint not loading properly")
    print(f"   • No face swapping integration")
    print(f"   • Missing balanced face swap from proven system")
    
    print(f"\n✅ FIXES APPLIED:")
    print(f"   • Use from_single_file() for proper checkpoint loading")
    print(f"   • Lower guidance_scale (7.5) for photorealistic results")
    print(f"   • RealisticVision-specific prompt engineering")
    print(f"   • Integrated balanced face swap system")
    print(f"   • Complete pipeline with quality validation")
    
    if os.path.exists(source_path) and os.path.exists(checkpoint_path):
        print(f"\n🧪 Testing fixed pipeline...")
        
        try:
            # Test the complete fixed pipeline
            result, metadata = fix_realistic_vision_issues(
                source_image_path=source_path,
                checkpoint_path=checkpoint_path,
                outfit_prompt="red evening dress",
                output_path="fixed_realistic_vision_result.jpg"
            )
            
            validation = metadata['validation']
            
            print(f"\n📊 RESULTS:")
            print(f"   Photorealistic: {validation['looks_photorealistic']}")
            print(f"   Single person: {validation['single_person']}")
            print(f"   Face quality: {validation['face_quality']:.2f}")
            print(f"   Face swap applied: {metadata['face_swap_applied']}")
            print(f"   Overall assessment: {validation['overall_assessment']}")
            
            if validation['looks_photorealistic'] and metadata['face_swap_applied']:
                print(f"\n🎉 SUCCESS! Both issues fixed:")
                print(f"   ✅ Photorealistic image (not painting-like)")
                print(f"   ✅ Face swap successfully applied")
                print(f"   ✅ RealisticVision features active")
            
        except Exception as e:
            print(f"❌ Test failed: {e}")
    
    else:
        print(f"\n⚠️ Files not found:")
        print(f"   Source: {source_path} - {os.path.exists(source_path)}")
        print(f"   Checkpoint: {checkpoint_path} - {os.path.exists(checkpoint_path)}")
    
    print(f"\n📋 USAGE:")
    print(f"""
# Fix both issues in one call
result, metadata = fix_realistic_vision_issues(
    source_image_path="your_source.jpg",
    checkpoint_path="realisticVisionV60B1_v51HyperVAE.safetensors",
    outfit_prompt="red evening dress"
)

# Check results
if metadata['validation']['looks_photorealistic']:
    print("✅ Photorealistic result achieved!")
    
if metadata['face_swap_applied']:
    print("✅ Face swap successfully applied!")
""")
    
    print(f"\n🎯 EXPECTED IMPROVEMENTS:")
    print(f"   • Photorealistic images instead of painting-like")
    print(f"   • RealisticVision single-person bias working") 
    print(f"   • Natural skin tones with face preservation")
    print(f"   • Proper checkpoint loading (no missing tensors)")
    print(f"   • Complete end-to-end transformation pipeline")