Add new zipnn scripts

zipnn_compress_file.py

@@ -0,0 +1,154 @@
+import os
+import subprocess
+import sys
+import argparse
+import time
+
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
+
+KB = 1024
+MB = 1024 * 1024
+GB = 1024 * 1024 * 1024
+
+
+def check_and_install_zipnn():
+    try:
+        import zipnn
+    except ImportError:
+        print("zipnn not found. Installing...")
+        subprocess.check_call(
+            [
+                sys.executable,
+                "-m",
+                "pip",
+                "install",
+                "zipnn",
+                "--upgrade",
+            ]
+        )
+        import zipnn
+
+
+def parse_streaming_chunk_size(
+    streaming_chunk_size,
+):
+    if str(streaming_chunk_size).isdigit():
+        final = int(streaming_chunk_size)
+    else:
+        size_value = int(streaming_chunk_size[:-2])
+        size_unit = streaming_chunk_size[-2].lower()
+
+        if size_unit == "k":
+            final = KB * size_value
+        elif size_unit == "m":
+            final = MB * size_value
+        elif size_unit == "g":
+            final = GB * size_value
+        else:
+            raise ValueError(f"Invalid size unit: {size_unit}. Use 'k', 'm', or 'g'.")
+
+    return final
+
+
+def compress_file(
+    input_file,
+    dtype="",
+    streaming_chunk_size=1048576,
+    delete=False,
+    force=False,
+):
+    import zipnn
+
+    streaming_chunk_size = parse_streaming_chunk_size(streaming_chunk_size)
+    full_path = input_file
+    if not os.path.exists(full_path):
+        print("File not found")
+        return
+    if delete:
+        print(f"Deleting {full_path}...")
+        os.remove(full_path)
+    else:
+        compressed_path = full_path + ".znn"
+        if not force and os.path.exists(compressed_path):
+            user_input = (
+                input(f"{compressed_path} already exists; overwrite (y/n)? ").strip().lower()
+            )
+            if user_input not in ("yes", "y"):
+                print(f"Skipping {full_path}...")
+                return
+        print(f"Compressing {full_path}...")
+        #
+        output_file = input_file + ".znn"
+        if dtype:
+            zpn = zipnn.ZipNN(
+                bytearray_dtype="float32",
+                is_streaming=True,
+                streaming_chunk_kb=streaming_chunk_size,
+            )
+        else:
+            zpn = zipnn.ZipNN(
+                is_streaming=True,
+                streaming_chunk_kb=streaming_chunk_size,
+            )
+        file_size_before = 0
+        file_size_after = 0
+        start_time = time.time()
+        with open(input_file, "rb") as infile, open(output_file, "wb") as outfile:
+            chunk = infile.read()
+            file_size_before += len(chunk)
+            compressed_chunk = zpn.compress(chunk)
+            if compressed_chunk:
+                file_size_after += len(compressed_chunk)
+                outfile.write(compressed_chunk)
+        end_time = time.time() - start_time
+        print(f"Compressed {input_file} to {output_file}")
+        print(
+            f"Original size:  {file_size_before/GB:.02f}GB size after compression: {file_size_after/GB:.02f}GB, Remaining size is {file_size_after/file_size_before*100:.02f}% of original, time: {end_time:.02f}"
+        )
+
+
+if __name__ == "__main__":
+    if len(sys.argv) < 2:
+        print("Usage: python compress_files.py <suffix>")
+        print("Example: python compress_files.py 'safetensors'")
+        sys.exit(1)
+
+    parser = argparse.ArgumentParser(description="Enter a file path to compress.")
+    parser.add_argument(
+        "input_file",
+        type=str,
+        help="Specify the path to the file to compress.",
+    )
+    parser.add_argument(
+        "--float32",
+        action="store_true",
+        help="A flag that triggers float32 compression",
+    )
+    parser.add_argument(
+        "--streaming_chunk_size",
+        type=str,
+        help="An optional streaming chunk size. The format is int (for size in Bytes) or int+KB/MB/GB. Default is 1MB",
+    )
+    parser.add_argument(
+        "--delete",
+        action="store_true",
+        help="A flag that triggers deletion of a single file instead of compression",
+    )
+    parser.add_argument(
+        "--force",
+        action="store_true",
+        help="A flag that forces overwriting when compressing.",
+    )
+    args = parser.parse_args()
+    optional_kwargs = {}
+    if args.float32:
+        optional_kwargs["dtype"] = 32
+    if args.streaming_chunk_size is not None:
+        optional_kwargs["streaming_chunk_size"] = args.streaming_chunk_size
+    if args.delete:
+        optional_kwargs["delete"] = args.delete
+    if args.force:
+        optional_kwargs["force"] = args.force
+
+    check_and_install_zipnn()
+    compress_file(args.input_file, **optional_kwargs)

