152 lines
5.1 KiB
Python
152 lines
5.1 KiB
Python
from pathlib import Path
|
|
|
|
from . import ensure_node_root_on_path
|
|
from .download_utils import validate_required_models
|
|
|
|
|
|
_MODEL_CACHE = {}
|
|
|
|
|
|
def _get_repo_root():
|
|
return Path(__file__).resolve().parents[1]
|
|
|
|
|
|
def _torch_load(path, torch):
|
|
try:
|
|
return torch.load(path, map_location="cpu", weights_only=True)
|
|
except TypeError:
|
|
return torch.load(path, map_location="cpu")
|
|
|
|
|
|
def _get_torch_device(torch):
|
|
try:
|
|
import comfy.model_management as model_management
|
|
|
|
return model_management.get_torch_device()
|
|
except Exception:
|
|
return torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
|
|
|
|
def _weight_dtype(torch, weight_dtype):
|
|
dtype_text = str(weight_dtype or "fp16").lower()
|
|
if dtype_text == "fp16":
|
|
return torch.float16
|
|
if dtype_text == "bf16":
|
|
return torch.bfloat16
|
|
if dtype_text == "fp32":
|
|
return torch.float32
|
|
raise ValueError(f"Unsupported HighSync weight_dtype: {weight_dtype}")
|
|
|
|
|
|
def _load_config(models_root, weight_dtype):
|
|
try:
|
|
from omegaconf import OmegaConf
|
|
except ImportError as exc:
|
|
raise RuntimeError("omegaconf is required to load HighSync config. Install requirements.txt first.") from exc
|
|
|
|
config_path = _get_repo_root() / "high_sync" / "configs" / "inference" / "default.yaml"
|
|
config = OmegaConf.load(str(config_path))
|
|
|
|
config.weight_dtype = str(weight_dtype)
|
|
config.denoising_unet_path = str(models_root / "denoising_unet-500.pth")
|
|
config.reference_unet_path = str(models_root / "reference_unet-500.pth")
|
|
config.base_model_path = str(models_root / "sd-image-variations-diffusers")
|
|
config.audio_model_path = str(models_root / "audio_processor" / "whisper_tiny.pt")
|
|
config.vae.model_path = str(models_root / "sd-vae-ft-mse")
|
|
config.mask_path = str(_get_repo_root() / "high_sync" / "masks" / "mask_64.png")
|
|
config.mask_512_path = str(_get_repo_root() / "high_sync" / "masks" / "mask_512.png")
|
|
|
|
return config
|
|
|
|
|
|
def _build_scheduler(config):
|
|
try:
|
|
from diffusers import DDIMScheduler
|
|
from omegaconf import OmegaConf
|
|
except ImportError as exc:
|
|
raise RuntimeError("diffusers and omegaconf are required for HighSync. Install requirements.txt first.") from exc
|
|
|
|
sched_kwargs = OmegaConf.to_container(config.noise_scheduler_kwargs, resolve=True)
|
|
if bool(config.enable_zero_snr):
|
|
sched_kwargs.update(
|
|
rescale_betas_zero_snr=True,
|
|
timestep_spacing="trailing",
|
|
prediction_type="v_prediction",
|
|
)
|
|
return DDIMScheduler(**sched_kwargs)
|
|
|
|
|
|
def load_highsync_model(models_dir, weight_dtype="fp16"):
|
|
ensure_node_root_on_path()
|
|
|
|
try:
|
|
import torch
|
|
from diffusers import AutoencoderKL
|
|
except ImportError as exc:
|
|
raise RuntimeError(
|
|
"torch and diffusers are required to load HighSync. Install this wrapper's requirements.txt first."
|
|
) from exc
|
|
|
|
models_root = validate_required_models(models_dir)
|
|
device = _get_torch_device(torch)
|
|
dtype = _weight_dtype(torch, weight_dtype)
|
|
cache_key = (str(models_root), str(weight_dtype).lower(), str(device))
|
|
|
|
if cache_key in _MODEL_CACHE:
|
|
return _MODEL_CACHE[cache_key]
|
|
|
|
config = _load_config(models_root, weight_dtype)
|
|
|
|
from omegaconf import OmegaConf
|
|
|
|
from high_sync.image_processor import ImageProcessor
|
|
from high_sync.src.models.unet_2d_condition import UNet2DConditionModel
|
|
from high_sync.src.models.unet_3d_echo import EchoUNet3DConditionModel
|
|
from high_sync.src.models.whisper.audio2feature import load_audio_model
|
|
from high_sync.src.pipelines.face_animate import FaceAnimatePipeline
|
|
|
|
scheduler = _build_scheduler(config)
|
|
|
|
vae = AutoencoderKL.from_pretrained(str(config.vae.model_path))
|
|
reference_unet = UNet2DConditionModel.from_pretrained(str(config.base_model_path), subfolder="unet")
|
|
denoising_unet = EchoUNet3DConditionModel.from_pretrained_2d(
|
|
str(config.base_model_path),
|
|
"",
|
|
subfolder="unet",
|
|
unet_additional_kwargs=OmegaConf.to_container(config.unet_additional_kwargs, resolve=True),
|
|
)
|
|
|
|
reference_unet.load_state_dict(_torch_load(str(config.reference_unet_path), torch))
|
|
denoising_unet.load_state_dict(_torch_load(str(config.denoising_unet_path), torch))
|
|
|
|
vae.requires_grad_(False)
|
|
reference_unet.requires_grad_(False)
|
|
denoising_unet.requires_grad_(False)
|
|
|
|
vae.eval()
|
|
reference_unet.eval()
|
|
denoising_unet.eval()
|
|
|
|
pipeline = FaceAnimatePipeline(
|
|
vae=vae,
|
|
reference_unet=reference_unet,
|
|
denoising_unet=denoising_unet,
|
|
scheduler=scheduler,
|
|
)
|
|
pipeline.to(device=device, dtype=dtype)
|
|
|
|
audio_processor = load_audio_model(model_path=str(config.audio_model_path), device=device)
|
|
|
|
model = {
|
|
"models_dir": str(models_root),
|
|
"device": device,
|
|
"weight_dtype": dtype,
|
|
"config": config,
|
|
"pipeline": pipeline,
|
|
"audio_processor": audio_processor,
|
|
"image_processor_cls": ImageProcessor,
|
|
"denoiser_model": None,
|
|
}
|
|
_MODEL_CACHE[cache_key] = model
|
|
return model
|