Update app.py

app.py

@@ -181,7 +181,6 @@ def create_temp_module(code: str) -> str:
 
 # Initialize GPU configuration
 setup_gpu_memory()
-
 class ModelManager:
     """Manages AI models and their configurations"""
     
@@ -203,29 +202,28 @@ class ModelManager:
             "image_processor": {
                 "model_id": "Salesforce/blip-image-captioning-base",
                 "processor": AutoProcessor,
-                "model": BlipForConditionalGeneration,  # Changed from AutoModel
+                "model": BlipForConditionalGeneration,
                 "kwargs": {
                     "cache_dir": str(cache_dir),
-                    "device_map": "auto"  # Add device mapping
-                   }
+                    "device_map": "auto"
+                }
             }
         }
 
     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 == "image_processor":
-            try:
-                processor = config["processor"].from_pretrained(
+            if model_type == "code_generator":
+                tokenizer = config["tokenizer"].from_pretrained(
                     config["model_id"],
                     **config["kwargs"]
                 )
@@ -233,14 +231,33 @@ 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")
-                
-            except Exception as e:
-                logger.error(f"Error loading {model_type} model: {e}")
-                raise ModelError(f"Failed to load {model_type} model: {e}")
+                self.loaded_models[model_type] = (model, tokenizer)
+
+            elif model_type == "image_processor":
+                try:
+                    processor = config["processor"].from_pretrained(
+                        config["model_id"],
+                        **config["kwargs"]
+                    )
+                    model = config["model"].from_pretrained(
+                        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")
+                    
+                except Exception as e:
+                    logger.error(f"Error loading {model_type} model: {e}")
+                    raise ModelError(f"Failed to load {model_type} model: {e}")
+
+            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)}")
+
     def unload_model(self, model_type: str):
         """Unload a model to free memory"""
         if model_type in self.loaded_models: