LlamaTor: Decentralized AI model sharing via BitTorrent for efficient, user-friendly distribution and collaboration.
import os
import json
import sys
import requests
import argparse
import logging
import datetime
from tqdm import tqdm
import signal

# Create the parser
parser = argparse.ArgumentParser(description="Download models from Hugging Face's model hub.")
# Add the arguments
parser.add_argument('--output_folder', type=str, required=True, help='The folder where the models will be saved.')
parser.add_argument('--models_file', type=str, required=True, help='The JSON file containing the models to download.')
parser.add_argument('--timeout', type=int, default=30, help='The timeout for the download process (in seconds).')
parser.add_argument('--verbose', action='store_true', help='Enable verbose output.')
# Parse the arguments
args = parser.parse_args()

def download_file(file_url, file_path, logger):
    response = None
    try:
        response = requests.get(file_url, allow_redirects=True, timeout=args.timeout, stream=True)
        if response.status_code == 200:
            os.makedirs(os.path.dirname(file_path), exist_ok=True)  # create the folder if it does not exist
            with open(file_path, 'wb') as f:
                total_size = int(response.headers.get('content-length', 0))
                logger.info(f"Downloading file {file_path} ({total_size/ 1e6} Mbytes)")
                with tqdm(total=total_size, unit='B', unit_scale=True, unit_divisor=1024, desc=file_path, miniters=1, file=sys.stdout) as t:
                    for data in response.iter_content(chunk_size=4096):
                        f.write(data)
                        t.update(len(data))
    except Exception as e:
        logger.error(f"Error downloading {file_url}: {e}")

def signal_handler(sig, frame):
    logging.info("Interrupted by user. Exiting...")
    sys.exit(0)

def main():
    # Load models from JSON file
    try:
        with open(args.models_file, 'r') as f:
            models = json.load(f)
    except FileNotFoundError:
        logging.error(f"Models file {args.models_file} not found.")
        sys.exit(1)

    # Check if there are any models
    if models:
        # Get the models
        for model in models:
            model_name = model['name']
            dir_name = model_name.replace('/', '_')  # replace '/' with '_'
            output_dir = os.path.join(args.output_folder, dir_name)  # specify the folder where you want to save the model

            # Set up logging for this model
            logger = logging.getLogger(model_name)
            logger.setLevel(logging.INFO)
            log_dir = 'logs'
            os.makedirs(log_dir, exist_ok=True)
            current_time = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
            file_handler = logging.FileHandler(os.path.join(log_dir, f"{model_name.replace('/', '_')}_{current_time}.log"))
            stream_handler = logging.StreamHandler()  # This handler will print messages to the console
            formatter = logging.Formatter('%(asctime)s - %(message)s', datefmt='%Y-%m-%d %H:%M:%S')
            file_handler.setFormatter(formatter)
            stream_handler.setFormatter(formatter)
            logger.addHandler(file_handler)
            logger.addHandler(stream_handler)
            # Download the files
            for branch in model['branches']:
                for file in branch['files']:
                    file_url = f"https://huggingface.co/{model_name}/resolve/{branch['name']}/{file}"
                    file_path = os.path.join(output_dir, branch['name'], file)
                    if os.path.exists(file_path):
                        logger.info(f"File {file_path} already exists. Skipping...")
                        continue
                    download_file(file_url, file_path, logger)
    else:
        logging.info("No models found.")


def download_repository(repo):
    model_name = repo['name']
    dir_name = model_name.replace('/', '_')  # replace '/' with '_'
    output_dir = os.path.join(args.output_folder, dir_name)  # specify the folder where you want to save the model

    # Set up logging for this model
    logger = logging.getLogger(model_name)
    logger.setLevel(logging.INFO)
    log_dir = 'logs'
    os.makedirs(log_dir, exist_ok=True)
    current_time = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
    file_handler = logging.FileHandler(os.path.join(log_dir, f"{model_name.replace('/', '_')}_{current_time}.log"))
    stream_handler = logging.StreamHandler()  # This handler will print messages to the console
    formatter = logging.Formatter('%(asctime)s - %(message)s', datefmt='%Y-%m-%d %H:%M:%S')
    file_handler.setFormatter(formatter)
    stream_handler.setFormatter(formatter)
    logger.addHandler(file_handler)
    logger.addHandler(stream_handler)

    # Download the files
    for branch in repo['branches']:
        for file in branch['files']:
            file_url = f"https://huggingface.co/{model_name}/resolve/{branch['name']}/{file}"
            file_path = os.path.join(output_dir, branch['name'], file)
            if os.path.exists(file_path):
                logger.info(f"File {file_path} already exists. Skipping...")
                continue
            download_file(file_url, file_path, logger)

