Update modules/model.py

modules/model.py

@@ -39,6 +39,20 @@ exists = lambda val: val is not None
 default = lambda val, d: val if exists(val) else d
 logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
 
+# from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.rescale_noise_cfg
+def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):
+    """
+    Rescale `noise_cfg` according to `guidance_rescale`. Based on findings of [Common Diffusion Noise Schedules and
+    Sample Steps are Flawed](https://arxiv.org/pdf/2305.08891.pdf). See Section 3.4
+    """
+    std_text = noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)), keepdim=True)
+    std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True)
+    # rescale the results from guidance (fixes overexposure)
+    noise_pred_rescaled = noise_cfg * (std_text / std_cfg)
+    # mix with the original results from guidance by factor guidance_rescale to avoid "plain looking" images
+    noise_cfg = guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * noise_cfg
+    return noise_cfg
+
 
 def get_attention_scores(attn, query, key, attention_mask=None):
 
@@ -528,6 +542,10 @@ class StableDiffusionPipeline(DiffusionPipeline):
             noise_pred = noise_pred_uncond + guidance_scale * (
                 noise_pred_text - noise_pred_uncond
             )
+
+            if guidance_rescale > 0.0:
+                noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=guidance_rescale)
+            
             return noise_pred
 
         sampler_args = self.get_sampler_extra_args_i2i(sigma_sched, sampler)
@@ -696,6 +714,10 @@ class StableDiffusionPipeline(DiffusionPipeline):
             noise_pred = noise_pred_uncond + guidance_scale * (
                 noise_pred_text - noise_pred_uncond
             )
+            
+            if guidance_rescale > 0.0:
+                noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=guidance_rescale)
+                
             return noise_pred
 
         extra_args = self.get_sampler_extra_args_t2i(