diff --git a/comfyui/custom_nodes/zone_lora/__init__.py b/comfyui/custom_nodes/zone_lora/__init__.py index b501b057..9a7ea380 100644 --- a/comfyui/custom_nodes/zone_lora/__init__.py +++ b/comfyui/custom_nodes/zone_lora/__init__.py @@ -1,4 +1,4 @@ -from comfy_api.latest import ComfyExtension +from comfy_api.latest import ComfyExtension, io from .cleanup_node import ZoneCleanupTrainingRun from .dataset_node import ZoneLoadTrainDataset @@ -11,16 +11,13 @@ install_all() -NODE_CLASS_MAPPINGS = { - 'ZoneIPAdapterFace': ZoneIPAdapterFace, -} - class ZoneLoraExtension(ComfyExtension): - async def get_node_list(self): + async def get_node_list(self) -> list[type[io.ComfyNode]]: return [ ZoneLoadTrainDataset, ZoneCleanupTrainingRun, + ZoneIPAdapterFace, ZoneLoadTrainFolder, ZoneProbeGradient, ZoneProbeLoss, diff --git a/comfyui/custom_nodes/zone_lora/face_node.py b/comfyui/custom_nodes/zone_lora/face_node.py index ac60dd7d..30aa63dc 100644 --- a/comfyui/custom_nodes/zone_lora/face_node.py +++ b/comfyui/custom_nodes/zone_lora/face_node.py @@ -11,6 +11,7 @@ from comfy.ldm.modules.attention import optimized_attention import comfy.model_management as model_management import comfy.utils +from comfy_api.latest import io PLUS_FACE = 'ip-adapter-plus-face_sdxl_vit-h.safetensors' @@ -267,32 +268,30 @@ def patch_sdxl(model, ipa, cond, uncond, weight: float): number += 1 -class ZoneIPAdapterFace: +class ZoneIPAdapterFace(io.ComfyNode): @classmethod - def INPUT_TYPES(cls): - names = folder_paths.get_filename_list('ipadapter') - if not names: - names = [PLUS_FACE] - return { - 'required': { - 'model': ('MODEL',), - 'clip_vision': ('CLIP_VISION',), - 'image': ('IMAGE',), - 'ipadapter_name': (names,), - 'strength': ( - 'FLOAT', - {'default': 0.4, 'min': 0.0, 'max': 2.0, 'step': 0.01}, - ), - } - } - - RETURN_TYPES = ('MODEL',) - FUNCTION = 'apply' - CATEGORY = 'zone' - - def apply(self, model, clip_vision, image, ipadapter_name, strength): + def define_schema(cls) -> io.Schema: + names = folder_paths.get_filename_list('ipadapter') or [PLUS_FACE] + return io.Schema( + node_id='ZoneIPAdapterFace', + display_name='Zone IP-Adapter Face', + category='zone', + inputs=[ + io.Model.Input('model'), + io.ClipVision.Input('clip_vision'), + io.Image.Input('image'), + io.Combo.Input('ipadapter_name', options=names), + io.Float.Input('strength', default=0.4, min=0.0, max=2.0, step=0.01), + ], + outputs=[io.Model.Output()], + ) + + @classmethod + def execute( + cls, model, clip_vision, image, ipadapter_name: str, strength: float + ) -> io.NodeOutput: if strength == 0: - return (model,) + return io.NodeOutput(model) path = folder_paths.get_full_path_or_raise('ipadapter', ipadapter_name) loaded = load_plus_face(path) if 'latents' not in loaded['image_proj']: @@ -314,4 +313,4 @@ def apply(self, model, clip_vision, image, ipadapter_name, strength): uncond = uncond.to(device, dtype=dtype) work = model.clone() patch_sdxl(work, ipa, cond, uncond, float(strength)) - return (work,) + return io.NodeOutput(work) diff --git a/comfyui/tests/test_graph_contract.py b/comfyui/tests/test_graph_contract.py index ac249179..066c0257 100644 --- a/comfyui/tests/test_graph_contract.py +++ b/comfyui/tests/test_graph_contract.py @@ -3,6 +3,7 @@ import asyncio import json import os +import re import subprocess import sys import tempfile @@ -19,6 +20,7 @@ ZONE_NODES = COMFYUI / 'custom_nodes' / 'zone_lora' TRAIN_SERVER = ROOT / 'runner/zone_comfy/src/train.rs' QUALITY_SERVER = ROOT / 'runner/zone_comfy/src/quality.rs' +PEOPLE_SERVER = ROOT / 'runner/zone_comfy/src/people.rs' if not INSTALL.is_dir(): raise RuntimeError('COMFYUI_INSTALL_DIR must point to the pinned ComfyUI checkout') @@ -165,6 +167,53 @@ def rust_graphs() -> dict[str, dict]: } +def resolve(value: object, constants: dict[str, str]) -> object: + if isinstance(value, Expression): + return constants.get(value.text, value) + if isinstance(value, list): + return [resolve(item, constants) for item in value] + if isinstance(value, dict): + return {name: resolve(item, constants) for name, item in value.items()} + return value + + +def people_strings() -> dict[str, str]: + return dict(re.findall(r'const (\w+): &str = "([^"]*)";', PEOPLE_SERVER.read_text())) + + +def people_numbers() -> dict[str, float]: + return { + name: float(value) + for name, value in re.findall(r'const (\w+): f64 = ([0-9.]+);', PEOPLE_SERVER.read_text()) + } + + +def rust_face_graph() -> dict: + source = PEOPLE_SERVER.read_text() + constants = people_strings() + start = source.find('fn inject_ipadapter(') + if start < 0: + raise AssertionError(f'{PEOPLE_SERVER}: inject_ipadapter is missing') + end = source.find('\n}\n', start) + if end <= start: + raise AssertionError(f'{PEOPLE_SERVER}: inject_ipadapter has no closing brace') + graph = { + '4': { + 'class_type': 'CheckpointLoaderSimple', + 'inputs': {'ckpt_name': 'model.safetensors'}, + } + } + for insert in re.finditer(r'object\.insert\(\s*(\w+)\.into\(\),\s*json!\(', source[start:end]): + key = insert.group(1) + if key not in constants: + raise AssertionError(f'{PEOPLE_SERVER}: node key {key} is not a &str constant') + node = RustLiteral(source, start + insert.end(), PEOPLE_SERVER).object() + graph[constants[key]] = resolve(node, constants) + face = graph[constants['IPADAPTER_NODE']] + face['inputs']['model'] = ['4', 0] + return graph + + def flux() -> train_lora.TrainingModel: return train_lora.TrainingModel(architecture='flux', checkpoint='flux.safetensors') @@ -371,6 +420,7 @@ def test_repository_extension_is_part_of_the_loaded_registry(self) -> None: 'ZoneCleanupTrainingRun', 'ZoneIPAdapterFace', 'ZoneLoadTrainDataset', + 'ZoneLoadTrainFolder', 'ZoneProbeGradient', 'ZoneProbeLoss', 'ZoneStageTrainingArtifact', @@ -378,6 +428,51 @@ def test_repository_extension_is_part_of_the_loaded_registry(self) -> None: } self.assertLessEqual(expected, set(self.nodes.NODE_CLASS_MAPPINGS)) + def test_rust_face_graph_matches_the_pinned_registry(self) -> None: + graph = rust_face_graph() + self.assertEqual( + sorted(node['class_type'] for node in graph.values()), + ['CLIPVisionLoader', 'CheckpointLoaderSimple', 'LoadImage', 'ZoneIPAdapterFace'], + ) + self.validate_graph('Rust face', graph) + + def test_rust_face_constants_fit_the_face_node(self) -> None: + node_class = self.nodes.NODE_CLASS_MAPPINGS['ZoneIPAdapterFace'] + inputs, _, _ = schema(node_class) + bounds = inputs['strength'][1] + strength = people_numbers()['IPADAPTER_STRENGTH'] + self.assertGreaterEqual(strength, bounds['min'], f'{PEOPLE_SERVER}: IPADAPTER_STRENGTH') + self.assertLessEqual(strength, bounds['max'], f'{PEOPLE_SERVER}: IPADAPTER_STRENGTH') + self.assertEqual( + people_strings()['IPADAPTER_FILE'], + sys.modules[node_class.__module__].PLUS_FACE, + f'{PEOPLE_SERVER}: IPADAPTER_FILE', + ) + + def test_face_node_at_zero_strength_returns_the_model_untouched(self) -> None: + from execution import _async_map_node_over_list + + node_class = self.nodes.NODE_CLASS_MAPPINGS['ZoneIPAdapterFace'] + model = object() + results = asyncio.run( + _async_map_node_over_list( + 'contract', + 'face', + node_class, + { + 'model': [model], + 'clip_vision': [None], + 'image': [None], + 'ipadapter_name': ['missing.safetensors'], + 'strength': [0.0], + }, + node_class.FUNCTION, + v3_data={}, + ) + ) + self.assertEqual(len(results), 1) + self.assertIs(results[0].result[0], model) + def test_sdxl_people_graphs_match_the_pinned_registry(self) -> None: for name in ( 'sdxl-api.json', diff --git a/comfyui/tests/test_setup_macos.py b/comfyui/tests/test_setup_macos.py index e0908ff3..9a826c21 100644 --- a/comfyui/tests/test_setup_macos.py +++ b/comfyui/tests/test_setup_macos.py @@ -202,6 +202,20 @@ def test_install_trainer_inherits_comfy_site_packages(self) -> None: text=True, ) self.assertNotEqual(mlx.returncode, 0, mlx.stdout + mlx.stderr) + for trainer in (".venv-train", ".venv-train-llm"): + trainer_site = subprocess.check_output( + [ + str(install / trainer / "bin" / "python"), + "-c", + "import sysconfig; print(sysconfig.get_path('purelib'))", + ], + text=True, + ).strip() + with self.subTest(venv=trainer): + self.assertFalse( + Path(trainer_site, "pip").exists(), + "a skipped pip install still bootstrapped pip into the trainer venv", + ) plist = (home / "Library/LaunchAgents/ai.zone.train.plist").read_text( encoding="utf-8" ) diff --git a/scripts/setup-comfyui-macos.sh b/scripts/setup-comfyui-macos.sh index e6007a2c..f541f790 100755 --- a/scripts/setup-comfyui-macos.sh +++ b/scripts/setup-comfyui-macos.sh @@ -160,7 +160,10 @@ install_trainer() { COMFY_PYTHON="$INSTALL_DIR/.venv/bin/python" TRAIN_VENV="$INSTALL_DIR/.venv-train" if [ ! -x "$TRAIN_VENV/bin/python" ]; then - "$COMFY_PYTHON" -m venv "$TRAIN_VENV" + "$COMFY_PYTHON" -m venv --without-pip "$TRAIN_VENV" + fi + if [ "${ZONE_TRAIN_SKIP_PIP:-}" != "1" ]; then + "$TRAIN_VENV/bin/python" -m ensurepip --upgrade fi # Nested venvs do not inherit the parent venv via --system-site-packages; # that flag only sees Homebrew's site. A .pth keeps Comfy's torch. @@ -181,11 +184,13 @@ install_trainer() { diffusers peft accelerate transformers safetensors pillow prodigyopt fi LLM_VENV="$INSTALL_DIR/.venv-train-llm" - if [ ! -x "$LLM_VENV/bin/python" ]; then - "$COMFY_PYTHON" -m venv "$LLM_VENV" + LLM_PYTHON="$LLM_VENV/bin/python" + if [ ! -x "$LLM_PYTHON" ]; then + "$COMFY_PYTHON" -m venv --without-pip "$LLM_VENV" fi if [ "${ZONE_TRAIN_SKIP_PIP:-}" != "1" ]; then - "$LLM_VENV/bin/python" -m pip install --disable-pip-version-check "mlx-lm[train]" + "$LLM_PYTHON" -m ensurepip --upgrade + "$LLM_PYTHON" -m pip install --disable-pip-version-check "mlx-lm[train]" fi mkdir -p "$HOME/Library/LaunchAgents" "$HOME/Library/Logs" PLIST="$HOME/Library/LaunchAgents/ai.zone.train.plist"