Update optimizer.py

optimizer.py

@@ -1,6 +1,6 @@
 """
 Ultra Supreme Optimizer - Main optimization engine for image analysis
-VERSIÓN MEJORADA - Usa el prompt completo de CLIP Interrogator
+VERSIÓN FLORENCE-2 - Usa Florence-2 en lugar de CLIP Interrogator
 """
 
 # IMPORTANT: spaces must be imported BEFORE torch or any CUDA-using library
@@ -14,7 +14,7 @@ from typing import Tuple, Dict, Any, Optional
 import torch
 import numpy as np
 from PIL import Image
-from clip_interrogator import Config, Interrogator
+from transformers import AutoProcessor, AutoModelForCausalLM
 
 from analyzer import UltraSupremeAnalyzer
 
@@ -25,12 +25,12 @@ class UltraSupremeOptimizer:
     """Main optimizer class for ultra supreme image analysis"""
     
     def __init__(self):
-        self.interrogator: Optional[Interrogator] = None
+        self.processor = None
+        self.model = None
         self.analyzer = UltraSupremeAnalyzer()
         self.usage_count = 0
         self.device = self._get_device()
         self.is_initialized = False
-        # NO inicializar modelo aquí - hacerlo lazy
         
     @staticmethod
     def _get_device() -> str:
@@ -43,31 +43,37 @@ class UltraSupremeOptimizer:
             return "cpu"
 
     def initialize_model(self) -> bool:
-        """Initialize the CLIP interrogator model"""
+        """Initialize Florence-2 model"""
         if self.is_initialized:
             return True
         
         try:
