generator.py

"""
Generation logic for Pixagram - Torch 2.1.1 + Depth Anything V2 optimized
"""
import torch
import numpy as np
import cv2
from PIL import Image
import torch.nn.functional as F
from torchvision import transforms

from config import *
from utils import *
from models import *


class RetroArtConverter:
    """Main retro art generator with torch 2.1.1 optimizations"""
    
    def __init__(self):
        self.device = device
        self.dtype = dtype
        self.models_loaded = {
            'custom_checkpoint': False,
            'lora': False,
            'instantid': False,
            'depth_detector': False,
            'ip_adapter': False
        }
        
        # Face analysis with CPU fallback
        self.face_app, self.face_detection_enabled = load_face_analysis()
        
        # Depth detector with Depth Anything V2 priority
        self.depth_detector, depth_success, self.depth_type = load_depth_detector()
        self.models_loaded['depth_detector'] = depth_success
        print(f"[DEPTH] Using: {self.depth_type}")
        
        # ControlNets
        controlnet_depth, self.controlnet_instantid, instantid_success = load_controlnets()
        self.controlnet_depth = controlnet_depth
        self.instantid_enabled = instantid_success
        self.models_loaded['instantid'] = instantid_success
        
        # Image encoder
        if self.instantid_enabled:
            self.image_encoder = load_image_encoder()
        else:
            self.image_encoder = None
        
        # Determine controlnets
        if self.instantid_enabled and self.controlnet_instantid is not None:
            controlnets = [self.controlnet_instantid, controlnet_depth]
        else:
            controlnets = controlnet_depth
        
        # SDXL pipeline
        self.pipe, checkpoint_success = load_sdxl_pipeline(controlnets)
        self.models_loaded['custom_checkpoint'] = checkpoint_success
        
        # LORA
        lora_success = load_lora(self.pipe)
        self.models_loaded['lora'] = lora_success
        
        # IP-Adapter
        if self.instantid_enabled 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:
            self.models_loaded['ip_adapter'] = False
            self.image_proj_model = None
        
        # Compel
        self.compel, self.use_compel = setup_compel(self.pipe)
        
        # LCM scheduler
        setup_scheduler(self.pipe)
        
        # TORCH 2.1.1: Apply optimizations (compile, etc.)
        optimize_pipeline(self.pipe)
        
        # Caption model
        self.caption_processor, self.caption_model, self.caption_enabled = load_caption_model()
        
        # CLIP skip
        set_clip_skip(self.pipe)
        
        self.using_multiple_controlnets = isinstance(controlnets, list)
        self._print_status()
        print("  [OK] Initialization complete")
    
    def _print_status(self):
        """Print model status"""
        print("\n=== MODEL STATUS ===")
        for model, loaded in self.models_loaded.items():
            status = "[OK]" if loaded else "[FALLBACK]"
            print(f"{model}: {status}")
        print("====================\n")
    
    def get_depth_map(self, image):
        """Generate depth map with Depth Anything V2 or fallback"""
        if self.depth_type == "depth_anything_v2" and self.depth_detector is not None:
            try:
                result = self.depth_detector(image)
                depth_image = result["depth"]
                # Convert to PIL if needed
                if not isinstance(depth_image, Image.Image):
                    depth_array = np.array(depth_image)
                    depth_image = Image.fromarray(depth_array)
                return depth_image
            except Exception as e:
                print(f"[WARNING] Depth Anything V2 failed: {e}, using fallback")
        
        if self.depth_type == "zoe" and self.depth_detector is not None:
            try:
                depth_image = self.depth_detector(image)
                return depth_image
            except Exception as e:
                print(f"[WARNING] Zoe failed: {e}, using grayscale")
        
        # Grayscale fallback
        gray = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2GRAY)
        depth_colored = cv2.cvtColor(gray, cv2.COLOR_GRAY2RGB)
        return Image.fromarray(depth_colored)
    
    def add_trigger_word(self, prompt):
        """Add trigger word if not present"""
        if TRIGGER_WORD.lower() not in prompt.lower():
            return f"{TRIGGER_WORD}, {prompt}"
        return prompt
    
    def extract_multi_scale_face(self, face_crop, face):
        """Multi-scale face extraction"""
        try:
            multi_scale_embeds = []
            for scale in MULTI_SCALE_FACTORS:
                w, h = face_crop.size
                scaled_size = (int(w * scale), int(h * scale))
                scaled_crop = face_crop.resize(scaled_size, Image.LANCZOS)
                scaled_crop = scaled_crop.resize((w, h), Image.LANCZOS)
                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)
            
            if len(multi_scale_embeds) > 0:
                averaged = np.mean(multi_scale_embeds, axis=0)
                averaged = averaged / np.linalg.norm(averaged)
                return averaged
            return face.normed_embedding
        except Exception as e:
            return face.normed_embedding
    
    def detect_face_quality(self, face):
        """Adaptive parameter adjustment"""
        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
            
            if face_size < ADAPTIVE_THRESHOLDS['small_face_size']:
                return ADAPTIVE_PARAMS['small_face'].copy()
            elif det_score < ADAPTIVE_THRESHOLDS['low_confidence']:
                return ADAPTIVE_PARAMS['low_confidence'].copy()
            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:
                    pass
            return None
        except:
            return None
    
    def validate_and_adjust_parameters(self, strength, guidance_scale, lora_scale, 
                                       identity_preservation, identity_control_scale,
                                       depth_control_scale, consistency_mode=True):
        """Parameter validation"""
        if consistency_mode:
            adjustments = []
            
            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}")
            
            if strength < 0.5:
                if identity_preservation < 1.3:
                    identity_preservation = 1.3
                if lora_scale > 0.9:
                    lora_scale = 0.9
            elif strength > 0.7:
                if identity_preservation > 1.0:
                    identity_preservation = 1.0
                if lora_scale < 1.2:
                    lora_scale = 1.2
            
            original_cfg = guidance_scale
            guidance_scale = max(1.0, min(guidance_scale, 1.5))
            
            if adjustments:
                print("  [OK] Applied adjustments")
        
        return strength, guidance_scale, lora_scale, identity_preservation, identity_control_scale, depth_control_scale
    
    def generate_caption(self, image, max_length=None, num_beams=None):
        """Generate caption"""
        if not self.caption_enabled or self.caption_model is None:
            return None
        
        if max_length is None:
            max_length = CAPTION_CONFIG['max_length']
        if num_beams is None:
            num_beams = CAPTION_CONFIG['num_beams']
        
        try:
            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)
            caption = self.caption_processor.decode(output[0], skip_special_tokens=True)
            return caption
        except Exception as e:
            return None
    
    def generate_retro_art(
        self,
        input_image,
        prompt="retro game character",
        negative_prompt="blurry, low quality",
        num_inference_steps=12,
        guidance_scale=1.0,
        depth_control_scale=0.8,
        identity_control_scale=0.85,
        lora_scale=1.0,
        identity_preservation=0.8,
        strength=0.75,
        enable_color_matching=False,
        consistency_mode=True,
        seed=-1
    ):
        """Generate retro art with torch 2.1.1 optimizations"""
        
        prompt = sanitize_text(prompt)
        negative_prompt = sanitize_text(negative_prompt)
        
        if consistency_mode:
            strength, guidance_scale, lora_scale, identity_preservation, identity_control_scale, depth_control_scale = \
                self.validate_and_adjust_parameters(
                    strength, guidance_scale, lora_scale, identity_preservation, 
                    identity_control_scale, depth_control_scale, consistency_mode
                )
        
        prompt = self.add_trigger_word(prompt)
        
        original_width, original_height = input_image.size
        target_width, target_height = calculate_optimal_size(original_width, original_height, RECOMMENDED_SIZES)
        
        resized_image = input_image.resize((int(target_width), int(target_height)), Image.LANCZOS)
        
        print("Generating depth map...")
        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)
        
        using_multiple_controlnets = self.using_multiple_controlnets
        face_kps_image = None
        face_embeddings = None
        face_crop_enhanced = None
        has_detected_faces = False
        face_bbox_original = None
        
        if using_multiple_controlnets and self.face_app is not None:
            print("Detecting faces...")
            img_array = cv2.cvtColor(np.array(resized_image), cv2.COLOR_RGB2BGR)
            faces = self.face_app.get(img_array)
            
            if len(faces) > 0:
                has_detected_faces = True
                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)
                
                # ENHANCED: Use new facial attributes extraction
                facial_attrs = get_facial_attributes(face)
                prompt = build_enhanced_prompt(prompt, facial_attrs, TRIGGER_WORD)
        
        if hasattr(self.pipe, 'set_adapters') and self.models_loaded['lora']:
            try:
                self.pipe.set_adapters(["retroart"], adapter_weights=[lora_scale])
            except:
                pass
        
        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()
        else:
            generator = torch.Generator(device=self.device).manual_seed(seed)
            actual_seed = seed
        
        pipe_kwargs["generator"] = generator
        
        if self.use_compel and self.compel is not None:
            try:
                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]
            except:
                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
        
        if using_multiple_controlnets and has_detected_faces and face_kps_image is not None:
            control_images = [face_kps_image, depth_image]
            conditioning_scales = [identity_control_scale, depth_control_scale]
            pipe_kwargs["control_image"] = control_images
            pipe_kwargs["controlnet_conditioning_scale"] = conditioning_scales
            
            if face_embeddings is not None and self.models_loaded.get('ip_adapter', False) and face_crop_enhanced is not None:
                with torch.no_grad():
                    insightface_embeds = torch.from_numpy(face_embeddings).to(
                        device=self.device, dtype=self.dtype
                    ).unsqueeze(0).unsqueeze(1)
                    
                    image_embeds = self.image_proj_model(insightface_embeds)
                
                boosted_scale = identity_preservation * IDENTITY_BOOST_MULTIPLIER
                
                pipe_kwargs["added_cond_kwargs"] = {"image_embeds": image_embeds, "time_ids": None}
                pipe_kwargs["cross_attention_kwargs"] = {"ip_adapter_scale": boosted_scale}
        else:
            if using_multiple_controlnets and not has_detected_faces:
                control_images = [depth_image, depth_image]
                conditioning_scales = [0.0, depth_control_scale]
                pipe_kwargs["control_image"] = control_images
                pipe_kwargs["controlnet_conditioning_scale"] = conditioning_scales
            else:
                pipe_kwargs["control_image"] = depth_image
                pipe_kwargs["controlnet_conditioning_scale"] = depth_control_scale
            
            if self.models_loaded.get('ip_adapter', False):
                dummy_embeds = torch.zeros(
                    (1, 4, self.pipe.unet.config.cross_attention_dim),
                    device=self.device, dtype=self.dtype
                )
                pipe_kwargs["added_cond_kwargs"] = {"image_embeds": dummy_embeds, "time_ids": None}
                pipe_kwargs["cross_attention_kwargs"] = {"ip_adapter_scale": 0.0}
        
        # TORCH 2.1.1: Use optimized attention backend
        print(f"Generating (steps={num_inference_steps}, cfg={guidance_scale}, strength={strength})...")
        
        if device == "cuda" and hasattr(torch.backends.cuda, 'sdp_kernel'):
            with torch.backends.cuda.sdp_kernel(
                enable_flash=True,
                enable_mem_efficient=True,
                enable_math=False
            ):
                result = self.pipe(**pipe_kwargs)
        else:
            result = self.pipe(**pipe_kwargs)
        
        generated_image = result.images[0]
        
        if enable_color_matching and has_detected_faces:
            try:
                if face_bbox_original is not None:
                    generated_image = enhanced_color_match(generated_image, resized_image, face_bbox=face_bbox_original)
                else:
                    generated_image = color_match(generated_image, resized_image, mode='mkl')
            except:
                pass
        elif enable_color_matching:
            try:
                generated_image = color_match(generated_image, resized_image, mode='mkl')
            except:
                pass
        
        return generated_image


print("[OK] Generator ready (Torch 2.1.1 + Depth Anything V2)")