mirror of
https://github.com/storytold/FineTrainers-Conditioning.git
synced 2026-10-09 00:09:45 +00:00
Sayak Paul
36 lines
1.5 KiB
Python
36 lines
1.5 KiB
Python
from pathlib import Path
|
|
from typing import Union
|
|
|
|
from accelerate.logging import get_logger
|
|
|
|
from ..constants import PRECOMPUTED_CONDITIONS_DIR_NAME, PRECOMPUTED_DIR_NAME, PRECOMPUTED_LATENTS_DIR_NAME
|
|
|
|
|
|
logger = get_logger("finetrainers")
|
|
|
|
|
|
def should_perform_precomputation(data_root: Union[str, Path]) -> bool:
|
|
if isinstance(data_root, str):
|
|
data_root = Path(data_root)
|
|
conditions_dir = data_root / PRECOMPUTED_DIR_NAME / PRECOMPUTED_CONDITIONS_DIR_NAME
|
|
latents_dir = data_root / PRECOMPUTED_DIR_NAME / PRECOMPUTED_LATENTS_DIR_NAME
|
|
if conditions_dir.exists() and latents_dir.exists():
|
|
num_files_conditions = len(list(conditions_dir.glob("*.pt")))
|
|
num_files_latents = len(list(latents_dir.glob("*.pt")))
|
|
if num_files_conditions != num_files_latents:
|
|
logger.warning(
|
|
f"Number of precomputed conditions ({num_files_conditions}) does not match number of precomputed latents ({num_files_latents})."
|
|
f"Cleaning up precomputed directories and re-running precomputation."
|
|
)
|
|
# clean up precomputed directories
|
|
for file in conditions_dir.glob("*.pt"):
|
|
file.unlink()
|
|
for file in latents_dir.glob("*.pt"):
|
|
file.unlink()
|
|
return True
|
|
if num_files_conditions > 0:
|
|
logger.info(f"Found {num_files_conditions} precomputed conditions and latents.")
|
|
return False
|
|
logger.info("Precomputed data not found. Running precomputation.")
|
|
return True
|