acecalisto3 commited on
Commit
6e9a206
·
verified ·
1 Parent(s): 7ada1ee

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +40 -23
app.py CHANGED
@@ -181,7 +181,6 @@ def create_temp_module(code: str) -> str:
181
 
182
  # Initialize GPU configuration
183
  setup_gpu_memory()
184
-
185
  class ModelManager:
186
  """Manages AI models and their configurations"""
187
 
@@ -203,29 +202,28 @@ class ModelManager:
203
  "image_processor": {
204
  "model_id": "Salesforce/blip-image-captioning-base",
205
  "processor": AutoProcessor,
206
- "model": BlipForConditionalGeneration, # Changed from AutoModel
207
  "kwargs": {
208
  "cache_dir": str(cache_dir),
209
- "device_map": "auto" # Add device mapping
210
- }
211
  }
212
  }
213
 
214
  def load_model(self, model_type: str):
215
- """Load a model by type"""
216
- try:
217
- if model_type not in self.model_configs:
218
- raise ModelError(f"Unknown model type: {model_type}")
219
 
220
- if model_type in self.loaded_models:
221
- return self.loaded_models[model_type]
222
 
223
- config = self.model_configs[model_type]
224
- logger.info(f"Loading {model_type} model...")
225
 
226
- if model_type == "image_processor":
227
- try:
228
- processor = config["processor"].from_pretrained(
229
  config["model_id"],
230
  **config["kwargs"]
231
  )
@@ -233,14 +231,33 @@ class ModelManager:
233
  config["model_id"],
234
  **config["kwargs"]
235
  )
236
- if torch.cuda.is_available():
237
- model = model.to("cuda")
238
- self.loaded_models[model_type] = (model, processor)
239
- logger.info(f"{model_type} model loaded successfully")
240
-
241
- except Exception as e:
242
- logger.error(f"Error loading {model_type} model: {e}")
243
- raise ModelError(f"Failed to load {model_type} model: {e}")
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
244
  def unload_model(self, model_type: str):
245
  """Unload a model to free memory"""
246
  if model_type in self.loaded_models:
 
181
 
182
  # Initialize GPU configuration
183
  setup_gpu_memory()
 
184
  class ModelManager:
185
  """Manages AI models and their configurations"""
186
 
 
202
  "image_processor": {
203
  "model_id": "Salesforce/blip-image-captioning-base",
204
  "processor": AutoProcessor,
205
+ "model": BlipForConditionalGeneration,
206
  "kwargs": {
207
  "cache_dir": str(cache_dir),
208
+ "device_map": "auto"
209
+ }
210
  }
211
  }
212
 
213
  def load_model(self, model_type: str):
214
+ """Load a model by type"""
215
+ try:
216
+ if model_type not in self.model_configs:
217
+ raise ModelError(f"Unknown model type: {model_type}")
218
 
219
+ if model_type in self.loaded_models:
220
+ return self.loaded_models[model_type]
221
 
222
+ config = self.model_configs[model_type]
223
+ logger.info(f"Loading {model_type} model...")
224
 
225
+ if model_type == "code_generator":
226
+ tokenizer = config["tokenizer"].from_pretrained(
 
227
  config["model_id"],
228
  **config["kwargs"]
229
  )
 
231
  config["model_id"],
232
  **config["kwargs"]
233
  )
234
+ self.loaded_models[model_type] = (model, tokenizer)
235
+
236
+ elif model_type == "image_processor":
237
+ try:
238
+ processor = config["processor"].from_pretrained(
239
+ config["model_id"],
240
+ **config["kwargs"]
241
+ )
242
+ model = config["model"].from_pretrained(
243
+ config["model_id"],
244
+ **config["kwargs"]
245
+ )
246
+ if torch.cuda.is_available():
247
+ model = model.to("cuda")
248
+ self.loaded_models[model_type] = (model, processor)
249
+ logger.info(f"{model_type} model loaded successfully")
250
+
251
+ except Exception as e:
252
+ logger.error(f"Error loading {model_type} model: {e}")
253
+ raise ModelError(f"Failed to load {model_type} model: {e}")
254
+
255
+ logger.info(f"{model_type} model loaded successfully")
256
+ return self.loaded_models[model_type]
257
+
258
+ except Exception as e:
259
+ raise ModelError(f"Error loading {model_type} model: {str(e)}")
260
+
261
  def unload_model(self, model_type: str):
262
  """Unload a model to free memory"""
263
  if model_type in self.loaded_models: