mirror of
https://github.com/rookiestar28/ComfyUI-OpenClaw.git
synced 2026-08-14 00:48:07 +00:00
refactor(r133): unify image encoding helper and converge node class names
This commit is contained in:
+12
-11
@@ -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 = {
|
||||
|
||||
@@ -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
@@ -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
|
||||
|
||||
@@ -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
@@ -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
|
||||
|
||||
@@ -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")
|
||||
@@ -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)
|
||||
|
||||
@@ -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()
|
||||
@@ -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
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user