120 lines
4.4 KiB
Python
120 lines
4.4 KiB
Python
from typing_extensions import override
|
|
|
|
from comfy_api.latest import ComfyExtension, io, _io
|
|
|
|
import os
|
|
|
|
import folder_paths
|
|
|
|
|
|
def list_dirs(folder_type: str) -> list[str]:
|
|
dirs = set([""])
|
|
for base in folder_paths.get_folder_paths(folder_type):
|
|
for root, subdirs, _ in os.walk(base):
|
|
for d in subdirs:
|
|
rel = os.path.relpath(os.path.join(root, d), base)
|
|
rel = rel.replace("\\", "/")
|
|
dirs.add(rel)
|
|
return sorted(dirs)
|
|
|
|
|
|
def list_files(folder_type: str, rel_dir: str) -> list[str]:
|
|
files: list[str] = []
|
|
allowed_ext = {".safetensors", ".ckpt", ".bin", ".pt", ".pth"}
|
|
for base in folder_paths.get_folder_paths(folder_type):
|
|
dir_path = os.path.join(base, rel_dir) if rel_dir else base
|
|
if not os.path.isdir(dir_path):
|
|
continue
|
|
for root, _, fnames in os.walk(dir_path):
|
|
for f in fnames:
|
|
name_low = f.lower()
|
|
ext = os.path.splitext(name_low)[1]
|
|
if ext not in allowed_ext:
|
|
continue
|
|
rel = os.path.relpath(os.path.join(root, f), dir_path)
|
|
rel = rel.replace("\\", "/")
|
|
files.append(rel if rel_dir == "" else os.path.join(rel_dir, rel).replace("\\", "/"))
|
|
return sorted(set(files))
|
|
|
|
|
|
def build_folder_file_input(folder_id: str, file_id: str, folder_type: str, tooltip: str | None = None) -> _io.DynamicCombo.Input:
|
|
options: list[_io.DynamicCombo.Option] = []
|
|
for rel_dir in list_dirs(folder_type):
|
|
file_options = list_files(folder_type, rel_dir)
|
|
if not file_options:
|
|
continue
|
|
inputs = [io.Combo.Input(file_id, options=file_options, tooltip=tooltip)]
|
|
options.append(_io.DynamicCombo.Option(rel_dir if rel_dir else "root", inputs))
|
|
if not options:
|
|
options.append(_io.DynamicCombo.Option("No files found", [io.Combo.Input(file_id, options=["<none>"], tooltip="No files found")]))
|
|
return _io.DynamicCombo.Input(folder_id, options=options, tooltip="Select folder")
|
|
|
|
|
|
def resolve_selected_path(folder_type: str, selection: dict, folder_id: str, file_id: str) -> str:
|
|
if not isinstance(selection, dict):
|
|
raise ValueError("Invalid selection")
|
|
rel_dir = selection.get(folder_id, "root")
|
|
if rel_dir == "root":
|
|
rel_dir = ""
|
|
file_rel = selection.get(file_id)
|
|
if not file_rel or file_rel == "<none>":
|
|
raise FileNotFoundError("No file selected")
|
|
rel_path = file_rel.replace("\\", "/")
|
|
for base in folder_paths.get_folder_paths(folder_type):
|
|
candidate = os.path.join(base, rel_path)
|
|
if os.path.exists(candidate):
|
|
return candidate
|
|
raise FileNotFoundError(f"Could not resolve path for selection: {rel_path}")
|
|
|
|
|
|
class UNETLoader(io.ComfyNode):
|
|
CATEGORY = "SK Loader/Advanced"
|
|
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
return io.Schema(
|
|
node_id="SK_UNETLoader",
|
|
display_name="[SK] Diffusion Model",
|
|
category="SK Loader/Advanced",
|
|
inputs=[
|
|
build_folder_file_input("unet_folder", "unet_name", "diffusion_models"),
|
|
io.Combo.Input(
|
|
"weight_dtype",
|
|
options=["default", "fp8_e4m3fn", "fp8_e4m3fn_fast", "fp8_e5m2"],
|
|
),
|
|
],
|
|
outputs=[
|
|
io.Model.Output(),
|
|
],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, unet_folder: dict, weight_dtype: str) -> io.NodeOutput:
|
|
import torch
|
|
import comfy.sd
|
|
|
|
model_options = {}
|
|
if weight_dtype == "fp8_e4m3fn":
|
|
model_options["dtype"] = torch.float8_e4m3fn
|
|
elif weight_dtype == "fp8_e4m3fn_fast":
|
|
model_options["dtype"] = torch.float8_e4m3fn
|
|
model_options["fp8_optimizations"] = True
|
|
elif weight_dtype == "fp8_e5m2":
|
|
model_options["dtype"] = torch.float8_e5m2
|
|
|
|
unet_path = resolve_selected_path("diffusion_models", unet_folder, "unet_folder", "unet_name")
|
|
model = comfy.sd.load_diffusion_model(unet_path, model_options=model_options)
|
|
return io.NodeOutput(model)
|
|
|
|
|
|
class DiffusionModelExtension(ComfyExtension):
|
|
@override
|
|
async def get_node_list(self) -> list[type[io.ComfyNode]]:
|
|
return [
|
|
UNETLoader,
|
|
]
|
|
|
|
|
|
async def comfy_entrypoint() -> DiffusionModelExtension:
|
|
return DiffusionModelExtension()
|