Add noise augmentation to hunyuan image refiner. (#9831)

This was missing and should help with colors being blown out.

Author: comfyanonymous

comfy/model_base.py

@@ -1437,6 +1437,7 @@ class HunyuanImage21Refiner(HunyuanImage21):
     def concat_cond(self, **kwargs):
         noise = kwargs.get("noise", None)
         image = kwargs.get("concat_latent_image", None)
+        noise_augmentation = kwargs.get("noise_augmentation", 0.0)
         device = kwargs["device"]
 
         if image is None:
@@ -1446,6 +1447,9 @@ def concat_cond(self, **kwargs):
             image = utils.common_upscale(image.to(device), noise.shape[-1], noise.shape[-2], "bilinear", "center")
             image = self.process_latent_in(image)
             image = utils.resize_to_batch_size(image, noise.shape[0])
+            if noise_augmentation > 0:
+                noise = torch.randn(image.shape, generator=torch.manual_seed(kwargs.get("seed", 0) - 10), dtype=image.dtype, device="cpu").to(image.device)
+                image = noise_augmentation * noise + (1.0 - noise_augmentation) * image
         return image
 
     def extra_conds(self, **kwargs):

comfy_extras/nodes_hunyuan.py

@@ -134,18 +134,18 @@ def INPUT_TYPES(s):
         return {"required": {"positive": ("CONDITIONING", ),
                              "negative": ("CONDITIONING", ),
                              "latent": ("LATENT", ),
+                             "noise_augmentation": ("FLOAT", {"default": 0.10, "min": 0.0, "max": 1.0, "step": 0.01}),
                              }}
 
     RETURN_TYPES = ("CONDITIONING", "CONDITIONING", "LATENT")
     RETURN_NAMES = ("positive", "negative", "latent")
 
     FUNCTION = "execute"
 
-    def execute(self, positive, negative, latent):
+    def execute(self, positive, negative, latent, noise_augmentation):
         latent = latent["samples"]
-
-        positive = node_helpers.conditioning_set_values(positive, {"concat_latent_image": latent})
-        negative = node_helpers.conditioning_set_values(negative, {"concat_latent_image": latent})
+        positive = node_helpers.conditioning_set_values(positive, {"concat_latent_image": latent, "noise_augmentation": noise_augmentation})
+        negative = node_helpers.conditioning_set_values(negative, {"concat_latent_image": latent, "noise_augmentation": noise_augmentation})
         out_latent = {}
         out_latent["samples"] = torch.zeros([latent.shape[0], 32, latent.shape[-3], latent.shape[-2], latent.shape[-1]], device=comfy.model_management.intermediate_device())
         return (positive, negative, out_latent)