63 lines
2.2 KiB
Python
63 lines
2.2 KiB
Python
# Copyright (c) 2025 Salvador E. Tropea
|
|
# Copyright (c) 2025 Instituto Nacional de Tecnología Industrial
|
|
# License: GPLv3
|
|
# Project: ComfyUI-AudioSeparation
|
|
# This helper is used to load a class from an arbitrary file
|
|
# Gemini 2.5 Pro code
|
|
import importlib
|
|
import logging
|
|
import os
|
|
import sys
|
|
from .misc import NODES_NAME
|
|
|
|
logger = logging.getLogger(f"{NODES_NAME}.load_class")
|
|
|
|
|
|
# Helper to dynamically import the target PyTorch model class
|
|
def import_model_class(location_string: str):
|
|
"""
|
|
Dynamically imports a PyTorch model class from a file path and class name.
|
|
|
|
The location_string is expected to be in the format:
|
|
'path/to/your/file.py:ClassName'
|
|
"""
|
|
module_dir = None
|
|
try:
|
|
# 1. Split the input string into a file path and a class name
|
|
filepath, class_name = location_string.split(':')
|
|
|
|
# Check if the file exists before proceeding
|
|
if not os.path.exists(filepath):
|
|
logger.error(f"File not found at '{filepath}'.")
|
|
sys.exit(1)
|
|
|
|
# 2. Get the directory and the module name from the file path
|
|
module_dir, module_file = os.path.split(filepath)
|
|
module_name = os.path.splitext(module_file)[0]
|
|
|
|
# Add the directory to sys.path to allow Python to find it
|
|
# Add it to the beginning to ensure it's checked first
|
|
sys.path.insert(0, module_dir)
|
|
|
|
# 3. Import the module
|
|
logger.info(f"Importing module '{module_name}' from '{module_dir}'...")
|
|
module = importlib.import_module(module_name)
|
|
|
|
# 4. Get the class from the imported module
|
|
model_class = getattr(module, class_name)
|
|
|
|
except (ValueError, ImportError, AttributeError, FileNotFoundError) as e:
|
|
logger.error(f"Could not import model class from '{location_string}'.")
|
|
logger.error("Please ensure the format is 'path/to/file.py:ClassName'.")
|
|
logger.error(f"Original error: {e}")
|
|
sys.exit(1)
|
|
|
|
finally:
|
|
# 5. Clean up by removing the path we added.
|
|
# This is crucial to avoid polluting the user's environment.
|
|
if module_dir is not None and module_dir in sys.path:
|
|
sys.path.pop(0)
|
|
|
|
logger.info(f"Successfully imported class '{class_name}'.")
|
|
return model_class
|