118 lines
4.0 KiB
Python
118 lines
4.0 KiB
Python
import json
|
|
import os
|
|
import time
|
|
|
|
import torch
|
|
from ..base import FlowyApiNode
|
|
from .schema_to_node import (
|
|
schema_to_comfyui_input_types,
|
|
get_return_type,
|
|
name_and_version,
|
|
inputs_that_need_arrays,
|
|
)
|
|
from ...api_key_manager import load_api_key
|
|
|
|
def create_comfyui_node(schema):
|
|
replicate_model, node_name = name_and_version(schema)
|
|
return_type = get_return_type(schema)
|
|
|
|
class ReplicateNode(FlowyApiNode):
|
|
@classmethod
|
|
def IS_CHANGED(cls, **kwargs):
|
|
return time.time() if kwargs["force_rerun"] else ""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return schema_to_comfyui_input_types(schema)
|
|
|
|
RETURN_TYPES = (
|
|
tuple(return_type.values())
|
|
if isinstance(return_type, dict)
|
|
else (return_type,)
|
|
)
|
|
CATEGORY = "Comflowy Replicate"
|
|
|
|
def get_model_type(self) -> str:
|
|
return "replicate"
|
|
|
|
def get_api_key(self) -> str:
|
|
"""获取 Comflowy API token"""
|
|
api_key = load_api_key()
|
|
if not api_key:
|
|
raise ValueError("Comflowy API key not found. Please set your API key first.")
|
|
return api_key
|
|
|
|
def prepare_payload(self, **kwargs):
|
|
# Remove force_rerun from kwargs
|
|
kwargs = {k: v for k, v in kwargs.items() if k != "force_rerun"}
|
|
|
|
# Handle array inputs
|
|
array_inputs = inputs_that_need_arrays(schema)
|
|
for input_name in array_inputs:
|
|
if input_name in kwargs:
|
|
if isinstance(kwargs[input_name], str):
|
|
kwargs[input_name] = (
|
|
[]
|
|
if kwargs[input_name] == ""
|
|
else kwargs[input_name].split("\n")
|
|
)
|
|
else:
|
|
kwargs[input_name] = [kwargs[input_name]]
|
|
|
|
# Process image and audio inputs
|
|
for key, value in kwargs.items():
|
|
if value is not None:
|
|
input_type = (
|
|
self.INPUT_TYPES()["required"].get(key, (None,))[0]
|
|
or self.INPUT_TYPES().get("optional", {}).get(key, (None,))[0]
|
|
)
|
|
if input_type == "IMAGE":
|
|
kwargs[key] = self.image_to_base64(value)
|
|
elif input_type == "AUDIO":
|
|
kwargs[key] = self.audio_to_base64(value)
|
|
|
|
# Remove empty optional inputs
|
|
optional_inputs = self.INPUT_TYPES().get("optional", {})
|
|
for key in list(kwargs.keys()):
|
|
if key in optional_inputs:
|
|
if isinstance(kwargs[key], torch.Tensor):
|
|
continue
|
|
elif not kwargs[key]:
|
|
del kwargs[key]
|
|
|
|
# 添加 API token 到 payload
|
|
return {
|
|
"replicate_model": replicate_model,
|
|
"input": kwargs,
|
|
"api_key": self.get_api_key() # 添加 API token
|
|
}
|
|
|
|
return node_name, ReplicateNode
|
|
|
|
|
|
def create_comfyui_nodes_from_schemas(schemas_dir):
|
|
nodes = {}
|
|
current_path = os.path.dirname(os.path.abspath(__file__))
|
|
schemas_dir_path = os.path.join(current_path, schemas_dir)
|
|
for schema_file in os.listdir(schemas_dir_path):
|
|
if schema_file.endswith(".json"):
|
|
with open(
|
|
os.path.join(schemas_dir_path, schema_file), "r", encoding="utf-8"
|
|
) as f:
|
|
schema = json.load(f)
|
|
node_name, node_class = create_comfyui_node(schema)
|
|
nodes[node_name] = node_class
|
|
return nodes
|
|
|
|
|
|
_cached_node_class_mappings = None
|
|
|
|
def get_node_class_mappings():
|
|
global _cached_node_class_mappings
|
|
if _cached_node_class_mappings is None:
|
|
_cached_node_class_mappings = create_comfyui_nodes_from_schemas("schemas")
|
|
return _cached_node_class_mappings
|
|
|
|
|
|
REPLICATE_NODE_CLASS_MAPPINGS = get_node_class_mappings()
|