Update optimizer.py

optimizer.py

@@ -30,9 +30,8 @@ class UltraSupremeOptimizer:
         self.usage_count = 0
         self.device = self._get_device()
         self.is_initialized = False
-        # Inicializar modelo inmediatamente
-        self.initialize_model()
-
+        # NO inicializar modelo aquí - hacerlo lazy
+        
     @staticmethod
     def _get_device() -> str:
         """Determine the best available device for computation"""
@@ -49,13 +48,13 @@ class UltraSupremeOptimizer:
             return True
         
         try:
-            # Configuración estándar sin forzar precisión
+            # Configuración para CPU inicialmente
             config = Config(
                 clip_model_name="ViT-L-14/openai",
                 download_cache=True,
                 chunk_size=2048,
                 quiet=True,
-                device="cpu"  # Inicializar en CPU
+                device="cpu"  # Siempre inicializar en CPU
             )
             
             self.interrogator = Interrogator(config)
@@ -63,7 +62,8 @@ class UltraSupremeOptimizer:
             
             # Clean up memory after initialization
             gc.collect()
-                
+            
+            logger.info("Model initialized successfully on CPU")
             return True
             
         except Exception as e:
@@ -86,8 +86,8 @@ class UltraSupremeOptimizer:
             if image.mode != 'RGB':
                 image = image.convert('RGB')
             
-            # Resize if too large - usar tamaño generoso para máxima calidad
-            max_size = 1024 if self.device != "cpu" else 768
+            # Resize if too large
+            max_size = 768  # Reducir tamaño para evitar problemas de memoria
             if image.size[0] > max_size or image.size[1] > max_size:
                 image.thumbnail((max_size, max_size), Image.Resampling.LANCZOS)
             
@@ -109,7 +109,7 @@ class UltraSupremeOptimizer:
             r',\s*artstation',
             r',\s*concept art',
             r',\s*digital art',
-            r',\s*by greg rutkowski',  # Remover artistas genéricos overused
+            r',\s*by greg rutkowski',
         ]
         
         cleaned_prompt = base_prompt
@@ -148,25 +148,101 @@ class UltraSupremeOptimizer:
         
         return final_prompt
 
-    @spaces.GPU
+    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"""
         try:
-            # Mover modelos a GPU sin forzar precisión
-            if self.device == "cuda":
-                # Configurar el dispositivo en el interrogator
-                self.interrogator.config.device = "cuda"
+            # Preparar modelos para GPU
+            self._prepare_models_for_gpu()
+            
+            # Usar autocast para manejar precisión mixta
+            with torch.cuda.amp.autocast(enabled=True, dtype=torch.float16):
+                # Convertir imagen a tensor y asegurar que esté en half precision
+                from torchvision import transforms
+                preprocess = transforms.Compose([
+                    transforms.Resize((224, 224)),
+                    transforms.ToTensor(),
+                    transforms.Normalize(mean=[0.48145466, 0.4578275, 0.40821073], 
+                                       std=[0.26862954, 0.26130258, 0.27577711]),
+                ])
                 
-                # Mover modelos a GPU manteniendo su precisión nativa
-                if hasattr(self.interrogator, 'clip_model') and self.interrogator.clip_model is not None:
-                    self.interrogator.clip_model = self.interrogator.clip_model.to("cuda")
-                    logger.info("CLIP model moved to GPU with native precision")
+                # Procesar imagen manualmente para controlar la precisión
+                image_tensor = preprocess(image).unsqueeze(0).half().to("cuda")
                 
-                if hasattr(self.interrogator, 'blip_model') and self.interrogator.blip_model is not None:
-                    self.interrogator.blip_model = self.interrogator.blip_model.to("cuda")
-                    logger.info("BLIP model moved to GPU with native precision")
+                # Ejecutar inferencias con manejo especial
+                full_prompt = self._safe_interrogate(image, 'interrogate')
+                clip_fast = self._safe_interrogate(image, 'interrogate_fast')
+                clip_classic = self._safe_interrogate(image, 'interrogate_classic')
             
-            # Ejecutar inferencias CLIP con precisión nativa
+            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")
+            
+            # Preparar modelos para CPU
+            self._prepare_models_for_cpu()
+            
+            # 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)
@@ -174,26 +250,8 @@ class UltraSupremeOptimizer:
             return full_prompt, clip_fast, clip_classic
             
         except Exception as e:
-            logger.error(f"CLIP inference error: {e}")
-            # Si falla en GPU, intentar en CPU
-            if self.device == "cuda":
-                logger.info("Falling back to CPU inference")
-                self.interrogator.config.device = "cpu"
-                
-                if hasattr(self.interrogator, 'clip_model') and self.interrogator.clip_model is not None:
-                    self.interrogator.clip_model = self.interrogator.clip_model.to("cpu")
-                
-                if hasattr(self.interrogator, 'blip_model') and self.interrogator.blip_model is not None:
-                    self.interrogator.blip_model = self.interrogator.blip_model.to("cpu")
-                
-                # Reintentar en CPU
-                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
-            else:
-                raise e
+            logger.error(f"CPU inference also failed: {e}")
+            return "Error: Failed to process image", "Error", "Error"
 
     def generate_ultra_supreme_prompt(self, image: Any) -> Tuple[str, str, int, Dict[str, int]]:
         """
