mirror of
https://git.datalinker.icu/comfyanonymous/ComfyUI
synced 2026-09-06 11:27:04 +08:00
use pydantic instead
This commit is contained in:
Terry Jia
parent
a21e49c18b
commit
37fcf3e8e8
@ -1,23 +1,17 @@
|
||||
import os
|
||||
from typing import Optional
|
||||
|
||||
import tomlkit
|
||||
import tomlkit.exceptions
|
||||
import tomllib
|
||||
|
||||
from pydantic import ValidationError
|
||||
import logging
|
||||
|
||||
from comfy_config.types import (
|
||||
ComfyConfig,
|
||||
License,
|
||||
Model,
|
||||
ProjectConfig,
|
||||
PyProjectConfig,
|
||||
URLs,
|
||||
)
|
||||
|
||||
"""
|
||||
Original implementation comes from https://github.com/Comfy-Org/comfy-cli/blob/2e36f33dd39ef43b5acf7d1fc5acc5e01be92360/comfy_cli/registry/config_parser.py#L146
|
||||
|
||||
Extract configuration from a custom node directory's pyproject.toml file.
|
||||
|
||||
This function reads and parses the pyproject.toml file in the specified directory
|
||||
@ -46,70 +40,42 @@ Example:
|
||||
>>> print(project_config.project.name) # "my_custom_node" or name from pyproject.toml
|
||||
>>> nodes.EXTENSION_WEB_DIRS[project_config.project.name] = js_dir
|
||||
"""
|
||||
def extract_node_configuration(
|
||||
path,
|
||||
) -> Optional[PyProjectConfig]:
|
||||
|
||||
|
||||
def extract_node_configuration(path) -> Optional[PyProjectConfig]:
|
||||
folder_name = os.path.basename(path)
|
||||
toml_path = os.path.join(path, "pyproject.toml")
|
||||
|
||||
path = os.path.join(path, "pyproject.toml")
|
||||
|
||||
if not os.path.isfile(path):
|
||||
logging.warning("No pyproject.toml file found in the current directory, will use custom node folder name as project name as default.")
|
||||
|
||||
project = ProjectConfig(
|
||||
name=folder_name,
|
||||
)
|
||||
|
||||
return PyProjectConfig(project=project)
|
||||
|
||||
with open(path, "r") as file:
|
||||
data = tomlkit.load(file)
|
||||
|
||||
project_data = data.get("project", {})
|
||||
urls_data = project_data.get("urls", {})
|
||||
comfy_data = data.get("tool", {}).get("comfy", {})
|
||||
|
||||
license_data = project_data.get("license", {})
|
||||
if isinstance(license_data, str):
|
||||
license = License(text=license_data)
|
||||
if not os.path.isfile(toml_path):
|
||||
logging.warning(
|
||||
'Warning: License should be in one of these two formats: license = {file = "LICENSE"} OR license = {text = "MIT License"}. Please check the documentation: https://docs.comfy.org/registry/specifications.'
|
||||
)
|
||||
elif isinstance(license_data, dict):
|
||||
if "file" in license_data or "text" in license_data:
|
||||
license = License(file=license_data.get("file", ""), text=license_data.get("text", ""))
|
||||
else:
|
||||
logging.warning(
|
||||
'Warning: License should be in one of these two formats: license = {file = "LICENSE"} OR license = {text = "MIT License"}. Please check the documentation: https://docs.comfy.org/registry/specifications.'
|
||||
)
|
||||
license = License()
|
||||
else:
|
||||
license = License()
|
||||
logging.warning(
|
||||
'Warning: License should be in one of these two formats: license = {file = "LICENSE"} OR license = {text = "MIT License"}. Please check the documentation: https://docs.comfy.org/registry/specifications.'
|
||||
)
|
||||
"No pyproject.toml file found in the current directory, will use custom node folder name as project name as default.")
|
||||
|
||||
project = ProjectConfig(
|
||||
name=project_data.get("name", ""),
|
||||
description=project_data.get("description", ""),
|
||||
version=project_data.get("version", ""),
|
||||
requires_python=project_data.get("requires-python", ""),
|
||||
dependencies=project_data.get("dependencies", []),
|
||||
license=license,
|
||||
urls=URLs(
|
||||
homepage=urls_data.get("Homepage", ""),
|
||||
documentation=urls_data.get("Documentation", ""),
|
||||
repository=urls_data.get("Repository", ""),
|
||||
issues=urls_data.get("Issues", ""),
|
||||
),
|
||||
)
|
||||
try:
|
||||
project = ProjectConfig(name=folder_name)
|
||||
return PyProjectConfig(project=project)
|
||||
except ValidationError as e:
|
||||
logging.error(f"Failed to create default configuration: {e}")
|
||||
return None
|
||||
|
||||
comfy = ComfyConfig(
|
||||
publisher_id=comfy_data.get("PublisherId", ""),
|
||||
display_name=comfy_data.get("DisplayName", ""),
|
||||
icon=comfy_data.get("Icon", ""),
|
||||
models=[Model(location=m["location"], model_url=m["model_url"]) for m in comfy_data.get("Models", [])],
|
||||
includes=comfy_data.get("includes", []),
|
||||
)
|
||||
try:
|
||||
with open(toml_path, "rb") as f:
|
||||
data = tomllib.load(f)
|
||||
except Exception as e:
|
||||
logging.error(f"Failed to read pyproject.toml: {e}")
|
||||
return None
|
||||
|
||||
return PyProjectConfig(project=project, tool_comfy=comfy)
|
||||
try:
|
||||
config_data = {
|
||||
"project": data.get("project", {}),
|
||||
"tool_comfy": data.get("tool", {}).get("comfy", {})
|
||||
}
|
||||
|
||||
return PyProjectConfig(**config_data)
|
||||
|
||||
except ValidationError as e:
|
||||
logging.error(f"Validation error while parsing configuration: {e}")
|
||||
logging.error(f"Validation details: {e.errors()}")
|
||||
return None
|
||||
except Exception as e:
|
||||
logging.error(f"Unexpected error while parsing configuration: {e}")
|
||||
return None
|
||||
|
||||
@ -1,12 +1,11 @@
|
||||
from dataclasses import dataclass, field
|
||||
from pydantic import BaseModel, Field
|
||||
from typing import List, Optional
|
||||
|
||||
# IMPORTANT: The type definitions specified in pyproject.toml for custom nodes
|
||||
# must remain synchronized with the corresponding files in the https://github.com/Comfy-Org/comfy-cli/blob/main/comfy_cli/registry/types.py.
|
||||
# Any changes to one must be reflected in the other to maintain consistency.
|
||||
|
||||
@dataclass
|
||||
class NodeVersion:
|
||||
class NodeVersion(BaseModel):
|
||||
changelog: str
|
||||
dependencies: List[str]
|
||||
deprecated: bool
|
||||
@ -15,8 +14,7 @@ class NodeVersion:
|
||||
download_url: str
|
||||
|
||||
|
||||
@dataclass
|
||||
class Node:
|
||||
class Node(BaseModel):
|
||||
id: str
|
||||
name: str
|
||||
description: str
|
||||
@ -24,57 +22,50 @@ class Node:
|
||||
license: Optional[str] = None
|
||||
icon: Optional[str] = None
|
||||
repository: Optional[str] = None
|
||||
tags: List[str] = field(default_factory=list)
|
||||
tags: List[str] = Field(default_factory=list)
|
||||
latest_version: Optional[NodeVersion] = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class PublishNodeVersionResponse:
|
||||
class PublishNodeVersionResponse(BaseModel):
|
||||
node_version: NodeVersion
|
||||
signedUrl: str
|
||||
|
||||
|
||||
@dataclass
|
||||
class URLs:
|
||||
homepage: str = ""
|
||||
documentation: str = ""
|
||||
repository: str = ""
|
||||
issues: str = ""
|
||||
class URLs(BaseModel):
|
||||
homepage: str = Field(default="", alias="Homepage")
|
||||
documentation: str = Field(default="", alias="Documentation")
|
||||
repository: str = Field(default="", alias="Repository")
|
||||
issues: str = Field(default="", alias="Issues")
|
||||
|
||||
|
||||
@dataclass
|
||||
class Model:
|
||||
class Model(BaseModel):
|
||||
location: str
|
||||
model_url: str
|
||||
|
||||
|
||||
@dataclass
|
||||
class ComfyConfig:
|
||||
publisher_id: str = ""
|
||||
display_name: str = ""
|
||||
icon: str = ""
|
||||
models: List[Model] = field(default_factory=list)
|
||||
includes: List[str] = field(default_factory=list)
|
||||
class ComfyConfig(BaseModel):
|
||||
publisher_id: str = Field(default="", alias="PublisherId")
|
||||
display_name: str = Field(default="", alias="DisplayName")
|
||||
icon: str = Field(default="", alias="Icon")
|
||||
models: List[Model] = Field(default_factory=list, alias="Models")
|
||||
includes: List[str] = Field(default_factory=list)
|
||||
|
||||
|
||||
@dataclass
|
||||
class License:
|
||||
class License(BaseModel):
|
||||
file: str = ""
|
||||
text: str = ""
|
||||
|
||||
|
||||
@dataclass
|
||||
class ProjectConfig:
|
||||
class ProjectConfig(BaseModel):
|
||||
name: str = ""
|
||||
description: str = ""
|
||||
version: str = "1.0.0"
|
||||
requires_python: str = ">= 3.9"
|
||||
dependencies: List[str] = field(default_factory=list)
|
||||
license: License = field(default_factory=License)
|
||||
urls: URLs = field(default_factory=URLs)
|
||||
requires_python: str = Field(default=">= 3.9", alias="requires-python")
|
||||
dependencies: List[str] = Field(default_factory=list)
|
||||
license: License = Field(default_factory=License)
|
||||
urls: URLs = Field(default_factory=URLs)
|
||||
|
||||
|
||||
@dataclass
|
||||
class PyProjectConfig:
|
||||
project: ProjectConfig = field(default_factory=ProjectConfig)
|
||||
tool_comfy: ComfyConfig = field(default_factory=ComfyConfig)
|
||||
class PyProjectConfig(BaseModel):
|
||||
project: ProjectConfig = Field(default_factory=ProjectConfig)
|
||||
tool_comfy: ComfyConfig = Field(default_factory=ComfyConfig)
|
||||
@ -24,4 +24,3 @@ spandrel
|
||||
soundfile
|
||||
av>=14.2.0
|
||||
pydantic~=2.0
|
||||
tomlkit
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user