-            # Configuración para CPU inicialmente
-            config = Config(
-                clip_model_name="ViT-L-14/openai",
-                download_cache=True,
-                chunk_size=2048,
-                quiet=True,
-                device="cpu"  # Siempre inicializar en CPU
+            logger.info("Loading Florence-2 model...")
+            
+            # Load Florence-2 base model (you can also use 'microsoft/Florence-2-large' for better quality)
+            model_id = "microsoft/Florence-2-base"
+            
+            self.processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
+            self.model = AutoModelForCausalLM.from_pretrained(
+                model_id, 
+                trust_remote_code=True,
+                torch_dtype=torch.float16 if self.device == "cuda" else torch.float32
             )
             
-            self.interrogator = Interrogator(config)
+            # Keep model on CPU initially
+            self.model = self.model.to("cpu")
+            self.model.eval()
+            
             self.is_initialized = True
             
             # Clean up memory after initialization
             gc.collect()
             
-            logger.info("Model initialized successfully on CPU")
+            logger.info("Florence-2 model initialized successfully")
             return True
             
         except Exception as e:
-            logger.error(f"Initialization error: {e}")
+            logger.error(f"Model initialization error: {e}")
             return False
 
     def optimize_image(self, image: Any) -> Optional[Image.Image]:
@@ -86,8 +92,8 @@ class UltraSupremeOptimizer:
             if image.mode != 'RGB':
                 image = image.convert('RGB')
             
-            # Resize if too large
-            max_size = 768  # Reducir tamaño para evitar problemas de memoria
+            # Florence-2 handles various sizes well, but let's be reasonable
+            max_size = 1024
             if image.size[0] > max_size or image.size[1] > max_size:
                 image.thumbnail((max_size, max_size), Image.Resampling.LANCZOS)
             
@@ -98,7 +104,7 @@ class UltraSupremeOptimizer:
             return None
 
     def apply_flux_rules(self, base_prompt: str) -> str:
-        """Aplica las reglas de Flux a un prompt base de CLIP Interrogator"""
+        """Aplica las reglas de Flux a un prompt base"""
         
         # Limpiar el prompt de elementos no deseados
         cleanup_patterns = [
@@ -148,110 +154,88 @@ class UltraSupremeOptimizer:
         
         return final_prompt
 
-    def _prepare_models_for_gpu(self):
-        """Prepara los modelos para GPU con la precisión correcta"""
-        try:
-            if hasattr(self.interrogator, 'caption_model'):
-                self.interrogator.caption_model = self.interrogator.caption_model.half().to("cuda")
-                
-            if hasattr(self.interrogator, 'clip_model'):
-                self.interrogator.clip_model = self.interrogator.clip_model.half().to("cuda")
-                
-            if hasattr(self.interrogator, 'blip_model'):
-                self.interrogator.blip_model = self.interrogator.blip_model.half().to("cuda")
-                
-            self.interrogator.config.device = "cuda"
-            logger.info("Models prepared for GPU with FP16")
-            
-        except Exception as e:
-            logger.error(f"Error preparing models for GPU: {e}")
-            raise
-
-    def _prepare_models_for_cpu(self):
-        """Prepara los modelos para CPU con float32"""
-        try:
-            if hasattr(self.interrogator, 'caption_model'):
-                self.interrogator.caption_model = self.interrogator.caption_model.float().to("cpu")
-                
-            if hasattr(self.interrogator, 'clip_model'):
-                self.interrogator.clip_model = self.interrogator.clip_model.float().to("cpu")
-                
-            if hasattr(self.interrogator, 'blip_model'):
-                self.interrogator.blip_model = self.interrogator.blip_model.float().to("cpu")
-                
-            self.interrogator.config.device = "cpu"
-            logger.info("Models prepared for CPU with FP32")
-            
-        except Exception as e:
-            logger.error(f"Error preparing models for CPU: {e}")
-            raise
-
     @spaces.GPU(duration=60)
-    def run_clip_inference(self, image: Image.Image) -> Tuple[str, str, str]:
-        """Solo la inferencia CLIP usa GPU"""
+    def run_florence_inference(self, image: Image.Image) -> Tuple[str, str, str]:
+        """Run Florence-2 inference on GPU"""
         try:
-            # NO usar half precision - mantener float32 para compatibilidad
-            if hasattr(self.interrogator, 'caption_model'):
-                self.interrogator.caption_model = self.interrogator.caption_model.to("cuda")
-                
-            if hasattr(self.interrogator, 'clip_model'):
-                self.interrogator.clip_model = self.interrogator.clip_model.to("cuda")
-                
-            if hasattr(self.interrogator, 'blip_model'):
-                self.interrogator.blip_model = self.interrogator.blip_model.to("cuda")
-                
-            self.interrogator.config.device = "cuda"
-            logger.info("Models moved to GPU with float32 (full precision)")
-            
-            # Ejecutar inferencias sin autocast para evitar problemas de half precision
-            full_prompt = self.interrogator.interrogate(image)
-            clip_fast = self.interrogator.interrogate_fast(image)
-            clip_classic = self.interrogator.interrogate_classic(image)
-            
-            return full_prompt, clip_fast, clip_classic
-            
-        except Exception as e:
-            logger.error(f"GPU inference error: {e}")
-            # Intentar en CPU como fallback
-            return self._run_cpu_inference(image)
-    
-    def _safe_interrogate(self, image: Image.Image, method: str) -> str:
-        """Ejecuta interrogate de forma segura manejando precisión"""
-        try:
-            # Temporalmente parchear el método de procesamiento de imagen
-            original_method = getattr(self.interrogator, method)
-            
-            # Ejecutar el método
-            result = original_method(image)
-            
-            return result
-            
-        except Exception as e:
-            logger.error(f"Error in {method}: {e}")
-            return f"Error processing with {method}"
-    
-    def _run_cpu_inference(self, image: Image.Image) -> Tuple[str, str, str]:
-        """Ejecuta inferencia en CPU como fallback"""
-        try:
-            logger.info("Running CPU inference as fallback")
+            # Move model to GPU
+            self.model = self.model.to("cuda")
+            logger.info("Florence-2 model moved to GPU")
+            
+            # Task prompts for different types of analysis
+            tasks = {
+                "detailed_caption": "<DETAILED_CAPTION>",
+                "more_detailed_caption": "<MORE_DETAILED_CAPTION>", 
+                "caption": "<CAPTION>",
+                "dense_region_caption": "<DENSE_REGION_CAPTION>"
+            }
+            
+            results = {}
+            
+            # Run different captioning tasks
+            for task_name, task_prompt in tasks.items():
+                try:
+                    inputs = self.processor(text=task_prompt, images=image, return_tensors="pt")
+                    inputs = {k: v.to("cuda") for k, v in inputs.items()}
+                    
+                    with torch.cuda.amp.autocast(dtype=torch.float16):
+                        generated_ids = self.model.generate(
+                            input_ids=inputs["input_ids"],
+                            pixel_values=inputs["pixel_values"],
+                            max_new_tokens=1024,
+                            num_beams=3,
+                            do_sample=False
+                        )
+                    
+                    generated_text = self.processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
+                    parsed = self.processor.post_process_generation(generated_text, task=task_prompt, image_size=(image.width, image.height))
+                    
+                    # Extract the caption from the parsed result
+                    if task_prompt in parsed:
+                        results[task_name] = parsed[task_prompt]
+                    else:
+                        # Sometimes the result is directly in the parsed output
+                        results[task_name] = str(parsed) if parsed else ""
+                        
+                except Exception as e:
+                    logger.warning(f"Error in {task_name}: {e}")
+                    results[task_name] = ""
+            
+            # Extract results
+            detailed_caption = results.get("detailed_caption", "")
+            more_detailed = results.get("more_detailed_caption", "")
+            caption = results.get("caption", "")
+            
+            # Combine for a comprehensive description
+            if more_detailed:
+                full_prompt = more_detailed
+            elif detailed_caption:
+                full_prompt = detailed_caption
+            else:
+                full_prompt = caption
             
-            # Preparar modelos para CPU
-            self._prepare_models_for_cpu()
+            # Use different levels as our three outputs
+            clip_fast = caption if caption else "A photograph"
+            clip_classic = detailed_caption if detailed_caption else full_prompt
+            clip_best = more_detailed if more_detailed else full_prompt
             
-            # Ejecutar en CPU sin autocast
-            full_prompt = self.interrogator.interrogate(image)
-            clip_fast = self.interrogator.interrogate_fast(image)
-            clip_classic = self.interrogator.interrogate_classic(image)
+            logger.info(f"Florence-2 captions generated successfully")
             
             return full_prompt, clip_fast, clip_classic
             
         except Exception as e:
-            logger.error(f"CPU inference also failed: {e}")
-            return "Error: Failed to process image", "Error", "Error"
+            logger.error(f"Florence-2 inference error: {e}")
+            # Move model back to CPU to free GPU memory
+            self.model = self.model.to("cpu")
+            raise e
+        finally:
+            # Always move model back to CPU after inference
+            self.model = self.model.to("cpu")
+            torch.cuda.empty_cache()
 
     def generate_ultra_supreme_prompt(self, image: Any) -> Tuple[str, str, int, Dict[str, int]]:
         """
-        Generate ultra supreme prompt from image usando el pipeline completo
+        Generate ultra supreme prompt from image usando Florence-2
         
         Returns:
             Tuple of (prompt, analysis_info, score, breakdown)
@@ -275,30 +259,30 @@ class UltraSupremeOptimizer:
             
             start_time = datetime.now()
             
-            logger.info("ULTRA SUPREME ANALYSIS - Starting complete pipeline with multi-model analysis")
+            logger.info("ULTRA SUPREME ANALYSIS - Starting with Florence-2")
             
-            # Ejecutar inferencia CLIP
-            full_prompt, clip_fast, clip_classic = self.run_clip_inference(image)
-            
-            # Verificar si hubo errores
-            if "Error" in full_prompt:
-                logger.warning("Using fallback prompt due to inference error")
+            # Ejecutar inferencia Florence-2
+            try:
+                full_prompt, caption_fast, caption_detailed = self.run_florence_inference(image)
+            except Exception as e:
+                logger.error(f"Florence-2 failed: {e}")
+                # Fallback básico
                 full_prompt = "A photograph"
-                clip_fast = "image"
-                clip_classic = "picture"
+                caption_fast = "image"
+                caption_detailed = "detailed image"
             
-            logger.info(f"CLIP complete prompt: {full_prompt[:100]}...")
+            logger.info(f"Florence-2 caption: {full_prompt[:100]}...")
             
-            # NUEVO: Ejecutar análisis ultra supremo con múltiples modelos
+            # Ejecutar análisis ultra supremo con múltiples modelos
             logger.info("Running multi-model ultra supreme analysis...")
             ultra_analysis = self.analyzer.ultra_supreme_analysis(
-                image, clip_fast, clip_classic, full_prompt
+                image, caption_fast, caption_detailed, full_prompt
             )
             
             # Construir prompt mejorado basado en análisis completo
             enhanced_prompt_parts = []
             
-            # Base prompt de CLIP
+            # Base prompt de Florence
             enhanced_prompt_parts.append(full_prompt)
             
             # Agregar información demográfica si está disponible
@@ -339,7 +323,7 @@ class UltraSupremeOptimizer:
             
             # Generate enhanced analysis report con datos de múltiples modelos
             analysis_info = self._generate_ultra_analysis_report(
-                ultra_analysis, score, breakdown, duration
+                ultra_analysis, score, breakdown, duration, "Florence-2"
             )
             
             return optimized_prompt, analysis_info, score, breakdown
@@ -348,58 +332,9 @@ class UltraSupremeOptimizer:
             logger.error(f"Ultra supreme generation error: {e}", exc_info=True)
             return f"❌ Error: {str(e)}", "Please try with a different image.", 0, {}
 
-    def _detect_style(self, prompt: str) -> str:
-        """Detecta el estilo principal del prompt"""
-        styles = {
-            "portrait": ["portrait", "person", "face", "headshot"],
-            "landscape": ["landscape", "mountain", "nature", "scenery"],
-            "street": ["street", "urban", "city"],
-            "artistic": ["artistic", "abstract", "conceptual"],
-            "dramatic": ["dramatic", "cinematic", "moody"]
-        }
-        
-        prompt_lower = prompt.lower()
-        for style_name, keywords in styles.items():
-            if any(keyword in prompt_lower for keyword in keywords):
-                return style_name
-        
-        return "general"
-
-    def _detect_subject(self, prompt: str) -> str:
-        """Detecta el sujeto principal del prompt"""
-        if not prompt:
-            return "Unknown"
-            
-        # Tomar las primeras palabras significativas
-        words = prompt.split(',')[0].split()
-        if len(words) > 3:
-            return ' '.join(words[:4])
-        return prompt.split(',')[0] if prompt else "Unknown"
-
-    def _calculate_score(self, optimized_prompt: str, base_prompt: str) -> int:
-        """Calcula el score basado en la calidad del prompt"""
-        score = 0
-        
-        # Base score por longitud y riqueza
-        score += min(len(base_prompt) // 10, 25)
-        
-        # Technical enhancement
-        if "Shot on" in optimized_prompt:
-            score += 25
-        
-        # Lighting quality
-        if "lighting" in optimized_prompt.lower():
-            score += 25
-        
-        # Professional quality
-        if any(word in optimized_prompt.lower() for word in ["professional", "masterful", "epic", "cinematic"]):
-            score += 25
-        
-        return min(score, 100)
-
     def _generate_ultra_analysis_report(self, analysis: Dict[str, Any], 
                                        score: int, breakdown: Dict[str, int], 
-                                       duration: float) -> str:
+                                       duration: float, caption_model: str = "Florence-2") -> str:
         """Generate ultra detailed analysis report with multi-model results"""
         
         device_used = "cuda" if torch.cuda.is_available() else "cpu"
@@ -457,9 +392,12 @@ class UltraSupremeOptimizer:
         # Intelligence metrics
         metrics = analysis["intelligence_metrics"]
         
+        # Caption info
+        caption_info = analysis.get("clip_best", "")[:150] + "..." if len(analysis.get("clip_best", "")) > 150 else analysis.get("clip_best", "")
+        
         analysis_info = f"""**🚀 ULTRA SUPREME MULTI-MODEL ANALYSIS COMPLETE**
-**Processing:** {gpu_status} • {duration:.1f}s • Multi-Model Pipeline
-**Ultra Score:** {score}/100 • Models: CLIP + DeepFace + MediaPipe + Transformers
+**Processing:** {gpu_status} • {duration:.1f}s • {caption_model} + Multi-Model Pipeline
+**Ultra Score:** {score}/100 • Models: {caption_model} + DeepFace + MediaPipe + Transformers
 
 **📊 BREAKDOWN:**
 • Prompt Quality: {breakdown.get('prompt_quality', 0)}/25
@@ -467,6 +405,9 @@ class UltraSupremeOptimizer:
 • Model Confidence: {breakdown.get('model_confidence', 0)}/25
 • Feature Richness: {breakdown.get('feature_richness', 0)}/25
 
+**📝 VISION-LANGUAGE ANALYSIS:**
+**{caption_model} Caption:** {caption_info}
+
 **🧠 DEEP ANALYSIS RESULTS:**
 
 **👤 DEMOGRAPHICS & IDENTITY:**
@@ -491,9 +432,9 @@ class UltraSupremeOptimizer:
 • **Technical Optimization:** {metrics['technical_optimization_score']}/100
 
 **✨ MULTI-MODEL ADVANTAGES:**
+✅ {caption_model}: State-of-the-art vision-language understanding
 ✅ DeepFace: Accurate age, gender, emotion detection
 ✅ MediaPipe: Body pose and gesture analysis
-✅ CLIP: Semantic understanding and context
 ✅ Transformers: Advanced emotion classification
 ✅ OpenCV: Robust face detection