Initial work to generate node from schema
This commit is contained in:
Paul
@@ -1,90 +1,157 @@
|
||||
import os
|
||||
import json
|
||||
import replicate
|
||||
|
||||
|
||||
class Llama3Replicate:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"prompt": (
|
||||
"STRING",
|
||||
{"default": "", "multiline": True, "dynamicPrompts": True},
|
||||
),
|
||||
"system_prompt": (
|
||||
"STRING",
|
||||
{
|
||||
"default": "You are a helpful assistant",
|
||||
"multiline": True,
|
||||
"dynamicPrompts": True,
|
||||
},
|
||||
),
|
||||
"top_p": (
|
||||
"FLOAT",
|
||||
{"default": 0.95, "max": 1.0, "min": -1.0},
|
||||
),
|
||||
"top_k": (
|
||||
"INT",
|
||||
{"default": 0, "min": -1},
|
||||
),
|
||||
"max_tokens": (
|
||||
"INT",
|
||||
{"default": 512, "min": 1},
|
||||
),
|
||||
"min_tokens": (
|
||||
"INT",
|
||||
{"default": 0, "min": 0},
|
||||
),
|
||||
"temperature": (
|
||||
"FLOAT",
|
||||
{"default": 0.7, "max": 5.0, "min": 0.0},
|
||||
),
|
||||
"length_penalty": (
|
||||
"FLOAT",
|
||||
{"default": 1.0, "max": 5.0, "min": 0.0},
|
||||
),
|
||||
"presence_penalty": (
|
||||
"FLOAT",
|
||||
{"default": 0},
|
||||
),
|
||||
"seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("STRING",)
|
||||
FUNCTION = "run_llama3_replicate"
|
||||
CATEGORY = "Replicate"
|
||||
|
||||
def run_llama3_replicate(
|
||||
self,
|
||||
top_p,
|
||||
top_k,
|
||||
prompt,
|
||||
system_prompt,
|
||||
max_tokens,
|
||||
min_tokens,
|
||||
temperature,
|
||||
length_penalty,
|
||||
presence_penalty,
|
||||
seed,
|
||||
):
|
||||
input = {
|
||||
"system_prompt": system_prompt,
|
||||
"prompt": prompt,
|
||||
"top_p": top_p,
|
||||
"top_k": top_k,
|
||||
"max_tokens": max_tokens,
|
||||
"min_tokens": min_tokens,
|
||||
"temperature": temperature,
|
||||
"length_penalty": length_penalty,
|
||||
"presence_penalty": presence_penalty,
|
||||
"seed": seed,
|
||||
}
|
||||
|
||||
output = replicate.run("meta/meta-llama-3-70b-instruct", input=input)
|
||||
output = "".join(output).strip()
|
||||
return (output,)
|
||||
def convert_type(openapi_type):
|
||||
type_mapping = {
|
||||
"string": "STRING",
|
||||
"integer": "INT",
|
||||
"number": "FLOAT",
|
||||
"boolean": "BOOL",
|
||||
}
|
||||
return type_mapping.get(openapi_type, "STRING")
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"Llama 3 Replicate": Llama3Replicate,
|
||||
}
|
||||
def resolve_schema(prop_data, schemas):
|
||||
if "$ref" in prop_data:
|
||||
ref_path = prop_data["$ref"].split("/")
|
||||
current = schemas
|
||||
for path in ref_path[1:]: # Skip the first '#' element
|
||||
current = current[path]
|
||||
return current
|
||||
return prop_data
|
||||
|
||||
|
||||
def convert_schema_to_comfyui(schema, schemas):
|
||||
input_types = {"required": {}}
|
||||
|
||||
for prop_name, prop_data in schema["properties"].items():
|
||||
prop_data = resolve_schema(prop_data, schemas)
|
||||
|
||||
if "allOf" in prop_data:
|
||||
prop_data = resolve_schema(prop_data["allOf"][0], schemas)
|
||||
|
||||
if "enum" in prop_data:
|
||||
input_type = prop_data["enum"]
|
||||
elif "type" in prop_data:
|
||||
input_type = convert_type(prop_data["type"])
|
||||
else:
|
||||
input_type = "STRING"
|
||||
|
||||
default_value = prop_data.get("default", "")
|
||||
|
||||
input_config = {"default": default_value}
|
||||
|
||||
if "minimum" in prop_data:
|
||||
input_config["min"] = prop_data["minimum"]
|
||||
if "maximum" in prop_data:
|
||||
input_config["max"] = prop_data["maximum"]
|
||||
|
||||
if prop_data.get("type") == "string" and prop_data.get("format") == "uri":
|
||||
input_config["multiline"] = True
|
||||
|
||||
if "prompt" in prop_name and prop_data.get("type") == "string":
|
||||
input_config["multiline"] = True
|
||||
|
||||
if "template" not in prop_name:
|
||||
input_config["dynamicPrompts"] = True
|
||||
|
||||
input_types["required"][prop_name] = (input_type, input_config)
|
||||
|
||||
# Reorder input_types to put prompt and negative_prompt first
|
||||
ordered_input_types = {"required": {}}
|
||||
for key in ["prompt", "negative_prompt"]:
|
||||
if key in input_types["required"]:
|
||||
ordered_input_types["required"][key] = input_types["required"][key]
|
||||
for key in list(input_types["required"].keys()):
|
||||
if "prompt" in key:
|
||||
ordered_input_types["required"][key] = input_types["required"][key]
|
||||
for key in list(input_types["required"].keys()):
|
||||
if key not in ordered_input_types["required"] and key != "seed":
|
||||
ordered_input_types["required"][key] = input_types["required"][key]
|
||||
if "seed" in input_types["required"]:
|
||||
ordered_input_types["required"]["seed"] = input_types["required"]["seed"]
|
||||
|
||||
return ordered_input_types
|
||||
|
||||
|
||||
def create_comfyui_node(schemas, model_info):
|
||||
author = model_info["owner"]
|
||||
name = model_info["name"]
|
||||
version = model_info["latest_version"]["id"]
|
||||
|
||||
replicate_model = f"{author}/{name}:{version}"
|
||||
node_name = f"Replicate {author}/{name}"
|
||||
input_schema = schemas["components"]["schemas"]["Input"]
|
||||
|
||||
class ReplicateToComfyUI:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return convert_schema_to_comfyui(input_schema, schemas)
|
||||
|
||||
RETURN_TYPES = ("STRING",)
|
||||
FUNCTION = "run_openapi_to_comfyui"
|
||||
CATEGORY = "Replicate"
|
||||
|
||||
def run_openapi_to_comfyui(self, **kwargs):
|
||||
print(f"Running {replicate_model} with {kwargs}")
|
||||
output = replicate.run(replicate_model, input=kwargs)
|
||||
output = "".join(output).strip()
|
||||
# print(f"Output: {output}")
|
||||
return (output,)
|
||||
|
||||
return node_name, ReplicateToComfyUI
|
||||
|
||||
|
||||
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") as f:
|
||||
schema = json.load(f)
|
||||
openapi_schema = schema["latest_version"]["openapi_schema"]
|
||||
model_info = schema
|
||||
node_name, node_class = create_comfyui_node(openapi_schema, model_info)
|
||||
nodes[node_name] = node_class
|
||||
return nodes
|
||||
|
||||
|
||||
# Create ComfyUI nodes for all schema files in the "schemas" directory
|
||||
comfyui_nodes = create_comfyui_nodes_from_schemas("schemas")
|
||||
|
||||
# Print the resulting node classes
|
||||
for schema_file, node_class in comfyui_nodes.items():
|
||||
print(f"Node class for {schema_file}:")
|
||||
print(node_class.INPUT_TYPES())
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = comfyui_nodes
|
||||
print(NODE_CLASS_MAPPINGS)
|
||||
|
||||
# # Load the schema
|
||||
# with open("schema.json", "r") as f:
|
||||
# schema = json.load(f)
|
||||
# openapi_schema = schema["latest_version"]["openapi_schema"]
|
||||
|
||||
# # Create the ComfyUI node
|
||||
# ComfyUINode = create_comfyui_node(openapi_schema, schema)
|
||||
|
||||
# # Print the resulting node class
|
||||
# print(ComfyUINode.INPUT_TYPES())
|
||||
|
||||
# # Create an instance of the node and pass in defaults
|
||||
# node_instance = ComfyUINode()
|
||||
# defaults = {
|
||||
# "seed": 0,
|
||||
# "image": "https://example.com/default_image.png",
|
||||
# "style": "3D",
|
||||
# "prompt": "a person",
|
||||
# "lora_scale": 1.0
|
||||
# }
|
||||
|
||||
# # Run the node with the defaults
|
||||
# result = node_instance.run_openapi_to_comfyui(**defaults)
|
||||
# print(result)
|
||||
|
||||
Reference in New Issue
Block a user