Update telebot/telebotTest/sd_models.py

telebot/telebotTest/sd_models.py

@@ -22,8 +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
-active_model_reloads = 0
-model_reload_lock = threading.Lock()
 
 
 model_dir = "Stable-diffusion"
@@ -451,8 +449,26 @@ 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()
-    return
+    # Add synchronization point before unloading
+    import time
+    import os
+    lock_file = "/kaggle/working/unload.lock"
+    
+    # Wait for previous unloading to complete
+    while os.path.exists(lock_file):
+        time.sleep(0.1)
+
+    # Create lock file
+    with open(lock_file, 'w') as f:
+        f.write("locked")
+
+    try:
+        memory_management.unload_all_models()
+        return
+    finally:
+        # Release lock
+        if os.path.exists(lock_file):
+            os.remove(lock_file)
 
 
 def apply_token_merging(sd_model, token_merging_ratio):
@@ -473,73 +489,56 @@ def apply_token_merging(sd_model, token_merging_ratio):
 
 @torch.inference_mode()
 def forge_model_reload():
-    global active_model_reloads
-    
-    # Increment counter of active reloads
-    with model_reload_lock:
-        active_model_reloads += 1
-    
-    try:
-        current_hash = str(model_data.forge_loading_parameters)
-
-        if model_data.forge_hash == current_hash:
-            return model_data.sd_model, False
+    current_hash = str(model_data.forge_loading_parameters)
 
-        # Only proceed with unloading/reloading if this is the last active reload
-        with model_reload_lock:
-            if active_model_reloads > 1:
-                return model_data.sd_model, False
+    if model_data.forge_hash == current_hash:
+        return model_data.sd_model, False
 
-        print('Loading Model: ' + str(model_data.forge_loading_parameters))
+    print('Loading Model: ' + str(model_data.forge_loading_parameters))
 
-        timer = Timer()
+    timer = Timer()
 
-        if model_data.sd_model:
-            model_data.sd_model = None
-            memory_management.unload_all_models()
-            memory_management.soft_empty_cache()
-            gc.collect()
+    if model_data.sd_model:
+        model_data.sd_model = None
+        memory_management.unload_all_models()
+        memory_management.soft_empty_cache()
+        gc.collect()
 
-        timer.record("unload existing model")
+    timer.record("unload existing model")
 
-        checkpoint_info = model_data.forge_loading_parameters['checkpoint_info']
+    checkpoint_info = model_data.forge_loading_parameters['checkpoint_info']
 
-        if checkpoint_info is None:
-            raise ValueError('You do not have any model! Please download at least one model in [models/Stable-diffusion].')
+    if checkpoint_info is None:
+        raise ValueError('You do not have any model! Please download at least one model in [models/Stable-diffusion].')
 
-        state_dict = checkpoint_info.filename
-        additional_state_dicts = model_data.forge_loading_parameters.get('additional_modules', [])
+    state_dict = checkpoint_info.filename
+    additional_state_dicts = model_data.forge_loading_parameters.get('additional_modules', [])
 
-        timer.record("cache state dict")
+    timer.record("cache state dict")
 
-        dynamic_args['forge_unet_storage_dtype'] = model_data.forge_loading_parameters.get('unet_storage_dtype', None)
-        dynamic_args['embedding_dir'] = cmd_opts.embeddings_dir
-        dynamic_args['emphasis_name'] = opts.emphasis
-        sd_model = forge_loader(state_dict, additional_state_dicts=additional_state_dicts)
-        timer.record("forge model load")
+    dynamic_args['forge_unet_storage_dtype'] = model_data.forge_loading_parameters.get('unet_storage_dtype', None)
+    dynamic_args['embedding_dir'] = cmd_opts.embeddings_dir
+    dynamic_args['emphasis_name'] = opts.emphasis
+    sd_model = forge_loader(state_dict, additional_state_dicts=additional_state_dicts)
+    timer.record("forge model load")
 
-        sd_model.extra_generation_params = {}
-        sd_model.comments = []
-        sd_model.sd_checkpoint_info = checkpoint_info
-        sd_model.filename = checkpoint_info.filename
-        sd_model.sd_model_hash = checkpoint_info.calculate_shorthash()
-        timer.record("calculate hash")
-
-        shared.opts.data["sd_checkpoint_hash"] = checkpoint_info.sha256
+    sd_model.extra_generation_params = {}
+    sd_model.comments = []
+    sd_model.sd_checkpoint_info = checkpoint_info
+    sd_model.filename = checkpoint_info.filename
+    sd_model.sd_model_hash = checkpoint_info.calculate_shorthash()
+    timer.record("calculate hash")
 
-        model_data.set_sd_model(sd_model)
+    shared.opts.data["sd_checkpoint_hash"] = checkpoint_info.sha256
 
-        script_callbacks.model_loaded_callback(sd_model)
+    model_data.set_sd_model(sd_model)
 
-        timer.record("scripts callbacks")
+    script_callbacks.model_loaded_callback(sd_model)
 
-        print(f"Model loaded in {timer.summary()}.")
+    timer.record("scripts callbacks")
 
-        model_data.forge_hash = current_hash
+    print(f"Model loaded in {timer.summary()}.")
 
-        return sd_model, True
+    model_data.forge_hash = current_hash
 
-    finally:
-        # Decrement counter when done
-        with model_reload_lock:
-            active_model_reloads -= 1
+    return sd_model, True