From c20ba34b3ea433bb8871e955db42dff258d2c850 Mon Sep 17 00:00:00 2001
From: Vladimir Mandic <mandic00@live.com>
Date: Fri, 20 Sep 2024 12:56:30 -0400
Subject: [PATCH] invalidate ui selection on load fail

---
 modules/model_flux.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/modules/model_flux.py b/modules/model_flux.py
index bb69683df..16f6bb8bb 100644
--- a/modules/model_flux.py
+++ b/modules/model_flux.py
@@ -175,6 +175,7 @@ def load_transformer(file_path): # triggered by opts.sd_unet change
         transformer = diffusers.FluxTransformer2DModel.from_single_file(file_path, **diffusers_load_config)
     if transformer is None:
         shared.log.error('Failed to load UNet model')
+        shared.opts.sd_unet = 'None'
     return transformer
 
 
@@ -196,9 +197,9 @@ def load_flux(checkpoint_info, diffusers_load_config): # triggered by opts.sd_ch
             from modules import sd_unet
             _transformer = load_transformer(sd_unet.unet_dict[shared.opts.sd_unet])
             if _transformer is not None:
-                sd_unet.loaded_unet = shared.opts.sd_unet
                 transformer = _transformer
             else:
+                shared.opts.sd_unet = 'None'
                 sd_unet.failed_unet.append(shared.opts.sd_unet)
         except Exception as e:
             shared.log.error(f"Loading FLUX: Failed to load UNet: {e}")
@@ -213,6 +214,8 @@ def load_flux(checkpoint_info, diffusers_load_config): # triggered by opts.sd_ch
             _text_encoder_2 = load_t5(name=shared.opts.sd_text_encoder, cache_dir=shared.opts.diffusers_dir)
             if _text_encoder_2 is not None:
                 text_encoder_2 = _text_encoder_2
+            else:
+                shared.opts.sd_text_encoder = 'None'
         except Exception as e:
             shared.log.error(f"Loading FLUX: Failed to load T5: {e}")
             shared.opts.sd_text_encoder = 'None'