forestcalled/text-generation-webui

1'''2Downloads models from Hugging Face to models/username_modelname.3 4Example:5python download-model.py facebook/opt-1.3b6 7'''8 9import argparse10import base6411import datetime12import hashlib13import json14import os15import re16import sys17from pathlib import Path18 19import requests20import tqdm21from requests.adapters import HTTPAdapter22from tqdm.contrib.concurrent import thread_map23 24base = "https://huggingface.co"25 26 27class ModelDownloader:28    def __init__(self, max_retries=5):29        self.session = requests.Session()30        if max_retries:31            self.session.mount('https://cdn-lfs.huggingface.co', HTTPAdapter(max_retries=max_retries))32            self.session.mount('https://huggingface.co', HTTPAdapter(max_retries=max_retries))33        if os.getenv('HF_USER') is not None and os.getenv('HF_PASS') is not None:34            self.session.auth = (os.getenv('HF_USER'), os.getenv('HF_PASS'))35        if os.getenv('HF_TOKEN') is not None:36            self.session.headers = {'authorization': f'Bearer {os.getenv("HF_TOKEN")}'}37 38    def sanitize_model_and_branch_names(self, model, branch):39        if model[-1] == '/':40            model = model[:-1]41 42        if model.startswith(base + '/'):43            model = model[len(base) + 1:]44 45        model_parts = model.split(":")46        model = model_parts[0] if len(model_parts) > 0 else model47        branch = model_parts[1] if len(model_parts) > 1 else branch48 49        if branch is None:50            branch = "main"51        else:52            pattern = re.compile(r"^[a-zA-Z0-9._-]+$")53            if not pattern.match(branch):54                raise ValueError(55                    "Invalid branch name. Only alphanumeric characters, period, underscore and dash are allowed.")56 57        return model, branch58 59    def get_download_links_from_huggingface(self, model, branch, text_only=False, specific_file=None):60        page = f"/api/models/{model}/tree/{branch}"61        cursor = b""62 63        links = []64        sha256 = []65        classifications = []66        has_pytorch = False67        has_pt = False68        has_gguf = False69        has_safetensors = False70        is_lora = False71        while True:72            url = f"{base}{page}" + (f"?cursor={cursor.decode()}" if cursor else "")73            r = self.session.get(url, timeout=10)74            r.raise_for_status()75            content = r.content76 77            dict = json.loads(content)78            if len(dict) == 0:79                break80 81            for i in range(len(dict)):82                fname = dict[i]['path']83                if specific_file not in [None, ''] and fname != specific_file:84                    continue85 86                if not is_lora and fname.endswith(('adapter_config.json', 'adapter_model.bin')):87                    is_lora = True88 89                is_pytorch = re.match(r"(pytorch|adapter|gptq)_model.*\.bin", fname)90                is_safetensors = re.match(r".*\.safetensors", fname)91                is_pt = re.match(r".*\.pt", fname)92                is_gguf = re.match(r'.*\.gguf', fname)93                is_tiktoken = re.match(r".*\.tiktoken", fname)94                is_tokenizer = re.match(r"(tokenizer|ice|spiece).*\.model", fname) or is_tiktoken95                is_text = re.match(r".*\.(txt|json|py|md)", fname) or is_tokenizer96                if any((is_pytorch, is_safetensors, is_pt, is_gguf, is_tokenizer, is_text)):97                    if 'lfs' in dict[i]:98                        sha256.append([fname, dict[i]['lfs']['oid']])99 100                    if is_text:101                        links.append(f"https://huggingface.co/{model}/resolve/{branch}/{fname}")102                        classifications.append('text')103                        continue104 105                    if not text_only:106                        links.append(f"https://huggingface.co/{model}/resolve/{branch}/{fname}")107                        if is_safetensors:108                            has_safetensors = True109                            classifications.append('safetensors')110                        elif is_pytorch:111                            has_pytorch = True112                            classifications.append('pytorch')113                        elif is_pt:114                            has_pt = True115                            classifications.append('pt')116                        elif is_gguf:117                            has_gguf = True118                            classifications.append('gguf')119 120            cursor = base64.b64encode(f'{{"file_name":"{dict[-1]["path"]}"}}'.encode()) + b':50'121            cursor = base64.b64encode(cursor)122            cursor = cursor.replace(b'=', b'%3D')123 124        # If both pytorch and safetensors are available, download safetensors only125        if (has_pytorch or has_pt) and has_safetensors:126            for i in range(len(classifications) - 1, -1, -1):127                if classifications[i] in ['pytorch', 'pt']:128                    links.pop(i)129 130        # For GGUF, try to download only the Q4_K_M if no specific file is specified.131        # If not present, exclude all GGUFs, as that's likely a repository with both132        # GGUF and fp16 files.133        if has_gguf and specific_file is None:134            has_q4km = False135            for i in range(len(classifications) - 1, -1, -1):136                if 'q4_k_m' in links[i].lower():137                    has_q4km = True138 139            if has_q4km:140                for i in range(len(classifications) - 1, -1, -1):141                    if 'q4_k_m' not in links[i].lower():142                        links.pop(i)143            else:144                for i in range(len(classifications) - 1, -1, -1):145                    if links[i].lower().endswith('.gguf'):146                        links.pop(i)147 148        is_llamacpp = has_gguf and specific_file is not None149        return links, sha256, is_lora, is_llamacpp150 151    def get_output_folder(self, model, branch, is_lora, is_llamacpp=False, base_folder=None):152        if base_folder is None:153            base_folder = 'models' if not is_lora else 'loras'154 155        # If the model is of type GGUF, save directly in the base_folder156        if is_llamacpp:157            return Path(base_folder)158 159        output_folder = f"{'_'.join(model.split('/')[-2:])}"160        if branch != 'main':161            output_folder += f'_{branch}'162 163        output_folder = Path(base_folder) / output_folder164        return output_folder165 166    def get_single_file(self, url, output_folder, start_from_scratch=False):167        filename = Path(url.rsplit('/', 1)[1])168        output_path = output_folder / filename169        headers = {}170        mode = 'wb'171        if output_path.exists() and not start_from_scratch:172 173            # Check if the file has already been downloaded completely174            r = self.session.get(url, stream=True, timeout=10)175            total_size = int(r.headers.get('content-length', 0))176            if output_path.stat().st_size >= total_size:177                return178 179            # Otherwise, resume the download from where it left off180            headers = {'Range': f'bytes={output_path.stat().st_size}-'}181            mode = 'ab'182 183        with self.session.get(url, stream=True, headers=headers, timeout=10) as r:184            r.raise_for_status()  # Do not continue the download if the request was unsuccessful185            total_size = int(r.headers.get('content-length', 0))186            block_size = 1024 * 1024  # 1MB187 188            tqdm_kwargs = {189                'total': total_size,190                'unit': 'iB',191                'unit_scale': True,192                'bar_format': '{l_bar}{bar}| {n_fmt:6}/{total_fmt:6} {rate_fmt:6}'193            }194 195            if 'COLAB_GPU' in os.environ:196                tqdm_kwargs.update({197                    'position': 0,198                    'leave': True199                })200 201            with open(output_path, mode) as f:202                with tqdm.tqdm(**tqdm_kwargs) as t:203                    count = 0204                    for data in r.iter_content(block_size):205                        t.update(len(data))206                        f.write(data)207                        if total_size != 0 and self.progress_bar is not None:208                            count += len(data)209                            self.progress_bar(float(count) / float(total_size), f"{filename}")210 211    def start_download_threads(self, file_list, output_folder, start_from_scratch=False, threads=4):212        thread_map(lambda url: self.get_single_file(url, output_folder, start_from_scratch=start_from_scratch), file_list, max_workers=threads, disable=True)213 214    def download_model_files(self, model, branch, links, sha256, output_folder, progress_bar=None, start_from_scratch=False, threads=4, specific_file=None, is_llamacpp=False):215        self.progress_bar = progress_bar216 217        # Create the folder and writing the metadata218        output_folder.mkdir(parents=True, exist_ok=True)219 220        if not is_llamacpp:221            metadata = f'url: https://huggingface.co/{model}\n' \222                       f'branch: {branch}\n' \223                       f'download date: {datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")}\n'224 225            sha256_str = '\n'.join([f'    {item[1]} {item[0]}' for item in sha256])226            if sha256_str:227                metadata += f'sha256sum:\n{sha256_str}'228 229            metadata += '\n'230            (output_folder / 'huggingface-metadata.txt').write_text(metadata)231 232        if specific_file:233            print(f"Downloading {specific_file} to {output_folder}")234        else:235            print(f"Downloading the model to {output_folder}")236 237        self.start_download_threads(links, output_folder, start_from_scratch=start_from_scratch, threads=threads)238 239    def check_model_files(self, model, branch, links, sha256, output_folder):240        # Validate the checksums241        validated = True242        for i in range(len(sha256)):243            fpath = (output_folder / sha256[i][0])244 245            if not fpath.exists():246                print(f"The following file is missing: {fpath}")247                validated = False248                continue249 250            with open(output_folder / sha256[i][0], "rb") as f:251                file_hash = hashlib.file_digest(f, "sha256").hexdigest()252                if file_hash != sha256[i][1]:253                    print(f'Checksum failed: {sha256[i][0]}  {sha256[i][1]}')254                    validated = False255                else:256                    print(f'Checksum validated: {sha256[i][0]}  {sha256[i][1]}')257 258        if validated:259            print('[+] Validated checksums of all model files!')260        else:261            print('[-] Invalid checksums. Rerun download-model.py with the --clean flag.')262 263 264if __name__ == '__main__':265 266    parser = argparse.ArgumentParser()267    parser.add_argument('MODEL', type=str, default=None, nargs='?')268    parser.add_argument('--branch', type=str, default='main', help='Name of the Git branch to download from.')269    parser.add_argument('--threads', type=int, default=4, help='Number of files to download simultaneously.')270    parser.add_argument('--text-only', action='store_true', help='Only download text files (txt/json).')271    parser.add_argument('--specific-file', type=str, default=None, help='Name of the specific file to download (if not provided, downloads all).')272    parser.add_argument('--output', type=str, default=None, help='The folder where the model should be saved.')273    parser.add_argument('--clean', action='store_true', help='Does not resume the previous download.')274    parser.add_argument('--check', action='store_true', help='Validates the checksums of model files.')275    parser.add_argument('--max-retries', type=int, default=5, help='Max retries count when get error in download time.')276    args = parser.parse_args()277 278    branch = args.branch279    model = args.MODEL280    specific_file = args.specific_file281 282    if model is None:283        print("Error: Please specify the model you'd like to download (e.g. 'python download-model.py facebook/opt-1.3b').")284        sys.exit()285 286    downloader = ModelDownloader(max_retries=args.max_retries)287    # Clean up the model/branch names288    try:289        model, branch = downloader.sanitize_model_and_branch_names(model, branch)290    except ValueError as err_branch:291        print(f"Error: {err_branch}")292        sys.exit()293 294    # Get the download links from Hugging Face295    links, sha256, is_lora, is_llamacpp = downloader.get_download_links_from_huggingface(model, branch, text_only=args.text_only, specific_file=specific_file)296 297    # Get the output folder298    output_folder = downloader.get_output_folder(model, branch, is_lora, is_llamacpp=is_llamacpp, base_folder=args.output)299 300    if args.check:301        # Check previously downloaded files302        downloader.check_model_files(model, branch, links, sha256, output_folder)303    else:304        # Download files305        downloader.download_model_files(model, branch, links, sha256, output_folder, specific_file=specific_file, threads=args.threads, is_llamacpp=is_llamacpp)306