diff --git a/__init__.py b/__init__.py index 65be3d4..0094628 100644 --- a/__init__.py +++ b/__init__.py @@ -2,21 +2,22 @@ import os import sys # Ensure this custom node root is on sys.path (ComfyUI loads modules by path, not package) -_MOLTBOT_ROOT = os.path.dirname(os.path.abspath(__file__)) -if _MOLTBOT_ROOT not in sys.path: - sys.path.insert(0, _MOLTBOT_ROOT) +_OPENCLAW_ROOT = os.path.dirname(os.path.abspath(__file__)) +if _OPENCLAW_ROOT not in sys.path: + sys.path.insert(0, _OPENCLAW_ROOT) if __package__: - from .nodes.batch_variants import MoltbotBatchVariants - from .nodes.image_to_prompt import MoltbotImageToPrompt - from .nodes.prompt_planner import MoltbotPromptPlanner - from .nodes.prompt_refiner import MoltbotPromptRefiner + from .nodes.batch_variants import OpenClawBatchVariants + from .nodes.image_to_prompt import OpenClawImageToPrompt + from .nodes.prompt_planner import OpenClawPromptPlanner + from .nodes.prompt_refiner import OpenClawPromptRefiner NODE_CLASS_MAPPINGS = { - "MoltbotPromptPlanner": MoltbotPromptPlanner, - "MoltbotBatchVariants": MoltbotBatchVariants, - "MoltbotImageToPrompt": MoltbotImageToPrompt, - "MoltbotPromptRefiner": MoltbotPromptRefiner, + # IMPORTANT: keep legacy mapping keys stable for existing workflows. + "MoltbotPromptPlanner": OpenClawPromptPlanner, + "MoltbotBatchVariants": OpenClawBatchVariants, + "MoltbotImageToPrompt": OpenClawImageToPrompt, + "MoltbotPromptRefiner": OpenClawPromptRefiner, } NODE_DISPLAY_NAME_MAPPINGS = { diff --git a/nodes/batch_variants.py b/nodes/batch_variants.py index f1d238d..a4066b5 100644 --- a/nodes/batch_variants.py +++ b/nodes/batch_variants.py @@ -15,7 +15,7 @@ except ImportError: logger = logging.getLogger("ComfyUI-OpenClaw.nodes.BatchVariants") -class MoltbotBatchVariants: +class OpenClawBatchVariants: """ Generates deterministic variants for batch processing. """ @@ -133,3 +133,7 @@ class MoltbotBatchVariants: params_list.append(json.dumps(validated.dict(), indent=2)) return (pos_list, neg_list, params_list) + + +# IMPORTANT: keep legacy class alias for existing imports and tests. +MoltbotBatchVariants = OpenClawBatchVariants diff --git a/nodes/image_to_prompt.py b/nodes/image_to_prompt.py index 93d60bb..75c9b87 100644 --- a/nodes/image_to_prompt.py +++ b/nodes/image_to_prompt.py @@ -1,20 +1,8 @@ -import base64 -import io -import json import logging -from typing import Any, Dict, List, Tuple - -try: - from PIL import Image # type: ignore -except ModuleNotFoundError: # pragma: no cover - Image = None # type: ignore - -try: - import numpy as np # type: ignore -except ModuleNotFoundError: # pragma: no cover - np = None # type: ignore +from typing import Any, Tuple try: + from ..services.image_utils import tensor_to_base64_png from ..services.llm_client import LLMClient from ..services.llm_output import ( extract_json_object, @@ -22,6 +10,7 @@ try: sanitize_string, ) except ImportError: + from services.image_utils import tensor_to_base64_png from services.llm_client import LLMClient from services.llm_output import ( extract_json_object, @@ -37,7 +26,7 @@ except ImportError: logger = logging.getLogger("ComfyUI-OpenClaw.nodes.ImageToPrompt") -class MoltbotImageToPrompt: +class OpenClawImageToPrompt: """ Experimental node that uses Vision LLM to generate prompt starters from an image. """ @@ -80,46 +69,9 @@ class MoltbotImageToPrompt: Convert ComfyUI tensor (Batch, H, W, C) to base64 PNG. Uses the first image in batch. """ - if Image is None: - raise RuntimeError( - "Pillow (PIL) is required for ImageToPrompt. Please install pillow." - ) - if np is None: - raise RuntimeError( - "numpy is required for ImageToPrompt. Please install numpy." - ) - # Tensor is typically [Batch, H, W, 3] float32 0..1 - # Take first image - if len(tensor_image.shape) == 4: - img_np = tensor_image[0] - else: - # Handle case where it might be single image [H, W, 3] - img_np = tensor_image - - # Check if tensor (convert to numpy if it is a torch tensor) - if hasattr(img_np, "cpu"): - img_np = img_np.cpu().numpy() - - # Convert to uint8 0..255 - img_np = np.clip(img_np * 255.0, 0, 255).astype(np.uint8) - - # To PIL - pil_img = Image.fromarray(img_np) - - # Resize if needed - width, height = pil_img.size - max_dim = max(width, height) - if max_dim > max_side: - scale = max_side / max_dim - new_w = int(width * scale) - new_h = int(height * scale) - pil_img = pil_img.resize((new_w, new_h), Image.Resampling.LANCZOS) - - # Bytes Metadata stripping (default save doesn't add much, but good practice) - buffered = io.BytesIO() - pil_img.save(buffered, format="PNG", optimize=True) - img_str = base64.b64encode(buffered.getvalue()).decode("utf-8") - return img_str + return tensor_to_base64_png( + tensor_image=tensor_image, max_side=max_side, context="ImageToPrompt" + ) def generate_prompt( self, image: Any, goal: str, detail_level: str, max_image_side: int @@ -186,3 +138,7 @@ Do not use markdown blocks. metrics.increment("errors") logger.error(f"Failed to generate prompt from image: {e}") raise e + + +# IMPORTANT: keep legacy class alias for existing imports and tests. +MoltbotImageToPrompt = OpenClawImageToPrompt diff --git a/nodes/prompt_planner.py b/nodes/prompt_planner.py index 443eae5..b893b12 100644 --- a/nodes/prompt_planner.py +++ b/nodes/prompt_planner.py @@ -16,7 +16,7 @@ except ImportError as e: logger = logging.getLogger("ComfyUI-OpenClaw.nodes.PromptPlanner") -class MoltbotPromptPlanner: +class OpenClawPromptPlanner: """ Experimental node that uses an LLM to plan the prompt and generation parameters. DELEGATES to services.planner.PlannerService (F8 Refactor). @@ -70,3 +70,7 @@ class MoltbotPromptPlanner: # Node expects params as JSON string return (positive, negative, json.dumps(params_dict, indent=2)) + + +# IMPORTANT: keep legacy class alias for existing imports and tests. +MoltbotPromptPlanner = OpenClawPromptPlanner diff --git a/nodes/prompt_refiner.py b/nodes/prompt_refiner.py index 9d0a1a7..e7468eb 100644 --- a/nodes/prompt_refiner.py +++ b/nodes/prompt_refiner.py @@ -1,22 +1,12 @@ -import base64 -import io import json import logging -from typing import Any, Dict, List, Tuple - -try: - import numpy as np # type: ignore -except ModuleNotFoundError: # pragma: no cover - np = None # type: ignore - -try: - from PIL import Image # type: ignore -except ModuleNotFoundError: # pragma: no cover - Image = None # type: ignore +from typing import Any, Tuple try: + from ..services.image_utils import tensor_to_base64_png from ..services.refiner import RefinerService except ImportError: + from services.image_utils import tensor_to_base64_png from services.refiner import RefinerService try: @@ -38,7 +28,7 @@ ALLOWED_PATCH_KEYS = { logger = logging.getLogger("ComfyUI-OpenClaw.nodes.PromptRefiner") -class MoltbotPromptRefiner: +class OpenClawPromptRefiner: """ Critiques and refines prompts/params based on a generated image and identified issues. DELEGATES to services.refiner.RefinerService (F21 Refactor). @@ -89,39 +79,10 @@ class MoltbotPromptRefiner: def _tensor_to_base64_png(self, tensor_image: Any, max_side: int) -> str: """ Convert ComfyUI tensor (Batch, H, W, C) to base64 PNG. - (Duplicated from ImageToPrompt for MVP robustness/isolation). """ - if Image is None: - raise RuntimeError( - "Pillow (PIL) is required for PromptRefiner. Please install pillow." - ) - if np is None: - raise RuntimeError( - "numpy is required for PromptRefiner. Please install numpy." - ) - if len(tensor_image.shape) == 4: - img_np = tensor_image[0] - else: - img_np = tensor_image - - if hasattr(img_np, "cpu"): - img_np = img_np.cpu().numpy() - - img_np = np.clip(img_np * 255.0, 0, 255).astype(np.uint8) - pil_img = Image.fromarray(img_np) - - width, height = pil_img.size - max_dim = max(width, height) - if max_dim > max_side: - scale = max_side / max_dim - new_w = int(width * scale) - new_h = int(height * scale) - pil_img = pil_img.resize((new_w, new_h), Image.Resampling.LANCZOS) - - buffered = io.BytesIO() - pil_img.save(buffered, format="PNG", optimize=True) - img_str = base64.b64encode(buffered.getvalue()).decode("utf-8") - return img_str + return tensor_to_base64_png( + tensor_image=tensor_image, max_side=max_side, context="PromptRefiner" + ) def refine_prompt( self, @@ -174,3 +135,7 @@ class MoltbotPromptRefiner: metrics.increment("errors") # Add metrics for service errors logger.error(f"Refiner Service failed: {e}") return (orig_positive, orig_negative, "{}", f"Error: {str(e)}") + + +# IMPORTANT: keep legacy class alias for existing imports and tests. +MoltbotPromptRefiner = OpenClawPromptRefiner diff --git a/services/image_utils.py b/services/image_utils.py new file mode 100644 index 0000000..c9e898c --- /dev/null +++ b/services/image_utils.py @@ -0,0 +1,48 @@ +import base64 +import io +from typing import Any + +# CRITICAL: keep optional imports at module-load time so test/loader paths +# without Pillow/numpy can still import modules and fail only on image use. +try: + from PIL import Image # type: ignore +except ModuleNotFoundError: # pragma: no cover + Image = None # type: ignore + +try: + import numpy as np # type: ignore +except ModuleNotFoundError: # pragma: no cover + np = None # type: ignore + + +def tensor_to_base64_png(tensor_image: Any, max_side: int, context: str) -> str: + """ + Convert ComfyUI IMAGE tensor ([B,H,W,C] or [H,W,C]) into base64 PNG. + """ + if Image is None: + raise RuntimeError( + f"Pillow (PIL) is required for {context}. Please install pillow." + ) + if np is None: + raise RuntimeError(f"numpy is required for {context}. Please install numpy.") + + img_np = tensor_image[0] if len(tensor_image.shape) == 4 else tensor_image + + if hasattr(img_np, "cpu"): + img_np = img_np.cpu().numpy() + + img_np = np.clip(img_np * 255.0, 0, 255).astype(np.uint8) + pil_img = Image.fromarray(img_np) + + width, height = pil_img.size + max_dim = max(width, height) + if max_dim > max_side: + scale = max_side / max_dim + new_w = int(width * scale) + new_h = int(height * scale) + resampling = getattr(Image, "Resampling", Image) + pil_img = pil_img.resize((new_w, new_h), resampling.LANCZOS) + + buffered = io.BytesIO() + pil_img.save(buffered, format="PNG", optimize=True) + return base64.b64encode(buffered.getvalue()).decode("utf-8") diff --git a/tests/test_comfyui_loader_import.py b/tests/test_comfyui_loader_import.py index 1848023..088ad6a 100644 --- a/tests/test_comfyui_loader_import.py +++ b/tests/test_comfyui_loader_import.py @@ -89,10 +89,16 @@ class TestComfyUICustomNodeLoaderImport(unittest.TestCase): self.assertTrue(hasattr(module, "NODE_CLASS_MAPPINGS")) self.assertIn("MoltbotPromptPlanner", module.NODE_CLASS_MAPPINGS) + planner_cls = module.NODE_CLASS_MAPPINGS["MoltbotPromptPlanner"] + self.assertEqual(planner_cls.__name__, "OpenClawPromptPlanner") # After import, `services.*` should be importable because __init__.py must self-heal sys.path. llm_mod = importlib.import_module("services.llm_client") self.assertTrue(hasattr(llm_mod, "LLMClient")) + planner_mod = importlib.import_module("nodes.prompt_planner") + self.assertIs( + planner_mod.MoltbotPromptPlanner, planner_mod.OpenClawPromptPlanner + ) finally: sys.path = old_sys_path sys.modules.pop(name, None) diff --git a/tests/test_node_class_aliases.py b/tests/test_node_class_aliases.py new file mode 100644 index 0000000..a3a8fc1 --- /dev/null +++ b/tests/test_node_class_aliases.py @@ -0,0 +1,22 @@ +import os +import sys +import unittest + +sys.path.append(os.getcwd()) + +from nodes.batch_variants import MoltbotBatchVariants, OpenClawBatchVariants +from nodes.image_to_prompt import MoltbotImageToPrompt, OpenClawImageToPrompt +from nodes.prompt_planner import MoltbotPromptPlanner, OpenClawPromptPlanner +from nodes.prompt_refiner import MoltbotPromptRefiner, OpenClawPromptRefiner + + +class TestNodeClassAliases(unittest.TestCase): + def test_legacy_aliases_resolve_to_openclaw_classes(self): + self.assertIs(MoltbotPromptPlanner, OpenClawPromptPlanner) + self.assertIs(MoltbotBatchVariants, OpenClawBatchVariants) + self.assertIs(MoltbotImageToPrompt, OpenClawImageToPrompt) + self.assertIs(MoltbotPromptRefiner, OpenClawPromptRefiner) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_planner.py b/tests/test_planner.py index 0f354b0..79ab8ae 100644 --- a/tests/test_planner.py +++ b/tests/test_planner.py @@ -77,8 +77,8 @@ class TestPromptPlanner(unittest.TestCase): self.assertEqual(params["width"], 1016) self.assertEqual(params["height"], 1024) # cfg/steps should be clamped by schema (exact limits in GenerationParams) - self.assertLessEqual(params["cfg"], 30.0) - self.assertLessEqual(params["steps"], 150) + self.assertEqual(params["cfg"], 30.0) + self.assertEqual(params["steps"], 100) if __name__ == "__main__": diff --git a/tests/test_vision.py b/tests/test_vision.py index c5c1a7a..74f9103 100644 --- a/tests/test_vision.py +++ b/tests/test_vision.py @@ -16,19 +16,23 @@ except ModuleNotFoundError: sys.path.append(os.getcwd()) try: - from nodes.image_to_prompt import MoltbotImageToPrompt + from nodes.image_to_prompt import MoltbotImageToPrompt, OpenClawImageToPrompt + from nodes.prompt_refiner import MoltbotPromptRefiner except ModuleNotFoundError: MoltbotImageToPrompt = None + OpenClawImageToPrompt = None + MoltbotPromptRefiner = None +from services.image_utils import tensor_to_base64_png from services.llm_client import LLMClient @unittest.skipIf( - (not NUMPY_AVAILABLE) or (MoltbotImageToPrompt is None), + (not NUMPY_AVAILABLE) or (OpenClawImageToPrompt is None), "numpy (and node deps) not available", ) class TestImageToPrompt(unittest.TestCase): def setUp(self): - self.node = MoltbotImageToPrompt() + self.node = OpenClawImageToPrompt() self.node.llm_client = MagicMock() def test_preprocessing_tensor_mock(self): @@ -100,6 +104,21 @@ class TestImageToPrompt(unittest.TestCase): self.assertEqual(tags, "sci-fi, neon") self.assertEqual(prompt, "Cyberpunk city with neon lights") + def test_shared_image_helper_encoding_parity(self): + """R133: wrappers in image nodes must preserve shared encoding behavior.""" + self.assertIs(MoltbotImageToPrompt, OpenClawImageToPrompt) + self.assertIsNotNone(MoltbotPromptRefiner) + + fake_tensor = np.zeros((1, 256, 128, 3), dtype=np.float32) + image_node_b64 = self.node._tensor_to_base64_png(fake_tensor, max_side=512) + helper_b64 = tensor_to_base64_png(fake_tensor, max_side=512, context="test") + refiner_b64 = MoltbotPromptRefiner()._tensor_to_base64_png( + fake_tensor, max_side=512 + ) + + self.assertEqual(image_node_b64, helper_b64) + self.assertEqual(image_node_b64, refiner_b64) + if __name__ == "__main__": unittest.main()