From f620f55e564ac146702b62228a069e4a6e943d04 Mon Sep 17 00:00:00 2001
From: catboxanon <122327233+catboxanon@users.noreply.github.com>
Date: Sat, 19 Oct 2024 07:47:56 -0400
Subject: [PATCH] Fallback to already detected prediction type if none
 applicable found

---
 backend/loader.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/backend/loader.py b/backend/loader.py
index a0460493..4f1f73c2 100644
--- a/backend/loader.py
+++ b/backend/loader.py
@@ -294,7 +294,7 @@ def forge_loader(sd, additional_state_dicts=None):
         'EDM': 'edm',
     }
     if 'scheduler' in huggingface_components and hasattr(huggingface_components['scheduler'], 'config') and 'prediction_type' in huggingface_components['scheduler'].config:
-        huggingface_components['scheduler'].config.prediction_type = prediction_types.get(estimated_config.model_type.name, 'epsilon')
+        huggingface_components['scheduler'].config.prediction_type = prediction_types.get(estimated_config.model_type.name, huggingface_components['scheduler'].config.prediction_type)
 
     for M in possible_models:
         if any(isinstance(estimated_config, x) for x in M.matched_guesses):