From d4d8ad406e7f6d6b4e13c297de3c9c3a728fb42d Mon Sep 17 00:00:00 2001
From: DenOfEquity <166248528+DenOfEquity@users.noreply.github.com>
Date: Thu, 24 Oct 2024 20:37:34 +0100
Subject: [PATCH] save (Hires) Distilled CFG to infotext (#2172)

changing hires model could lead to not writing
---
 modules/processing.py | 12 +++++++++---
 1 file changed, 9 insertions(+), 3 deletions(-)

diff --git a/modules/processing.py b/modules/processing.py
index f7f7c0cb..acc2718b 100644
--- a/modules/processing.py
+++ b/modules/processing.py
@@ -737,7 +737,9 @@ def create_infotext(p, all_prompts, all_seeds, all_subseeds, comments=None, iter
         "CFG scale": p.cfg_scale
     }
 
-    if p.sd_model.use_distilled_cfg_scale:
+    # if hires fix was used, p.firstpass_use_distilled_cfg_scale is appropriately set, otherwise it doesn't exist
+    firstpass_use_distilled_cfg_scale = getattr(p,'firstpass_use_distilled_cfg_scale', p.sd_model.use_distilled_cfg_scale)
+    if firstpass_use_distilled_cfg_scale:
         generation_params['Distilled CFG Scale'] = p.distilled_cfg_scale
 
     noise_source_type = get_noise_source_type()
@@ -1305,8 +1307,7 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
             self.extra_generation_params["Hires negative prompt"] = get_hr_negative_prompt
 
             self.extra_generation_params["Hires CFG Scale"] = self.hr_cfg
-            if shared.sd_model.use_distilled_cfg_scale:
-                self.extra_generation_params['Hires Distilled CFG Scale'] = self.hr_distilled_cfg
+            self.extra_generation_params["Hires Distilled CFG Scale"] = None  # set after potential hires model load
 
             self.extra_generation_params["Hires schedule type"] = None  # to be set in sd_samplers_kdiffusion.py
 
@@ -1407,6 +1408,8 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
 
             if self.hr_checkpoint_name and self.hr_checkpoint_name != 'Use same checkpoint':
                 if self.hr_checkpoint_name != fp_checkpoint:
+                    self.firstpass_use_distilled_cfg_scale = self.sd_model.use_distilled_cfg_scale
+
                     main_entry.checkpoint_change(self.hr_checkpoint_name, save=False, refresh=False)
                     reload = True
             
@@ -1418,6 +1421,9 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
                     main_entry.modules_change(fp_additional_modules, save=False, refresh=False)
                     main_entry.checkpoint_change(fp_checkpoint, save=False)
 
+            if self.sd_model.use_distilled_cfg_scale:
+                self.extra_generation_params['Hires Distilled CFG Scale'] = self.hr_distilled_cfg
+
         return self.sample_hr_pass(samples, decoded_samples, seeds, subseeds, subseed_strength, prompts)
 
     def sample_hr_pass(self, samples, decoded_samples, seeds, subseeds, subseed_strength, prompts):