small fixs

SUC-DriverOld · Sep 16, 2024 · 00fa813 · 00fa813
1 parent eaf6744
commit 00fa813
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 4 deletions.
diff --git 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
@@ -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":