Ying Cai
da2aefeb9a
- Added a new method for validating inputs based on type definitions, including support for custom enum values. - Updated input handling to manage override values for enum fields, allowing users to connect custom values. - Enhanced schema parsing to include `allow_custom_enum` for specific fields, enabling custom value connections in the UI. - Bumped project version to 0.0.13 to reflect these updates.
370 lines
13 KiB
Python
370 lines
13 KiB
Python
"""
|
|
Base node class for Sunra ComfyUI nodes.
|
|
|
|
This module provides the base class that all dynamically generated
|
|
Sunra nodes inherit from.
|
|
"""
|
|
|
|
import time
|
|
import hashlib
|
|
import json
|
|
import io
|
|
from typing import Dict, Any, Tuple, List
|
|
import torch
|
|
|
|
from .type_converters import (
|
|
process_image_input,
|
|
process_image_output,
|
|
process_video_input,
|
|
process_video_output,
|
|
process_audio_input,
|
|
process_audio_output,
|
|
process_audio_path_output,
|
|
process_3d_output,
|
|
validate_value,
|
|
)
|
|
|
|
|
|
class SunraBaseNode:
|
|
"""
|
|
Base class for all Sunra ComfyUI nodes.
|
|
|
|
This class handles common functionality like:
|
|
- Input validation and processing
|
|
- API communication
|
|
- Output processing
|
|
- Caching and change detection
|
|
"""
|
|
|
|
def __init__(self, schema: Dict[str, Any], api_client):
|
|
"""
|
|
Initialize the node with schema and API client.
|
|
|
|
Args:
|
|
schema: The node schema definition
|
|
api_client: The Sunra API client instance
|
|
"""
|
|
self.schema = schema
|
|
self.api_client = api_client
|
|
self.model_id = schema["model_id"]
|
|
self.api_endpoint = schema["api_endpoint"]
|
|
self.metadata = schema.get("metadata", {})
|
|
|
|
@classmethod
|
|
def create_from_schema(cls, schema: Dict[str, Any], api_client):
|
|
"""
|
|
Create a dynamic node class from schema.
|
|
|
|
Args:
|
|
schema: The complete node schema
|
|
api_client: The Sunra API client instance
|
|
|
|
Returns:
|
|
A new node class configured according to the schema
|
|
"""
|
|
from .schema_to_node import schema_to_comfyui_node_spec
|
|
|
|
# Parse schema to get node specifications
|
|
node_spec = schema_to_comfyui_node_spec(schema)
|
|
|
|
# Create the dynamic class
|
|
class DynamicSunraNode(cls):
|
|
# Set class attributes from schema
|
|
CATEGORY = node_spec["category"]
|
|
FUNCTION = node_spec["function"]
|
|
OUTPUT_NODE = node_spec.get("output_node", False)
|
|
|
|
# Store schema and api_client as class attributes
|
|
_schema = schema
|
|
_api_client = api_client
|
|
|
|
def __init__(self):
|
|
"""Initialize without arguments for ComfyUI compatibility."""
|
|
super().__init__(self._schema, self._api_client)
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
"""Return input types for ComfyUI."""
|
|
return node_spec["input_types"]
|
|
|
|
RETURN_TYPES = node_spec["return_types"]
|
|
RETURN_NAMES = node_spec["return_names"]
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, **kwargs):
|
|
"""
|
|
Determine if the node should be re-executed.
|
|
|
|
Uses force_rerun flag and input hashing for caching.
|
|
"""
|
|
if kwargs.get("force_rerun", False):
|
|
return time.time()
|
|
|
|
# Create hash of all inputs
|
|
input_hash = cls._create_input_hash(kwargs)
|
|
return input_hash
|
|
|
|
@classmethod
|
|
def _create_input_hash(cls, inputs: Dict[str, Any]) -> str:
|
|
"""Create a hash of inputs for caching."""
|
|
# Convert inputs to a hashable format
|
|
hashable_inputs = {}
|
|
for key, value in inputs.items():
|
|
if isinstance(value, torch.Tensor):
|
|
# Hash tensor data
|
|
hashable_inputs[key] = hashlib.sha256(
|
|
value.cpu().numpy().tobytes()
|
|
).hexdigest()[:16]
|
|
elif isinstance(value, (list, dict)):
|
|
# Convert to JSON string
|
|
hashable_inputs[key] = json.dumps(value, sort_keys=True)
|
|
else:
|
|
hashable_inputs[key] = str(value)
|
|
|
|
# Create final hash
|
|
content = json.dumps(hashable_inputs, sort_keys=True)
|
|
return hashlib.sha256(content.encode()).hexdigest()[:16]
|
|
|
|
@classmethod
|
|
def VALIDATE_INPUTS(cls, **kwargs):
|
|
"""
|
|
Validate inputs according to their type definitions.
|
|
"""
|
|
# Get all input definitions
|
|
all_inputs = {}
|
|
input_types = cls.INPUT_TYPES()
|
|
all_inputs.update(input_types.get("required", {}))
|
|
all_inputs.update(input_types.get("optional", {}))
|
|
|
|
for name, value in kwargs.items():
|
|
# Skip internal parameters and override inputs
|
|
if name == "force_rerun" or name.endswith("_override"):
|
|
continue
|
|
|
|
if name in all_inputs:
|
|
input_def = all_inputs[name]
|
|
|
|
# Check if this is an enum field
|
|
if isinstance(input_def[0], list) and value not in input_def[0]:
|
|
return f"Value '{value}' not in allowed values for {name}"
|
|
|
|
return True
|
|
|
|
# Return the class itself, not an instance
|
|
return DynamicSunraNode
|
|
|
|
def execute(self, **kwargs) -> Tuple:
|
|
"""
|
|
Main execution function called by ComfyUI.
|
|
|
|
Args:
|
|
**kwargs: Input parameters from ComfyUI
|
|
|
|
Returns:
|
|
Tuple of outputs according to RETURN_TYPES
|
|
"""
|
|
try:
|
|
# Remove internal parameters
|
|
kwargs = self._clean_inputs(kwargs)
|
|
|
|
# Validate inputs
|
|
validated_inputs = self._validate_inputs(kwargs)
|
|
|
|
# Process inputs for API
|
|
processed_inputs = self._process_inputs(validated_inputs)
|
|
|
|
# Call Sunra API
|
|
response = self._call_api(processed_inputs)
|
|
|
|
# Process outputs
|
|
outputs = self._process_outputs(response)
|
|
|
|
# Return as tuple
|
|
return tuple(outputs)
|
|
|
|
except Exception as e:
|
|
# Log error and re-raise
|
|
print(f"Error in {self.model_id}: {str(e)}")
|
|
raise
|
|
|
|
def _clean_inputs(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
|
|
"""Remove internal parameters like force_rerun."""
|
|
cleaned = inputs.copy()
|
|
cleaned.pop("force_rerun", None)
|
|
return cleaned
|
|
|
|
def _validate_inputs(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
|
|
"""Validate inputs according to schema."""
|
|
validated = {}
|
|
|
|
# Get input definitions from schema
|
|
all_inputs = {}
|
|
all_inputs.update(self.schema["inputs"].get("required", {}))
|
|
all_inputs.update(self.schema["inputs"].get("optional", {}))
|
|
|
|
for name, value in inputs.items():
|
|
if name in all_inputs:
|
|
config = all_inputs[name]
|
|
validated[name] = validate_value(value, config["type"], config)
|
|
else:
|
|
# Pass through unknown inputs
|
|
validated[name] = value
|
|
|
|
return validated
|
|
|
|
def _process_inputs(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
|
|
"""Process inputs for API call."""
|
|
processed = {}
|
|
|
|
# Get all input definitions
|
|
all_inputs = {
|
|
**self.schema["inputs"].get("required", {}),
|
|
**self.schema["inputs"].get("optional", {}),
|
|
}
|
|
|
|
# First, handle override values for enum fields with custom value support
|
|
override_map = {}
|
|
for name, value in inputs.items():
|
|
if name.endswith("_override") and value is not None and value != "None":
|
|
base_name = name[:-9] # Remove "_override" suffix
|
|
override_map[base_name] = value
|
|
|
|
for name, value in inputs.items():
|
|
# Skip override inputs - they're handled above
|
|
if name.endswith("_override"):
|
|
continue
|
|
|
|
# Skip None values and special "None" strings
|
|
if value is None or (isinstance(value, str) and value == "None"):
|
|
continue
|
|
|
|
# Check if this field has an override value
|
|
if name in override_map:
|
|
value = override_map[name]
|
|
|
|
input_config = all_inputs.get(name, {})
|
|
input_type = input_config.get("type", "string")
|
|
api_name = input_config.get("api_field_name", name)
|
|
|
|
# Process based on type
|
|
if input_type in ("image", "mask"):
|
|
processed[api_name] = process_image_input(value)
|
|
elif input_type == "video":
|
|
processed[api_name] = process_video_input(value)
|
|
elif input_type == "audio":
|
|
processed[api_name] = process_audio_input(value)
|
|
elif input_type == "boolean":
|
|
processed[api_name] = bool(value)
|
|
else:
|
|
processed[api_name] = value
|
|
|
|
return processed
|
|
|
|
def _call_api(self, inputs: Dict[str, Any]) -> Any:
|
|
"""
|
|
Call Sunra API with processed inputs.
|
|
|
|
Args:
|
|
inputs: Processed input parameters
|
|
|
|
Returns:
|
|
API response
|
|
"""
|
|
# Use the API client to make the call
|
|
return self.api_client.call(self.api_endpoint, inputs)
|
|
|
|
def _process_outputs(self, response: Any) -> List[Any]:
|
|
"""
|
|
Process API response to ComfyUI outputs.
|
|
|
|
Args:
|
|
response: Raw API response
|
|
|
|
Returns:
|
|
List of outputs matching RETURN_TYPES
|
|
"""
|
|
output_schema = self.schema.get("outputs", {})
|
|
return_names = self.__class__.RETURN_NAMES
|
|
return_types = self.__class__.RETURN_TYPES
|
|
|
|
outputs = []
|
|
audio_data_cache = {} # Cache to store audio data for path outputs
|
|
|
|
for i, name in enumerate(return_names):
|
|
# Check if this is an automatically added audio_path output
|
|
if name.endswith("_path") and i > 0:
|
|
audio_name = name[:-5] # Remove '_path' suffix
|
|
# Check if previous output was the corresponding audio
|
|
if i > 0 and return_names[i-1] == audio_name and return_types[i-1] == "AUDIO":
|
|
# Use cached audio file path
|
|
if audio_name in audio_data_cache:
|
|
outputs.append(audio_data_cache[audio_name])
|
|
else:
|
|
outputs.append("") # Default empty path
|
|
continue
|
|
|
|
output_config = output_schema.get(name, {})
|
|
output_type = output_config.get("type", "string")
|
|
|
|
# Extract value from response
|
|
if isinstance(response, dict):
|
|
value = response.get(name, self._get_default_output(output_type))
|
|
elif len(output_schema) == 1: # Use original schema length, not return_names
|
|
value = response # Single output, response is the value
|
|
else:
|
|
value = self._get_default_output(output_type)
|
|
|
|
# Process and convert to appropriate type
|
|
if output_type == "audio" and name not in audio_data_cache:
|
|
# Process audio and get path from global variable
|
|
audio_dict = self._convert_output_value(value, output_type)
|
|
audio_path = process_audio_path_output(value)
|
|
outputs.append(audio_dict)
|
|
audio_data_cache[name] = audio_path
|
|
else:
|
|
outputs.append(self._convert_output_value(value, output_type))
|
|
|
|
return outputs
|
|
|
|
def _convert_output_value(self, value: Any, output_type: str) -> Any:
|
|
"""Convert output value to the appropriate type."""
|
|
if value is None:
|
|
return self._get_default_output(output_type)
|
|
|
|
type_converters = {
|
|
"image": process_image_output,
|
|
"mask": process_image_output,
|
|
"video": process_video_output,
|
|
"audio": process_audio_output,
|
|
"3d": process_3d_output,
|
|
"model": process_3d_output, # Alternative name for 3D outputs
|
|
"string": str,
|
|
"integer": int,
|
|
"float": float,
|
|
"boolean": bool,
|
|
}
|
|
|
|
converter = type_converters.get(output_type)
|
|
return converter(value) if converter else value
|
|
|
|
def _get_default_output(self, output_type: str) -> Any:
|
|
"""Get default value for an output type."""
|
|
if output_type == "video":
|
|
from comfy_api.input_impl.video_types import VideoFromFile
|
|
# Create empty video bytes
|
|
empty_video = io.BytesIO()
|
|
return VideoFromFile(empty_video)
|
|
|
|
defaults = {
|
|
"image": torch.zeros(1, 3, 512, 512),
|
|
"mask": torch.zeros(1, 1, 512, 512),
|
|
"string": "",
|
|
"integer": 0,
|
|
"float": 0.0,
|
|
"boolean": False,
|
|
"audio": {"waveform": torch.zeros(1, 1, 44100), "sample_rate": 44100}, # Empty 1-second mono audio
|
|
"3d": "", # Empty path for 3D models
|
|
"model": "", # Alternative name for 3D outputs
|
|
}
|
|
return defaults.get(output_type, None)
|