diff --git a/nodes.py b/nodes.py index eb92d4b..84e8ef3 100644 --- a/nodes.py +++ b/nodes.py @@ -3,10 +3,12 @@ import torch import comfy.sample import comfy.samplers import comfy.utils +import comfy.model_sampling + import latent_preview -def wan_ksampler(model_high_noise, model_low_noise, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent, boundary = 0.875, denoise=1.0, disable_noise=False, start_step=None, last_step=None, force_full_denoise=False): +def wan_ksampler(model_high_noise, model_low_noise, seed, steps, cfgs, sampler_name, scheduler, positive, negative, latent, boundary = 0.875, denoise=1.0, disable_noise=False, start_step=None, last_step=None, force_full_denoise=False): # boundary is .9 for i2v, .875 for t2v latent_image = latent["samples"] @@ -32,9 +34,12 @@ def wan_ksampler(model_high_noise, model_low_noise, seed, steps, cfg, sampler_na sampler = comfy.samplers.KSampler(model_high_noise, steps=steps, device=model_high_noise.load_device, sampler=sampler_name, scheduler=scheduler, denoise=denoise, model_options=model_high_noise.model_options) sigmas = sampler.sigmas # why are timesteps 0-1000? - timesteps = [model_high_noise.get_model_object("model_sampling").timestep(sigma)/1000 for sigma in sigmas.tolist()] + sampling = model_high_noise.get_model_object("model_sampling") + timesteps = [sampling.timestep(sigma)/1000 for sigma in sigmas.tolist()] + print(sigmas.tolist()) + print(timesteps) switching_step = steps - for (i,t) in enumerate(timesteps): + for (i,t) in enumerate(timesteps[1:]): if t < boundary: switching_step = i break @@ -47,7 +52,7 @@ def wan_ksampler(model_high_noise, model_low_noise, seed, steps, cfg, sampler_na callback = latent_preview.prepare_callback(model_high_noise, steps) end_step = min(last_step,switching_step) latent_image = comfy.sample.fix_empty_latent_channels(model_high_noise, latent_image) - latent_image = comfy.sample.sample(model_high_noise, noise, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, + latent_image = comfy.sample.sample(model_high_noise, noise, steps, cfgs[0], sampler_name, scheduler, positive, negative, latent_image, denoise=denoise, disable_noise=end_wth_low or disable_noise, start_step=start_step, last_step=end_step, force_full_denoise=end_wth_low or force_full_denoise, noise_mask=noise_mask, callback=callback, disable_pbar=disable_pbar, seed=seed) @@ -57,7 +62,7 @@ def wan_ksampler(model_high_noise, model_low_noise, seed, steps, cfg, sampler_na callback = latent_preview.prepare_callback(model_low_noise, steps) begin_step = max(start_step, switching_step) latent_image = comfy.sample.fix_empty_latent_channels(model_low_noise, latent_image) - latent_image = comfy.sample.sample(model_low_noise, noise, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, + latent_image = comfy.sample.sample(model_low_noise, noise, steps, cfgs[1], sampler_name, scheduler, positive, negative, latent_image, denoise=denoise, disable_noise=disable_noise, start_step=begin_step, last_step=last_step, force_full_denoise=force_full_denoise, noise_mask=noise_mask, callback=callback, disable_pbar=disable_pbar, seed=seed) @@ -65,6 +70,20 @@ def wan_ksampler(model_high_noise, model_low_noise, seed, steps, cfg, sampler_na out["samples"] = latent_image return (out, ) +def set_shift(model,sigma_shift): + + sampling_base = comfy.model_sampling.ModelSamplingDiscreteFlow + sampling_type = comfy.model_sampling.CONST + + class ModelSamplingAdvanced(sampling_base, sampling_type): + pass + + model_sampling = ModelSamplingAdvanced(model.model.model_config) + model_sampling.set_parameters(shift=sigma_shift, multiplier=1000) + model.add_object_patch("model_sampling", model_sampling) + model.add_object_patch("model_sampling", model_sampling) + return model + class WanMoeKSampler: @classmethod def INPUT_TYPES(s): @@ -72,12 +91,14 @@ class WanMoeKSampler: "required": { "model_high_noise": ("MODEL", {"tooltip": "The first expert of the model used for denoising the input latent."}), "model_low_noise": ("MODEL", {"tooltip": "The second expert of the model used for denoising the input latent."}), - "boundary": ("FLOAT", {"default:": 0.875, "min": 0.0, "max": 1.0, "step": 0.001, "round": 0.001,"tooltip": "Boundary (or t_moe): Timestep (not to be confused with denoising step) at which models should be swapped. Recommended values: 0.875 for t2v, 0.9 for i2v"}), + "boundary": ("FLOAT", {"default": 0.875, "min": 0.0, "max": 1.0, "step": 0.001, "round": 0.001,"tooltip": "Boundary (or t_moe): Timestep (not to be confused with denoising step) at which models should be switched. Recommended values: 0.875 for t2v, 0.9 for i2v"}), "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "control_after_generate": True, "tooltip": "The random seed used for creating the noise."}), "steps": ("INT", {"default": 20, "min": 1, "max": 10000, "tooltip": "The number of steps used in the denoising process."}), - "cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step":0.1, "round": 0.01, "tooltip": "The Classifier-Free Guidance scale balances creativity and adherence to the prompt. Higher values result in images more closely matching the prompt however too high values will negatively impact quality."}), + "cfg_high_noise": ("FLOAT", {"default": 4.0, "min": 0.0, "max": 100.0, "step":0.1, "round": 0.01, "tooltip": "The Classifier-Free Guidance scale balances creativity and adherence to the prompt. Higher values result in images more closely matching the prompt however too high values will negatively impact quality."}), + "cfg_low_noise": ("FLOAT", {"default": 3.0, "min": 0.0, "max": 100.0, "step":0.1, "round": 0.01, "tooltip": "The Classifier-Free Guidance scale balances creativity and adherence to the prompt. Higher values result in images more closely matching the prompt however too high values will negatively impact quality."}), "sampler_name": (comfy.samplers.KSampler.SAMPLERS, {"tooltip": "The algorithm used when sampling, this can affect the quality, speed, and style of the generated output."}), "scheduler": (comfy.samplers.KSampler.SCHEDULERS, {"tooltip": "The scheduler controls how noise is gradually removed to form the image."}), + "sigma_shift": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step":0.01, "tooltip": "Same purpose as the a shift parameter in the ModelSamplingSD3 node (same value applied to both models)"}), "positive": ("CONDITIONING", {"tooltip": "The conditioning describing the attributes you want to include in the image."}), "negative": ("CONDITIONING", {"tooltip": "The conditioning describing the attributes you want to exclude from the image."}), "latent_image": ("LATENT", {"tooltip": "The latent image to denoise."}), @@ -92,8 +113,11 @@ class WanMoeKSampler: CATEGORY = "sampling" DESCRIPTION = "Uses the provided model, positive and negative conditioning to denoise the latent image." - def sample(self, model_high_noise, model_low_noise, boundary, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=1.0): - return wan_ksampler(model_high_noise, model_low_noise, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image,boundary=boundary, denoise=denoise) + def sample(self, model_high_noise, model_low_noise, boundary, seed, steps, cfg_high_noise, cfg_low_noise, sampler_name, scheduler, sigma_shift, positive, negative, latent_image, denoise=1.0): + model_high_noise = set_shift(model_high_noise, sigma_shift) + model_low_noise = set_shift(model_low_noise, sigma_shift) + + return wan_ksampler(model_high_noise, model_low_noise, seed, steps, (cfg_high_noise, cfg_low_noise), sampler_name, scheduler, positive, negative, latent_image,boundary=boundary, denoise=denoise) class WanMoeKSamplerAdvanced: @classmethod @@ -101,13 +125,15 @@ class WanMoeKSamplerAdvanced: return {"required": {"model_high_noise": ("MODEL", {"tooltip": "The first expert of the model used for denoising the input latent."}), "model_low_noise": ("MODEL", {"tooltip": "The second expert of the model used for denoising the input latent."}), - "boundary":("FLOAT", {"default:": 0.875, "min": 0.0, "max": 1.0, "step": 0.001, "round":0.001,"tooltip": "Boundary (or t_moe): Timestep (not to be confused with denoising step) at which models should be swapped. Recommended values: 0.875 for t2v, 0.9 for i2v"}), + "boundary": ("FLOAT", {"default": 0.875, "min": 0.0, "max": 1.0, "step": 0.001, "round":0.001}), "add_noise": (["enable", "disable"], ), "noise_seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "control_after_generate": True}), "steps": ("INT", {"default": 20, "min": 1, "max": 10000}), - "cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step":0.1, "round": 0.01}), + "cfg_high_noise": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step":0.1, "round": 0.01}), + "cfg_low_noise": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step":0.1, "round": 0.01}), "sampler_name": (comfy.samplers.KSampler.SAMPLERS, ), "scheduler": (comfy.samplers.KSampler.SCHEDULERS, ), + "sigma_shift": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step":0.01, "tooltip": "Same purpose as the a shift parameter in the ModelSamplingSD3 node (same value applied to both models)"}), "positive": ("CONDITIONING", ), "negative": ("CONDITIONING", ), "latent_image": ("LATENT", ), @@ -122,11 +148,13 @@ class WanMoeKSamplerAdvanced: CATEGORY = "sampling" - def sample(self, model_high_noise, model_low_noise, boundary, add_noise, noise_seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, start_at_step, end_at_step, return_with_leftover_noise, denoise=1.0): + def sample(self, model_high_noise, model_low_noise, boundary, add_noise, noise_seed, steps, cfg_high_noise, cfg_low_noise, sampler_name, scheduler, sigma_shift, positive, negative, latent_image, start_at_step, end_at_step, return_with_leftover_noise, denoise=1.0): + model_high_noise = set_shift(model_high_noise, sigma_shift) + model_low_noise = set_shift(model_low_noise, sigma_shift) force_full_denoise = True if return_with_leftover_noise == "enable": force_full_denoise = False disable_noise = False if add_noise == "disable": disable_noise = True - return wan_ksampler(model_high_noise, model_low_noise, noise_seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, boundary=boundary, denoise=denoise, disable_noise=disable_noise, start_step=start_at_step, last_step=end_at_step, force_full_denoise=force_full_denoise) + return wan_ksampler(model_high_noise, model_low_noise, noise_seed, steps, (cfg_high_noise, cfg_low_noise), sampler_name, scheduler, positive, negative, latent_image, boundary=boundary, denoise=denoise, disable_noise=disable_noise, start_step=start_at_step, last_step=end_at_step, force_full_denoise=force_full_denoise)