Fix: In ondemand mode, fully unload current model before loading new one

- In ondemand mode (no --load-all or --loadswap specified), when a new model
  is requested, the current model in VRAM is now fully unloaded before loading
  the new one. This ensures clean model switching.
- Added cleanup logic to both /v1/chat/completions and /v1/completions endpoints
- Added same logic to image generation endpoints (diffusers and sd.cpp paths)
- Cleanup includes: model cleanup, gc.collect(), torch.cuda.empty_cache()

codai/api/images.py

@@ -272,15 +272,87 @@ async def create_image_generation(request: ImageGenerationRequest, http_request:
     model_key = f"image:{model_to_use}"
     pipeline = multi_model_manager.get_model(model_key)
     
-    # Try to load if not cached
-    if pipeline is None:
-        # FIRST: Clean up any loaded text models to free VRAM
-        # This must happen BEFORE attempting to load any new model
-        print("Cleaning up text models to free VRAM for image generation...")
+    # Check what model is currently in VRAM (if any)
+    current_image_model = None
+    for key in multi_model_manager.models.keys():
+        if key.startswith("image:"):
+            current_image_model = key
+            break
+    
+    # Check if the legacy model_manager has a model loaded
+    has_legacy_model = False
+    try:
+        from codai.models.manager import model_manager
+        if model_manager.backend is not None:
+            has_legacy_model = True
+    except:
+        pass
+    
+    # Determine if we need to fully unload current model before loading new one
+    # In "ondemand" mode (default - no --load-all, no --loadswap), 
+    # always unload current model before loading a different one
+    needs_full_unload = (mode == "ondemand" and (current_image_model is not None or has_legacy_model))
+    
+    # Try to load if not cached OR if we need to unload and reload
+    if pipeline is None or needs_full_unload:
+        # If in ondemand mode and a different model is already loaded, do FULL cleanup first
+        if needs_full_unload:
+            print(f"In ondemand mode - fully unloading current model before loading new one...")
+            
+            # FIRST: Clean up ALL loaded models (text, image, audio, etc.)
+            print("Full cleanup: Removing all models from VRAM...")
+            
+            # Cleanup all models in multi_model_manager
+            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()
+                        elif hasattr(model_to_cleanup, 'model') and model_to_cleanup.model is not None:
+                            if hasattr(model_to_cleanup.model, 'cleanup'):
+                                model_to_cleanup.model.cleanup()
+                    except Exception as e:
+                        print(f"Warning during cleanup of '{key}': {e}")
+                    del multi_model_manager.models[key]
+            
+            # Also unload the legacy model_manager backend if loaded
+            try:
+                if model_manager.backend is not None:
+                    print("Unloading legacy model_manager from VRAM...")
+                    if hasattr(model_manager.backend, 'unload'):
+                        model_manager.backend.unload()
+                    elif hasattr(model_manager.backend, 'model') and model_manager.backend.model is not None:
+                        if hasattr(model_manager.backend.model, 'unload'):
+                            model_manager.backend.model.unload()
+                    elif hasattr(model_manager.backend, 'cleanup'):
+                        model_manager.backend.cleanup()
+                    model_manager.backend = None
+            except Exception as e:
+                print(f"Warning during legacy model cleanup: {e}")
+            
+            # 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)
+        else:
+            # Non-ondemand mode or no model loaded: just clean up non-image models
+            print("Cleaning up non-image models to free VRAM for image generation...")
             for key in list(multi_model_manager.models.keys()):
                 if key.startswith("image:"):
                     continue
-            # Unload any other model (text, audio, etc.) to free VRAM
                 model_to_cleanup = multi_model_manager.models.get(key)
                 if model_to_cleanup is not None:
                     print(f"Unloading '{key}' from VRAM to make room for image model")
@@ -304,7 +376,6 @@ async def create_image_generation(request: ImageGenerationRequest, http_request:
                     elif hasattr(model_manager.backend, 'model') and model_manager.backend.model is not None:
                         if hasattr(model_manager.backend.model, 'unload'):
                             model_manager.backend.model.unload()
-                # Reset backend to None
                     model_manager.backend = None
             except Exception as e:
                 print(f"Warning during main model cleanup: {e}")
@@ -560,8 +631,62 @@ async def create_image_generation(request: ImageGenerationRequest, http_request:
     
     # If no cached image model found, need to load one - first cleanup any existing models
     if sd_model is None:
+        # Check if we need to do full cleanup (ondemand mode with any model loaded)
+        current_any_model = len(multi_model_manager.models) > 0
+        has_legacy_model = False
+        try:
+            from codai.models.manager import model_manager
+            has_legacy_model = model_manager.backend is not None
+        except:
+            pass
+        
+        needs_full_unload = (mode == "ondemand" and (current_any_model or has_legacy_model))
+        
         # Check if there's a text model loaded and unload it to free VRAM
         # Cleanup ALL models except the one we're about to load
+        if needs_full_unload:
+            print(f"In ondemand mode - fully unloading current model before loading new one (sd.cpp)...")
+            
+            # 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()
+                        elif hasattr(model_to_cleanup, 'model') and model_to_cleanup.model is not None:
+                            if hasattr(model_to_cleanup.model, 'cleanup'):
+                                model_to_cleanup.model.cleanup()
+                    except Exception as e:
+                        print(f"Warning during cleanup of '{key}': {e}")
+                    del multi_model_manager.models[key]
+            
+            # Also cleanup legacy model
+            try:
+                from codai.models.manager import model_manager
+                if model_manager.backend is not None:
+                    print("Unloading legacy model_manager from VRAM...")
+                    if hasattr(model_manager.backend, 'unload'):
+                        model_manager.backend.unload()
+                    elif hasattr(model_manager.backend, 'cleanup'):
+                        model_manager.backend.cleanup()
+                    model_manager.backend = None
+            except:
+                pass
+            
+            # 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()
+            except:
+                pass
+        else:
+            # Non-ondemand or no model loaded: cleanup non-image models only
             for key in list(multi_model_manager.models.keys()):
                 # Skip the image model we'll be loading (if we find it later)
                 # For now, cleanup all other models

codai/api/text.py

@@ -295,6 +295,68 @@ async def chat_completions(request: ChatCompletionRequest, http_request: Request
     # Get the model for this request
     requested_model = request.model
     
+    # Get load mode to determine if we need to unload other models first
+    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))
+    
+    # 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...")
+        
+        # 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]
+        
+        # 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()
+            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)
     
@@ -1688,6 +1750,58 @@ async def completions(request: CompletionRequest):
     # Get the model for this request
     requested_model = request.model
     
+    # Get load mode to determine if we need to unload other models first
+    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]
+        
+        # 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()
+            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()
+        except:
+            pass
+    
     # Try to get the appropriate model
     mm = multi_model_manager.get_model_for_request(requested_model)