73 lines
2.5 KiB
Python
73 lines
2.5 KiB
Python
import torch
|
|
import os
|
|
import sys
|
|
from pathlib import Path
|
|
EXTENSION_PATH = Path(__file__).parent
|
|
sys.path.insert(0, str(EXTENSION_PATH.resolve()))
|
|
|
|
import torch
|
|
from scipy.ndimage import gaussian_filter
|
|
import comfy.model_management as model_management
|
|
from torch.hub import download_url_to_file, get_dir
|
|
from urllib.parse import urlparse
|
|
from motiondiff_modules.mogen.utils.plot_utils import recover_from_ric
|
|
|
|
HF_PREFIX = "https://huggingface.co/spaces/mingyuan/ReMoDiffuse/resolve/main/"
|
|
|
|
def motion_temporal_filter(motion, sigma=1):
|
|
motion = motion.reshape(motion.shape[0], -1)
|
|
for i in range(motion.shape[1]):
|
|
motion[:, i] = gaussian_filter(motion[:, i], sigma=sigma, mode="nearest")
|
|
return motion.reshape(motion.shape[0], -1, 3)
|
|
|
|
def get_motion_length(motion_data):
|
|
return motion_data["motion"].shape[0]
|
|
|
|
def to_cpu(x):
|
|
if isinstance(x, torch.Tensor):
|
|
return x.detach().cpu()
|
|
return x
|
|
|
|
def to_gpu(x):
|
|
if isinstance(x, torch.Tensor):
|
|
return x.to(model_management.get_torch_device())
|
|
return x
|
|
|
|
def load_file_from_url(url, model_dir=None, progress=True, file_name=None):
|
|
"""Load file form http url, will download models if necessary.
|
|
|
|
Ref:https://github.com/1adrianb/face-alignment/blob/master/face_alignment/utils.py
|
|
|
|
Args:
|
|
url (str): URL to be downloaded.
|
|
model_dir (str): The path to save the downloaded model. Should be a full path. If None, use pytorch hub_dir.
|
|
Default: None.
|
|
progress (bool): Whether to show the download progress. Default: True.
|
|
file_name (str): The downloaded file name. If None, use the file name in the url. Default: None.
|
|
|
|
Returns:
|
|
str: The path to the downloaded file.
|
|
"""
|
|
if model_dir is None: # use the pytorch hub_dir
|
|
hub_dir = get_dir()
|
|
model_dir = os.path.join(hub_dir, 'checkpoints')
|
|
|
|
os.makedirs(model_dir, exist_ok=True)
|
|
|
|
parts = urlparse(url)
|
|
filename = os.path.basename(parts.path)
|
|
if file_name is not None:
|
|
filename = file_name
|
|
cached_file = os.path.abspath(os.path.join(model_dir, filename))
|
|
if not os.path.exists(cached_file):
|
|
print(f'Downloading: "{url}" to {cached_file}\n')
|
|
download_url_to_file(url, cached_file, hash_prefix=None, progress=progress)
|
|
return cached_file
|
|
|
|
def motion_data_to_joints(motion_data_tensor):
|
|
joint = recover_from_ric(motion_data_tensor, 22).numpy()
|
|
joint = motion_temporal_filter(joint, sigma=2.5)
|
|
return joint
|
|
|
|
__all__ = ["motion_temporal_filter", "get_motion_length", "to_cpu", "to_gpu", "motion_data_to_joints"]
|