From fbbb56f6ca1db73fd48eb76c020c15bf9e32b00c Mon Sep 17 00:00:00 2001
From: Vladimir Mandic <mandic00@live.com>
Date: Sat, 1 Jul 2023 16:12:38 -0400
Subject: [PATCH] dissalow ckpt option

---
 CHANGELOG.md         | 2 ++
 modules/sd_models.py | 6 +++++-
 modules/shared.py    | 1 +
 3 files changed, 8 insertions(+), 1 deletion(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 7b1a68eb7..e60a29277 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -4,8 +4,10 @@
 
 Small quality-of-life updates and bugfixes:
 
+- add option to disallow usage of ckpt checkpoints
 - change lora and lyco dir without server restart
 - additional filename template fields: `uuid`, `seq`, `image_hash`  
+- image toolbar is now shown only when image is present
 - image `Zip` button gone and its not optional setting that applies to standard `Save` button
 - folder `Show` button is present only when working on localhost,  
   otherwise its replaced with `Copy` that places image URLs on clipboard so they can be used in other apps
diff --git a/modules/sd_models.py b/modules/sd_models.py
index b9977ccb2..30d5dfdb8 100644
--- a/modules/sd_models.py
+++ b/modules/sd_models.py
@@ -111,7 +111,8 @@ def list_models():
     checkpoints_list.clear()
     checkpoint_aliases.clear()
     if shared.backend == shared.Backend.ORIGINAL:
-        model_list = modelloader.load_models(model_path=os.path.join(models_path, 'Stable-diffusion'), model_url=None, command_path=shared.opts.ckpt_dir, ext_filter=[".ckpt", ".safetensors"], download_name=None, ext_blacklist=[".vae.ckpt", ".vae.safetensors"])
+        ext_filter=[".safetensors"] if shared.opts.sd_disable_ckpt else [".ckpt", ".safetensors"]
+        model_list = modelloader.load_models(model_path=os.path.join(models_path, 'Stable-diffusion'), model_url=None, command_path=shared.opts.ckpt_dir, ext_filter=ext_filter, download_name=None, ext_blacklist=[".vae.ckpt", ".vae.safetensors"])
     else:
         global model_path # pylint: disable=global-statement
         model_path = os.path.join(models_path, 'Diffusers')
@@ -306,6 +307,9 @@ def read_state_dict(checkpoint_file, map_location=None): # pylint: disable=unuse
         pl_sd = None
         with progress.open(checkpoint_file, 'rb', description=f'Loading weights: [cyan]{checkpoint_file}', auto_refresh=True) as f:
             _, extension = os.path.splitext(checkpoint_file)
+            if extension.lower() == ".ckpt" and shared.opts.sd_disable_ckpt:
+                shared.log.warning(f"Checkpoint loading disabled: {checkpoint_file}")
+                return None
             if shared.opts.stream_load:
                 if extension.lower() == ".safetensors":
                     # shared.log.debug('Model weights loading: type=safetensors mode=buffered')
diff --git a/modules/shared.py b/modules/shared.py
index 8d68e7fe7..36854d33a 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -298,6 +298,7 @@ options_templates.update(options_section(('sd', "Stable Diffusion"), {
     "prompt_attention": OptionInfo("Full parser", "Prompt attention parser", gr.Radio, lambda: {"choices": ["Full parser", "Compel parser", "A1111 parser", "Fixed attention"] }),
     "prompt_mean_norm": OptionInfo(True, "Prompt attention mean normalization"),
     "comma_padding_backtrack": OptionInfo(20, "Prompt padding for long prompts", gr.Slider, {"minimum": 0, "maximum": 74, "step": 1 }),
+    "sd_disable_ckpt": OptionInfo(False, "Disallow usage of checkpoints in ckpt format"),
     "sd_backend": OptionInfo("Original", "Stable Diffusion backend (experimental)", gr.Radio, lambda: {"choices": ["Original", "Diffusers"] }),
 }))