Update telebot/telebotTest/sd_models.py

telebot/telebotTest/sd_models.py

@@ -22,7 +22,6 @@ from backend.loader import forge_loader
 from backend import memory_management
 from backend.args import dynamic_args
 from backend.utils import load_torch_file
-from modules import shared
 
 
 model_dir = "Stable-diffusion"
@@ -450,7 +449,8 @@ def reload_model_weights(sd_model=None, info=None, forced_reload=False):
 
 
 def unload_model_weights(sd_model=None, info=None):
-    memory_management.unload_all_models(api_mode=shared.cmd_opts.api)
+    device_id = sd_model.device.index if sd_model else None
+    memory_management.unload_all_models(device_id)
     return