add epsilon scaling

ldm_patched/modules/samplers.py

@@ -360,6 +360,15 @@ def calc_cond_uncond_batch(model, cond, uncond, x_in, timestep, model_options):
 
 import math
 
+def _epsilon_scaling_post_cfg(args, scaling_factor: float):
+    """Scale the denoised prediction by adjusting the predicted noise."""
+    denoised = args["denoised"]
+    x = args["input"]
+    noise_pred = x - denoised
+    scaled_noise_pred = noise_pred / scaling_factor
+    return x - scaled_noise_pred
+
+
 def cfg_function(model, cond_pred, uncond_pred, cond_scale, x, timestep, model_options={}, cond=None, uncond=None):
     edit_strength = sum((item.get('strength', 1) for item in cond))
 
@@ -389,6 +398,14 @@ def cfg_function(model, cond_pred, uncond_pred, cond_scale, x, timestep, model_o
                 "sigma": timestep, "model_options": model_options, "input": x}
         cfg_result = fn(args)
 
+    if getattr(shared.opts, "epsilon_scaling_enabled", False):
+        scaling_factor = getattr(shared.opts, "epsilon_scaling_factor", 1.005)
+        if math.isclose(scaling_factor, 0.0):
+            scaling_factor = 1e-9
+        args = {"denoised": cfg_result, "cond": cond, "uncond": uncond, "cond_scale": cond_scale, "model": model, "uncond_denoised": uncond_pred, "cond_denoised": cond_pred,
+                "sigma": timestep, "model_options": model_options, "input": x}
+        cfg_result = _epsilon_scaling_post_cfg(args, scaling_factor)
+
     return cfg_result
 
 #The main sampling function shared by all the samplers

modules/shared_options.py

@@ -250,6 +250,11 @@
     ),
 }))
 
+options_templates.update(options_section(('epsilon_scaling', "Epsilon Scaling", "sd"), {
+    "epsilon_scaling_enabled": OptionInfo(False, "Enable epsilon scaling").info("Scale the guided prediction to mitigate exposure bias during sampling."),
+    "epsilon_scaling_factor": OptionInfo(1.005, "Epsilon scaling factor", gr.Slider, {"minimum": 0.5, "maximum": 1.5, "step": 0.001}).info("Noise scaling multiplier from 'Elucidating the Exposure Bias in Diffusion Models'."),
+}))
+
 options_templates.update(options_section(('img2img', "img2img", "sd"), {
     "inpainting_mask_weight": OptionInfo(1.0, "Inpainting conditioning mask strength", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01}, infotext='Conditional mask weight'),
     "initial_noise_multiplier": OptionInfo(1.0, "Noise multiplier for img2img", gr.Slider, {"minimum": 0.0, "maximum": 1.5, "step": 0.001}, infotext='Noise multiplier'),