xinyun/ComfyUI
1
0
Fork 0
mirror of https://git.datalinker.icu/comfyanonymous/ComfyUI synced 2026-08-12 22:03:33 +08:00
Code Issues Packages Projects Releases Wiki Activity

remove step tracking code

Browse Source
This commit is contained in:
KimbingNg 2025-09-25 17:55:15 +08:00
parent f23bac8108
commit fd20999994
1 changed files with 10 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

25
comfy_extras/nodes_hunyuan.py
View File

@ -1,3 +1,4 @@
import math
import nodes import nodes
import node_helpers import node_helpers
import torch import torch
@ -91,9 +92,9 @@ class AdaptiveProjectedGuidance:
self.use_original_formulation = use_original_formulation self.use_original_formulation = use_original_formulation
self.momentum_buffer = None self.momentum_buffer = None
def __call__(self, pred_cond: torch.Tensor, pred_uncond=None, step=None) -> torch.Tensor: def __call__(self, pred_cond: torch.Tensor, pred_uncond=None, is_first_step=False) -> torch.Tensor:
if step == 0 and self.adaptive_projected_guidance_momentum is not None: if is_first_step and self.adaptive_projected_guidance_momentum is not None:
self.momentum_buffer = MomentumBuffer(self.adaptive_projected_guidance_momentum) self.momentum_buffer = MomentumBuffer(self.adaptive_projected_guidance_momentum)
pred = normalized_guidance_apg( pred = normalized_guidance_apg(
@ -160,42 +161,36 @@ class HunyuanMixModeAPG:
general_start_t = model_sampling.percent_to_sigma(general_start_percent) general_start_t = model_sampling.percent_to_sigma(general_start_percent)
ocr_start_t = model_sampling.percent_to_sigma(ocr_start_percent) ocr_start_t = model_sampling.percent_to_sigma(ocr_start_percent)
step_tracker = {"step": 0}
def hunyuan_apg_outer_sample_wrapper(executor, *args, **kwargs):
step_tracker['step'] = 0
return executor(*args, **kwargs)
def cfg_function(args): def cfg_function(args):
sigma = args["sigma"].to(torch.float32) sigma = args["sigma"].to(torch.float32)
sigma = sigma[:, None, None, None] is_first_step = math.isclose(sigma.item(), args['model_options']['transformer_options']['sample_sigmas'][0].item())
cond = args["cond"] cond = args["cond"]
uncond = args["uncond"] uncond = args["uncond"]
cond_scale = args["cond_scale"] cond_scale = args["cond_scale"]
step = step_tracker['step'] sigma = sigma[:, None, None, None]
step_tracker['step'] += 1
if not has_quoted_text: if not has_quoted_text:
if sigma[0] <= general_start_t: if sigma[0] <= general_start_t:
modified_cond = general_apg(cond / sigma, uncond / sigma, step) modified_cond = general_apg(cond / sigma, uncond / sigma, is_first_step=is_first_step)
return modified_cond * sigma return modified_cond * sigma
else: else:
if cond_scale > 1: if cond_scale > 1:
_ = general_apg(cond / sigma, uncond / sigma, step) # track momentum _ = general_apg(cond / sigma, uncond / sigma, is_first_step=is_first_step) # track momentum
return uncond + (cond - uncond) * cond_scale return uncond + (cond - uncond) * cond_scale
else: else:
if sigma[0] <= ocr_start_t: if sigma[0] <= ocr_start_t:
modified_cond = ocr_apg(cond / sigma, uncond / sigma, step) modified_cond = ocr_apg(cond / sigma, uncond / sigma, is_first_step=is_first_step)
return modified_cond * sigma return modified_cond * sigma
else: else:
if cond_scale > 1: if cond_scale > 1:
_ = ocr_apg(cond / sigma, uncond / sigma, step) # track momentum _ = ocr_apg(cond / sigma, uncond / sigma, is_first_step=is_first_step) # track momentum
return uncond + (cond - uncond) * cond_scale return uncond + (cond - uncond) * cond_scale
return cond return cond
m.add_wrapper_with_key(comfy.patcher_extension.WrappersMP.OUTER_SAMPLE, "hunyuan_apg", hunyuan_apg_outer_sample_wrapper)
m.set_model_sampler_cfg_function(cfg_function, disable_cfg1_optimization=True) m.set_model_sampler_cfg_function(cfg_function, disable_cfg1_optimization=True)
return (m,) return (m,)