File size: 1,747 Bytes
f2dd2b8 | 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 | import argparse
import os
import shutil
from huggingface_hub import snapshot_download
available_models = [
"marigold_appearance/finetuned",
"marigold_appearance/pretrained",
"marigold_lighting/finetuned",
"marigold_lighting/pretrained",
"rgbx/finetuned",
"rgbx/pretrained"
]
def main():
parser = argparse.ArgumentParser()
parser.add_argument(
"--model",
type=str,
default="rgbx/finetuned",
choices=available_models,
help="Select model to download (default: rgbx/finetuned)"
)
parser.add_argument(
"--local_dir",
type=str,
default="checkpoint",
help="Directory to save the model"
)
args = parser.parse_args()
LOCAL_DIR = args.local_dir
selected_model = args.model
if os.path.exists(LOCAL_DIR):
if os.path.abspath(LOCAL_DIR) in ["/", os.path.expanduser("~")]:
raise ValueError("Refusing to delete critical directory.")
print(f"Removing existing directory: {LOCAL_DIR}")
shutil.rmtree(LOCAL_DIR)
print(f"Downloading model: {selected_model}")
snapshot_download(
repo_id="GDAOSU/olbedo",
allow_patterns=f"{selected_model}/*",
local_dir=LOCAL_DIR,
local_dir_use_symlinks=False,
)
src = os.path.join(LOCAL_DIR, *selected_model.split("/"))
for name in os.listdir(src):
shutil.move(
os.path.join(src, name),
os.path.join(LOCAL_DIR, name)
)
top_level_folder = selected_model.split("/")[0]
shutil.rmtree(os.path.join(LOCAL_DIR, top_level_folder), ignore_errors=True)
shutil.rmtree(os.path.join(LOCAL_DIR, ".cache"), ignore_errors=True)
if __name__ == "__main__":
main()
|