Create models.py

models.py

@@ -0,0 +1,362 @@
+"""
+Model management for FLUX Prompt Optimizer
+Handles Florence-2 and Bagel model integration
+"""
+
+import logging
+import requests
+import spaces
+import torch
+from typing import Optional, Dict, Any, Tuple
+from PIL import Image
+from transformers import AutoProcessor, AutoModelForCausalLM
+
+from config import MODEL_CONFIG, get_device_config
+from utils import clean_memory, safe_execute
+
+logger = logging.getLogger(__name__)
+
+
+class BaseImageAnalyzer:
+    """Base class for image analysis models"""
+    
+    def __init__(self):
+        self.model = None
+        self.processor = None
+        self.device_config = get_device_config()
+        self.is_initialized = False
+        
+    def initialize(self) -> bool:
+        """Initialize the model"""
+        raise NotImplementedError
+        
+    def analyze_image(self, image: Image.Image) -> Tuple[str, Dict[str, Any]]:
+        """Analyze image and return description"""
+        raise NotImplementedError
+        
+    def cleanup(self) -> None:
+        """Clean up model resources"""
+        if self.model is not None:
+            del self.model
+            self.model = None
+        if self.processor is not None:
+            del self.processor
+            self.processor = None
+        clean_memory()
+
+
+class Florence2Analyzer(BaseImageAnalyzer):
+    """Florence-2 model for image analysis"""
+    
+    def __init__(self):
+        super().__init__()
+        self.config = MODEL_CONFIG["florence2"]
+        
+    def initialize(self) -> bool:
+        """Initialize Florence-2 model"""
+        if self.is_initialized:
+            return True
+            
+        try:
+            logger.info("Initializing Florence-2 model...")
+            
+            model_id = self.config["model_id"]
+            
+            # Load processor
+            self.processor = AutoProcessor.from_pretrained(
+                model_id, 
+                trust_remote_code=self.config["trust_remote_code"]
+            )
+            
+            # Load model
+            self.model = AutoModelForCausalLM.from_pretrained(
+                model_id,
+                trust_remote_code=self.config["trust_remote_code"],
+                torch_dtype=self.config["torch_dtype"] if self.device_config["use_gpu"] else torch.float32
+            )
+            
+            # Move to appropriate device
+            if self.device_config["use_gpu"]:
+                self.model = self.model.to(self.device_config["device"])
+            else:
+                self.model = self.model.to("cpu")
+                
+            self.model.eval()
+            self.is_initialized = True
+            
+            logger.info(f"Florence-2 initialized on {self.device_config['device']}")
+            return True
+            
+        except Exception as e:
+            logger.error(f"Florence-2 initialization failed: {e}")
+            self.cleanup()
+            return False
+    
+    @spaces.GPU(duration=60)
+    def _gpu_inference(self, image: Image.Image, task_prompt: str) -> str:
+        """Run inference on GPU with spaces decorator"""
+        try:
+            # Move model to GPU for inference
+            if self.device_config["use_gpu"]:
+                self.model = self.model.to("cuda")
+            
+            # Prepare inputs
+            inputs = self.processor(text=task_prompt, images=image, return_tensors="pt")
+            
+            # Move inputs to device
+            device = "cuda" if self.device_config["use_gpu"] else self.device_config["device"]
+            inputs = {k: v.to(device) for k, v in inputs.items()}
+            
+            # Generate response
+            with torch.no_grad():
+                if self.device_config["use_gpu"]:
+                    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=self.config["max_new_tokens"],
+                            num_beams=3,
+                            do_sample=False
+                        )
+                else:
+                    generated_ids = self.model.generate(
+                        input_ids=inputs["input_ids"],
+                        pixel_values=inputs["pixel_values"],
+                        max_new_tokens=self.config["max_new_tokens"],
+                        num_beams=3,
+                        do_sample=False
+                    )
+            
+            # Decode response
+            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 caption
+            if task_prompt in parsed:
+                return parsed[task_prompt]
+            else:
+                return str(parsed) if parsed else ""
+                
+        except Exception as e:
+            logger.error(f"Florence-2 GPU inference failed: {e}")
+            return ""
+        finally:
+            # Move model back to CPU to free GPU memory
+            if self.device_config["use_gpu"]:
+                self.model = self.model.to("cpu")
+                clean_memory()
+    
+    def analyze_image(self, image: Image.Image) -> Tuple[str, Dict[str, Any]]:
+        """Analyze image using Florence-2"""
+        if not self.is_initialized:
+            success = self.initialize()
+            if not success:
+                return "Model initialization failed", {"error": "Florence-2 not available"}
+        
+        try:
+            # Define analysis tasks
+            tasks = {
+                "detailed": "<DETAILED_CAPTION>",
+                "more_detailed": "<MORE_DETAILED_CAPTION>",
+                "caption": "<CAPTION>"
+            }
+            
+            results = {}
+            
+            # Run analysis for each task
+            for task_name, task_prompt in tasks.items():
+                if self.device_config["use_gpu"]:
+                    result = self._gpu_inference(image, task_prompt)
+                else:
+                    result = self._cpu_inference(image, task_prompt)
+                results[task_name] = result
+            
+            # Choose best result
+            if results["more_detailed"]:
+                main_description = results["more_detailed"]
+            elif results["detailed"]:
+                main_description = results["detailed"]
+            else:
+                main_description = results["caption"] or "A photograph"
+            
+            # Prepare metadata
+            metadata = {
+                "model": "Florence-2",
+                "device": self.device_config["device"],
+                "all_results": results,
+                "confidence": 0.85  # Florence-2 generally reliable
+            }
+            
+            logger.info(f"Florence-2 analysis complete: {len(main_description)} chars")
+            return main_description, metadata
+            
+        except Exception as e:
+            logger.error(f"Florence-2 analysis failed: {e}")
+            return "Analysis failed", {"error": str(e)}
+    
+    def _cpu_inference(self, image: Image.Image, task_prompt: str) -> str:
+        """Run inference on CPU"""
+        try:
+            inputs = self.processor(text=task_prompt, images=image, return_tensors="pt")
+            
+            with torch.no_grad():
+                generated_ids = self.model.generate(
+                    input_ids=inputs["input_ids"],
+                    pixel_values=inputs["pixel_values"],
+                    max_new_tokens=self.config["max_new_tokens"],
+                    num_beams=2,  # Reduced for CPU
+                    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)
+            )
+            
+            if task_prompt in parsed:
+                return parsed[task_prompt]
+            else:
+                return str(parsed) if parsed else ""
+                
+        except Exception as e:
+            logger.error(f"Florence-2 CPU inference failed: {e}")
+            return ""
+
+
+class BagelAnalyzer(BaseImageAnalyzer):
+    """Bagel-7B model analyzer via API"""
+    
+    def __init__(self):
+        super().__init__()
+        self.config = MODEL_CONFIG["bagel"]
+        self.session = requests.Session()
+        
+    def initialize(self) -> bool:
+        """Initialize Bagel analyzer (API-based)"""
+        try:
+            # Test API connectivity
+            test_response = self.session.get(
+                self.config["api_url"],
+                timeout=self.config["timeout"]
+            )
+            
+            if test_response.status_code == 200:
+                self.is_initialized = True
+                logger.info("Bagel API connection established")
+                return True
+            else:
+                logger.error(f"Bagel API not accessible: {test_response.status_code}")
+                return False
+                
+        except Exception as e:
+            logger.error(f"Bagel initialization failed: {e}")
+            return False
+    
+    def analyze_image(self, image: Image.Image) -> Tuple[str, Dict[str, Any]]:
+        """Analyze image using Bagel-7B API"""
+        if not self.is_initialized:
+            success = self.initialize()
+            if not success:
+                return "Bagel API not available", {"error": "API connection failed"}
+        
+        try:
+            # Convert image to base64 or prepare for API call
+            # Note: This is a placeholder - actual implementation would depend on Bagel API format
+            
+            # For now, return a placeholder response
+            # In real implementation, you would:
+            # 1. Convert image to required format
+            # 2. Make API call to Bagel endpoint
+            # 3. Parse response
+            
+            description = "Detailed image analysis via Bagel-7B (API implementation needed)"
+            metadata = {
+                "model": "Bagel-7B",
+                "method": "API",
+                "confidence": 0.8
+            }
+            
+            logger.info("Bagel analysis complete (placeholder)")
+            return description, metadata
+            
+        except Exception as e:
+            logger.error(f"Bagel analysis failed: {e}")
+            return "Analysis failed", {"error": str(e)}
+
+
+class ModelManager:
+    """Manager for handling multiple analysis models"""
+    
+    def __init__(self, preferred_model: str = None):
+        self.preferred_model = preferred_model or MODEL_CONFIG["primary_model"]
+        self.analyzers = {}
+        self.current_analyzer = None
+        
+    def get_analyzer(self, model_name: str = None) -> Optional[BaseImageAnalyzer]:
+        """Get or create analyzer for specified model"""
+        model_name = model_name or self.preferred_model
+        
+        if model_name not in self.analyzers:
+            if model_name == "florence2":
+                self.analyzers[model_name] = Florence2Analyzer()
+            elif model_name == "bagel":
+                self.analyzers[model_name] = BagelAnalyzer()
+            else:
+                logger.error(f"Unknown model: {model_name}")
+                return None
+        
+        return self.analyzers[model_name]
+    
+    def analyze_image(self, image: Image.Image, model_name: str = None) -> Tuple[str, Dict[str, Any]]:
+        """Analyze image with specified or preferred model"""
+        analyzer = self.get_analyzer(model_name)
+        if analyzer is None:
+            return "No analyzer available", {"error": "Model not found"}
+        
+        success, result = safe_execute(analyzer.analyze_image, image)
+        if success:
+            return result
+        else:
+            return "Analysis failed", {"error": result}
+    
+    def cleanup_all(self) -> None:
+        """Clean up all model resources"""
+        for analyzer in self.analyzers.values():
+            analyzer.cleanup()
+        self.analyzers.clear()
+        clean_memory()
+
+
+# Global model manager instance
+model_manager = ModelManager()
+
+
+def analyze_image(image: Image.Image, model_name: str = None) -> Tuple[str, Dict[str, Any]]:
+    """
+    Convenience function for image analysis
+    
+    Args:
+        image: PIL Image to analyze
+        model_name: Optional model name ("florence2" or "bagel")
+        
+    Returns:
+        Tuple of (description, metadata)
+    """
+    return model_manager.analyze_image(image, model_name)
+
+
+# Export main components
+__all__ = [
+    "BaseImageAnalyzer",
+    "Florence2Analyzer", 
+    "BagelAnalyzer",
+    "ModelManager",
+    "model_manager",
+    "analyze_image"
+]