zipnn_compress_path.py

@@ -0,0 +1,253 @@
+import os
+import subprocess
+import sys
+import argparse
+from concurrent.futures import (
+    ProcessPoolExecutor,
+    as_completed,
+)
+from zipnn_compress_file import compress_file
+import zipnn
+
+sys.path.append(
+    os.path.abspath(
+        os.path.join(
+            os.path.dirname(__file__), ".."
+        )
+    )
+)
+
+
+KB = 1024
+MB = 1024 * 1024
+GB = 1024 * 1024 * 1024
+
+
+def check_and_install_zipnn():
+    try:
+        import zipnn
+    except ImportError:
+        print("zipnn not found. Installing...")
+        subprocess.check_call(
+            [
+                sys.executable,
+                "-m",
+                "pip",
+                "install",
+                "zipnn",
+                "--upgrade",
+            ]
+        )
+        import zipnn
+
+
+def parse_streaming_chunk_size(
+    streaming_chunk_size,
+):
+    if str(streaming_chunk_size).isdigit():
+        final = int(streaming_chunk_size)
+    else:
+        size_value = int(
+            streaming_chunk_size[:-2]
+        )
+        size_unit = streaming_chunk_size[
+            -2
+        ].lower()
+
+        if size_unit == "k":
+            final = KB * size_value
+        elif size_unit == "m":
+            final = MB * size_value
+        elif size_unit == "g":
+            final = GB * size_value
+        else:
+            raise ValueError(
+                f"Invalid size unit: {size_unit}. Use 'k', 'm', or 'g'."
+            )
+
+    return final
+
+
+def compress_files_with_suffix(
+    suffix,
+    dtype="",
+    streaming_chunk_size=1048576,
+    path=".",
+    delete=False,
+    r=False,
+    force=False,
+    max_processes=1,
+):
+
+    file_list = []
+    streaming_chunk_size = (
+        parse_streaming_chunk_size(
+            streaming_chunk_size
+        )
+    )
+    directories_to_search = (
+        os.walk(path)
+        if r
+        else [(path, [], os.listdir(path))]
+    )
+    files_found = False
+    for root, _, files in directories_to_search:
+        for file_name in files:
+            if file_name.endswith(suffix):
+                compressed_path = (
+                    file_name + ".znn"
+                )
+                if not force and os.path.exists(
+                    compressed_path
+                ):
+                    user_input = (
+                        input(
+                            f"{compressed_path} already exists; overwrite (y/n)? "
+                        )
+                        .strip()
+                        .lower()
+                    )
+                    if user_input not in (
+                        "y",
+                        "yes",
+                    ):
+                        print(
+                            f"Skipping {file_name}..."
+                        )
+                        continue
+                files_found = True
+                full_path = os.path.join(
+                    root, file_name
+                )
+                file_list.append(full_path)
+
+    with ProcessPoolExecutor(
+        max_workers=max_processes
+    ) as executor:
+        future_to_file = {
+            executor.submit(
+                compress_file,
+                file,
+                dtype,
+                streaming_chunk_size,
+                delete,
+                True,
+            ): file
+            for file in file_list[:max_processes]
+        }
+        file_list = file_list[max_processes:]
+        while future_to_file:
+            for future in as_completed(
+                future_to_file
+            ):
+                file = future_to_file.pop(future)
+
+                try:
+                    future.result()
+                except Exception as exc:
+                    print(
+                        f"File {file} generated an exception: {exc}"
+                    )
+
+                if file_list:
+                    next_file = file_list.pop(0)
+                    future_to_file[
+                        executor.submit(
+                            compress_file,
+                            next_file,
+                            dtype,
+                            streaming_chunk_size,
+                            delete,
+                            True,
+                        )
+                    ] = next_file
+
+    if not files_found:
+        print(
+            f"No files with the suffix '{suffix}' found."
+        )
+
+
+if __name__ == "__main__":
+    if len(sys.argv) < 2:
+        print(
+            "Usage: python compress_files.py <suffix>"
+        )
+        print(
+            "Example: python compress_files.py 'safetensors'"
+        )
+        sys.exit(1)
+
+    parser = argparse.ArgumentParser(
+        description="Enter a suffix to compress, (optional) dtype, (optional) streaming chunk size, (optional) path to files."
+    )
+    parser.add_argument(
+        "suffix",
+        type=str,
+        help="Specify the file suffix to compress all files with that suffix. If a single file name is provided, only that file will be compressed.",
+    )
+    parser.add_argument(
+        "--float32",
+        action="store_true",
+        help="A flag that triggers float32 compression",
+    )
+    parser.add_argument(
+        "--streaming_chunk_size",
+        type=str,
+        help="An optional streaming chunk size. The format is int (for size in Bytes) or int+KB/MB/GB. Default is 1MB",
+    )
+    parser.add_argument(
+        "--path",
+        type=str,
+        help="Path to files to compress",
+    )
+    parser.add_argument(
+        "--delete",
+        action="store_true",
+        help="A flag that triggers deletion of a single file instead of compression",
+    )
+    parser.add_argument(
+        "-r",
+        action="store_true",
+        help="A flag that triggers recursive search on all subdirectories",
+    )
+    parser.add_argument(
+        "--recursive",
+        action="store_true",
+        help="A flag that triggers recursive search on all subdirectories",
+    )
+    parser.add_argument(
+        "--force",
+        action="store_true",
+        help="A flag that forces overwriting when compressing.",
+    )
+    parser.add_argument(
+        "--max_processes",
+        type=int,
+        help="The amount of maximum processes.",
+    )
+    args = parser.parse_args()
+    optional_kwargs = {}
+    if args.float32:
+        optional_kwargs["dtype"] = 32
+    if args.streaming_chunk_size is not None:
+        optional_kwargs[
+            "streaming_chunk_size"
+        ] = args.streaming_chunk_size
+    if args.path is not None:
+        optional_kwargs["path"] = args.path
+    if args.delete:
+        optional_kwargs["delete"] = args.delete
+    if args.r or args.recursive:
+        optional_kwargs["r"] = args.r
+    if args.force:
+        optional_kwargs["force"] = args.force
+    if args.max_processes:
+        optional_kwargs["max_processes"] = (
+            args.max_processes
+        )
+
+    check_and_install_zipnn()
+    compress_files_with_suffix(
+        args.suffix, **optional_kwargs
+    )

