From 00fa813601f3ff77c954bdc439eca01f90818e57 Mon Sep 17 00:00:00 2001
From: Sucial <75163782+SUC-DriverOld@users.noreply.github.com>
Date: Tue, 17 Sep 2024 01:16:40 +0800
Subject: [PATCH] small fixs

---
 tools/webUI_for_clouds/webUI_for_clouds.py | 7 +++++--
 webUI.py                                   | 7 +++++--
 2 files changed, 10 insertions(+), 4 deletions(-)

diff --git a/tools/webUI_for_clouds/webUI_for_clouds.py b/tools/webUI_for_clouds/webUI_for_clouds.py
index 28c8224..6c67e94 100644
--- a/tools/webUI_for_clouds/webUI_for_clouds.py
+++ b/tools/webUI_for_clouds/webUI_for_clouds.py
@@ -507,14 +507,17 @@ def run_inference_flow(input_folder, store_dir, preset_name, force_cpu, output_f
             break
         if i == 0:
             input_to_use = input_folder
-        elif i < len(model_list.keys()) - 1 and i > 0:
+        if i < len(model_list.keys()) - 1 and i > 0:
             if input_to_use != input_folder:
                 shutil.rmtree(input_to_use)
             input_to_use = tmp_store_dir
             tmp_store_dir = f"{TEMP_PATH}/inferflow_step{i+1}_output"
-        elif i == len(model_list.keys()) - 1:
+        if i == len(model_list.keys()) - 1:
             input_to_use = tmp_store_dir
             tmp_store_dir = store_dir
+        if len(model_list.keys()) == 1:
+            input_to_use = input_folder
+            tmp_store_dir = store_dir
         model_name = model_list[step]["model_name"]
         print(f"===== Step {i+1}: Running inference using {model_name} =====")
         if model_list[step]["model_type"] == "UVR_VR_Models":
diff --git a/webUI.py b/webUI.py
index 32718d5..abe6666 100644
--- a/webUI.py
+++ b/webUI.py
@@ -720,14 +720,17 @@ def run_inference_flow(input_folder, store_dir, preset_name, force_cpu, output_f
             break
         if i == 0:
             input_to_use = input_folder
-        elif i < len(model_list.keys()) - 1 and i > 0:
+        if i < len(model_list.keys()) - 1 and i > 0:
             if input_to_use != input_folder:
                 shutil.rmtree(input_to_use)
             input_to_use = tmp_store_dir
             tmp_store_dir = f"{TEMP_PATH}/inferflow_step{i+1}_output"
-        elif i == len(model_list.keys()) - 1:
+        if i == len(model_list.keys()) - 1:
             input_to_use = tmp_store_dir
             tmp_store_dir = store_dir
+        if len(model_list.keys()) == 1:
+            input_to_use = input_folder
+            tmp_store_dir = store_dir
         model_name = model_list[step]["model_name"]
         console.print(f"[yellow]Step {i+1}: Running inference using {model_name}", style="yellow", justify='center')
         if model_list[step]["model_type"] == "UVR_VR_Models":