diff --git a/configs/config.py b/configs/config.py index 9b1028c..a330fb5 100644 --- a/configs/config.py +++ b/configs/config.py @@ -78,11 +78,6 @@ class Config: exe = sys.executable or "python" parser = argparse.ArgumentParser() parser.add_argument("--port", type=int, default=7865, help="Listen port") - parser.add_argument( - "--share", - action="store_true", - help="If set, create a shareable public link." - ) parser.add_argument("--pycmd", type=str, default=exe, help="Python command") parser.add_argument("--colab", action="store_true", help="Launch in colab") parser.add_argument( diff --git a/tools/download_model.py b/tools/download_model.py new file mode 100644 index 0000000..41bc029 --- /dev/null +++ b/tools/download_model.py @@ -0,0 +1,73 @@ +import os +import requests +import zipfile +import argparse +from pathlib import Path +import tempfile +import shutil + +def download_file(url: str, local_path: Path): + """Download a file from a URL to a local path.""" + if local_path.exists(): + print(f"[SKIP] {local_path} already exists.") + return + local_path.parent.mkdir(parents=True, exist_ok=True) + print(f"[DOWNLOAD] {url} -> {local_path}") + with requests.get(url, stream=True) as r: + r.raise_for_status() + with open(local_path, "wb") as f: + for chunk in r.iter_content(chunk_size=8192): + f.write(chunk) + print(f"[DONE] Downloaded {local_path}") + +def extract_and_move(zip_path: Path, logs_dir: Path, weights_dir: Path): + """Extract a zip file and move .index files to logs and .pth files to assets/weights.""" + logs_dir.mkdir(parents=True, exist_ok=True) + weights_dir.mkdir(parents=True, exist_ok=True) + with tempfile.TemporaryDirectory() as tmpdirname: + tmpdir = Path(tmpdirname) + print(f"[EXTRACT] Extracting {zip_path} to temporary directory {tmpdir}") + with zipfile.ZipFile(zip_path, "r") as zip_ref: + zip_ref.extractall(tmpdir) + # Walk through all extracted files recursively. + for extracted_file in tmpdir.rglob("*"): + if extracted_file.is_file(): + if extracted_file.suffix == ".index": + dest = logs_dir / extracted_file.name + print(f"[MOVE] Moving {extracted_file} to {dest}") + shutil.move(str(extracted_file), str(dest)) + elif extracted_file.suffix == ".pth": + dest = weights_dir / extracted_file.name + print(f"[MOVE] Moving {extracted_file} to {dest}") + shutil.move(str(extracted_file), str(dest)) + print("[CLEANUP] Extraction complete.") + +def main(): + parser = argparse.ArgumentParser( + description="Download a model zip file, extract it, and place .index files in ./logs and .pth files in ./assets/weights." + ) + parser.add_argument("url", type=str, help="URL of the zip file to download.") + args = parser.parse_args() + + url = args.url + zip_filename = Path(url).name + zip_path = Path(zip_filename) + + # Download the zip archive. + download_file(url, zip_path) + + # Define destination folders. + logs_dir = Path("./logs") + weights_dir = Path("./assets/weights") + + # Extract the archive and move files accordingly. + extract_and_move(zip_path, logs_dir, weights_dir) + + # Delete the zip file. + if zip_path.exists(): + print(f"[DELETE] Removing {zip_path}") + zip_path.unlink() + print("[COMPLETE] Model download and extraction complete.") + +if __name__ == "__main__": + main()