fix custom_node import

This commit is contained in:
Tung Nguyen
2023-08-30 09:27:25 +07:00
parent 2848b82d6b
commit 627b24305f
2 changed files with 73 additions and 53 deletions
+68 -50
View File
@@ -1,4 +1,5 @@
import os
import sys
from typing import Dict
import folder_paths
@@ -10,15 +11,80 @@ from .patch import patch_all, unpatch_all
custom_nodes = folder_paths.get_folder_paths("custom_nodes")
efficieny_dir_names = ["Efficiency", "efficiency-nodes-comfyui"]
NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS = {}, {}
class KSamplerWithSharpness(KSampler):
@classmethod
def INPUT_TYPES(cls):
inputs = KSampler.INPUT_TYPES()
inputs["optional"] = {
"sharpness": (
"FLOAT",
{"default": 2.0, "min": 0.0, "max": 100.0, "step": 0.01},
)
}
return inputs
CATEGORY = "Art Venture/Sampling"
def sample(self, *args, sharpness=2.0, **kwargs):
patch.sharpness = sharpness
patch_all()
results = super().sample(*args, **kwargs)
unpatch_all()
return results
class KSamplerAdvancedWithSharpness(KSamplerAdvanced):
@classmethod
def INPUT_TYPES(cls):
inputs = KSamplerAdvanced.INPUT_TYPES()
inputs["optional"] = {
"sharpness": (
"FLOAT",
{"default": 2.0, "min": 0.0, "max": 100.0, "step": 0.01},
)
}
return inputs
CATEGORY = "Art Venture/Sampling"
def sample(self, *args, sharpness=2.0, **kwargs):
patch.sharpness = sharpness
patch_all()
results = super().sample(*args, **kwargs)
unpatch_all()
return results
NODE_CLASS_MAPPINGS = {
"Fooocus_KSampler": KSamplerWithSharpness,
"Fooocus_KSamplerAdvanced": KSamplerAdvancedWithSharpness,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"Fooocus_KSampler": "KSampler Fooocus",
"Fooocus_KSamplerAdvanced": "KSampler Adv. Fooocus",
}
try:
efficieny_path = None
for custom_node in custom_nodes:
custom_node = (
custom_node
if not os.path.islink(custom_node)
else os.readlink(custom_node)
)
for efficiency_dir in efficieny_dir_names:
if efficiency_dir in os.listdir(custom_node):
efficieny_path = os.path.join(custom_node, efficiency_dir)
efficieny_path = os.path.abspath(
os.path.join(custom_node, efficiency_dir)
)
if custom_node not in sys.path:
sys.path.append(custom_node)
break
if efficieny_path is None:
@@ -32,50 +98,6 @@ try:
TSC_KSampler = efficieny_nodes["KSampler (Efficient)"]
TSC_KSamplerAdvanced = efficieny_nodes["KSampler Adv. (Efficient)"]
class KSamplerWithSharpness(KSampler):
@classmethod
def INPUT_TYPES(cls):
inputs = KSampler.INPUT_TYPES()
inputs["optional"] = {
"sharpness": (
"FLOAT",
{"default": 2.0, "min": 0.0, "max": 100.0, "step": 0.01},
)
}
return inputs
CATEGORY = "Art Venture/Sampling"
def sample(self, *args, sharpness=2.0, **kwargs):
patch.sharpness = sharpness
patch_all()
results = super().sample(*args, **kwargs)
unpatch_all()
return results
class KSamplerAdvancedWithSharpness(KSamplerAdvanced):
@classmethod
def INPUT_TYPES(cls):
inputs = KSamplerAdvanced.INPUT_TYPES()
inputs["optional"] = {
"sharpness": (
"FLOAT",
{"default": 2.0, "min": 0.0, "max": 100.0, "step": 0.01},
)
}
return inputs
CATEGORY = "Art Venture/Sampling"
def sample(self, *args, sharpness=2.0, **kwargs):
patch.sharpness = sharpness
patch_all()
results = super().sample(*args, **kwargs)
unpatch_all()
return results
class KSamplerEfficientWithSharpness(TSC_KSampler):
@classmethod
def INPUT_TYPES(cls):
@@ -118,16 +140,12 @@ try:
NODE_CLASS_MAPPINGS.update(
{
"Fooocus_KSampler": KSamplerWithSharpness,
"Fooocus_KSamplerAdvanced": KSamplerAdvancedWithSharpness,
"Fooocus_KSamplerEfficient": KSamplerEfficientWithSharpness,
"Fooocus_KSamplerEfficientAdvanced": KSamplerEfficientAdvancedWithSharpness,
}
)
NODE_DISPLAY_NAME_MAPPINGS.update(
{
"Fooocus_KSampler": "KSampler Fooocus",
"Fooocus_KSamplerAdvanced": "KSampler Adv. Fooocus",
"Fooocus_KSamplerEfficient": "KSampler Efficient Fooocus",
"Fooocus_KSamplerEfficientAdvanced": "KSampler Adv. Efficient Fooocus",
}
+5 -3
View File
@@ -168,10 +168,12 @@ def load_module(module_path):
module_name = os.path.basename(module_path)
if os.path.isfile(module_path):
module_spec = importlib.util.spec_from_file_location(module_name, module_path)
module_spec = importlib.util.spec_from_file_location(module_name, module_path)
else:
module_spec = importlib.util.spec_from_file_location(module_name, os.path.join(module_path, "__init__.py"))
module_spec = importlib.util.spec_from_file_location(
module_name, os.path.join(module_path, "__init__.py")
)
module = importlib.util.module_from_spec(module_spec)
module_spec.loader.exec_module(module)