From a189f3e53e81dd2ffac100044129b8f648a7bf25 Mon Sep 17 00:00:00 2001
From: layerdiffusion <19834515+lllyasviel@users.noreply.github.com>
Date: Thu, 8 Aug 2024 01:22:27 -0700
Subject: [PATCH] distilled cfg scale info text

---
 backend/diffusion_engine/base.py |  1 +
 backend/diffusion_engine/flux.py |  2 ++
 modules/processing.py            | 10 ++++++++--
 modules/ui.py                    |  4 ++--
 4 files changed, 13 insertions(+), 4 deletions(-)

diff --git a/backend/diffusion_engine/base.py b/backend/diffusion_engine/base.py
index 1a697232..a7d72ef2 100644
--- a/backend/diffusion_engine/base.py
+++ b/backend/diffusion_engine/base.py
@@ -32,6 +32,7 @@ class ForgeDiffusionEngine:
         self.tiling_enabled = False
 
         self.first_stage_model = None  # set this so that you can change VAE in UI
+        self.use_distilled_cfg_scale = False
 
         # WebUI Dirty Legacy
         self.is_sd1 = False
diff --git a/backend/diffusion_engine/flux.py b/backend/diffusion_engine/flux.py
index 46845c9c..010022ad 100644
--- a/backend/diffusion_engine/flux.py
+++ b/backend/diffusion_engine/flux.py
@@ -66,6 +66,8 @@ class Flux(ForgeDiffusionEngine):
         self.forge_objects_original = self.forge_objects.shallow_copy()
         self.forge_objects_after_applying_lora = self.forge_objects.shallow_copy()
 
+        self.use_distilled_cfg_scale = True
+
         # WebUI Legacy
         self.first_stage_model = vae.first_stage_model
 
diff --git a/modules/processing.py b/modules/processing.py
index 4d321bf7..2a156d97 100644
--- a/modules/processing.py
+++ b/modules/processing.py
@@ -722,7 +722,13 @@ def create_infotext(p, all_prompts, all_seeds, all_subseeds, comments=None, iter
         "Steps": p.steps,
         "Sampler": p.sampler_name,
         "Schedule type": p.scheduler,
-        "CFG scale": p.cfg_scale,
+        "CFG scale": p.cfg_scale
+    }
+
+    if p.sd_model.use_distilled_cfg_scale:
+        generation_params['Distilled CFG Scale'] = p.distilled_cfg_scale
+
+    generation_params.update({
         "Image CFG scale": getattr(p, 'image_cfg_scale', None),
         "Seed": p.all_seeds[0] if use_main_prompt else all_seeds[index],
         "Face restoration": opts.face_restoration_model if p.restore_faces else None,
@@ -748,7 +754,7 @@ def create_infotext(p, all_prompts, all_seeds, all_subseeds, comments=None, iter
         **p.extra_generation_params,
         "Version": program_version() if opts.add_version_to_infotext else None,
         "User": p.user if opts.add_user_name_to_info else None,
-    }
+    })
 
     for key, value in generation_params.items():
         try:
diff --git a/modules/ui.py b/modules/ui.py
index 13b3d245..49e49f4d 100644
--- a/modules/ui.py
+++ b/modules/ui.py
@@ -460,7 +460,7 @@ def create_ui():
                 PasteField(toprow.prompt, "Prompt", api="prompt"),
                 PasteField(toprow.negative_prompt, "Negative prompt", api="negative_prompt"),
                 PasteField(cfg_scale, "CFG scale", api="cfg_scale"),
-                PasteField(distilled_cfg_scale, "Distilled CFG scale", api="distilled_cfg_scale"),
+                PasteField(distilled_cfg_scale, "Distilled CFG Scale", api="distilled_cfg_scale"),
                 PasteField(width, "Size-1", api="width"),
                 PasteField(height, "Size-2", api="height"),
                 PasteField(batch_size, "Batch size", api="batch_size"),
@@ -828,7 +828,7 @@ def create_ui():
                 (toprow.prompt, "Prompt"),
                 (toprow.negative_prompt, "Negative prompt"),
                 (cfg_scale, "CFG scale"),
-                (distilled_cfg_scale, "Distilled CFG scale"),
+                (distilled_cfg_scale, "Distilled CFG Scale"),
                 (image_cfg_scale, "Image CFG scale"),
                 (width, "Size-1"),
                 (height, "Size-2"),