import os
import shutil

MODEL_NAME = "THUDM/CogVideoX-5b"
HF_CACHE_DIR = "/root/.cache/huggingface/hub/"
hash_code = None

for cache_model_name in os.listdir(HF_CACHE_DIR):
    flag = False
    for model_name_str_split in MODEL_NAME.split("/"):
        if model_name_str_split in cache_model_name:
            flag = True
        else:
            flag = False
            break
    if flag:
        if hash_code is None:
            full_path = os.path.join(HF_CACHE_DIR, cache_model_name, "snapshots",
                                     os.listdir(os.path.join(HF_CACHE_DIR, cache_model_name, "snapshots"))[0])
        else:
            full_path = os.path.join(HF_CACHE_DIR, cache_model_name, "snapshots", hash_code)
        os.makedirs(MODEL_NAME, exist_ok=True)
        for filename in os.listdir(full_path):
            print("copy: ", os.path.join(full_path, filename), " to ", os.path.join(MODEL_NAME, filename))
            if os.path.isdir(os.path.join(full_path, filename)):
                shutil.copytree(os.path.join(full_path, filename), os.path.join(MODEL_NAME, filename))
            elif os.path.isfile(os.path.join(full_path, filename)):
                shutil.copy(os.path.join(full_path, filename), os.path.join(MODEL_NAME, filename))
            else:
                raise Exception("path is not a file or directory: ", os.path.join(full_path, filename))