refactor(r133): unify image encoding helper and converge node class names

This commit is contained in:
rookiestar28
2026-03-01 01:15:40 +08:00
parent 9aef5eb2d4
commit 80ae3a83a0
10 changed files with 144 additions and 119 deletions
+12 -11
View File
@@ -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 = {
+5 -1
View File
@@ -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
+11 -55
View File
@@ -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
+5 -1
View File
@@ -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
+11 -46
View File
@@ -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
+48
View File
@@ -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")
+6
View File
@@ -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)
+22
View File
@@ -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()
+2 -2
View File
@@ -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__":
+22 -3
View File
@@ -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()