@@ -203,9 +261,10 @@ class UltraSupremeOptimizer:
             Tuple of (prompt, analysis_info, score, breakdown)
         """
         try:
-            # Verificar que el modelo esté inicializado
+            # Inicializar modelo si no está inicializado
             if not self.is_initialized:
-                return "❌ Model initialization failed.", "Please refresh and try again.", 0, {}
+                if not self.initialize_model():
+                    return "❌ Model initialization failed.", "Please refresh and try again.", 0, {}
             
             # Validate input
             if image is None:
@@ -222,17 +281,24 @@ class UltraSupremeOptimizer:
             
             logger.info("ULTRA SUPREME ANALYSIS - Starting pipeline")
             
-            # Ejecutar inferencia CLIP en GPU
+            # Ejecutar inferencia CLIP
             full_prompt, clip_fast, clip_classic = self.run_clip_inference(image)
             
-            logger.info(f"Prompt completo de CLIP Interrogator: {full_prompt}")
-            logger.info(f"Análisis Fast: {clip_fast}")
-            logger.info(f"Análisis Classic: {clip_classic}")
+            # Verificar si hubo errores
+            if "Error" in full_prompt:
+                logger.warning("Using fallback prompt due to inference error")
+                full_prompt = "A photograph"
+                clip_fast = "image"
+                clip_classic = "picture"
             
-            # 3. Aplicar reglas de Flux al prompt completo
+            logger.info(f"Prompt completo: {full_prompt[:100]}...")
+            logger.info(f"Fast: {clip_fast[:50]}...")
+            logger.info(f"Classic: {clip_classic[:50]}...")
+            
+            # Aplicar reglas de Flux al prompt completo
             optimized_prompt = self.apply_flux_rules(full_prompt)
             
-            # 4. Crear análisis para el reporte (simplificado)
+            # Crear análisis para el reporte
             analysis_summary = {
                 "base_prompt": full_prompt,
                 "clip_fast": clip_fast,
@@ -242,7 +308,7 @@ class UltraSupremeOptimizer:
                 "detected_subject": self._detect_subject(full_prompt)
             }
             
-            # 5. Calcular score basado en la riqueza del prompt
+            # Calcular score
             score = self._calculate_score(optimized_prompt, full_prompt)
             breakdown = {
                 "base_quality": min(len(full_prompt) // 10, 25),
@@ -268,7 +334,7 @@ class UltraSupremeOptimizer:
             return optimized_prompt, analysis_info, score, breakdown
             
         except Exception as e:
-            logger.error(f"Ultra supreme generation error: {e}")
+            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:
@@ -281,19 +347,23 @@ class UltraSupremeOptimizer:
             "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):
+            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]
+        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"""
@@ -321,8 +391,9 @@ class UltraSupremeOptimizer:
                                   duration: float) -> str:
         """Generate detailed analysis report"""
         
-        gpu_status = "⚡ ZeroGPU" if torch.cuda.is_available() else "💻 CPU"
-        precision_info = "Native Model Precision" if torch.cuda.is_available() else "CPU Processing"
+        device_used = "cuda" if torch.cuda.is_available() else "cpu"
+        gpu_status = "⚡ ZeroGPU" if device_used == "cuda" else "💻 CPU"
+        precision_info = "Half Precision (FP16)" if device_used == "cuda" else "Full Precision (FP32)"
         
         # Extraer información clave
         detected_style = analysis.get("detected_style", "general").title()
@@ -337,8 +408,8 @@ class UltraSupremeOptimizer:
 **🧠 INTELLIGENT DETECTION:**
 - **Detected Style:** {detected_style}
 - **Main Subject:** {detected_subject}
-- **Precision:** Using native model precision for optimal performance
-- **Quality:** Maximum resolution processing (1024px)
+- **Precision:** Using {precision_info} for optimal performance
+- **Quality:** Maximum resolution processing (768px)
 
 **📊 CLIP INTERROGATOR ANALYSIS:**
 - **Base Prompt:** {base_prompt_preview}
@@ -346,9 +417,9 @@ class UltraSupremeOptimizer:
 - **Classic Analysis:** {analysis.get('clip_classic', '')[:80]}...
 
 **⚡ OPTIMIZATION APPLIED:**
-- ✅ Native precision inference for stability
-- ✅ GPU acceleration when available
-- ✅ Automatic fallback to CPU if needed
+- ✅ Mixed precision handling for stability
+- ✅ Automatic GPU/CPU fallback
+- ✅ Memory-efficient processing
 - ✅ Added professional camera specifications
 - ✅ Enhanced lighting descriptions
 - ✅ Applied Flux-specific optimizations