68 lines
3.1 KiB
Python
68 lines
3.1 KiB
Python
# multi_controlnet_loader.py
|
|
import os
|
|
import folder_paths
|
|
from comfy import model_management
|
|
import comfy.controlnet
|
|
import comfy.utils
|
|
|
|
class MultiControlNetLoader:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
# ControlNetモデルリストを取得
|
|
controlnet_list = folder_paths.get_filename_list("controlnet")
|
|
|
|
return {
|
|
"required": {
|
|
"checkpoint": ("MODEL", {"default": None}),
|
|
"sd_cn_name": (controlnet_list, {"default": controlnet_list[0] if controlnet_list else ""}),
|
|
"xl_cn_name": (controlnet_list, {"default": controlnet_list[0] if controlnet_list else ""}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("CONTROL_NET",)
|
|
FUNCTION = "switch_and_load"
|
|
CATEGORY = "ControlNet"
|
|
|
|
def switch_and_load(self, checkpoint, sd_cn_name, xl_cn_name):
|
|
"""
|
|
checkpoint 名に "sdxl" を含む場合は SDXL 用モデルを、それ以外は SD 用モデルをロードします。
|
|
"""
|
|
try:
|
|
# checkpoint 名で SD or SDXL を判定
|
|
# checkpointがModelPatcherオブジェクトの場合の処理
|
|
if hasattr(checkpoint, 'model') and hasattr(checkpoint.model, 'model_config'):
|
|
model_config = checkpoint.model.model_config
|
|
# SDXL判定 - model_configからSDXLかどうかを判定
|
|
is_sdxl = (
|
|
hasattr(model_config, 'unet_config') and
|
|
model_config.unet_config.get('model_channels', 320) == 320 and
|
|
model_config.unet_config.get('adm_in_channels', None) is not None
|
|
) or str(model_config).lower().find('sdxl') != -1
|
|
else:
|
|
# フォールバック: 文字列として名前をチェック
|
|
name = str(checkpoint).lower()
|
|
is_sdxl = "sdxl" in name
|
|
|
|
# 使用するControlNetモデル名を決定
|
|
target_name = xl_cn_name if is_sdxl else sd_cn_name
|
|
|
|
# ControlNetモデルのパスを取得
|
|
controlnet_path = folder_paths.get_full_path("controlnet", target_name)
|
|
if controlnet_path is None:
|
|
raise FileNotFoundError(f"ControlNet model '{target_name}' not found in controlnet folders")
|
|
|
|
# ControlNetモデルをロード
|
|
controlnet = comfy.controlnet.load_controlnet(controlnet_path)
|
|
|
|
print(f"MultiControlNetLoader: Loaded {'SDXL' if is_sdxl else 'SD'} ControlNet model: {target_name}")
|
|
|
|
return (controlnet,)
|
|
|
|
except Exception as e:
|
|
print(f"Error in MultiControlNetLoader: {str(e)}")
|
|
print(f"Attempted to load: {'SDXL' if 'is_sdxl' in locals() and is_sdxl else 'SD'} model with name: {target_name if 'target_name' in locals() else 'Unknown'}")
|
|
raise e
|
|
|
|
# ノードクラスマッピング
|
|
NODE_CLASS_MAPPINGS = {"MultiControlNetLoader": MultiControlNetLoader}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {"MultiControlNetLoader": "Multi ControlNet Loader"} |