From 073cc3bda2e061c9f528f3bd814d9f1dd2a16dfa Mon Sep 17 00:00:00 2001
From: Karl Hudgell <karl.hudgell@bjss.com>
Date: Fri, 25 Apr 2025 13:16:02 +0100
Subject: [PATCH] update flux logic

---
 lib.py                 | 49 +++++++++++++++++++++++++-----------------
 user_config.cfg.sample |  5 ++++-
 2 files changed, 33 insertions(+), 21 deletions(-)

diff --git a/lib.py b/lib.py
index 400f3ed..89fbff7 100644
--- a/lib.py
+++ b/lib.py
@@ -249,27 +249,36 @@ def create_image(prompt: str | None = None) -> None:
     """Main function for generating images."""
     if prompt is None:
         prompt = create_prompt_on_openwebui(user_config["comfyui"]["prompt"])
-    if prompt:
-        logging.info(f"Generated prompt: {prompt}")  # Log generated prompt
-        save_prompt(prompt)
-        if user_config["comfyui"]["FLUX"]:
-            generate_image(
-                file_name="image",
-                comfy_prompt=prompt,
-                workflow_path="./FLUX.json",
-                prompt_node="Positive Prompt T5",
-                seed_node="Seed",
-                seed_param="seed",
-                save_node="CivitAI Image Saver",
-                save_param="filename",
-                model_node="CivitAI Image Saver",
-                model_param="modelname",
-            )
-        else:
-            generate_image("image", prompt)
-        print(f"Image generation started with prompt: {prompt}")
-    else:
+
+    if not prompt:
         logging.error("No prompt generated.")
+        return
+    save_prompt(prompt)
+
+    use_flux = user_config["comfyui"].get("USE_FLUX", False)
+    only_flux = user_config["comfyui"].get("ONLY_FLUX", False)
+
+    selected_workflow = "SDXL"
+    if use_flux:
+        selected_workflow = "FLUX" if only_flux else random.choice(["FLUX", "SDXL"])
+
+    if selected_workflow == "FLUX":
+        generate_image(
+            file_name="image",
+            comfy_prompt=prompt,
+            workflow_path="./FLUX.json",
+            prompt_node="Positive Prompt T5",
+            seed_node="Seed",
+            seed_param="seed",
+            save_node="CivitAI Image Saver",
+            save_param="filename",
+            model_node="CivitAI Image Saver",
+            model_param="modelname",
+        )
+    else:
+        generate_image("image", prompt)
+
+    logging.info(f"{selected_workflow} generation started with prompt: {prompt}")
 
 
 def get_prompt_from_png(path):
diff --git a/user_config.cfg.sample b/user_config.cfg.sample
index a024ac2..b4787c3 100644
--- a/user_config.cfg.sample
+++ b/user_config.cfg.sample
@@ -11,7 +11,10 @@ output_dir = ./output/
 prompt = "Generate a random detailed prompt for stable diffusion."
 width = 1568
 height = 672
-FLUX = False
+topics = 
+
+USE_FLUX = False
+ONLY_FLUX = False
 
 [comfyui:flux]
 models = flux1-dev-Q4_0.gguf,flux1-schnell-Q4_0.gguf