diff --git a/libs/comfyui.py b/libs/comfyui.py
index 144862a..a33992c 100644
--- a/libs/comfyui.py
+++ b/libs/comfyui.py
@@ -102,7 +102,7 @@ def generate_image(
 
         # Conditionally set model if node and param are provided
         if model_node and model_param:
-            if user_config["comfyui"].get("FLUX"):
+            if "FLUX" in workflow_path:
                 valid_models = user_config["comfyui:flux"]["models"].split(",")
             else:
                 available_model_list = user_config["comfyui"]["models"].split(",")