diff --git a/nodes_sampler.py b/nodes_sampler.py index 22c9e5d..3c4c1ef 100644 --- a/nodes_sampler.py +++ b/nodes_sampler.py @@ -1297,6 +1297,12 @@ class WanVideoSampler: extra_channel_latents_input = extra_channel_latents.to(z) z = torch.cat([z, extra_channel_latents_input]) + if scheduler.lower() == "rcm": + c_in = 1 / (torch.cos(timestep) + torch.sin(timestep)) + c_noise = (torch.sin(timestep) / (torch.cos(timestep) + torch.sin(timestep))) * 1000 + z = z * c_in + timestep = c_noise + base_params = { 'x': [z], # latent 'y': [image_cond_input] if image_cond_input is not None else None, # image cond @@ -2956,6 +2962,8 @@ class WanVideoSampler: # callback_latent = (latent_model_input[:,:-phantom_latents.shape[1]].to(device) - noise_pred[:,:-phantom_latents.shape[1]].to(device) * t.to(device) / 1000).detach() elif humo_reference_count > 0: callback_latent = (latent_model_input[:,:-humo_reference_count].to(device) - noise_pred[:,:-humo_reference_count].to(device) * t.to(device) / 1000).detach() + elif scheduler == "rcm": + callback_latent = (latent_model_input.to(device) - noise_pred.to(device) * t.to(device)).detach() else: callback_latent = (latent_model_input.to(device) - noise_pred.to(device) * t.to(device) / 1000).detach() callback(idx, callback_latent.permute(1,0,2,3), None, len(timesteps)) diff --git a/wanvideo/schedulers/__init__.py b/wanvideo/schedulers/__init__.py index 6a9b276..8df1749 100644 --- a/wanvideo/schedulers/__init__.py +++ b/wanvideo/schedulers/__init__.py @@ -6,6 +6,7 @@ from .flowmatch_pusa import FlowMatchSchedulerPusa from .flowmatch_res_multistep import FlowMatchSchedulerResMultistep from .scheduling_flow_match_lcm import FlowMatchLCMScheduler from .fm_sa_ode import FlowMatchSAODEStableScheduler +from .fm_rcm import rCMFlowMatchScheduler from ...utils import log try: @@ -26,7 +27,8 @@ scheduler_list = [ "flowmatch_distill", "flowmatch_pusa", "multitalk", - "sa_ode_stable" + "sa_ode_stable", + "rcm" ] def get_scheduler(scheduler, steps, start_step, end_step, shift, device, transformer_dim=5120, flowedit_args=None, denoise_strength=1.0, sigmas=None, log_timesteps=False, **kwargs): @@ -105,6 +107,10 @@ def get_scheduler(scheduler, steps, start_step, end_step, shift, device, transfo elif "sa_ode_stable" in scheduler: sample_scheduler = FlowMatchSAODEStableScheduler(shift=shift, **kwargs) sample_scheduler.set_timesteps(steps, device=device, sigmas=sigmas[:-1].tolist() if sigmas is not None else None) + elif 'rcm' in scheduler: + sample_scheduler = rCMFlowMatchScheduler() + sample_scheduler.set_timesteps(steps, sigma_max=120) + if timesteps is None: timesteps = sample_scheduler.timesteps diff --git a/wanvideo/schedulers/fm_rcm.py b/wanvideo/schedulers/fm_rcm.py new file mode 100644 index 0000000..46ed703 --- /dev/null +++ b/wanvideo/schedulers/fm_rcm.py @@ -0,0 +1,40 @@ +# based on https://github.com/NVlabs/rcm +import torch +import math + +class rCMFlowMatchScheduler(): + + def __init__(self, num_inference_steps=4, num_train_timesteps=1000, sigma_max=120): + self.num_train_timesteps = num_train_timesteps + self.sigma_max = sigma_max + self.step_index = 0 + + self.set_timesteps(num_inference_steps, sigma_max=sigma_max) + + def set_timesteps(self, num_inference_steps=4, sigma_max=120): + mid_t = [1.5, 1.4, 1.0][:num_inference_steps - 1] + self.timesteps = torch.tensor([math.atan(sigma_max), *mid_t], dtype=torch.float32) + self.sigmas = torch.tensor([math.atan(sigma_max), *mid_t, 0], dtype=torch.float32) + self.step_index = 0 + + def step(self, model_output, timestep, sample, generator): + + c_skip = 1 / (torch.cos(timestep) + torch.sin(timestep)) + c_out = -1 * torch.sin(timestep) / (torch.cos(timestep) + torch.sin(timestep)) + + # Get next timestep + if self.step_index + 1 < len(self.sigmas): + t_next = self.sigmas[self.step_index + 1] + else: + t_next = torch.tensor(0.0) + + x = c_skip * sample + c_out * model_output + if t_next > 1e-5: + x = torch.cos(t_next) * x + torch.sin(t_next) * torch.randn( + *x.shape, + dtype=torch.float32, + device=torch.device("cpu"), + generator=generator, + ).to(x) + self.step_index += 1 + return x