jaidaken
f09734b0ee
Some checks failed
Python Linting / Run Ruff (push) Has been cancelled
Python Linting / Run Pylint (push) Has been cancelled
Full Comfy CI Workflow Runs / test-stable (12.1, , linux, 3.10, [self-hosted Linux], stable) (push) Has been cancelled
Full Comfy CI Workflow Runs / test-stable (12.1, , linux, 3.11, [self-hosted Linux], stable) (push) Has been cancelled
Full Comfy CI Workflow Runs / test-stable (12.1, , linux, 3.12, [self-hosted Linux], stable) (push) Has been cancelled
Full Comfy CI Workflow Runs / test-unix-nightly (12.1, , linux, 3.11, [self-hosted Linux], nightly) (push) Has been cancelled
Execution Tests / test (macos-latest) (push) Has been cancelled
Execution Tests / test (ubuntu-latest) (push) Has been cancelled
Execution Tests / test (windows-latest) (push) Has been cancelled
Test server launches without errors / test (push) Has been cancelled
Unit Tests / test (macos-latest) (push) Has been cancelled
Unit Tests / test (ubuntu-latest) (push) Has been cancelled
Unit Tests / test (windows-2022) (push) Has been cancelled
Includes 30 custom nodes committed directly, 7 Civitai-exclusive loras stored via Git LFS, and a setup script that installs all dependencies and downloads HuggingFace-hosted models on vast.ai. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
71 lines
2.9 KiB
Python
71 lines
2.9 KiB
Python
import os
|
|
import folder_paths
|
|
from . import subcore
|
|
from . import utils
|
|
import logging
|
|
|
|
|
|
def update_model_paths(model_path):
|
|
utils.add_folder_path_and_extensions("ultralytics_bbox", [os.path.join(model_path, "ultralytics", "bbox")], folder_paths.supported_pt_extensions)
|
|
utils.add_folder_path_and_extensions("ultralytics_segm", [os.path.join(model_path, "ultralytics", "segm")], folder_paths.supported_pt_extensions)
|
|
utils.add_folder_path_and_extensions("ultralytics", [os.path.join(model_path, "ultralytics")], folder_paths.supported_pt_extensions)
|
|
logging.info(f'[Impact Subpack] ultralytics_bbox: {", ".join(folder_paths.folder_names_and_paths["ultralytics_bbox"][0])}')
|
|
logging.info(f'[Impact Subpack] ultralytics_segm: {", ".join(folder_paths.folder_names_and_paths["ultralytics_segm"][0])}')
|
|
|
|
update_model_paths(folder_paths.models_dir)
|
|
if 'download_model_base' in folder_paths.folder_names_and_paths:
|
|
update_model_paths(folder_paths.get_folder_paths('download_model_base')[0])
|
|
|
|
|
|
class UltralyticsDetectorProvider:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
bboxs = ["bbox/"+x for x in folder_paths.get_filename_list("ultralytics_bbox")]
|
|
segms = ["segm/"+x for x in folder_paths.get_filename_list("ultralytics_segm")]
|
|
return {"required": {"model_name": (bboxs + segms, )}}
|
|
RETURN_TYPES = ("BBOX_DETECTOR", "SEGM_DETECTOR")
|
|
FUNCTION = "doit"
|
|
|
|
CATEGORY = "ImpactPack"
|
|
|
|
def doit(self, model_name):
|
|
model_path = folder_paths.get_full_path("ultralytics", model_name)
|
|
|
|
if model_path is None:
|
|
if model_name.startswith('bbox/'):
|
|
model_path = folder_paths.get_full_path("ultralytics_bbox", model_name[5:])
|
|
elif model_name.startswith('segm/'):
|
|
model_path = folder_paths.get_full_path("ultralytics_segm", model_name[5:])
|
|
|
|
if model_path is None:
|
|
logging.error(f"[Impact Subpack] model file '{model_name}' is not found in one of the following directories:")
|
|
|
|
cands = []
|
|
cands.extend(folder_paths.get_folder_paths("ultralytics"))
|
|
if model_name.startswith('bbox/'):
|
|
cands.extend(folder_paths.get_folder_paths("ultralytics_bbox"))
|
|
elif model_name.startswith('segm/'):
|
|
cands.extend(folder_paths.get_folder_paths("ultralytics_segm"))
|
|
|
|
formatted_cands = "\n\t".join(cands)
|
|
logging.error(f'\t{formatted_cands}\n')
|
|
|
|
raise ValueError(f"[Impact Subpack] model file '{model_name}' is not found.")
|
|
|
|
model = subcore.load_yolo(model_path)
|
|
|
|
if model_name.startswith("bbox"):
|
|
return subcore.UltraBBoxDetector(model), subcore.NO_SEGM_DETECTOR()
|
|
else:
|
|
return subcore.UltraBBoxDetector(model), subcore.UltraSegmDetector(model)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"UltralyticsDetectorProvider": UltralyticsDetectorProvider
|
|
}
|
|
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
|
|
}
|