diff --git a/comfy/model_base.py b/comfy/model_base.py index 7cbad81e7..4edd63667 100644 --- a/comfy/model_base.py +++ b/comfy/model_base.py @@ -2267,43 +2267,12 @@ class QwenImage(BaseModel): class JoyImage(BaseModel): # The noise latent and every reference latent are concatenated as a token sequence inside the - # transformer. A single-reference edit is just the len(ref_latents) == 1 case. The built-in CFG - # guidance rescale is installed from here. + # transformer. A single-reference edit is just the len(ref_latents) == 1 case. The required CFG + # guidance rescale is applied by the JoyImageGuidanceRescale node (comfy_extras/nodes_joyimage.py). def __init__(self, model_config, model_type=ModelType.FLOW, device=None): super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.joyimage.model.JoyImageTransformer3DModel) self.memory_usage_factor_conds = ("ref_latents",) - @staticmethod - def _guidance_rescale_cfg(args): - # CFG combine + per-row L2 rescale in eps-space (guidance rescale). - cond = args["cond"] - uncond = args["uncond"] - cond_scale = args["cond_scale"] - comb = uncond + cond_scale * (cond - uncond) - cond_norm = torch.norm(cond, dim=1, keepdim=True) - comb_norm = torch.norm(comb, dim=1, keepdim=True) - return comb * (cond_norm / comb_norm.clamp_min(1e-6)) - - def _ensure_guidance_rescale_installed(self): - # Self-install the hard-wired guidance rescale once the patcher binds (sd.py doesn't expose a hook - # for this; doing it here keeps the edit confined to model_base.py). Idempotent; refuses to install - # if a different sampler_cfg_function is already present (e.g. a CFGNorm node) so the user's - # override does not silently shadow JoyImage's required rescale. - patcher = self.current_patcher - if patcher is None: - return - existing = patcher.model_options.get("sampler_cfg_function", None) - if existing is JoyImage._guidance_rescale_cfg: - return - if existing is not None: - raise RuntimeError( - "JoyImage requires its built-in CFG guidance-rescale function " - "(comb * cond_norm / comb_norm); an external sampler_cfg_function " - "(e.g. CFGNorm) is already installed and would override it. " - "Remove the external function before sampling JoyImage." - ) - patcher.set_model_sampler_cfg_function(JoyImage._guidance_rescale_cfg) - def extra_conds(self, **kwargs): out = super().extra_conds(**kwargs) cross_attn = kwargs.get("cross_attn", None) @@ -2336,7 +2305,6 @@ class JoyImage(BaseModel): if c_concat is not None: raise ValueError("JoyImage does not support c_concat / noise_concat conditioning") transformer_options = transformer_options.copy() - self._ensure_guidance_rescale_installed() sigma = t xc = self.model_sampling.calculate_input(sigma, x) context = c_crossattn diff --git a/comfy_extras/nodes_joyimage.py b/comfy_extras/nodes_joyimage.py index 72c7f3b7f..00f7b072d 100644 --- a/comfy_extras/nodes_joyimage.py +++ b/comfy_extras/nodes_joyimage.py @@ -1,5 +1,6 @@ import node_helpers import comfy.utils +import torch from typing_extensions import override from comfy_api.latest import ComfyExtension, io @@ -144,12 +145,50 @@ class TextEncodeJoyImageEditPlus(io.ComfyNode): return io.NodeOutput(conditioning, resized_images[-1]) +class JoyImageGuidanceRescale(io.ComfyNode): + """CFG combine + per-token L2 norm rescale required by JoyImageEdit. + + Wire this onto the model before sampling: JoyImageEdit's diffusers pipeline + rescales the combined noise prediction back to the conditional branch's norm + (comb * ||cond|| / ||comb||), the same rescale CFGNorm's pre_cfg branch does. + """ + + @classmethod + def define_schema(cls): + return io.Schema( + node_id="JoyImageGuidanceRescale", + category="model/patch", + inputs=[ + io.Model.Input("model"), + ], + outputs=[ + io.Model.Output(), + ], + ) + + @classmethod + def execute(cls, model) -> io.NodeOutput: + def guidance_rescale(args): + cond = args["cond"] + uncond = args["uncond"] + cond_scale = args["cond_scale"] + comb = uncond + cond_scale * (cond - uncond) + cond_norm = torch.norm(cond, dim=1, keepdim=True) + comb_norm = torch.norm(comb, dim=1, keepdim=True) + return comb * (cond_norm / comb_norm.clamp_min(1e-6)) + + m = model.clone() + m.set_model_sampler_cfg_function(guidance_rescale) + return io.NodeOutput(m) + + class JoyImageExtension(ComfyExtension): @override async def get_node_list(self) -> list[type[io.ComfyNode]]: return [ TextEncodeJoyImageEdit, TextEncodeJoyImageEditPlus, + JoyImageGuidanceRescale, ]