diff --git a/ai_frame_image_server.py b/ai_frame_image_server.py index ddbf030..571b2bf 100644 --- a/ai_frame_image_server.py +++ b/ai_frame_image_server.py @@ -100,14 +100,14 @@ def create() -> str: str: Redirect to the main page or a JSON response. """ prompt = request.form.get("prompt") if request.method == "POST" else None - model = request.form.get("model") if request.method == "POST" else None + model = request.form.get("model") if request.method == "POST" else "Random" if prompt is None: prompt = create_prompt_on_openwebui(user_config["comfyui"]["prompt"]) def create_image_in_background(): - create_image(prompt) + create_image(prompt, model) threading.Thread(target=create_image_in_background).start() return render_template('image_queued.html', prompt=prompt) diff --git a/libs/comfyui.py b/libs/comfyui.py index 5f97e78..32ea369 100644 --- a/libs/comfyui.py +++ b/libs/comfyui.py @@ -71,6 +71,7 @@ def generate_image( save_param: str = "filename_prefix", model_node: Optional[str] = "Load Checkpoint", model_param: Optional[str] = "ckpt_name", + model: Optional[str] = "None", ) -> None: """Generates an image using the Comfy API with configurable workflow settings.""" try: @@ -100,21 +101,7 @@ def generate_image( user_config["comfyui"]["height"], ) - # Conditionally set model if node and param are provided - if model_node and model_param: - if "FLUX" in workflow_path: - valid_models = user_config["comfyui:flux"]["models"].split(",") - else: - available_model_list = user_config["comfyui"]["models"].split(",") - valid_models = list( - set(get_available_models()) & set(available_model_list) - ) - - if not valid_models: - raise Exception("No valid models available.") - - model = random.choice(valid_models) - wf.set_node_param(model_node, model_param, model) + wf.set_node_param(model_node, model_param, model) # Generate image logging.debug(f"Generating image: {file_name}") @@ -136,24 +123,29 @@ def generate_image( raise -def create_image(prompt: str | None = None) -> None: - """Main function for generating images.""" +def create_image(prompt: str | None = None, model: str = "Random") -> None: + """Generate an image with a chosen workflow (Random, FLUX*, or SDXL*).""" + if prompt is None: prompt = create_prompt_on_openwebui(user_config["comfyui"]["prompt"]) if not prompt: logging.error("No prompt generated.") return + save_prompt(prompt) - - use_flux = json.loads((user_config["comfyui"].get("FLUX", False)).lower()) - only_flux = json.loads((user_config["comfyui"].get("ONLY_FLUX", False)).lower()) - - selected_workflow = "SDXL" - if use_flux: - selected_workflow = "FLUX" if only_flux else random.choice(["FLUX", "SDXL"]) - + use_flux = json.loads(user_config["comfyui"].get("FLUX", "false").lower()) + only_flux = json.loads(user_config["comfyui"].get("ONLY_FLUX", "false").lower()) + if model == "Random": + selected_workflow = "FLUX" if (use_flux and (only_flux or random.choice([True, False]))) else "SDXL" + elif "flux" in model.lower(): + selected_workflow = "FLUX" + else: + selected_workflow = "SDXL" if selected_workflow == "FLUX": + if model == "Random": + valid_models = user_config["comfyui:flux"]["models"].split(",") + model = random.choice(valid_models) generate_image( file_name="image", comfy_prompt=prompt, @@ -165,9 +157,14 @@ def create_image(prompt: str | None = None) -> None: save_param="filename", model_node="CivitAI Image Saver", model_param="modelname", + model=model ) - else: - generate_image("image", prompt) + else: # SDXL + if model == "Random": + available_model_list = user_config["comfyui"]["models"].split(",") + valid_models = list(set(get_available_models()) & set(available_model_list)) + model = random.choice(valid_models) + generate_image("image", comfy_prompt=prompt, model=model) logging.info(f"{selected_workflow} generation started with prompt: {prompt}")