import torch import numpy as np from types import MethodType from functools import partial from typing import Any try: import torchdiffeq except ImportError: print("Warning: torchdiffeq is not installed. The advanced ODE solvers will not be available.") print("Please install it by running: pip install torchdiffeq") torchdiffeq = None import comfy.model_sampling import comfy.samplers import comfy.sample from comfy.ldm.modules.diffusionmodules.util import make_beta_schedule from tqdm.auto import trange, tqdm import latent_preview # Diff2Flow Core Logic def _d2f_extract_into_tensor(self, a: torch.Tensor, t: torch.Tensor, x_shape: tuple[int, ...]) -> torch.Tensor: b, *_ = t.shape t = t.to(a.device) t = t.clamp(0, a.shape[-1] - 1) left_idx = t.long() right_idx = (left_idx + 1).clamp(max=a.shape[-1] - 1) left_val = a.gather(-1, left_idx) right_val = a.gather(-1, right_idx) t_ = t - left_idx.float() out = left_val * (1 - t_) + right_val * t_ return out.reshape(b, *((1,) * (len(x_shape) - 1))) def _d2f_convert_fm_t_to_dm_t(self, t: torch.Tensor) -> torch.Tensor: rectified_alphas_cumprod_full = self.df_rectified_alphas_cumprod_full.clone().to(t.device) rectified_alphas_cumprod_full = torch.flip(rectified_alphas_cumprod_full, [0]) right_index = torch.searchsorted(rectified_alphas_cumprod_full, t, right=True) left_index = right_index - 1 right_value = rectified_alphas_cumprod_full[right_index] left_value = rectified_alphas_cumprod_full[left_index] dm_t = left_index.float() + (t - left_value) / (right_value - left_value) dm_t = self.model_sampling.num_timesteps - dm_t return dm_t def _d2f_convert_fm_xt_to_dm_xt(self, fm_xt: torch.Tensor, fm_t: torch.Tensor) -> torch.Tensor: scale = self.df_sqrt_alphas_cumprod_full + self.df_sqrt_one_minus_alphas_cumprod_full dm_t = self._df_convert_fm_t_to_dm_t(fm_t) dm_t_left_index = torch.floor(dm_t) dm_t_right_index = torch.ceil(dm_t) max_idx = scale.shape[-1] - 1 dm_t_left_index = dm_t_left_index.clamp(0, max_idx).long() dm_t_right_index = dm_t_right_index.clamp(0, max_idx).long() dm_t_left_value = scale[dm_t_left_index] dm_t_right_value = scale[dm_t_right_index] scale_t = dm_t_left_value + (dm_t - dm_t_left_index.float()) * (dm_t_right_value - dm_t_left_value) scale_t = scale_t.view(-1, 1, 1, 1) dm_xt = fm_xt * scale_t return dm_xt def _d2f_predict_start_from_z_and_v(self, x_t: torch.Tensor, t: torch.Tensor, v: torch.Tensor) -> torch.Tensor: return ( self._df_extract_into_tensor(self.df_sqrt_alphas_cumprod, t, x_t.shape) * x_t - self._df_extract_into_tensor(self.df_sqrt_one_minus_alphas_cumprod, t, x_t.shape) * v ) def _d2f_predict_eps_from_z_and_v(self, x_t: torch.Tensor, t: torch.Tensor, v: torch.Tensor) -> torch.Tensor: return ( self._df_extract_into_tensor(self.df_sqrt_alphas_cumprod, t, x_t.shape) * v + self._df_extract_into_tensor(self.df_sqrt_one_minus_alphas_cumprod, t, x_t.shape) * x_t ) def _d2f_predict_start_from_eps(self, x_t: torch.Tensor, t: torch.Tensor, noise: torch.Tensor) -> torch.Tensor: return ( self._df_extract_into_tensor(self.df_sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t - self._df_extract_into_tensor(self.df_sqrt_recipm1_alphas_cumprod, t, x_t.shape) * noise ) def _d2f_get_vector_field_from_v(self, v: torch.Tensor, x_t: torch.Tensor, t: torch.Tensor) -> torch.Tensor: z_pred = self._df_predict_start_from_z_and_v(x_t, t, v) eps_pred = self._df_predict_eps_from_z_and_v(x_t, t, v) vector_field = z_pred - eps_pred return vector_field def _d2f_get_vector_field_from_eps(self, noise: torch.Tensor, x_t: torch.Tensor, t: torch.Tensor) -> torch.Tensor: z_pred = self._df_predict_start_from_eps(x_t, t, noise) eps_pred = noise vector_field = z_pred - eps_pred return vector_field @torch.no_grad() def _d2f_get_diff2flow_velocity(self, fm_x: torch.Tensor, fm_t: torch.Tensor, **kwargs: Any) -> torch.Tensor: """ The main inference/sampling function. Predicts the vector field (Diff2Flow velocity) using the underlying UNet model. """ dm_t_continuous = self._df_convert_fm_t_to_dm_t(fm_t) dm_x = self._df_convert_fm_xt_to_dm_xt(fm_x, fm_t) model_device = self.device model_dtype = self.get_dtype() unet_kwargs = kwargs.copy() if 'c_crossattn' in unet_kwargs: unet_kwargs['context'] = unet_kwargs.pop('c_crossattn') # Cast inputs to the correct device and dtype for the model dm_x = dm_x.to(model_device, model_dtype) dm_t_continuous = dm_t_continuous.to(model_device) # Timesteps are usually float32 for key, value in unet_kwargs.items(): if hasattr(value, "to") and key not in ['transformer_options', 'control']: unet_kwargs[key] = value.to(model_device, model_dtype) model_pred = self.diffusion_model( dm_x, dm_t_continuous, **unet_kwargs ) model_pred = model_pred.float() # Handle NaN prediction (as in original code) if torch.isnan(model_pred).any(): model_pred[torch.isnan(model_pred)] = 0 dm_x_float = dm_x.float() dm_t_float = dm_t_continuous.float() vector_field = None if isinstance(self.model_sampling, comfy.model_sampling.V_PREDICTION): vector_field = self._df_get_vector_field_from_v(model_pred, dm_x_float, dm_t_float) elif isinstance(self.model_sampling, comfy.model_sampling.EPS): vector_field = self._df_get_vector_field_from_eps(model_pred, dm_x_float, dm_t_float) else: raise NotImplementedError(f"Diff2Flow is not implemented for prediction type: {type(self.model_sampling)}") return vector_field def enable_diff2flow(model): """ Modifies the model in-place to enable training and inference with the Diff2Flow methodology. The logic based on the paper: "Diff2Flow: Training Flow Matching Models via Diffusion Model Alignment" Modified from: https://github.com/CompVis/diff2flow/blob/33239aa0c02c554ee0b3fff5c5f0167a8dabdf6a/diff2flow/flow_obj.py ADAPTED FOR COMFYUI. """ if hasattr(model, 'get_diff2flow_velocity'): return # Already enabled device = model.device sampling_settings = model.model_config.sampling_settings beta_schedule_name = sampling_settings.get("beta_schedule", "linear") if beta_schedule_name == "scaled_linear": beta_schedule_name = "linear" betas_tensor = make_beta_schedule( beta_schedule_name, model.model_sampling.num_timesteps, linear_start=model.model_sampling.linear_start, linear_end=model.model_sampling.linear_end ) betas = betas_tensor.cpu().numpy() alphas = 1. - betas alphas_cumprod = np.cumprod(alphas, axis=0) alphas_cumprod_full = np.append(1., alphas_cumprod) to_torch = partial(torch.tensor, dtype=torch.float32, device=device) # Register all constants as buffers instead of simple attributes model.register_buffer('df_betas', to_torch(betas)) model.register_buffer('df_alphas_cumprod', to_torch(alphas_cumprod)) model.register_buffer('df_alphas_cumprod_full', to_torch(alphas_cumprod_full)) model.register_buffer('df_sqrt_alphas_cumprod', to_torch(np.sqrt(alphas_cumprod))) model.register_buffer('df_sqrt_one_minus_alphas_cumprod', to_torch(np.sqrt(1. - alphas_cumprod))) model.register_buffer('df_sqrt_alphas_cumprod_full', to_torch(np.sqrt(alphas_cumprod_full))) model.register_buffer('df_sqrt_one_minus_alphas_cumprod_full', to_torch(np.sqrt(1. - alphas_cumprod_full))) model.register_buffer('df_sqrt_recip_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod))) model.register_buffer('df_sqrt_recipm1_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod - 1))) # These are calculated from tensors that will now be buffers rectified_alphas = model.df_sqrt_alphas_cumprod_full / (model.df_sqrt_alphas_cumprod_full + model.df_sqrt_one_minus_alphas_cumprod_full) rectified_sqrt_alphas = model.df_sqrt_one_minus_alphas_cumprod_full / (model.df_sqrt_alphas_cumprod_full + model.df_sqrt_one_minus_alphas_cumprod_full) model.register_buffer('df_rectified_alphas_cumprod_full', rectified_alphas) model._df_extract_into_tensor = MethodType(_d2f_extract_into_tensor, model) model._df_convert_fm_t_to_dm_t = MethodType(_d2f_convert_fm_t_to_dm_t, model) model._df_convert_fm_xt_to_dm_xt = MethodType(_d2f_convert_fm_xt_to_dm_xt, model) model._df_predict_start_from_z_and_v = MethodType(_d2f_predict_start_from_z_and_v, model) model._df_predict_eps_from_z_and_v = MethodType(_d2f_predict_eps_from_z_and_v, model) model._df_predict_start_from_eps = MethodType(_d2f_predict_start_from_eps, model) model._df_get_vector_field_from_v = MethodType(_d2f_get_vector_field_from_v, model) model._df_get_vector_field_from_eps = MethodType(_d2f_get_vector_field_from_eps, model) model.get_diff2flow_velocity = MethodType(_d2f_get_diff2flow_velocity, model) # ODE Solver Logic class ODEFunction: def __init__(self, model, t_min, t_max, n_steps, is_adaptive, extra_args=None, callback=None): self.model = model self.extra_args = {} if extra_args is None else extra_args self.callback = callback self.t_min = t_min.item() self.t_max = t_max.item() self.n_steps = n_steps self.is_adaptive = is_adaptive self.step = 0 self.last_denoised = None if is_adaptive: self.pbar = tqdm(total=100, desc="solve", unit="%", leave=False, position=1) else: self.pbar = tqdm(total=n_steps, desc="solve", leave=False, position=1) def __call__(self, t, y): if t <= 1e-5: return torch.zeros_like(y) # Add batch dimension for the model y_in = y.unsqueeze(0) # prepare timestep for the model timestep_in = t.unsqueeze(0) if t.ndim == 0 else t # Call the CFGGuider denoised = self.model(y_in, timestep_in, **self.extra_args) self.last_denoised = denoised # Calculate derivative derivative = (y_in - denoised) / t # Remove batch dimension for torchdiffeq return derivative.squeeze(0) def _callback(self, t0, y0, step): if self.callback is not None: denoised_for_preview = self.last_denoised if self.last_denoised is not None else y0.unsqueeze(0) self.callback({ "x": y0.unsqueeze(0), "i": step, "sigma": t0, "sigma_hat": t0, "denoised": denoised_for_preview }) def callback_step(self, t0, y0, dt): if self.is_adaptive: return self._callback(t0, y0, self.step) self.pbar.update(1) self.step += 1 def callback_accept_step(self, t0, y0, dt): if not self.is_adaptive: return progress = (self.t_max - t0.item()) / (self.t_max - self.t_min) self._callback(t0, y0, round((self.n_steps - 1) * progress)) new_step = round(100 * progress) self.pbar.update(new_step - self.step) self.step = new_step def reset(self): self.step = 0 self.pbar.reset() self.last_denoised = None class ODESampler: def __init__(self, solver, rtol, atol, max_steps): if torchdiffeq is None: raise ImportError("torchdiffeq is not installed. Please install it to use ODE solvers.") self.solver = solver self.rtol = rtol self.atol = atol self.max_steps = max_steps self.ADAPTIVE_SOLVERS = { "dopri8", "dopri5", "bosh3", "fehlberg2", "adaptive_heun" } @torch.no_grad() def __call__(self, model, x, sigmas, extra_args=None, callback=None, disable=None): is_adaptive = self.solver in self.ADAPTIVE_SOLVERS t_max = sigmas[0] # Use the provided sigma schedule t_min = sigmas[-1] t = torch.stack([t_max, t_min]) if is_adaptive else sigmas.to(x.device) # The model passed here from KSampler is a wrapper, we need the inner CFGGuider cfg_guider_model = model.inner_model # The ODE function needs a model that returns denoised x0. The CFG guider does exactly that. # We must clean up extra_args to remove arguments not expected by the CFGGuider's predict function. ode_extra_args = extra_args.copy() ode_extra_args.pop('denoise_mask', None) ode = ODEFunction(cfg_guider_model, t_min, t_max, len(sigmas), is_adaptive, ode_extra_args, callback) samples = torch.empty_like(x) for i in trange(x.shape[0], desc=self.solver, disable=disable): ode.reset() samples[i] = torchdiffeq.odeint(ode, x[i], t, rtol=self.rtol, atol=self.atol, method=self.solver, options={"max_num_steps": self.max_steps})[-1] return samples # Node Definition class Diff2FlowODESampler: # Build solver list dynamically if torchdiffeq is not None: ADAPTIVE_SOLVERS = { "dopri8", "dopri5", "bosh3", "fehlberg2", "adaptive_heun" } FIXED_SOLVERS = { "euler", "midpoint", "rk4", "heun3", "explicit_adams", "implicit_adams" } ALL_ODE_SOLVERS = sorted(list(ADAPTIVE_SOLVERS | FIXED_SOLVERS)) SOLVERS = ALL_ODE_SOLVERS else: raise ImportError("torchdiffeq is not installed. Please install it to use this node.") @classmethod def INPUT_TYPES(s): inputs = { "required": { "model": ("MODEL",), "solver": (s.SOLVERS, {"default": "euler"}), "steps": ("INT", {"default": 20, "min": 1, "max": 10000}), "cfg": ("FLOAT", {"default": 5, "min": 0.0, "max": 100.0, "step":0.1, "round": 0.01}), "denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), "positive": ("CONDITIONING", ), "negative": ("CONDITIONING", ), "latent_image": ("LATENT", ), "scheduler": (comfy.samplers.SCHEDULER_NAMES, {"default": "simple"}), } } # The optional inputs for advanced ODE solvers if torchdiffeq is not None: inputs["optional"] = { "log_relative_tolerance": ("FLOAT", { "min": -7, "max": 0, "default": -2.5, "step": 0.1 }), "log_absolute_tolerance": ("FLOAT", { "min": -7, "max": 0, "default": -3.5, "step": 0.1 }), "max_steps": ("INT", { "min": 1, "max": 500, "default": 50, "step": 1 }), } return inputs RETURN_TYPES = ("LATENT",) FUNCTION = "sample" CATEGORY = "sampling/custom_sampling" def sample(self, model, solver, steps, cfg, denoise, seed, positive, negative, latent_image, disable_pbar=False, **kwargs): # Clone model and enable Diff2Flow model_patched = model.clone() enable_diff2flow(model_patched.model) # Determine which sampler to use if torchdiffeq is None: raise ImportError("torchdiffeq is not installed. Please install it to use this solver.") log_relative_tolerance = kwargs.get("log_relative_tolerance", -2.5) log_absolute_tolerance = kwargs.get("log_absolute_tolerance", -3.5) max_steps_ode = kwargs.get("max_steps", 50) rtol = 10 ** log_relative_tolerance atol = 10 ** log_absolute_tolerance sampler_function = ODESampler(solver, rtol, atol, max_steps_ode) sampler = comfy.samplers.KSAMPLER(sampler_function) # Prepare sigmas based on steps and denoise value # This standardizes the process for all samplers model_sampling = model_patched.model.model_sampling scheduler = kwargs.get("scheduler", "simple") if denoise == 1.0: sigmas = comfy.samplers.calculate_sigmas(model_sampling, scheduler, steps) else: total_steps = int(steps / denoise) if denoise > 0 else steps sigmas_full = comfy.samplers.calculate_sigmas(model_sampling, scheduler, total_steps) sigmas = sigmas_full[-(steps + 1):] # Prepare noise and call the main sample function noise = comfy.sample.prepare_noise(latent_image['samples'], seed) # Create the callback function that the UI uses for previews callback = latent_preview.prepare_callback(model, steps) samples = comfy.samplers.sample(model_patched, noise, positive, negative, cfg, device=model_patched.load_device, sampler=sampler, sigmas=sigmas.to(model_patched.load_device), model_options=model_patched.model_options, latent_image=latent_image['samples'], denoise_mask=None, # ODE samplers don't support masks well callback=callback, disable_pbar=disable_pbar, seed=seed) # Format and return the output out_latent = latent_image.copy() out_latent["samples"] = samples return (out_latent,) NODE_CLASS_MAPPINGS = { "Diff2FlowODESampler": Diff2FlowODESampler, } NODE_DISPLAY_NAME_MAPPINGS = { "Diff2FlowODESampler": "Diff2Flow ODE KSampler", }