Merge pull request #12909 from AUTOMATIC1111/Action-to-calculate-all-…

…SD-checkpoint-hashes Action to calculate all SD checkpoint hashes
AUTOMATIC1111 · Sep 9, 2023 · fccde0c · fccde0c
2 parents 3ca4655 + 348c602
commit fccde0c
Showing 1 changed file with 19 additions and 0 deletions.
diff --git a/modules/ui_settings.py b/modules/ui_settings.py
@@ -5,6 +5,7 @@
 from modules.shared import opts
 from modules.ui_components import FormRow
 from modules.ui_gradio_extensions import reload_javascript
+from concurrent.futures import ThreadPoolExecutor, as_completed
 
 
 def get_value_for_setting(key):
@@ -175,6 +176,9 @@ def create_ui(self, loadsave, dummy_component):
                     with gr.Row():
                         unload_sd_model = gr.Button(value='Unload SD checkpoint to free VRAM', elem_id="sett_unload_sd_model")
                         reload_sd_model = gr.Button(value='Reload the last SD checkpoint back into VRAM', elem_id="sett_reload_sd_model")
+                    with gr.Row():
+                        calculate_all_checkpoint_hash = gr.Button(value='Calculate hash for all checkpoint', elem_id="calculate_all_checkpoint_hash")
+                        calculate_all_checkpoint_hash_threads = gr.Number(value=1, label="Number of parallel calculations", elem_id="calculate_all_checkpoint_hash_threads", precision=0, minimum=1)
 
                 with gr.TabItem("Licenses", id="licenses", elem_id="settings_tab_licenses"):
                     gr.HTML(shared.html("licenses.html"), elem_id="licenses")
@@ -241,6 +245,21 @@ def check_file(x):
                 outputs=[sysinfo_check_output],
             )
 
+            def calculate_all_checkpoint_hash_fn(max_thread):
+                checkpoints_list = sd_models.checkpoints_list.values()
+                with ThreadPoolExecutor(max_workers=max_thread) as executor:
+                    futures = [executor.submit(checkpoint.calculate_shorthash) for checkpoint in checkpoints_list]
+                    completed = 0
+                    for _ in as_completed(futures):
+                        completed += 1
+                        print(f"{completed} / {len(checkpoints_list)} ")
+                    print("Finish calculating hash for all checkpoints")
+
+            calculate_all_checkpoint_hash.click(
+                fn=calculate_all_checkpoint_hash_fn,
+                inputs=[calculate_all_checkpoint_hash_threads],
+            )
+
         self.interface = settings_interface
 
     def add_quicksettings(self):