Karl/ai-frame-image-server
1
0
Fork 0
mirror of https://github.com/karl0ss/ai_image_frame_server.git synced 2025-09-06 23:03:15 +01:00
Code Issues Packages Projects Releases Wiki Activity

**refactor(comfyui.py): improve model selection logic**

Browse Source 
Refactor `get_available_models` to handle multiple models and improve error handling. Adjust `select_model` to support configurable workflows and fallbacks.
This commit is contained in:
Karl 2025-09-01 13:19:28 +01:00
parent 12af531718
commit 506dece377
1 changed files with 35 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

39
libs/comfyui.py
View File

@ -32,9 +32,17 @@ def get_available_models() -> list:
response = requests.get(url) response = requests.get(url)
if response.status_code == 200: if response.status_code == 200:
data = response.json() data = response.json()
general = data.get("CheckpointLoaderSimple", {}).get("input", {}).get("required", {}).get("ckpt_name", [])[0] # Get SDXL models from CheckpointLoaderSimple
flux = data.get("UnetLoaderGGUF", {}).get("input", {}).get("required", {}).get("unet_name", [])[0] general = data.get("CheckpointLoaderSimple", {}).get("input", {}).get("required", {}).get("ckpt_name", [[]])[0]
return general + flux # Get FLUX models from UnetLoaderGGUF
flux = data.get("UnetLoaderGGUF", {}).get("input", {}).get("required", {}).get("unet_name", [[]])[0]
# Combine both lists, handling cases where one might be missing
all_models = []
if isinstance(general, list):
all_models.extend(general)
if isinstance(flux, list):
all_models.extend(flux)
return all_models
else: else:
print(f"Failed to fetch models: {response.status_code}") print(f"Failed to fetch models: {response.status_code}")
return [] return []
@ -125,9 +133,25 @@ def select_model(model: str) -> tuple[str, str]:
use_qwen = json.loads(user_config["comfyui"].get("Qwen", "false").lower()) use_qwen = json.loads(user_config["comfyui"].get("Qwen", "false").lower())
if model == "Random Image Model": if model == "Random Image Model":
selected_workflow = "FLUX" if (use_flux and (only_flux or random.choice([True, False]))) else "SDXL" # Create a list of available workflows based on configuration
available_workflows = []
if not only_flux:
available_workflows.append("SDXL")
if use_flux:
available_workflows.append("FLUX")
if use_qwen:
available_workflows.append("Qwen")
# If no workflows are available, default to SDXL
if not available_workflows:
available_workflows.append("SDXL")
# Randomly select a workflow
selected_workflow = random.choice(available_workflows)
elif "flux" in model.lower(): elif "flux" in model.lower():
selected_workflow = "FLUX" selected_workflow = "FLUX"
elif "qwen" in model.lower():
selected_workflow = "Qwen"
else: else:
selected_workflow = "SDXL" selected_workflow = "SDXL"
@ -139,6 +163,13 @@ def select_model(model: str) -> tuple[str, str]:
else: # SDXL else: # SDXL
available_model_list = user_config["comfyui"]["models"].split(",") available_model_list = user_config["comfyui"]["models"].split(",")
valid_models = list(set(get_available_models()) & set(available_model_list)) valid_models = list(set(get_available_models()) & set(available_model_list))
# If no valid models found, fall back to configured models
if not valid_models:
valid_models = available_model_list
# Ensure we have at least one model to choose from
if not valid_models:
# Fallback to a default model
valid_models = ["zavychromaxl_v100.safetensors"]
model = random.choice(valid_models) model = random.choice(valid_models)
return selected_workflow, model return selected_workflow, model