Skip to content

C
coderai
Commit 00775972 authored by Your Name's avatar Your Name
Browse files
Options

Fix: Centralize model unloading - properly handle all model types in ondemand mode 

- Added unload_all_models() to MultiModelManager that handles ALL model types:
  ModelManager, diffusers pipelines, sd.cpp StableDiffusion, and any other objects
- Text endpoints now properly unload image models before loading text models
- Image endpoints now properly unload text models before loading image models
- The rule: in ondemand mode, if the model in VRAM differs from the requested
  model (regardless of type), fully unload before loading the new one
- Includes gc.collect(), torch.cuda.empty_cache(), and 1s settle delay
parent 7d838962
Expand all Show whitespace changes
Inline Side-by-side
Showing with 150 additions and 283 deletions
codai/api/images.py
View file @ 00775972
This diff is collapsed.
codai/api/text.py
View file @ 00775972
......@@ -299,63 +299,36 @@ async def chat_completions(request: ChatCompletionRequest, http_request: Request
from codai.api.state import get_load_mode
load_mode = get_load_mode()
# Check if there's an image model already loaded in VRAM
current_image_model = None
for key in multi_model_manager.models.keys():
if key.startswith("image:"):
current_image_model = key
break
# Check if legacy model_manager has a model loaded
has_legacy_model = model_manager.backend is not None
# In ondemand mode, if any model (text, image, etc.) is already loaded and we're requesting a different model,
# we should unload the current model first to free VRAM
needs_full_unload = (load_mode == "ondemand" and (current_image_model is not None or has_legacy_model))
# In ondemand mode (no --load-all or --loadswap), if ANY model is loaded in VRAM
# and it's different from what we need, fully unload it first to free VRAM
if load_mode == "ondemand":
has_any_model = len(multi_model_manager.models) > 0 or model_manager.backend is not None
if has_any_model:
# Check if the requested model is already loaded (no need to unload)
already_loaded = False
if requested_model and requested_model in multi_model_manager.models:
already_loaded = True
elif multi_model_manager.default_model and (
not requested_model or requested_model == "default" or
requested_model == multi_model_manager.default_model
):
if multi_model_manager.default_model in multi_model_manager.models:
already_loaded = True
# If we're requesting a text model and there's an image model loaded, unload it first
if needs_full_unload:
print(f"In ondemand mode - fully unloading current model before loading text model...")
if not already_loaded:
print(f"In ondemand mode - fully unloading current model(s) before loading text model '{requested_model}'...")
# Full cleanup: remove all models from VRAM
for key in list(multi_model_manager.models.keys()):
model_to_cleanup = multi_model_manager.models.get(key)
if model_to_cleanup is not None:
print(f"Unloading '{key}' from VRAM...")
try:
if hasattr(model_to_cleanup, 'cleanup') and callable(getattr(model_to_cleanup, 'cleanup')):
model_to_cleanup.cleanup()
except Exception as e:
print(f"Warning during cleanup of '{key}': {e}")
del multi_model_manager.models[key]
# Use centralized unload method
multi_model_manager.unload_all_models()
# Also cleanup legacy model_manager
if model_manager.backend is not None:
print("Unloading legacy model_manager from VRAM...")
try:
if hasattr(model_manager.backend, 'unload'):
model_manager.backend.unload()
elif hasattr(model_manager.backend, 'cleanup'):
model_manager.backend.cleanup()
model_manager.cleanup()
except Exception as e:
print(f"Warning during legacy model cleanup: {e}")
model_manager.backend = None
# Force garbage collection and clear CUDA cache
import gc
gc.collect()
try:
import torch
if torch.cuda.is_available():
torch.cuda.synchronize()
torch.cuda.empty_cache()
print("CUDA cache cleared")
except:
pass
# Add delay to let VRAM settle
import time
time.sleep(1)
# Try to get the appropriate model
mm = multi_model_manager.get_model_for_request(requested_model)
......@@ -1754,51 +1727,27 @@ async def completions(request: CompletionRequest):
from codai.api.state import get_load_mode
load_mode = get_load_mode()
# Check if there's an image model already loaded in VRAM
current_image_model = None
for key in multi_model_manager.models.keys():
if key.startswith("image:"):
current_image_model = key
break
# In ondemand mode, if any model is already loaded, unload it first
needs_full_unload = (load_mode == "ondemand" and current_image_model is not None)
if needs_full_unload:
print(f"In ondemand mode - fully unloading current model before loading text model...")
# Full cleanup
for key in list(multi_model_manager.models.keys()):
model_to_cleanup = multi_model_manager.models.get(key)
if model_to_cleanup is not None:
print(f"Unloading '{key}' from VRAM...")
try:
if hasattr(model_to_cleanup, 'cleanup') and callable(getattr(model_to_cleanup, 'cleanup')):
model_to_cleanup.cleanup()
except Exception as e:
print(f"Warning during cleanup of '{key}': {e}")
del multi_model_manager.models[key]
# In ondemand mode, if ANY model is loaded and it's different from what we need, unload first
if load_mode == "ondemand":
has_any_model = len(multi_model_manager.models) > 0 or model_manager.backend is not None
if has_any_model:
already_loaded = False
if requested_model and requested_model in multi_model_manager.models:
already_loaded = True
elif multi_model_manager.default_model and (
not requested_model or requested_model == "default" or
requested_model == multi_model_manager.default_model
):
if multi_model_manager.default_model in multi_model_manager.models:
already_loaded = True
# Also cleanup legacy model_manager
if not already_loaded:
print(f"In ondemand mode - fully unloading current model(s) before loading text model '{requested_model}'...")
multi_model_manager.unload_all_models()
if model_manager.backend is not None:
print("Unloading legacy model_manager from VRAM...")
try:
if hasattr(model_manager.backend, 'unload'):
model_manager.backend.unload()
elif hasattr(model_manager.backend, 'cleanup'):
model_manager.backend.cleanup()
except:
pass
model_manager.backend = None
# Force garbage collection and clear CUDA cache
import gc
gc.collect()
try:
import torch
if torch.cuda.is_available():
torch.cuda.synchronize()
torch.cuda.empty_cache()
model_manager.cleanup()
except:
pass
......
codai/models/manager.py
View file @ 00775972
......@@ -711,6 +711,78 @@ class MultiModelManager:
# Model not found - try to load it as a new model
return self._load_model_by_name(requested_model)
def unload_all_models(self):
"""
Fully unload ALL models from VRAM. Used in ondemand mode when switching
between different model types (e.g., text -> image or image -> text).
This handles all model types:
- ModelManager instances (have cleanup() method)
- Diffusers pipelines (need to be moved to CPU and deleted)
- stable-diffusion-cpp StableDiffusion instances
- Any other model objects
"""
print("=== FULL VRAM CLEANUP: Unloading all models ===")
for key in list(self.models.keys()):
model_obj = self.models.get(key)
if model_obj is None:
continue
print(f"Unloading '{key}' from VRAM...")
try:
# Method 1: ModelManager with cleanup()
if hasattr(model_obj, 'cleanup') and callable(getattr(model_obj, 'cleanup')):
model_obj.cleanup()
# Method 2: Diffusers pipeline (has 'to' method to move to CPU)
elif hasattr(model_obj, 'to') and callable(getattr(model_obj, 'to')):
try:
model_obj.to('cpu')
except:
pass
del model_obj
# Method 3: Object with 'model' attribute (e.g., wrapper)
elif hasattr(model_obj, 'model') and model_obj.model is not None:
if hasattr(model_obj.model, 'cleanup'):
model_obj.model.cleanup()
elif hasattr(model_obj.model, 'to'):
try:
model_obj.model.to('cpu')
except:
pass
del model_obj
# Method 4: Just delete it
else:
del model_obj
except Exception as e:
print(f"Warning during cleanup of '{key}': {e}")
# Remove from dict
if key in self.models:
del self.models[key]
# Reset tracking state
self.current_model_key = None
self.active_in_vram = None
# Force garbage collection
for _ in range(3):
gc.collect()
# Clear CUDA cache if available
try:
import torch
if torch.cuda.is_available():
torch.cuda.synchronize()
torch.cuda.empty_cache()
print("CUDA cache cleared")
except:
pass
# Small delay to let GPU memory settle
time.sleep(1)
print("=== FULL VRAM CLEANUP: Complete ===")
def add_model(self, key: str, manager: ModelManager):
"""Add a model manager for a specific key."""
self.models[key] = manager
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment