From a78ce0a3ca4521a5661042c47e3dd56eee7a1eeb Mon Sep 17 00:00:00 2001
From: cool-bigdogs-tshirt
 <131823432+cool-bigdogs-tshirt@users.noreply.github.com>
Date: Thu, 27 Apr 2023 12:04:14 -0500
Subject: [PATCH] xyz override for latent upscaler fallback

---
 modules/processing.py | 5 +++--
 scripts/xyz_grid.py   | 9 +++++++++
 2 files changed, 12 insertions(+), 2 deletions(-)

diff --git a/modules/processing.py b/modules/processing.py
index c3ea4b2b8..a83d9fa12 100644
--- a/modules/processing.py
+++ b/modules/processing.py
@@ -970,9 +970,10 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
         shared.state.nextjob()
 
         img2img_sampler_name = self.sampler_name
-        if self.sampler_name in ['PLMS']:
+        force_latent_upscaler = shared.opts.data.get('xyz_fallback_sampler')
+        if self.sampler_name in ['PLMS'] or force_latent_upscaler is not None:
             # PLMS does not support img2img, use fallback instead
-            img2img_sampler_name = shared.opts.fallback_sampler
+            img2img_sampler_name = force_latent_upscaler or shared.opts.fallback_sampler
         self.sampler = sd_samplers.create_sampler(img2img_sampler_name, self.sd_model)
 
         samples = samples[:, :, self.truncate_y//2:samples.shape[2]-(self.truncate_y+1)//2, self.truncate_x//2:samples.shape[3]-(self.truncate_x+1)//2]
diff --git a/scripts/xyz_grid.py b/scripts/xyz_grid.py
index 52ae1c6e1..9a5a67241 100644
--- a/scripts/xyz_grid.py
+++ b/scripts/xyz_grid.py
@@ -128,6 +128,14 @@ def apply_styles(p: StableDiffusionProcessingTxt2Img, x: str, _):
     p.styles.extend(x.split(','))
 
 
+def apply_fallback(p, x, xs):
+    sampler_name = sd_samplers.samplers_map.get(x.lower(), None)
+    if sampler_name is None:
+        raise RuntimeError(f"Unknown sampler: {x}")
+
+    opts.data["xyz_fallback_sampler"] = sampler_name
+
+
 def apply_uni_pc_order(p, x, xs):
     opts.data["uni_pc_order"] = min(x, p.steps - 1)
 
@@ -220,6 +228,7 @@ axis_options = [
     AxisOption("Clip skip", int, apply_clip_skip),
     AxisOption("Denoising", float, apply_field("denoising_strength")),
     AxisOptionTxt2Img("Hires upscaler", str, apply_field("hr_upscaler"), choices=lambda: [*shared.latent_upscale_modes, *[x.name for x in shared.sd_upscalers]]),
+    AxisOptionTxt2Img("Fallback latent upscaler sampler", str, apply_fallback, format_value=format_value, confirm=confirm_samplers, choices=lambda: [x.name for x in sd_samplers.samplers]),
     AxisOptionImg2Img("Cond. Image Mask Weight", float, apply_field("inpainting_mask_weight")),
     AxisOption("VAE", str, apply_vae, cost=0.7, choices=lambda: list(sd_vae.vae_dict)),
     AxisOption("Styles", str, apply_styles, choices=lambda: list(shared.prompt_styles.styles)),