fix

app.py

@@ -17,8 +17,6 @@ import functools
 
 app = FastAPI()
 
-# Comprehensive model registry
-
 # Comprehensive model registry
 MODELS = {
     "SDXL-Base": {
@@ -158,7 +156,11 @@ class ModelContext:
         self.model_name = model_name
         
     def __enter__(self):
-        return model_manager.load_model(self.model_name)
+        pipeline = model_manager.load_model(self.model_name)
+        if hasattr(pipeline, 'reset_device_map'):
+            pipeline.reset_device_map()
+        pipeline.enable_model_cpu_offload()
+        return pipeline
         
     def __exit__(self, exc_type, exc_val, exc_tb):
         if exc_type is not None: