change model context

app.py

@@ -1,4 +1,5 @@
 import spaces
+from accelerate import dispatch_model
 from fastapi import FastAPI, HTTPException, UploadFile, File
 from typing import Optional, Dict, Any
 import torch
@@ -123,8 +124,8 @@ class ModelManager:
                 if hasattr(self.current_pipeline, 'enable_xformers_memory_efficient_attention'):
                     self.current_pipeline.enable_xformers_memory_efficient_attention()
                 
-                if self._device == "cuda":
-                    self.current_pipeline.enable_model_cpu_offload()
+                # if self._device == "cuda":
+                #     self.current_pipeline.enable_model_cpu_offload()
                     
                 self.current_model = model_name
                 
@@ -159,7 +160,8 @@ class ModelContext:
         pipeline = model_manager.load_model(self.model_name)
         if hasattr(pipeline, 'reset_device_map'):
             pipeline.reset_device_map()
-        pipeline.enable_model_cpu_offload()
+        # Enable automatic device mapping strategy
+        dispatch_model(pipeline)
         return pipeline
         
     def __exit__(self, exc_type, exc_val, exc_tb):