376 lines
11 KiB
Python
376 lines
11 KiB
Python
"""Model discovery, loading, and GPU residency for NormalCrafter.
|
|
|
|
NormalCrafter uses its own ComfyUI model category: ``normalcrafter``.
|
|
All path discovery goes through ComfyUI's ``folder_paths`` registry, so paths
|
|
registered through ``extra_model_paths.yaml`` are visible automatically.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
import gc
|
|
import logging
|
|
import os
|
|
|
|
import folder_paths
|
|
import torch
|
|
from diffusers import AutoencoderKLTemporalDecoder, EulerDiscreteScheduler
|
|
from huggingface_hub import snapshot_download
|
|
from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection
|
|
|
|
from .unet import NormalCrafterUNet
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
MODEL_CATEGORY = "normalcrafter"
|
|
NORMALCRAFTER_REPO = "Yanrui95/NormalCrafter"
|
|
SVD_REPO = "stabilityai/stable-video-diffusion-img2vid-xt"
|
|
|
|
HF_PREFIX = "HF: "
|
|
LOCAL_PREFIX = "local: "
|
|
|
|
|
|
def _register_model_category() -> None:
|
|
"""Ensure ComfyUI knows the default ``models/normalcrafter`` location.
|
|
|
|
If ``extra_model_paths.yaml`` already registered one or more ``normalcrafter``
|
|
paths, they stay in ComfyUI's ordering and the standard ComfyUI model path is
|
|
simply added as another search location.
|
|
"""
|
|
|
|
default_path = os.path.join(folder_paths.models_dir, MODEL_CATEGORY)
|
|
|
|
try:
|
|
existing = folder_paths.get_folder_paths(MODEL_CATEGORY)
|
|
except KeyError:
|
|
existing = []
|
|
|
|
if default_path in existing:
|
|
return
|
|
|
|
# Current ComfyUI supports ``is_default``. The fallback keeps compatibility
|
|
# with older versions whose function accepted only two arguments.
|
|
try:
|
|
folder_paths.add_model_folder_path(
|
|
MODEL_CATEGORY,
|
|
default_path,
|
|
is_default=not existing,
|
|
)
|
|
except TypeError:
|
|
folder_paths.add_model_folder_path(MODEL_CATEGORY, default_path)
|
|
|
|
|
|
_register_model_category()
|
|
|
|
|
|
def _model_roots() -> list[str]:
|
|
"""Return all NormalCrafter roots currently registered in ComfyUI."""
|
|
|
|
roots = folder_paths.get_folder_paths(MODEL_CATEGORY)
|
|
# Preserve ComfyUI ordering while removing duplicate spellings of a path.
|
|
seen: set[str] = set()
|
|
result: list[str] = []
|
|
for root in roots:
|
|
normalized = os.path.normcase(os.path.abspath(root))
|
|
if normalized in seen:
|
|
continue
|
|
seen.add(normalized)
|
|
result.append(os.path.abspath(root))
|
|
return result
|
|
|
|
|
|
def _is_within_registered_root(path: str) -> bool:
|
|
candidate = os.path.normcase(os.path.abspath(path))
|
|
for root in _model_roots():
|
|
registered = os.path.normcase(os.path.abspath(root))
|
|
try:
|
|
if os.path.commonpath((candidate, registered)) == registered:
|
|
return True
|
|
except ValueError:
|
|
# Different Windows drives cannot share a common path.
|
|
continue
|
|
return False
|
|
|
|
|
|
def _local_model_dirs(required_subfolders: tuple[str, ...]) -> list[str]:
|
|
"""Find compatible model directories across every ComfyUI-registered root."""
|
|
|
|
results: list[str] = []
|
|
|
|
for root in _model_roots():
|
|
if not os.path.isdir(root):
|
|
continue
|
|
|
|
try:
|
|
names = sorted(os.listdir(root), key=str.lower)
|
|
except OSError:
|
|
continue
|
|
|
|
for name in names:
|
|
path = os.path.join(root, name)
|
|
if not os.path.isdir(path):
|
|
continue
|
|
if all(
|
|
os.path.isdir(os.path.join(path, subfolder))
|
|
for subfolder in required_subfolders
|
|
):
|
|
results.append(f"{LOCAL_PREFIX}{path}")
|
|
|
|
return results
|
|
|
|
|
|
def normalcrafter_sources() -> list[str]:
|
|
"""Dropdown values suitable for NormalCrafter UNet + VAE loading."""
|
|
|
|
return [
|
|
f"{HF_PREFIX}{NORMALCRAFTER_REPO}",
|
|
*_local_model_dirs(("unet", "vae")),
|
|
]
|
|
|
|
|
|
def base_model_sources() -> list[str]:
|
|
"""Dropdown values suitable for SVD CLIP + scheduler loading."""
|
|
|
|
return [
|
|
f"{HF_PREFIX}{SVD_REPO}",
|
|
*_local_model_dirs(("image_encoder", "feature_extractor", "scheduler")),
|
|
]
|
|
|
|
|
|
def _existing_repo_dir(repo_id: str) -> str | None:
|
|
"""Return an already-downloaded repo from any registered model root."""
|
|
|
|
folder_name = repo_id.replace("/", "--")
|
|
for root in _model_roots():
|
|
candidate = os.path.join(root, folder_name)
|
|
if os.path.isdir(candidate):
|
|
return candidate
|
|
return None
|
|
|
|
|
|
def _download_root() -> str:
|
|
"""Choose the first writable NormalCrafter root in ComfyUI's own order."""
|
|
|
|
errors: list[str] = []
|
|
for root in _model_roots():
|
|
try:
|
|
os.makedirs(root, exist_ok=True)
|
|
if os.access(root, os.W_OK):
|
|
return root
|
|
errors.append(f"{root} (not writable)")
|
|
except OSError as exc:
|
|
errors.append(f"{root} ({exc})")
|
|
|
|
details = "; ".join(errors) or "no paths registered"
|
|
raise RuntimeError(
|
|
"No writable ComfyUI 'normalcrafter' model path is available: " + details
|
|
)
|
|
|
|
|
|
def _download_repo(repo_id: str, required_subfolders: tuple[str, ...]) -> str:
|
|
"""Download only required components into a registered NormalCrafter root."""
|
|
|
|
existing = _existing_repo_dir(repo_id)
|
|
if existing is not None:
|
|
return existing
|
|
|
|
target = os.path.join(_download_root(), repo_id.replace("/", "--"))
|
|
allow_patterns = [f"{subfolder}/**" for subfolder in required_subfolders]
|
|
|
|
snapshot_download(
|
|
repo_id=repo_id,
|
|
local_dir=target,
|
|
allow_patterns=allow_patterns,
|
|
)
|
|
return target
|
|
|
|
|
|
def _resolve_source(source: str, required_subfolders: tuple[str, ...]) -> str:
|
|
"""Resolve a dropdown value to a validated local model directory."""
|
|
|
|
if source.startswith(HF_PREFIX):
|
|
path = _download_repo(source[len(HF_PREFIX) :], required_subfolders)
|
|
elif source.startswith(LOCAL_PREFIX):
|
|
path = source[len(LOCAL_PREFIX) :]
|
|
if not _is_within_registered_root(path):
|
|
raise ValueError(
|
|
"Selected model path is not inside a ComfyUI-registered "
|
|
f"'{MODEL_CATEGORY}' directory: {path}"
|
|
)
|
|
else:
|
|
raise ValueError(f"Unsupported model source: {source}")
|
|
|
|
if not os.path.isdir(path):
|
|
raise FileNotFoundError(f"NormalCrafter model folder not found: {path}")
|
|
|
|
missing = [
|
|
subfolder
|
|
for subfolder in required_subfolders
|
|
if not os.path.isdir(os.path.join(path, subfolder))
|
|
]
|
|
if missing:
|
|
raise FileNotFoundError(
|
|
f"Model folder {path} is missing required subfolders: {', '.join(missing)}"
|
|
)
|
|
|
|
return path
|
|
|
|
|
|
def _clear_cuda_cache() -> None:
|
|
gc.collect()
|
|
if torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
|
|
|
|
def _auto_dtype(device: torch.device) -> torch.dtype:
|
|
return torch.float16 if device.type == "cuda" else torch.float32
|
|
|
|
|
|
@dataclass
|
|
class NormalCrafterModels:
|
|
"""Own the heavy modules and keep GPU residency explicit."""
|
|
|
|
unet: NormalCrafterUNet
|
|
vae: AutoencoderKLTemporalDecoder
|
|
image_encoder: CLIPVisionModelWithProjection
|
|
feature_extractor: CLIPImageProcessor
|
|
scheduler: EulerDiscreteScheduler
|
|
dtype: torch.dtype
|
|
|
|
def _modules(self) -> dict[str, torch.nn.Module]:
|
|
return {
|
|
"image_encoder": self.image_encoder,
|
|
"vae": self.vae,
|
|
"unet": self.unet,
|
|
}
|
|
|
|
def activate(self, name: str, device: torch.device) -> torch.nn.Module:
|
|
"""Keep only one heavy component on the inference device."""
|
|
|
|
modules = self._modules()
|
|
if name not in modules:
|
|
raise ValueError(f"Unknown NormalCrafter component: {name}")
|
|
|
|
cpu = torch.device("cpu")
|
|
for module_name, module in modules.items():
|
|
target = device if module_name == name else cpu
|
|
try:
|
|
current = next(module.parameters()).device
|
|
except StopIteration:
|
|
current = cpu
|
|
|
|
if current != target:
|
|
module.to(target)
|
|
|
|
_clear_cuda_cache()
|
|
return modules[name]
|
|
|
|
def offload_all(self) -> None:
|
|
"""Return every heavy module to CPU and release unused CUDA blocks."""
|
|
|
|
for module in self._modules().values():
|
|
module.to("cpu")
|
|
_clear_cuda_cache()
|
|
|
|
|
|
def _enable_memory_efficient_attention(unet: NormalCrafterUNet) -> None:
|
|
try:
|
|
unet.enable_xformers_memory_efficient_attention()
|
|
logger.info("NormalCrafter: xFormers attention enabled")
|
|
except Exception as exc:
|
|
logger.info("NormalCrafter: using Diffusers default attention (%s)", exc)
|
|
|
|
|
|
def load_models(
|
|
*,
|
|
device: torch.device,
|
|
normalcrafter_source: str,
|
|
base_model_source: str,
|
|
) -> NormalCrafterModels:
|
|
"""Load the selected model components on CPU.
|
|
|
|
UNet and CLIP use FP16 on CUDA systems and FP32 on CPU. The temporal VAE is
|
|
intentionally loaded in FP32, matching its force-upcast inference behavior.
|
|
Inference moves CLIP, VAE and UNet to the GPU one stage at a time.
|
|
"""
|
|
|
|
normalcrafter_path = _resolve_source(
|
|
normalcrafter_source,
|
|
("unet", "vae"),
|
|
)
|
|
base_model_path = _resolve_source(
|
|
base_model_source,
|
|
("image_encoder", "feature_extractor", "scheduler"),
|
|
)
|
|
|
|
dtype = _auto_dtype(device)
|
|
common = {
|
|
"torch_dtype": dtype,
|
|
"low_cpu_mem_usage": True,
|
|
"local_files_only": True,
|
|
}
|
|
|
|
unet = NormalCrafterUNet.from_pretrained(
|
|
normalcrafter_path,
|
|
subfolder="unet",
|
|
**common,
|
|
)
|
|
|
|
# The temporal VAE declares force_upcast and is evaluated in FP32 in the
|
|
# reference NormalCrafter pipeline. Load it in FP32 from the start instead of
|
|
# repeatedly calling vae.to(dtype=...) during inference. Recent Diffusers
|
|
# versions explicitly warn against those runtime dtype casts because selected
|
|
# VAE modules may need to remain in FP32.
|
|
vae = AutoencoderKLTemporalDecoder.from_pretrained(
|
|
normalcrafter_path,
|
|
subfolder="vae",
|
|
torch_dtype=torch.float32,
|
|
low_cpu_mem_usage=True,
|
|
local_files_only=True,
|
|
)
|
|
|
|
encoder_kwargs = dict(common)
|
|
if dtype == torch.float16:
|
|
encoder_kwargs["variant"] = "fp16"
|
|
|
|
try:
|
|
image_encoder = CLIPVisionModelWithProjection.from_pretrained(
|
|
base_model_path,
|
|
subfolder="image_encoder",
|
|
**encoder_kwargs,
|
|
)
|
|
except (OSError, ValueError, TypeError):
|
|
encoder_kwargs.pop("variant", None)
|
|
image_encoder = CLIPVisionModelWithProjection.from_pretrained(
|
|
base_model_path,
|
|
subfolder="image_encoder",
|
|
**encoder_kwargs,
|
|
)
|
|
|
|
feature_extractor = CLIPImageProcessor.from_pretrained(
|
|
base_model_path,
|
|
subfolder="feature_extractor",
|
|
local_files_only=True,
|
|
)
|
|
scheduler = EulerDiscreteScheduler.from_pretrained(
|
|
base_model_path,
|
|
subfolder="scheduler",
|
|
local_files_only=True,
|
|
)
|
|
|
|
unet.eval()
|
|
vae.eval()
|
|
image_encoder.eval()
|
|
_enable_memory_efficient_attention(unet)
|
|
|
|
models = NormalCrafterModels(
|
|
unet=unet,
|
|
vae=vae,
|
|
image_encoder=image_encoder,
|
|
feature_extractor=feature_extractor,
|
|
scheduler=scheduler,
|
|
dtype=dtype,
|
|
)
|
|
models.offload_all()
|
|
return models
|