Merge 66bbdefcc0 into b08debceca

comfy/k_diffusion/sampling.py

@@ -1521,70 +1521,132 @@ def sample_gradient_estimation_cfg_pp(model, x, sigmas, extra_args=None, callbac
     return sample_gradient_estimation(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, ge_gamma=ge_gamma, cfg_pp=True)
 
 
+# Extended Reverse-Time SDE solver (VP ER-SDE-Solver-3). arXiv: https://arxiv.org/abs/2309.06169.
+# Code reference for initial implementation: https://github.com/QinpengCui/ER-SDE-Solver/blob/main/er_sde_solver.py.
 @torch.no_grad()
-def sample_er_sde(model, x, sigmas, extra_args=None, callback=None, disable=None, s_noise=1.0, noise_sampler=None, noise_scaler=None, max_stage=3):
-    """Extended Reverse-Time SDE solver (VP ER-SDE-Solver-3). arXiv: https://arxiv.org/abs/2309.06169.
-    Code reference: https://github.com/QinpengCui/ER-SDE-Solver/blob/main/er_sde_solver.py.
-    """
+def sample_er_sde(
+    model,
+    x: torch.Tensor,
+    sigmas: torch.Tensor,
+    extra_args=None,
+    callback=None,
+    disable=None,
+    eta: float = 1.0,
+    s_noise: float = 1.0,
+    noise_sampler=None,
+    noise_scaler=None,
+    max_stage: int = 3,
+    num_integration_points: int = 200,
+    scaling_power: float = 0.3,
+    scaling_constant: float = 10.0,
+    interpolation_function=torch.lerp,
+    # One of default, ersde or sde.
+    solver_type: str = "default",
+) -> torch.Tensor:
+    model_sampling = model.inner_model.model_patcher.get_model_object("model_sampling")
     extra_args = {} if extra_args is None else extra_args
     seed = extra_args.get("seed", None)
-    noise_sampler = default_noise_sampler(x, seed=seed) if noise_sampler is None else noise_sampler
-    s_noise = s_noise * getattr(model.inner_model.model_patcher.get_model_object('model_sampling'), "noise_scale", 1.0)
-    s_in = x.new_ones([x.shape[0]])
+    eta = max(0.0, eta)
+    if eta > 0:
+        s_noise = s_noise * getattr(model_sampling, "noise_scale", 1.0)
+        if noise_sampler is None:
+            noise_sampler = default_noise_sampler(x, seed=seed)
 
-    def default_er_sde_noise_scaler(x):
-        return x * ((x ** 0.3).exp() + 10.0)
+    s_in = x.new_ones(x.shape[:1])
 
-    noise_scaler = default_er_sde_noise_scaler if noise_scaler is None else noise_scaler
-    num_integration_points = 200.0
-    point_indice = torch.arange(0, num_integration_points, dtype=torch.float32, device=x.device)
+    if solver_type not in {"default", "sde", "ersde"}:
+        raise ValueError("Bad solver_type, must be one of default, ersde or sde")
+    if noise_scaler is None:
+        if solver_type == "sde":
+
+            def noise_scaler(val_x: torch.Tensor) -> torch.Tensor:
+                return val_x ** (eta + 1)
+
+        else:  # default or ersde.
+            solver_type = "ersde"
+
+            def noise_scaler(val_x: torch.Tensor) -> torch.Tensor:
+                rho_sde = val_x * ((val_x**scaling_power).exp_() + scaling_constant)
+                squared_scale = (1.0 - eta**2) * (val_x**2) + (eta**2) * (rho_sde**2)
+                return squared_scale.clamp_min_(1e-09).sqrt_()
+
+    elif solver_type == "default":
+        solver_type = "sde"
+
+    point_indice = torch.arange(
+        0, num_integration_points, dtype=x.dtype, device=x.device
+    )
 
-    model_sampling = model.inner_model.model_patcher.get_model_object("model_sampling")
     sigmas = offset_first_sigma_for_snr(sigmas, model_sampling)
     half_log_snrs = sigma_to_half_log_snr(sigmas, model_sampling)
-    er_lambdas = half_log_snrs.neg().exp()  # er_lambda_t = sigma_t / alpha_t
+    er_lambdas = half_log_snrs.neg().exp_()  # er_lambda_t = sigma_t / alpha_t
 
-    old_denoised = None
-    old_denoised_d = None
+    old_denoised = old_denoised_d = None
 
     for i in trange(len(sigmas) - 1, disable=disable):
-        denoised = model(x, sigmas[i] * s_in, **extra_args)
+        sigma, sigma_next = sigmas[i : i + 2]
+        denoised = model(x, sigma * s_in, **extra_args)
         if callback is not None:
-            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
+            callback(
+                {
+                    "x": x,
+                    "i": i,
+                    "sigma": sigma,
+                    "sigma_hat": sigma,
+                    "denoised": denoised,
+                }
+            )
+        if sigma_next <= 0:
+            return denoised
+
         stage_used = min(max_stage, i + 1)
-        if sigmas[i + 1] == 0:
-            x = denoised
+        er_lambda_s, er_lambda_t = er_lambdas[i], er_lambdas[i + 1]
+
+        alpha_s = sigma / er_lambda_s
+        alpha_t = sigma_next / er_lambda_t
+        rho_sde_s = noise_scaler(er_lambda_s)
+        rho_sde_t = noise_scaler(er_lambda_t)
+        r_alpha = alpha_t / alpha_s
+        r_SDE = rho_sde_t / rho_sde_s
+        if solver_type == "sde":
+            r, r_sq = r_SDE, r_SDE**2
         else:
-            er_lambda_s, er_lambda_t = er_lambdas[i], er_lambdas[i + 1]
-            alpha_s = sigmas[i] / er_lambda_s
-            alpha_t = sigmas[i + 1] / er_lambda_t
-            r_alpha = alpha_t / alpha_s
-            r = noise_scaler(er_lambda_t) / noise_scaler(er_lambda_s)
+            r_ODE = er_lambda_t / er_lambda_s
+            r_sq = interpolation_function(r_ODE**2, r_SDE**2, eta**2).clamp_min_(0.0)
+            r = r_sq.sqrt()
 
-            # Stage 1 Euler
-            x = r_alpha * r * x + alpha_t * (1 - r) * denoised
+        # Stage 1 Euler
+        x = r_alpha * r * x + alpha_t * (1 - r) * denoised
 
-            if stage_used >= 2:
-                dt = er_lambda_t - er_lambda_s
-                lambda_step_size = -dt / num_integration_points
-                lambda_pos = er_lambda_t + point_indice * lambda_step_size
-                scaled_pos = noise_scaler(lambda_pos)
+        if stage_used >= 2:
+            dt = er_lambda_t - er_lambda_s
+            lambda_step_size = -dt / num_integration_points
+            lambda_pos = er_lambda_t + point_indice * lambda_step_size
+            scaled_pos = noise_scaler(lambda_pos)
 
-                # Stage 2
-                s = torch.sum(1 / scaled_pos) * lambda_step_size
-                denoised_d = (denoised - old_denoised) / (er_lambda_s - er_lambdas[i - 1])
-                x = x + alpha_t * (dt + s * noise_scaler(er_lambda_t)) * denoised_d
+            # Stage 2
+            s = (1 / scaled_pos).sum() * lambda_step_size
+            denoised_d = (denoised - old_denoised) / (er_lambda_s - er_lambdas[i - 1])
+            x += alpha_t * (dt + s * rho_sde_t) * denoised_d
 
-                if stage_used >= 3:
-                    # Stage 3
-                    s_u = torch.sum((lambda_pos - er_lambda_s) / scaled_pos) * lambda_step_size
-                    denoised_u = (denoised_d - old_denoised_d) / ((er_lambda_s - er_lambdas[i - 2]) / 2)
-                    x = x + alpha_t * ((dt ** 2) / 2 + s_u * noise_scaler(er_lambda_t)) * denoised_u
-                old_denoised_d = denoised_d
-
-            if s_noise > 0:
-                x = x + alpha_t * noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * (er_lambda_t ** 2 - er_lambda_s ** 2 * r ** 2).sqrt().nan_to_num(nan=0.0)
+            if stage_used >= 3:
+                # Stage 3
+                s_u = ((lambda_pos - er_lambda_s) / scaled_pos).sum() * lambda_step_size
+                denoised_u = (denoised_d - old_denoised_d) / (
+                    (er_lambda_s - er_lambdas[i - 2]) / 2
+                )
+                x += alpha_t * ((dt**2) / 2 + s_u * rho_sde_t) * denoised_u
+            old_denoised_d = denoised_d
         old_denoised = denoised
+
+        if eta <= 0:
+            continue
+
+        # When r approaches 0.0, noise_coeff approaches er_lambda_t (maximum possible added noise).
+        noise_coeff = (
+            (er_lambda_t**2 - er_lambda_s**2 * r_sq).sqrt_().nan_to_num_(nan=0.0)
+        )
+        x += alpha_t * noise_sampler(sigma, sigma_next) * s_noise * noise_coeff
     return x
 
 

comfy/text_encoders/llama.py

@@ -937,22 +937,41 @@ class BaseGenerate:
             return torch.argmax(logits, dim=-1, keepdim=True)
 
         # Sampling mode
-        if repetition_penalty != 1.0:
-            for i in range(logits.shape[0]):
-                for token_id in set(token_history):
-                    logits[i, token_id] *= repetition_penalty if logits[i, token_id] < 0 else 1/repetition_penalty
-
-        if presence_penalty is not None and presence_penalty != 0.0:
-            for i in range(logits.shape[0]):
-                for token_id in set(token_history):
-                    logits[i, token_id] -= presence_penalty
+        if len(token_history) > 0 and (repetition_penalty != 1.0 or (presence_penalty is not None and presence_penalty != 0.0)):
+            token_ids = torch.tensor(list(set(token_history)), device=logits.device)
+            token_logits = logits[:, token_ids]
+            if repetition_penalty != 1.0:
+                token_logits = torch.where(token_logits < 0, token_logits * repetition_penalty, token_logits / repetition_penalty)
+            if presence_penalty is not None and presence_penalty != 0.0:
+                token_logits = token_logits - presence_penalty
+            logits[:, token_ids] = token_logits
 
         if temperature != 1.0:
             logits = logits / temperature
 
         if top_k > 0:
-            indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
-            logits[indices_to_remove] = torch.finfo(logits.dtype).min
+            top_k = min(top_k, logits.shape[-1])
+            logits, top_indices = torch.topk(logits, top_k)
+
+            if min_p > 0.0:
+                probs_before_filter = torch.nn.functional.softmax(logits, dim=-1)
+                top_probs, _ = probs_before_filter.max(dim=-1, keepdim=True)
+                min_threshold = min_p * top_probs
+                indices_to_remove = probs_before_filter < min_threshold
+                logits[indices_to_remove] = torch.finfo(logits.dtype).min
+
+            if top_p < 1.0:
+                sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+                cumulative_probs = torch.cumsum(torch.nn.functional.softmax(sorted_logits, dim=-1), dim=-1)
+                sorted_indices_to_remove = cumulative_probs > top_p
+                sorted_indices_to_remove[..., 0] = False
+                indices_to_remove = torch.zeros_like(logits, dtype=torch.bool)
+                indices_to_remove.scatter_(1, sorted_indices, sorted_indices_to_remove)
+                logits[indices_to_remove] = torch.finfo(logits.dtype).min
+
+            probs = torch.nn.functional.softmax(logits, dim=-1)
+            next_token = torch.multinomial(probs, num_samples=1, generator=generator)
+            return top_indices.gather(1, next_token)
 
         if min_p > 0.0:
             probs_before_filter = torch.nn.functional.softmax(logits, dim=-1)

comfy_extras/nodes_custom_sampler.py

@@ -584,40 +584,102 @@ class SamplerDPMAdaptative(io.ComfyNode):
 
 class SamplerER_SDE(io.ComfyNode):
     @classmethod
-    def define_schema(cls):
+    def define_schema(cls) -> io.Schema:
         return io.Schema(
             node_id="SamplerER_SDE",
+            search_aliases=["sde", "er_sde", "ersde"],
             category="model/sampling/samplers",
             inputs=[
                 io.Combo.Input("solver_type", options=["ER-SDE", "Reverse-time SDE", "ODE"]),
-                io.Int.Input("max_stage", default=3, min=1, max=3, advanced=True),
-                io.Float.Input("eta", default=1.0, min=0.0, max=100.0, step=0.01, round=False, tooltip="Stochastic strength of reverse-time SDE.\nWhen eta=0, it reduces to deterministic ODE. This setting doesn't apply to ER-SDE solver type.", advanced=True),
-                io.Float.Input("s_noise", default=1.0, min=0.0, max=100.0, step=0.01, round=False, advanced=True),
+                io.Int.Input(
+                    "max_stage",
+                    default=3,
+                    min=1,
+                    max=3,
+                    advanced=True,
+                    tooltip="Controls the number of stages the sampler uses. Stages: 1 - only uses the current step (Euler). 2 - Uses history from the previous step to improve accuracy. 3 - Uses two previous steps.",
+                ),
+                io.Float.Input(
+                    "eta",
+                    default=1.0,
+                    min=0.0,
+                    max=100.0,
+                    step=0.01,
+                    advanced=True,
+                    tooltip="Stochastic strength. Only has an effect when solver_type is not ODE.",
+                ),
+                io.Float.Input(
+                    "s_noise",
+                    default=1.0,
+                    min=-100.0,
+                    max=100.0,
+                    step=0.01,
+                    advanced=True,
+                    tooltip="SDE noise multiplier. Only has an effect when solver_type is not ODE.",
+                ),
+                io.Int.Input(
+                    "integration_points",
+                    default=200,
+                    min=1,
+                    max=10000,
+                    advanced=True,
+                    tooltip="More integration points improves accuracy with diminishing returns. The default is a good compromise. Only applies to the ER-SDE solver type.",
+                ),
+                io.Float.Input(
+                    "scaling_power",
+                    default=0.3,
+                    min=0.0,
+                    max=0.7,
+                    step=0.01,
+                    advanced=True,
+                    tooltip="Controls the exponent used for ER-SDE steps. Lower values make the sampler act more like a linear solver. Values above 0.5 may cause numerical overflow. Only has an effect when ETA is non-zero.",
+                ),
+                io.Float.Input(
+                    "scaling_constant",
+                    default=10.0,
+                    min=-0.99,
+                    max=100.0,
+                    step=0.1,
+                    advanced=True,
+                    tooltip="Constant value used for ER-SDE steps. Higher values cause the sampler to transition its stable, linear mode earlier while lower values will delay the transition.",
+                ),
             ],
-            outputs=[io.Sampler.Output()]
+            outputs=[io.Sampler.Output()],
         )
 
     @classmethod
-    def execute(cls, solver_type, max_stage, eta, s_noise) -> io.NodeOutput:
-        if solver_type == "ODE" or (solver_type == "Reverse-time SDE" and eta == 0):
-            eta = 0
-            s_noise = 0
-
-        def reverse_time_sde_noise_scaler(x):
-            return x ** (eta + 1)
-
-        if solver_type == "ER-SDE":
-            # Use the default one in sample_er_sde()
-            noise_scaler = None
+    def execute(
+        cls,
+        *,
+        solver_type: str,
+        max_stage: int,
+        eta: float,
+        s_noise: float,
+        integration_points: int,
+        scaling_power: float,
+        scaling_constant: float,
+    ) -> io.NodeOutput:
+        if solver_type == "ODE":
+            eta = s_noise = 0.0
+            solver_type = "sde"
+        elif solver_type == "Reverse-time SDE":
+            solver_type = "sde"
         else:
-            noise_scaler = reverse_time_sde_noise_scaler
-
-        sampler_name = "er_sde"
-        sampler = comfy.samplers.ksampler(sampler_name, {"s_noise": s_noise, "noise_scaler": noise_scaler, "max_stage": max_stage})
+            solver_type = "ersde"
+        sampler = comfy.samplers.ksampler(
+            "er_sde",
+            {
+                "solver_type": solver_type,
+                "eta": eta,
+                "s_noise": s_noise,
+                "max_stage": max_stage,
+                "num_integration_points": integration_points,
+                "scaling_power": scaling_power,
+                "scaling_constant": scaling_constant,
+            }
+        )
         return io.NodeOutput(sampler)
 
-    get_sampler = execute
-
 
 class SamplerSASolver(io.ComfyNode):
     @classmethod

requirements.txt

@@ -1,6 +1,6 @@
 comfyui-frontend-package==1.45.20
 comfyui-workflow-templates==0.11.2
-comfyui-embedded-docs==0.5.6
+comfyui-embedded-docs==0.5.7
 torch
 torchsde
 torchvision