Extend ER-SDE noise scaler by scaling h(t) (#15428)

comfy_extras/nodes_custom_sampler.py

@@ -591,7 +591,7 @@ class SamplerER_SDE(io.ComfyNode):
             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("eta", default=1.0, min=0.0, max=10.0, step=0.01, round=False, tooltip="Stochastic strength of SDEs.\nWhen eta=0, they reduce to deterministic ODE.\nLarge eta may cause invalid outputs. If this occurs, try decreasing this value.", advanced=True),
                 io.Float.Input("s_noise", default=1.0, min=0.0, max=100.0, step=0.01, round=False, advanced=True),
             ],
             outputs=[io.Sampler.Output()]
@@ -599,21 +599,35 @@ class SamplerER_SDE(io.ComfyNode):
 
     @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
+        # Extend existing noise scalers phi(x) with eta-controlled noise scalers:
+        #   psi(x) = x**(1-eta) * phi(x)**eta
+        # where eta is constant and directly scales the h^2(t) contribution.
 
-        def reverse_time_sde_noise_scaler(x):
+        def er_sde_noise_scaler(x: torch.Tensor) -> torch.Tensor:
+            return x * ((x ** 0.3).exp() + 10.0) ** eta
+
+        def reverse_time_sde_noise_scaler(x: torch.Tensor) -> torch.Tensor:
             return x ** (eta + 1)
 
-        if solver_type == "ER-SDE":
-            # Use the default one in sample_er_sde()
-            noise_scaler = None
-        else:
-            noise_scaler = reverse_time_sde_noise_scaler
+        def ode_noise_scaler(x: torch.Tensor) -> torch.Tensor:
+            return x
+
+        solver_scalers = {
+            "ER-SDE": er_sde_noise_scaler,
+            "Reverse-time SDE": reverse_time_sde_noise_scaler,
+            "ODE": ode_noise_scaler,
+        }
+
+        if solver_type == "ODE" or eta == 0:
+            s_noise = 0.0
+            solver_type = "ODE"
+        noise_scaler = solver_scalers[solver_type]
 
         sampler_name = "er_sde"
-        sampler = comfy.samplers.ksampler(sampler_name, {"s_noise": s_noise, "noise_scaler": noise_scaler, "max_stage": max_stage})
+        sampler = comfy.samplers.ksampler(
+            sampler_name,
+            {"s_noise": s_noise, "noise_scaler": noise_scaler, "max_stage": max_stage},
+        )
         return io.NodeOutput(sampler)
 
     get_sampler = execute