File size: 1,302 Bytes
759dfe0 |
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 |
#! /usr/bin/env python3
import os
import click
from huggingface_hub import snapshot_download
# Based off of Kijai's script
@click.command()
@click.argument('output_dir', required=True)
def download_weights(output_dir):
repo_id = "genmo/mochi-1-preview"
model = "dit.safetensors"
decoder = "decoder.safetensors"
encoder = "encoder.safetensors"
if not os.path.exists(output_dir):
print(f"Creating output directory: {output_dir}")
os.makedirs(output_dir, exist_ok=True)
def download_file(repo_id, output_dir, filename, description):
file_path = os.path.join(output_dir, filename)
if not os.path.exists(file_path):
print(f"Downloading mochi {description} to: {file_path}")
snapshot_download(
repo_id=repo_id,
allow_patterns=[f"*{filename}*"],
local_dir=output_dir,
local_dir_use_symlinks=False,
)
else:
print(f"{description} already exists in: {file_path}")
assert os.path.exists(file_path)
download_file(repo_id, output_dir, decoder, "decoder")
download_file(repo_id, output_dir, encoder, "encoder")
download_file(repo_id, output_dir, model, "model")
if __name__ == "__main__":
download_weights()
|