File size: 3,085 Bytes
22b4354 4b6dc3c 22b4354 2b7fd54 d1ad5e8 71b7d31 6058d6e a83a3d9 e492152 04b83e7 f2aa528 ce9c587 cb528cc f2a34e7 cb528cc 6f578f7 ccfe1b7 6f578f7 2444752 f2aa528 6736291 22b4354 a83a3d9 75ca5c9 a83a3d9 84ca156 a83a3d9 22b4354 71b7d31 2b7fd54 d1ad5e8 f2aa528 a83a3d9 e492152 6058d6e 22b4354 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 |
import functools
import os
PROJECT_ROOT_DIR = os.path.normpath(os.path.join(os.path.abspath(os.path.dirname(__file__)), "..", ".."))
MODELS_DIR = os.path.join(PROJECT_ROOT_DIR, "models")
MODELS_ANIMAL_DIR = os.path.join(MODELS_DIR, "animal")
MODELS_REAL_ESRGAN_DIR = os.path.join(MODELS_DIR, "RealESRGAN")
OUTPUTS_DIR = os.path.join(PROJECT_ROOT_DIR, "outputs")
OUTPUTS_VIDEOS_DIR = os.path.join(OUTPUTS_DIR, "videos")
TEMP_DIR = os.path.join(OUTPUTS_DIR, "temp")
TEMP_VIDEO_FRAMES_DIR = os.path.join(TEMP_DIR, "video_frames")
TEMP_VIDEO_OUT_FRAMES_DIR = os.path.join(TEMP_VIDEO_FRAMES_DIR, "out")
EXP_OUTPUT_DIR = os.path.join(OUTPUTS_DIR, "exp_data")
MODEL_CONFIG = os.path.join(PROJECT_ROOT_DIR, "modules", "config", "models.yaml")
MODEL_PATHS = {
"appearance_feature_extractor": os.path.join(MODELS_DIR, "appearance_feature_extractor.safetensors"),
"motion_extractor": os.path.join(MODELS_DIR, "motion_extractor.safetensors"),
"warping_module": os.path.join(MODELS_DIR, "warping_module.safetensors"),
"spade_generator": os.path.join(MODELS_DIR, "spade_generator.safetensors"),
"stitching_retargeting_module": os.path.join(MODELS_DIR, "stitching_retargeting_module.safetensors"),
"face_yolov8n": os.path.join(MODELS_DIR, "face_yolov8n.pt")
}
MODEL_ANIMAL_PATHS = {
"appearance_feature_extractor": os.path.join(MODELS_ANIMAL_DIR, "appearance_feature_extractor.safetensors"),
"motion_extractor": os.path.join(MODELS_ANIMAL_DIR, "motion_extractor.safetensors"),
"warping_module": os.path.join(MODELS_ANIMAL_DIR, "warping_module.safetensors"),
"spade_generator": os.path.join(MODELS_ANIMAL_DIR, "spade_generator.safetensors"),
"stitching_retargeting_module": os.path.join(MODELS_ANIMAL_DIR, "stitching_retargeting_module.safetensors"),
# Just animal detection model not the face, needs better model
"yolo_v5s_animal_det": os.path.join(MODELS_ANIMAL_DIR, "yolo_v5s_animal_det.n2x")
}
MODEL_REAL_ESRGAN_PATH = {
"realesr-general-x4v3": os.path.join(MODELS_REAL_ESRGAN_DIR, "realesr-general-x4v3.pth")
}
MASK_TEMPLATES = os.path.join(PROJECT_ROOT_DIR, "modules", "utils", "resources", "mask_template.png")
I18N_YAML_PATH = os.path.join(PROJECT_ROOT_DIR, "i18n", "translation.yaml")
def get_auto_incremental_file_path(dir_path: str, extension: str, prefix: str = ""):
counter = len(os.listdir(dir_path))
while True:
if prefix:
filename = f"{prefix}_{counter:05d}.{extension}"
else:
filename = f"{counter:05d}.{extension}"
full_path = os.path.join(dir_path, filename)
if not os.path.exists(full_path):
full_path = os.path.normpath(full_path)
return full_path
counter += 1
@functools.lru_cache
def init_dirs():
for dir_path in [
MODELS_DIR,
MODELS_ANIMAL_DIR,
MODELS_REAL_ESRGAN_DIR,
OUTPUTS_DIR,
EXP_OUTPUT_DIR,
TEMP_DIR,
TEMP_VIDEO_FRAMES_DIR,
TEMP_VIDEO_OUT_FRAMES_DIR,
OUTPUTS_VIDEOS_DIR
]:
os.makedirs(dir_path, exist_ok=True)
init_dirs()
|