Removed IF safety checker and watermarker, improved UI

MMqd · Jul 23, 2023 · 15dd794 · 15dd794
1 parent 177eb41
commit 15dd794
Show file tree

Hide file tree

Showing 4 changed files with 43 additions and 28 deletions.
diff --git a/scripts/abstract_model.py b/scripts/abstract_model.py
@@ -111,15 +111,15 @@ class AbstractModel():
     def __init__(self):
         self.cache_dir = os.path.join(os.path.join(script_path, 'models'), '')
 
-    def load_pipeline(self, pipe_name: str, pipeline: DiffusionPipeline, pretrained_model_name_or_path, move_to_cuda = True):
+    def load_pipeline(self, pipe_name: str, pipeline: DiffusionPipeline, pretrained_model_name_or_path, move_to_cuda = True, kwargs={}):
         pipe = getattr(self, pipe_name, None)
 
         if not isinstance(pipe, pipeline) or pipe is None:
             if pipe is not None:
                 pipe = None
                 gc.collect()
                 devices.torch_gc()
-            pipe = pipeline.from_pretrained(pretrained_model_name_or_path, variant="fp16", torch_dtype=torch.float16, cache_dir=self.cache_dir)#, scheduler=dpm)
+            pipe = pipeline.from_pretrained(pretrained_model_name_or_path, variant="fp16", torch_dtype=torch.float16, cache_dir=self.cache_dir, **kwargs)#, scheduler=dpm)
             if move_to_cuda:
                 pipe.to("cuda")
             else:

diff --git a/scripts/deepfloydif.py b/scripts/deepfloydif.py
@@ -28,11 +28,13 @@ class IFModel(AbstractModel):
     pipe = None
 
     def __init__(self):
+        self.stageI_model = "XL"
+        self.stageII_model = "L"
         self.cache_dir = os.path.join(os.path.join(script_path, 'models'), 'IF')
 
     def load_encoder(self):
         try:
-            self.pipe = self.load_pipeline("pipe", IFPipeline, "DeepFloyd/IF-I-XL-v1.0")
+            self.pipe = self.load_pipeline("pipe", IFPipeline, f"DeepFloyd/IF-I-{self.stageI_model}-v1.0", {"safety_checker": None, "watermarker": None})
         except FileNotFoundError as fe:
             errors.print_error_explanation(f'File {fe.filename} not found. Did you forget the Hugging Face token?')
 

diff --git a/scripts/deepfloydif_script.py b/scripts/deepfloydif_script.py
@@ -56,26 +56,31 @@ def title(self):
 
     def ui(self, is_img2img):
         gr.Markdown("To save VRAM unload the Stable Diffusion Model")
+        unload_sd_model = None
+        reload_sd_model = None
+        unload_k_model = None
 
-        unload_sd_model = gr.Button("Unload Stable Diffusion Model")
-        unload_sd_model.click(unload_model)
-        reload_sd_model = gr.Button("Reload Stable Diffusion Model")
-        reload_sd_model.click(reload_model)
-
-        unload_k_model = gr.Button("Unload IF Model")
-        unload_k_model.click(unload_if_model)
         with gr.Row():
-            unload_sd_model
-            reload_sd_model
-            unload_k_model
+            unload_sd_model = gr.Button("Unload Stable Diffusion Model")
+            unload_sd_model.click(unload_model)
+            reload_sd_model = gr.Button("Reload Stable Diffusion Model")
+            reload_sd_model.click(reload_model)
+            unload_k_model = gr.Button("Unload IF Model")
+            unload_k_model.click(unload_if_model)
+
+        stageI_model = None
+        stageII_model = None
+        with gr.Row():
+            stageI_model = gr.inputs.Dropdown(label="Stage I Model Type", choices=["M", "L", "XL"], default="XL")
+            stageII_model = gr.inputs.Dropdown(label="Stage II Model Type", choices=["None", "M", "L"], default="L")
 
         token_textbox = gr.inputs.Textbox(label="Hugging Face Token", type="password")
 
-        inputs = [token_textbox]
+        inputs = [token_textbox, stageI_model, stageII_model]
 
         return inputs
 
-    def run(self, p, token) -> Processed:
+    def run(self, p, token, stageI_model, stageII_model) -> Processed:
         p.sampler_name = "DDPM"
         p.init_image = getattr(p, 'init_images', None)
         p.extra_generation_params["Script"] = self.title()
@@ -84,6 +89,8 @@ def run(self, p, token) -> Processed:
 
         if shared.if_model is None:
             shared.if_model = IFModel()
+            shared.if_model.stageI_model = stageI_model
+            shared.if_model.stageII_model = stageII_model
             if token != "":
                 login(token=token)
 

diff --git a/scripts/kandinsky_script.py b/scripts/kandinsky_script.py
@@ -56,23 +56,29 @@ def title(self):
     def ui(self, is_img2img):
         gr.Markdown("To save VRAM unload the Stable Diffusion Model")
 
-        unload_sd_model = gr.Button("Unload Stable Diffusion Model")
-        unload_sd_model.click(unload_model)
-        reload_sd_model = gr.Button("Reload Stable Diffusion Model")
-        reload_sd_model.click(reload_model)
+        unload_sd_model = None
+        reload_sd_model = None
+        unload_k_model = None
+
+        prior_inference_steps = None
+        prior_cfg_scale = None
+
+        with gr.Row():
+            unload_sd_model = gr.Button("Unload Stable Diffusion Model")
+            unload_sd_model.click(unload_model)
+            reload_sd_model = gr.Button("Reload Stable Diffusion Model")
+            reload_sd_model.click(reload_model)
+            unload_k_model = gr.Button("Unload Kandinsky Model")
+            unload_k_model.click(unload_kandinsky_model)
 
-        unload_k_model = gr.Button("Unload Kandinsky Model")
-        unload_k_model.click(unload_kandinsky_model)
         with gr.Row():
-            unload_sd_model
-            reload_sd_model
-            unload_k_model
-        prior_inference_steps = gr.inputs.Slider(minimum=2, maximum=1024, step=1, label="Prior Inference Steps", default=128)
-        prior_cfg_scale = gr.inputs.Slider(minimum=1, maximum=20, step=0.5, label="Prior CFG Scale", default=4)
+            prior_inference_steps = gr.inputs.Slider(minimum=2, maximum=1024, step=1, label="Prior Inference Steps", default=128)
+            prior_cfg_scale = gr.inputs.Slider(minimum=1, maximum=20, step=0.5, label="Prior CFG Scale", default=4)
 
         with gr.Accordion("Image Mixing", open=False):
-            img1_strength = gr.inputs.Slider(minimum=-2, maximum=2, label="Interpolate Image 1 Strength", default=0.5)
-            img2_strength = gr.inputs.Slider(minimum=-2, maximum=2, label="Interpolate Image 2 Strength (image below)", default=0.5)
+            with gr.Row():
+                img1_strength = gr.inputs.Slider(minimum=-2, maximum=2, label="Interpolate Image 1 Strength", default=0.5)
+                img2_strength = gr.inputs.Slider(minimum=-2, maximum=2, label="Interpolate Image 2 Strength (image below)", default=0.5)
             extra_image = gr.inputs.Image()
 
         inputs = [extra_image, prior_inference_steps, prior_cfg_scale, img1_strength, img2_strength]