if __name__ == "__main__":
    signal.signal(signal.SIGINT, signal_handler)
    main()

download.py
Nondzu/LlamaTor
import argparse
from huggingface_hub import HfApi, ModelFilter
from tqdm import tqdm
import json
import time
import datetime
import requests
import datetime as dt
from dateutil.parser import parse

class ModelProcessor:
    def __init__(self, args):
        self.args = args
        self.api = HfApi()
        self.model_filter = ModelFilter(author=args.user) if args.user else None

    def fetch_and_filter_models(self):
        all_models = list(self.api.list_models(filter=self.model_filter))
        print(f"Total models count: {len(all_models)}")
        filtered_models = (
            [
                m
                for m in all_models
                if any(f.lower() in m.modelId.lower() for f in self.args.filter)
                and m.author.lower() == self.args.user.lower()
            ]
            if self.args.filter
            else all_models
        )
        if self.args.sort == "lastModified":
            filtered_models.sort(key=lambda x: x.lastModified, reverse=True)
        elif self.args.sort == "name":
            filtered_models.sort(key=lambda x: x.modelId.lower())
        if self.args.limit is not None:
            filtered_models = filtered_models[: self.args.limit]
        print("Number of models after filtering:", len(filtered_models))
        time.sleep(1)
        return filtered_models

    def process_models(self, filtered_models):
        now = dt.datetime.now().date()
        repo_table = []
        for model in tqdm(filtered_models, desc="Processing models"):
            try:
                last_modified = parse(model.lastModified).date()
                if (now - last_modified).days > self.args.age:
                    print(
                        f"Removed outdated repo: {(model.modelId)} : last update: {last_modified}"
                    )  # Added debug print info
                    continue
                repo_data = {
                    "name": model.modelId.replace("/", "#")
                    .replace("_", "#")
                    .replace("-", "#"),
                    "original_name": model.modelId,
                    "branches": [],
                }
                git_refs = self.api.list_repo_refs(model.modelId, repo_type="model")
                branches = [b for b in git_refs.branches if not b.name.startswith(".git")]
                repo_data["last_update"] = model.lastModified
                non_empty_branch_found = False
                for branch in branches:
                    try:
                        files = self.api.list_repo_files(model.modelId, revision=branch.name)
                        if files:
                            non_empty_branch_found = True
                            repo_data["branches"].append(
                                {
                                    "name": branch.name,
                                    "files": [
                                        f
                                        for f in files
                                        if not f.endswith(".gitattributes")
                                        and not f.startswith(".git")
                                    ],
                                }
                            )
                    except Exception as e:
                        print(f"Error fetching files for model {model.modelId} on branch {branch.name}: {e}")
                if non_empty_branch_found:
                    repo_table.append(repo_data)
                else:
                    print(
                        f"Empty repo detected and skipped: {model.modelId}"
                    )  # Print info about empty repo
            except Exception as e:
                print(f"Error fetching branches for model {model.modelId}: {e}")
        return repo_table

    def get_existing_torrents(self, url):
        response = requests.get(url)
        files = response.json()
        file_list = []
        for file in files:
            if file["type"] == "file":
                file_name = (
                    file["name"]
                    .replace(".torrent", "")
                    .replace("/", "#")
                    .replace("_", "#")
                    .replace("-", "#")
                )
                if not file_name.startswith(".gitattributes"):
                    file_list.append(file_name)
        return file_list

    def filter_and_output_repos(self, repo_table, existing_torrents):
        removed_repos = []
        duplicated_repos = []
        output_repos = []
        for repo in repo_table:
            if repo["name"] not in existing_torrents:
                output_repos.append(repo)
            else:
                if self.args.rd:
                    removed_repos.append(repo)
                    print(
                        f"Removed duplicated repo: {repo['original_name']}"
                    )  # Added debug print info
                else:
                    duplicated_repos.append(repo)
        with open(self.args.filename + "-models.json", "w") as f:
            json.dump(
                [{**repo, "name": repo["original_name"]} for repo in output_repos],
                f,
                indent=4,
            )
        if removed_repos:
            with open(self.args.filename + "removed_repos-models.json", "w") as f:
                json.dump(
                    {
                        "removed_repos": [
                            {"name": repo["original_name"]} for repo in removed_repos
                        ],
                        "duplicated_repos": [
                            {"name": repo["original_name"]} for repo in duplicated_repos
                        ],
                    },
                    f,
                    indent=4,
                )

