diff --git a/nodes.py b/nodes.py index 32694af..95f7981 100644 --- a/nodes.py +++ b/nodes.py @@ -6,7 +6,7 @@ class uncondZeroNode: return {"required": { "model": ("MODEL",), "scale": ("FLOAT", {"default": 1, "min": 0.0, "max": 10.0, "step": 0.01, "round": 0.01}), - "method":(["uncond_zero","rescale_cfg"],), + "method":(["uncond_zero_v1","uncond_zero_v2","uncond_zero_v3"], {"default": "uncond_zero_v3"},), }} RETURN_TYPES = ("MODEL",) FUNCTION = "patch" @@ -19,10 +19,27 @@ class uncondZeroNode: x_orig = args["input"] x_orig -= x_orig.mean() cond -= cond.mean() - return x_orig - (cond / cond.std() ** .5) * scale - + return x_orig - cond / cond.std() ** .5 * scale + + def uncond_zero_v2(args): + cond = args["cond_denoised"] + x_orig = args["input"] + cond -= cond.mean() + result = torch.zeros_like(x_orig) + + for b in range(len(x_orig)): + for c in range(len(x_orig[b])): + x_orig[b][c] -= x_orig[b][c].mean() + cond_c_mean = cond[b][c].mean() + cond[b][c] -= cond_c_mean + result[b][c] = x_orig[b][c] - cond[b][c] / cond[b][c].std() ** .5 * scale + cond_c_mean + + return result + + new_scale = 1 / (model.model.latent_format.scale_factor * 8) - #Taken and adapted from comfy_extras/nodes_model_advanced + + #Taken and modified from comfy_extras/nodes_model_advanced def rescale_cfg(args): x_orig = args["input"] x_orig -= x_orig.mean() @@ -30,7 +47,7 @@ class uncondZeroNode: cond = args["cond_denoised"] cond -= cond.mean() - cond = x_orig - (cond / cond.std() ** .5) * scale + cond = x_orig - cond / cond.std() ** .5 sigma = args["sigma"] sigma = sigma.view(sigma.shape[:1] + (1,) * (cond.ndim - 1)) @@ -41,7 +58,7 @@ class uncondZeroNode: cond = ((x - (x_orig - cond)) * (sigma ** 2 + 1.0) ** 0.5) / (sigma) #rescalecfg - x_cfg = uncond + new_scale * (cond - uncond) + x_cfg = uncond + new_scale * scale * (cond - uncond) ro_pos = torch.std(cond, dim=(1,2,3), keepdim=True) ro_cfg = torch.std(x_cfg, dim=(1,2,3), keepdim=True) @@ -50,5 +67,5 @@ class uncondZeroNode: return x_orig - (x - x_rescaled * sigma / (sigma * sigma + 1.0) ** 0.5) m = model.clone() - m.set_model_sampler_cfg_function({"uncond_zero":uncond_zero,"rescale_cfg":rescale_cfg}[method]) + m.set_model_sampler_cfg_function({"uncond_zero_v1":uncond_zero,"uncond_zero_v2":uncond_zero_v2,"uncond_zero_v3":rescale_cfg}[method]) return (m, ) \ No newline at end of file