zipnn_decompress_file.py

@@ -0,0 +1,84 @@
+import os
+import subprocess
+import sys
+import argparse
+
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
+
+
+def check_and_install_zipnn():
+    try:
+        import zipnn
+    except ImportError:
+        print("zipnn not found. Installing...")
+        subprocess.check_call([sys.executable, "-m", "pip", "install", "zipnn"])
+        import zipnn
+
+
+def decompress_file(input_file, dtype="", delete=False, force=False):
+    import zipnn
+
+    if not input_file.endswith(".znn"):
+        raise ValueError("Input file does not have the '.znn' suffix")
+
+    if os.path.exists(input_file):
+        if delete:
+            print(f"Deleting {input_file}...")
+            os.remove(input_file)
+        else:
+            decompressed_path = input_file[:-4]
+            if not force and os.path.exists(decompressed_path):
+
+                user_input = (
+                    input(f"{decompressed_path} already exists; overwrite (y/n)? ").strip().lower()
+                )
+
+                if user_input not in ("yes", "y"):
+                    print(f"Skipping {input_file}...")
+                    return
+            print(f"Decompressing {input_file}...")
+
+            output_file = input_file[:-4]
+
+            if dtype:
+                zpn = zipnn.ZipNN(is_streaming=True, bytearray_dtype="float32")
+            else:
+                zpn = zipnn.ZipNN(is_streaming=True)
+
+            with open(input_file, "rb") as infile, open(output_file, "wb") as outfile:
+                d_data = b""
+                chunk = infile.read()
+                d_data += zpn.decompress(chunk)
+                outfile.write(d_data)
+                print(f"Decompressed {input_file} to {output_file}")
+
+    else:
+        print(f"Error: The file {input_file} does not exist.")
+
+
+if __name__ == "__main__":
+    check_and_install_zipnn()
+
+    parser = argparse.ArgumentParser(description="Enter a file path to decompress.")
+    parser.add_argument("input_file", type=str, help="Specify the path to the file to decompress.")
+    parser.add_argument(
+        "--float32", action="store_true", help="A flag that triggers float32 compression."
+    )
+    parser.add_argument(
+        "--delete",
+        action="store_true",
+        help="A flag that triggers deletion of a single compressed file instead of decompression",
+    )
+    parser.add_argument(
+        "--force", action="store_true", help="A flag that forces overwriting when decompressing."
+    )
+    args = parser.parse_args()
+    optional_kwargs = {}
+    if args.float32:
+        optional_kwargs["dtype"] = 32
+    if args.delete:
+        optional_kwargs["delete"] = args.delete
+    if args.force:
+        optional_kwargs["force"] = args.force
+
+    decompress_file(args.input_file, **optional_kwargs)

