|
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"), |
|
|
|
"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() |
|
|