Merge branch 'master' of github.com:comfyanonymous/ComfyUI

comfy/cmd/latent_preview.py

@@ -14,6 +14,15 @@ import logging
 
 MAX_PREVIEW_RESOLUTION = 512
 
+
+def preview_to_image(latent_image):
+    latents_ubyte = (((latent_image + 1.0) / 2.0).clamp(0, 1)  # change scale from -1..1 to 0..1
+                        .mul(0xFF)  # to 0..255
+                        ).to(device="cpu", dtype=torch.uint8, non_blocking=model_management.device_supports_non_blocking(latent_image.device))
+
+    return Image.fromarray(latents_ubyte.numpy())
+
+
 class LatentPreviewer:
     def decode_latent_to_preview(self, x0):
         pass
@@ -22,17 +31,14 @@ class LatentPreviewer:
         preview_image = self.decode_latent_to_preview(x0)
         return ("JPEG", preview_image, MAX_PREVIEW_RESOLUTION)
 
+
 class TAESDPreviewerImpl(LatentPreviewer):
     def __init__(self, taesd):
         self.taesd = taesd
 
     def decode_latent_to_preview(self, x0):
-        x_sample = self.taesd.decode(x0[:1])[0].detach()
+        x_sample = self.taesd.decode(x0[:1])[0].movedim(0, 2)
-        x_sample = 255. * torch.clamp((x_sample + 1.0) / 2.0, min=0.0, max=1.0)
+        return preview_to_image(x_sample)
-        x_sample = np.moveaxis(x_sample.to(device="cpu", dtype=torch.uint8, non_blocking=model_management.device_supports_non_blocking(x_sample.device)).numpy(), 0, 2)
-
-        preview_image = Image.fromarray(x_sample)
-        return preview_image
 
 
 class Latent2RGBPreviewer(LatentPreviewer):
@@ -42,13 +48,7 @@ class Latent2RGBPreviewer(LatentPreviewer):
     def decode_latent_to_preview(self, x0):
         self.latent_rgb_factors = self.latent_rgb_factors.to(dtype=x0.dtype, device=x0.device)
         latent_image = x0[0].permute(1, 2, 0) @ self.latent_rgb_factors
-
+        return preview_to_image(latent_image)
-        latents_ubyte = (((latent_image + 1) / 2)
-                            .clamp(0, 1)  # change scale from -1..1 to 0..1
-                            .mul(0xFF)  # to 0..255
-                            ).to(device="cpu", dtype=torch.uint8, non_blocking=model_management.device_supports_non_blocking(latent_image.device))
-
-        return Image.fromarray(latents_ubyte.numpy())
 
 
 def get_previewer(device, latent_format):
@@ -88,6 +88,7 @@ def prepare_callback(model, steps, x0_output_dict=None):
     previewer = get_previewer(model.load_device, model.model.latent_format)
 
     pbar = utils.ProgressBar(steps)
+
     def callback(step, x0, x, total_steps):
         if x0_output_dict is not None:
             x0_output_dict["x0"] = x0

comfy/model_management.py

@@ -964,6 +964,7 @@ def unload_all_models():
 
 
 def resolve_lowvram_weight(weight, model, key):  # TODO: remove
+    print("WARNING: The comfy.model_management.resolve_lowvram_weight function will be removed soon, please stop using it.")
     return weight
 
 

comfy/model_patcher.py

@@ -88,9 +88,7 @@ class ModelPatcher(ModelManageable):
     def model_size(self):
         if self.size > 0:
             return self.size
-        model_sd = self.model.state_dict()
         self.size = model_management.module_size(self.model)
-        self.model_keys = set(model_sd.keys())
         return self.size
 
     def clone(self):
@@ -102,7 +100,6 @@ class ModelPatcher(ModelManageable):
 
         n.object_patches = self.object_patches.copy()
         n.model_options = copy.deepcopy(self.model_options)
-        n.model_keys = self.model_keys
         n.backup = self.backup
         n.object_patches_backup = self.object_patches_backup
         return n
@@ -222,8 +219,9 @@ class ModelPatcher(ModelManageable):
 
     def add_patches(self, patches, strength_patch=1.0, strength_model=1.0):
         p = set()
+        model_sd = self.model.state_dict()
         for k in patches:
-            if k in self.model_keys:
+            if k in model_sd:
                 p.add(k)
                 current_patches = self.patches.get(k, [])
                 current_patches.append((strength_patch, patches[k], strength_model))

comfy_extras/nodes/nodes_perpneg.py

@@ -59,12 +59,39 @@ class Guider_PerpNeg(samplers.CFGGuider):
         self.neg_scale = neg_scale
 
     def predict_noise(self, x, timestep, model_options={}, seed=None):
+        # in CFGGuider.predict_noise, we call sampling_function(), which uses cfg_function() to compute pos & neg
+        # but we'd rather do a single batch of sampling pos, neg, and empty, so we call calc_cond_batch([pos,neg,empty]) directly
+
         positive_cond = self.conds.get("positive", None)
         negative_cond = self.conds.get("negative", None)
         empty_cond = self.conds.get("empty_negative_prompt", None)
 
-        out = samplers.calc_cond_batch(self.inner_model, [negative_cond, positive_cond, empty_cond], x, timestep, model_options)
+        (noise_pred_pos, noise_pred_neg, noise_pred_empty) = \
-        return perp_neg(x, out[1], out[0], out[2], self.neg_scale, self.cfg)
+            samplers.calc_cond_batch(self.inner_model, [positive_cond, negative_cond, empty_cond], x, timestep, model_options)
+        cfg_result = perp_neg(x, noise_pred_pos, noise_pred_neg, noise_pred_empty, self.neg_scale, self.cfg)
+
+        # normally this would be done in cfg_function, but we skipped
+        # that for efficiency: we can compute the noise predictions in
+        # a single call to calc_cond_batch() (rather than two)
+        # so we replicate the hook here
+        for fn in model_options.get("sampler_post_cfg_function", []):
+            args = {
+                "denoised": cfg_result,
+                "cond": positive_cond,
+                "uncond": negative_cond,
+                "model": self.inner_model,
+                "uncond_denoised": noise_pred_neg,
+                "cond_denoised": noise_pred_pos,
+                "sigma": timestep,
+                "model_options": model_options,
+                "input": x,
+                # not in the original call in samplers.py:cfg_function, but made available for future hooks
+                "empty_cond": empty_cond,
+                "empty_cond_denoised": noise_pred_empty,}
+            # todo: is this supposed to be ** spread?
+            cfg_result = fn(args)
+
+        return cfg_result
 
 class PerpNegGuider:
     @classmethod