zipnn_decompress_path.py

@@ -0,0 +1,184 @@
+import os
+import sys
+import argparse
+import subprocess
+import zipnn
+from concurrent.futures import (
+    ProcessPoolExecutor,
+    as_completed,
+)
+from zipnn_decompress_file import (
+    decompress_file,
+)
+
+sys.path.append(
+    os.path.abspath(
+        os.path.join(
+            os.path.dirname(__file__),
+            "..",
+        )
+    )
+)
+
+
+def check_and_install_zipnn():
+    try:
+        import zipnn
+    except ImportError:
+        print("zipnn not found. Installing...")
+        subprocess.check_call(
+            [
+                sys.executable,
+                "-m",
+                "pip",
+                "install",
+                "zipnn",
+            ]
+        )
+        import zipnn
+
+
+def decompress_zpn_files(
+    dtype="",
+    path=".",
+    delete=False,
+    force=False,
+    max_processes=1,
+):
+
+    file_list = []
+    directories_to_search = [
+        (
+            path,
+            [],
+            os.listdir(path),
+        )
+    ]
+    for (
+        root,
+        _,
+        files,
+    ) in directories_to_search:
+        for file_name in files:
+            if file_name.endswith(".znn"):
+                decompressed_path = file_name[:-4]
+                if not force and os.path.exists(
+                    decompressed_path
+                ):
+                    user_input = (
+                        input(
+                            f"{decompressed_path} already exists; overwrite (y/n)? "
+                        )
+                        .strip()
+                        .lower()
+                    )
+                    if user_input not in (
+                        "y",
+                        "yes",
+                    ):
+                        print(
+                            f"Skipping {file_name}..."
+                        )
+                        continue
+                full_path = os.path.join(
+                    root,
+                    file_name,
+                )
+                file_list.append(full_path)
+
+    with ProcessPoolExecutor(
+        max_workers=max_processes
+    ) as executor:
+        for file in file_list[:max_processes]:
+            future_to_file = {
+                executor.submit(
+                    decompress_file,
+                    file,
+                    dtype,
+                    delete,
+                    True,
+                ): file
+                for file in file_list[
+                    :max_processes
+                ]
+            }
+
+            file_list = file_list[max_processes:]
+            while future_to_file:
+
+                for future in as_completed(
+                    future_to_file
+                ):
+                    file = future_to_file.pop(
+                        future
+                    )
+                    try:
+                        future.result()
+                    except Exception as exc:
+                        print(
+                            f"File {file} generated an exception: {exc}"
+                        )
+
+                    if file_list:
+                        next_file = file_list.pop(
+                            0
+                        )
+                        future_to_file[
+                            executor.submit(
+                                decompress_file,
+                                next_file,
+                                dtype,
+                                delete,
+                                True,
+                            )
+                        ] = next_file
+                        #
+
+
+if __name__ == "__main__":
+    check_and_install_zipnn()
+
+    parser = argparse.ArgumentParser(
+        description="Compresses all .znn files. (optional) dtype."
+    )
+    parser.add_argument(
+        "--float32",
+        action="store_true",
+        help="A flag that triggers float32 compression.",
+    )
+    parser.add_argument(
+        "--path",
+        type=str,
+        help="Path to folder of files to decompress. If left empty, checks current folder.",
+    )
+    parser.add_argument(
+        "--delete",
+        action="store_true",
+        help="A flag that triggers deletion of a single compressed file instead of decompression",
+    )
+    parser.add_argument(
+        "--force",
+        action="store_true",
+        help="A flag that forces overwriting when decompressing.",
+    )
+    parser.add_argument(
+        "--max_processes",
+        type=int,
+        help="The amount of maximum processes.",
+    )
+    args = parser.parse_args()
+    optional_kwargs = {}
+    if args.float32:
+        optional_kwargs["dtype"] = 32
+    if args.path is not None:
+        optional_kwargs["path"] = args.path
+    if args.delete:
+        optional_kwargs["delete"] = args.delete
+    if args.force:
+        optional_kwargs["force"] = args.force
+    if args.max_processes:
+        optional_kwargs["max_processes"] = (
+            args.max_processes
+        )
+
+    decompress_zpn_files(**optional_kwargs)