Update app.py

app.py

@@ -23,7 +23,7 @@ from sentence_transformers import SentenceTransformer
 import faiss
 import numpy as np
 from PIL import Image
-
+from transformers import BlipForConditionalGeneration
 
 
 
@@ -213,29 +213,19 @@ class ModelManager:
         }
 
     def load_model(self, model_type: str):
-        """Load a model by type"""
-        try:
-            if model_type not in self.model_configs:
-                raise ModelError(f"Unknown model type: {model_type}")
+    """Load a model by type"""
+    try:
+        if model_type not in self.model_configs:
+            raise ModelError(f"Unknown model type: {model_type}")
 
-            if model_type in self.loaded_models:
-                return self.loaded_models[model_type]
+        if model_type in self.loaded_models:
+            return self.loaded_models[model_type]
 
-            config = self.model_configs[model_type]
-            logger.info(f"Loading {model_type} model...")
+        config = self.model_configs[model_type]
+        logger.info(f"Loading {model_type} model...")
 
-            if model_type == "code_generator":
-                tokenizer = config["tokenizer"].from_pretrained(
-                    config["model_id"],
-                    **config["kwargs"]
-                )
-                model = config["model"].from_pretrained(
-                    config["model_id"],
-                    **config["kwargs"]
-                )
-                self.loaded_models[model_type] = (model, tokenizer)
-
-            elif model_type == "image_processor":
+        if model_type == "image_processor":
+            try:
                 processor = config["processor"].from_pretrained(
                     config["model_id"],
                     **config["kwargs"]
@@ -244,14 +234,14 @@ class ModelManager:
                     config["model_id"],
                     **config["kwargs"]
                 )
+                if torch.cuda.is_available():
+                    model = model.to("cuda")
                 self.loaded_models[model_type] = (model, processor)
-
-            logger.info(f"{model_type} model loaded successfully")
-            return self.loaded_models[model_type]
-
-        except Exception as e:
-            raise ModelError(f"Error loading {model_type} model: {str(e)}")
-
+                logger.info(f"{model_type} model loaded successfully")
+                
+            except Exception as e:
+                logger.error(f"Error loading {model_type} model: {e}")
+                raise ModelError(f"Failed to load {model_type} model: {e}")
     def unload_model(self, model_type: str):
         """Unload a model to free memory"""
         if model_type in self.loaded_models:
@@ -309,21 +299,27 @@ class MultimodalRAG:
             logger.error(f"Error saving template embeddings: {e}")
 
     def encode_image(self, image: Image.Image) -> np.ndarray:
-        """Encode image using BLIP"""
-        try:
-            model, processor = self.model_manager.load_model("image_processor")
-            
-            inputs = processor(images=image, return_tensors="pt")
-            with torch.no_grad():
-                image_features = model.get_image_features(**inputs)
-            
-            return image_features.detach().numpy()
+    """Encode image using BLIP"""
+    try:
+        model, processor = self.model_manager.load_model("image_processor")
+        
+        # Process image
+        inputs = processor(images=image, return_tensors="pt").to(model.device)
+        
+        # Get image features using the proper method
+        with torch.no_grad():
+            outputs = model.get_image_features(**inputs)
+            image_features = outputs.last_hidden_state.mean(dim=1)  # Average pooling
+        
+        return image_features.cpu().numpy()
             
-        except Exception as e:
-            raise ModelError(f"Error encoding image: {str(e)}")
-
+    except Exception as e:
+        logger.error(f"Error encoding image: {str(e)}")
+        raise ModelError(f"Error encoding image: {str(e)}")
+        
     def encode_text(self, text: str) -> np.ndarray:
         """Encode text using sentence-transformers"""
+        
         try:
             return self.text_encoder.encode(text)
         except Exception as e: