Refactor: Replace dynamic relative imports with model class registry for cross-platform compatibility. Minor: Adjust ASCII logo width for better screen fit

This commit is contained in:
Adrien Toupet
2025-09-20 01:53:25 -04:00
parent e89a8b00a5
commit d0d24a18cb
5 changed files with 37 additions and 50 deletions
+2 -8
View File
@@ -5,10 +5,7 @@ __object__:
dit:
model:
__object__:
path:
- "custom_nodes.ComfyUI-SeedVR2_VideoUpscaler.src.models.dit_3b.nadit"
- "ComfyUI.custom_nodes.ComfyUI-SeedVR2_VideoUpscaler.src.models.dit_3b.nadit"
- "src.models.dit_3b.nadit"
path: "dit_3b.nadit"
name: "NaDiT"
args: "as_params"
vid_in_channels: 33
@@ -48,10 +45,7 @@ ema:
vae:
model:
__object__:
path:
- "custom_nodes.ComfyUI-SeedVR2_VideoUpscaler.src.models.video_vae_v3.modules.attn_video_vae"
- "ComfyUI.custom_nodes.ComfyUI-SeedVR2_VideoUpscaler.src.models.video_vae_v3.modules.attn_video_vae"
- "src.models.video_vae_v3.modules.attn_video_vae"
path: "video_vae_v3.modules.attn_video_vae"
name: "VideoAutoencoderKLWrapper"
args: "as_params"
freeze_encoder: False
+2 -8
View File
@@ -5,10 +5,7 @@ __object__:
dit:
model:
__object__:
path:
- "custom_nodes.ComfyUI-SeedVR2_VideoUpscaler.src.models.dit_7b.nadit"
- "ComfyUI.custom_nodes.ComfyUI-SeedVR2_VideoUpscaler.src.models.dit_7b.nadit"
- "src.models.dit_7b.nadit"
path: "dit_7b.nadit"
name: "NaDiT"
args: "as_params"
vid_in_channels: 33
@@ -45,10 +42,7 @@ ema:
vae:
model:
__object__:
path:
- "custom_nodes.ComfyUI-SeedVR2_VideoUpscaler.src.models.video_vae_v3.modules.attn_video_vae"
- "ComfyUI.custom_nodes.ComfyUI-SeedVR2_VideoUpscaler.src.models.video_vae_v3.modules.attn_video_vae"
- "src.models.video_vae_v3.modules.attn_video_vae"
path: "video_vae_v3.modules.attn_video_vae"
name: "VideoAutoencoderKLWrapper"
args: "as_params"
freeze_encoder: False
+13 -26
View File
@@ -19,6 +19,7 @@ Configuration utility functions
import importlib
from typing import Any, Callable, List, Union
from omegaconf import DictConfig, ListConfig, OmegaConf
from ..utils.model_registry import MODEL_CLASSES
try:
OmegaConf.register_new_resolver("eval", eval)
@@ -87,40 +88,26 @@ def resolve_inheritance(config: Union[DictConfig, ListConfig]) -> Any:
return config
def import_item(path: Union[str, List[str]], name: str) -> Any:
def import_item(path: str, name: str) -> Any:
"""
Import a python item with fallback support.
Import a python item, checking model registry first.
Args:
path: Single path string or list of paths to try (fallback order)
path: Module path
name: Class/function name to import
Returns:
Imported object
Example:
import_item("path.to.file", "MyClass") -> MyClass
import_item(["path1.to.file", "path2.to.file"], "MyClass") -> MyClass (first working path)
"""
if isinstance(path, str):
# Single path - original behavior
"""
# Simple lookup with path as key
if path in MODEL_CLASSES:
return MODEL_CLASSES[path]
# Fallback to dynamic import for everything else
try:
return getattr(importlib.import_module(path), name)
elif isinstance(path, (list, ListConfig)):
# Multiple paths - try each until one works
last_error = None
for single_path in path:
try:
return getattr(importlib.import_module(single_path), name)
except ImportError as e:
last_error = e
continue
# If we get here, none of the paths worked
raise ImportError(f"Could not import '{name}' from any of the paths: {path}. Last error: {last_error}")
else:
raise ValueError(f"Path must be string or list of strings, got: {type(path)}")
except (ImportError, AttributeError) as e:
raise ImportError(f"Could not import '{name}' from '{path}': {e}")
def create_object(config: DictConfig) -> Any:
+7 -7
View File
@@ -222,13 +222,13 @@ class SeedVR2:
debug.start_timer("total_execution", force=True)
debug.log("", category="none", force=True)
debug.log(" ╔══════════════════════════════════════════════════════════════╗", category="none", force=True)
debug.log(" ║ ███████ ███████ ███████ ██████ ██ ██ ██████ ██████ ║", category="none", force=True)
debug.log(" ║ ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ ║", category="none", force=True)
debug.log(" ║ ███████ █████ █████ ██ ██ ██ ██ ██████ █████ ║", category="none", force=True)
debug.log(" ║ ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ ║", category="none", force=True)
debug.log(" ║ ███████ ███████ ███████ ██████ ████ ██ ██ ███████ ║", category="none", force=True)
debug.log(" ╚══════════════════════════════════════════════════════════════╝", category="none", force=True)
debug.log(" ╔══════════════════════════════════════════════════════════╗", category="none", force=True)
debug.log(" ║ ███████ ███████ ███████ ██████ ██ ██ ██████ ███████ ║", category="none", force=True)
debug.log(" ║ ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ ║", category="none", force=True)
debug.log(" ║ ███████ █████ █████ ██ ██ ██ ██ ██████ █████ ║", category="none", force=True)
debug.log(" ║ ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ ║", category="none", force=True)
debug.log(" ║ ███████ ███████ ███████ ██████ ████ ██ ██ ███████ ║", category="none", force=True)
debug.log(" ╚══════════════════════════════════════════════════════════╝", category="none", force=True)
debug.log("", category="none", force=True)
debug.log("━━━━━━━━━ Model Preparation ━━━━━━━━━", category="none")
+13 -1
View File
@@ -3,10 +3,22 @@ Model Registry for SeedVR2
Central registry for model definitions, repositories, and metadata
"""
from typing import Dict, List, Optional
from typing import Dict, List, Optional, Any
from dataclasses import dataclass
from .constants import SEEDVR2_MODEL_TYPE, is_supported_model_file, get_base_cache_dir
# Model class imports using relative imports
from ..models.dit_3b.nadit import NaDiT as NaDiT3B
from ..models.dit_7b.nadit import NaDiT as NaDiT7B
from ..models.video_vae_v3.modules.attn_video_vae import VideoAutoencoderKLWrapper
# Model classes - simple registry with clear keys
MODEL_CLASSES = {
"dit_3b.nadit": NaDiT3B,
"dit_7b.nadit": NaDiT7B,
"video_vae_v3.modules.attn_video_vae": VideoAutoencoderKLWrapper,
}
@dataclass
class ModelInfo:
"""Model metadata"""