145 lines
6.2 KiB
Python
145 lines
6.2 KiB
Python
import torch
|
|
import os
|
|
from comfy import model_management
|
|
import folder_paths
|
|
|
|
from .. import logger # Import from parent __init__.py
|
|
from .utils import (
|
|
load_file_from_url_comfy, ARCH_REGISTRY,
|
|
KEEP_MODEL_CONFIGS, FACELIB_MODEL_URLS, FACELIB_DEST_DIR
|
|
)
|
|
|
|
try:
|
|
from wm_facelib.utils.face_restoration_helper import FaceRestoreHelper
|
|
except ImportError as e:
|
|
logger.error(f"Critical Import Error in keep_model_loader.py (facelib): {e}")
|
|
class FaceRestoreHelper: pass
|
|
|
|
class KEEPModelPack:
|
|
def __init__(self, keep_net, face_helper, bg_upscale_model, face_upscale_model, model_type_str):
|
|
self.keep_net = keep_net
|
|
self.face_helper = face_helper
|
|
self.bg_upscale_model = bg_upscale_model
|
|
self.face_upscale_model = face_upscale_model
|
|
self.model_type_str = model_type_str
|
|
self.device = model_management.get_torch_device()
|
|
self.offload_device = model_management.unet_offload_device()
|
|
|
|
def load_device(self):
|
|
"""Move all models to the main compute device."""
|
|
if self.keep_net is not None:
|
|
self.keep_net.to(self.device)
|
|
# The spandrel model descriptor has a 'model' attribute that needs to be moved
|
|
if self.bg_upscale_model is not None:
|
|
self.bg_upscale_model.model.to(self.device)
|
|
if self.face_upscale_model is not None:
|
|
self.face_upscale_model.model.to(self.device)
|
|
|
|
if self.face_helper is not None:
|
|
self.face_helper.device = self.device
|
|
if hasattr(self.face_helper, 'face_detector'):
|
|
self.face_helper.face_detector.to(self.device)
|
|
if hasattr(self.face_helper, 'face_parse'):
|
|
self.face_helper.face_parse.to(self.device)
|
|
|
|
def offload(self):
|
|
"""Move all models to the offload device (CPU)."""
|
|
if self.keep_net is not None:
|
|
self.keep_net.to(self.offload_device)
|
|
if self.bg_upscale_model is not None:
|
|
self.bg_upscale_model.model.to(self.offload_device)
|
|
if self.face_upscale_model is not None:
|
|
self.face_upscale_model.model.to(self.offload_device)
|
|
|
|
if self.face_helper is not None:
|
|
self.face_helper.device = self.offload_device
|
|
if hasattr(self.face_helper, 'face_detector'):
|
|
self.face_helper.face_detector.to(self.offload_device)
|
|
if hasattr(self.face_helper, 'face_parse'):
|
|
self.face_helper.face_parse.to(self.offload_device)
|
|
|
|
model_management.soft_empty_cache()
|
|
|
|
class KEEPModelLoader:
|
|
def __init__(self):
|
|
self.device = model_management.get_torch_device()
|
|
self.offload_device = model_management.unet_offload_device()
|
|
self.loaded_models = {}
|
|
|
|
def load_keep_model_pack(self, model_type_str, detection_model_str,
|
|
bg_upscale_model=None, face_upscale_model=None):
|
|
|
|
bg_upscaler_present = bg_upscale_model is not None
|
|
face_upscaler_present = face_upscale_model is not None
|
|
cache_key = (model_type_str, detection_model_str, bg_upscaler_present, face_upscaler_present)
|
|
|
|
if cache_key in self.loaded_models:
|
|
logger.debug(f"Returning cached base models for {cache_key}")
|
|
cached_pack = self.loaded_models[cache_key]
|
|
model_pack = KEEPModelPack(
|
|
cached_pack.keep_net,
|
|
cached_pack.face_helper,
|
|
bg_upscale_model,
|
|
face_upscale_model,
|
|
model_type_str
|
|
)
|
|
return model_pack
|
|
|
|
if model_type_str not in KEEP_MODEL_CONFIGS:
|
|
raise ValueError(f"Unknown KEEP model type: {model_type_str}")
|
|
|
|
config_entry = KEEP_MODEL_CONFIGS[model_type_str]
|
|
|
|
keep_model_arch_class = ARCH_REGISTRY.get('KEEP')
|
|
if keep_model_arch_class is None:
|
|
raise RuntimeError("KEEP architecture class not found in ARCH_REGISTRY.")
|
|
|
|
net = keep_model_arch_class(**config_entry['architecture']).to(self.offload_device)
|
|
|
|
ckpt_path = load_file_from_url_comfy(
|
|
url=config_entry['url'],
|
|
model_dir_name=config_entry['dest_dir'],
|
|
file_name=os.path.basename(config_entry['url']),
|
|
sha256=config_entry.get('sha256')
|
|
)
|
|
checkpoint = torch.load(ckpt_path, map_location=self.offload_device, weights_only=True)
|
|
|
|
state_dict_key = 'params_ema' if 'params_ema' in checkpoint else 'params'
|
|
state_dict = checkpoint.get(state_dict_key, checkpoint)
|
|
|
|
converted_state_dict = {}
|
|
needs_conversion = any('cross_fuse' in k or 'fuse_convs_dict' in k for k in state_dict.keys())
|
|
if needs_conversion:
|
|
for k, v in state_dict.items():
|
|
new_k = k
|
|
if 'cross_fuse' in k: new_k = k.replace('cross_fuse', 'cfa')
|
|
if 'fuse_convs_dict' in k: new_k = k.replace('fuse_convs_dict', 'cft')
|
|
converted_state_dict[new_k] = v
|
|
state_dict = converted_state_dict
|
|
|
|
net.load_state_dict(state_dict, strict=True)
|
|
net.eval()
|
|
logger.debug(f"KEEP model '{model_type_str}' loaded onto {self.offload_device}.")
|
|
|
|
for fname, (url, sha) in FACELIB_MODEL_URLS.items():
|
|
load_file_from_url_comfy(url=url, model_dir_name=FACELIB_DEST_DIR, file_name=fname, sha256=sha)
|
|
|
|
facelib_actual_paths = os.path.join(folder_paths.models_dir, FACELIB_DEST_DIR)
|
|
os.makedirs(facelib_actual_paths, exist_ok=True)
|
|
|
|
try:
|
|
face_helper = FaceRestoreHelper(
|
|
upscale_factor=1, face_size=512, crop_ratio=(1, 1),
|
|
det_model=detection_model_str, save_ext='png', use_parse=True, device=self.offload_device,
|
|
model_rootpath=facelib_actual_paths
|
|
)
|
|
except Exception as e:
|
|
logger.error(f"Error initializing FaceRestoreHelper: {e}")
|
|
raise
|
|
|
|
model_pack = KEEPModelPack(net, face_helper, bg_upscale_model, face_upscale_model, model_type_str)
|
|
|
|
base_pack_for_cache = KEEPModelPack(net, face_helper, None, None, model_type_str)
|
|
self.loaded_models[cache_key] = base_pack_for_cache
|
|
|
|
return model_pack |