KSamplerProgress fix and enhance

fix: ommit_start_latent - doesn't work if it is false improve: add ommit_final_latent modified: move progress latent to cpu
ltdrdata · Aug 4, 2024 · 09ae2ea · 09ae2ea
1 parent eab0b95
commit 09ae2ea
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 11 deletions.
diff --git a/__init__.py b/__init__.py
@@ -7,7 +7,7 @@
 
 import importlib
 
-version_code = [0, 82, 6]
+version_code = [0, 82, 7]
 version_str = f"V{version_code[0]}.{version_code[1]}" + (f'.{version_code[2]}' if len(version_code) > 2 else '')
 print(f"### Loading: ComfyUI-Inspire-Pack ({version_str})")
 

diff --git a/inspire/sampler_nodes.py b/inspire/sampler_nodes.py
@@ -24,6 +24,7 @@ def INPUT_TYPES(s):
                     "noise_mode": (["GPU(=A1111)", "CPU"],),
                     "interval": ("INT", {"default": 1, "min": 1, "max": 10000}),
                     "omit_start_latent": ("BOOLEAN", {"default": True, "label_on": "True", "label_off": "False"}),
+                    "omit_final_latent": ("BOOLEAN", {"default": False, "label_on": "True", "label_off": "False"}),
                     },
                 "optional": {
                     "scheduler_func_opt": ("SCHEDULER_FUNC",),
@@ -36,27 +37,29 @@ def INPUT_TYPES(s):
     RETURN_NAMES = ("latent", "progress_latent")
 
     @staticmethod
-    def doit(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise, noise_mode, interval, omit_start_latent, scheduler_func_opt=None):
+    def doit(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise, noise_mode,
+             interval, omit_start_latent, omit_final_latent, scheduler_func_opt=None):
         adv_steps = int(steps / denoise)
 
         if omit_start_latent:
             result = []
         else:
             result = [latent_image['samples']]
 
-        result = []
-
         def progress_callback(step, x0, x, total_steps):
             if (total_steps-1) != step and step % interval != 0:
                 return
 
             x = model.model.process_latent_out(x)
-            x = x.to(model_management.intermediate_device())
+            x = x.cpu()
             result.append(x)
 
         latent_image, noise = a1111_compat.KSamplerAdvanced_inspire.sample(model, True, seed, adv_steps, cfg, sampler_name, scheduler, positive, negative, latent_image, (adv_steps-steps),
                                                                            adv_steps, noise_mode, False, callback=progress_callback, scheduler_func_opt=scheduler_func_opt)
 
+        if not omit_final_latent:
+            result.append(latent_image['samples'].cpu())
+
         if len(result) > 0:
             result = torch.cat(result)
             result = {'samples': result}
@@ -86,6 +89,7 @@ def INPUT_TYPES(s):
                     "return_with_leftover_noise": ("BOOLEAN", {"default": False, "label_on": "enable", "label_off": "disable"}),
                     "interval": ("INT", {"default": 1, "min": 1, "max": 10000}),
                     "omit_start_latent": ("BOOLEAN", {"default": False, "label_on": "True", "label_off": "False"}),
+                    "omit_final_latent": ("BOOLEAN", {"default": False, "label_on": "True", "label_off": "False"}),
                     },
                 "optional": {
                     "prev_progress_latent_opt": ("LATENT",),
@@ -100,26 +104,29 @@ def INPUT_TYPES(s):
     RETURN_TYPES = ("LATENT", "LATENT")
     RETURN_NAMES = ("latent", "progress_latent")
 
-    def doit(self, model, add_noise, noise_seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, start_at_step, end_at_step,
-             noise_mode, return_with_leftover_noise, interval, omit_start_latent, prev_progress_latent_opt=None, scheduler_func_opt=None):
+    def doit(self, model, add_noise, noise_seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image,
+             start_at_step, end_at_step, noise_mode, return_with_leftover_noise, interval, omit_start_latent, omit_final_latent,
+             prev_progress_latent_opt=None, scheduler_func_opt=None):
+
         if omit_start_latent:
             result = []
         else:
             result = [latent_image['samples']]
 
-        result = []
-
         def progress_callback(step, x0, x, total_steps):
             if (total_steps-1) != step and step % interval != 0:
                 return
 
             x = model.model.process_latent_out(x)
-            x = x.to(model_management.intermediate_device())
+            x = x.cpu()
             result.append(x)
 
         latent_image, noise = a1111_compat.KSamplerAdvanced_inspire.sample(model, add_noise, noise_seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, start_at_step, end_at_step,
                                                                            noise_mode, False, callback=progress_callback, scheduler_func_opt=scheduler_func_opt)
 
+        if not omit_final_latent:
+            result.append(latent_image['samples'].cpu())
+
         if len(result) > 0:
             result = torch.cat(result)
             result = {'samples': result}

diff --git a/pyproject.toml b/pyproject.toml
@@ -1,7 +1,7 @@
 [project]
 name = "comfyui-inspire-pack"
 description = "This extension provides various nodes to support Lora Block Weight and the Impact Pack. Provides many easily applicable regional features and applications for Variation Seed."
-version = "0.82.6"
+version = "0.82.7"
 license = { file = "LICENSE" }
 dependencies = ["matplotlib", "cachetools"]