""" 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": "", "more_detailed": "", "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" ]