def parse_arguments():
    parser = argparse.ArgumentParser()
    parser.add_argument("--user", type=str, default=None)
    parser.add_argument("--filter", type=str, nargs="+", default=[])
    parser.add_argument("--age", type=int, default=30)  # in days
    parser.add_argument(
        "--sort", type=str, default="lastModified", choices=["lastModified", "name"]
    )
    parser.add_argument("--limit", type=int, default=None)
    parser.add_argument(
        "--rd",
        "--remove-duplicates",
        action="store_true",
        help="Enable removing duplicate repositories",
    )
    parser.add_argument("--filename", type=str, default=f"string_model_{datetime.date.today().strftime('%d%m%Y')}")
    return parser.parse_args()

def main():
    args = parse_arguments()
    if not (args.user or args.filter):
        raise ValueError("At least one of --user or --filter must be provided")
    processor = ModelProcessor(args)
    filtered_models = processor.fetch_and_filter_models()
    repo_table = processor.process_models(filtered_models)
    existing_torrents = processor.get_existing_torrents(
        "https://api.github.com/repos/Nondzu/LlamaTor/contents/torrents?ref=torrents"
    )
    processor.filter_and_output_repos(repo_table, existing_torrents)
if __name__ == "__main__":
    main()
generate.py
Nondzu/LlamaTor
import os
import subprocess
import time
import argparse
import fnmatch

def generate_torrents(input_folder, output_folder, filter_pattern="*AWQ*"):
    trackers = [
        "http://atrack.pow7.com/announce",
        "udp://explodie.org:6969/announce",
        "udp://p4p.arenabg.com:1337/announce",
        "http://pow7.com:80/announce",
        "udp://tracker.opentrackr.org:1337/announce",
        "udp://tracker.tiny-vps.com:6969/announce",
        "http://tracker2.itzmx.com:6961/announce",
        "udp://open.stealth.si:80/announce",
        "http://p4p.arenabg.com:1337/announce",
        "http://tracker.opentrackr.org:1337/announce"
    ]

    folders = [f for f in os.listdir(input_folder) if os.path.isdir(os.path.join(input_folder, f)) and fnmatch.fnmatch(f, filter_pattern)]
    total_folders = len(folders)


    start_time = time.time()

    for i, folder in enumerate(folders, start=1):
        output_file = os.path.join(output_folder, f"{folder}.torrent")
        command = ["mktorrent", "-l", "19", "-p", "-o", output_file]

        for tracker in trackers:
            command.extend(["-a", tracker])

        command.append(os.path.join(input_folder, folder))

        print(f"Processing {i} of {total_folders}: {folder}")

        folder_start_time = time.time()
        subprocess.run(command)
        folder_end_time = time.time()

        print(f"Time taken for {folder}: {folder_end_time - folder_start_time} seconds")

    end_time = time.time()

    print(f"Total time taken: {end_time - start_time} seconds")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Generate torrent files for AI models.")
    parser.add_argument('--input_folder', type=str, required=True, help='The folder containing the downloaded models.')
    parser.add_argument('--output_folder', type=str, required=True, help='The folder where the generated torrent files will be saved.')
    parser.add_argument('--filter', type=str, default="*", help="Filter pattern for folder names (e.g., *AWQ*). If you want to create torrents for all folders, leave this blank.")
    args = parser.parse_args()

    generate_torrents(args.input_folder, args.output_folder, filter_pattern=args.filter)

make_torrent.py
Nondzu/LlamaTor