ai-toolkit/toolkit/civitai.py

from toolkit.paths import MODELS_PATH
import requests
import os
import json
import tqdm


class ModelCache:
    def __init__(self):
        self.raw_cache = {}
        self.cache_path = os.path.join(MODELS_PATH, '.ai_toolkit_cache.json')
        if os.path.exists(self.cache_path):
            with open(self.cache_path, 'r') as f:
                all_cache = json.load(f)
                if 'models' in all_cache:
                    self.raw_cache = all_cache['models']
                else:
                    self.raw_cache = all_cache

    def get_model_path(self, model_id: int, model_version_id: int = None):
        if str(model_id) not in self.raw_cache:
            return None
        if model_version_id is None:
            # get latest version
            model_version_id = max([int(x) for x in self.raw_cache[str(model_id)].keys()])
            if model_version_id is None:
                return None
            model_path = self.raw_cache[str(model_id)][str(model_version_id)]['model_path']
            # check if model path exists
            if not os.path.exists(model_path):
                # remove version from cache
                del self.raw_cache[str(model_id)][str(model_version_id)]
                self.save()
                return None
            return model_path
        else:
            if str(model_version_id) not in self.raw_cache[str(model_id)]:
                return None
            model_path = self.raw_cache[str(model_id)][str(model_version_id)]['model_path']
            # check if model path exists
            if not os.path.exists(model_path):
                # remove version from cache
                del self.raw_cache[str(model_id)][str(model_version_id)]
                self.save()
                return None
            return model_path

    def update_cache(self, model_id: int, model_version_id: int, model_path: str):
        if str(model_id) not in self.raw_cache:
            self.raw_cache[str(model_id)] = {}
        if str(model_version_id) not in self.raw_cache[str(model_id)]:
            self.raw_cache[str(model_id)][str(model_version_id)] = {}
        self.raw_cache[str(model_id)][str(model_version_id)] = {
            'model_path': model_path
        }
        self.save()

    def save(self):
        if not os.path.exists(os.path.dirname(self.cache_path)):
            os.makedirs(os.path.dirname(self.cache_path), exist_ok=True)
        all_cache = {'models': {}}
        if os.path.exists(self.cache_path):
            # load it first
            with open(self.cache_path, 'r') as f:
                all_cache = json.load(f)

        all_cache['models'] = self.raw_cache

        with open(self.cache_path, 'w') as f:
            json.dump(all_cache, f, indent=2)


def get_model_download_info(model_id: int, model_version_id: int = None):
    # curl https://civitai.com/api/v1/models?limit=3&types=TextualInversion \
    # -H "Content-Type: application/json" \
    # -X GET
    print(
        f"Getting model info for model id: {model_id}{f' and version id: {model_version_id}' if model_version_id is not None else ''}")
    endpoint = f"https://civitai.com/api/v1/models/{model_id}"

    # get the json
    response = requests.get(endpoint)
    response.raise_for_status()
    model_data = response.json()

    model_version = None

    # go through versions and get the top one if one is not set
    for version in model_data['modelVersions']:
        if model_version_id is not None:
            if str(version['id']) == str(model_version_id):
                model_version = version
                break
        else:
            # get first version
            model_version = version
            break

    if model_version is None:
        raise ValueError(
            f"Could not find a model version for model id: {model_id}{f' and version id: {model_version_id}' if model_version_id is not None else ''}")

    model_file = None
    # go through files and prefer fp16 safetensors
    # "metadata": {
    #   "fp": "fp16",
    #   "size": "pruned",
    #   "format": "SafeTensor"
    # },
    # todo check pickle scans and skip if not good
    # try to get fp16 safetensor
    for file in model_version['files']:
        if file['metadata']['fp'] == 'fp16' and file['metadata']['format'] == 'SafeTensor':
            model_file = file
            break

    if model_file is None:
        # try to get primary
        for file in model_version['files']:
            if file['primary']:
                model_file = file
                break

    if model_file is None:
        # try to get any safetensor
        for file in model_version['files']:
            if file['metadata']['format'] == 'SafeTensor':
                model_file = file
                break

    if model_file is None:
        # try to get any fp16
        for file in model_version['files']:
            if file['metadata']['fp'] == 'fp16':
                model_file = file
                break

    if model_file is None:
        # try to get any
        for file in model_version['files']:
            model_file = file
            break

    if model_file is None:
        raise ValueError(f"Could not find a model file to download for model id: {model_id}")

    return model_file, model_version['id']


def get_model_path_from_url(url: str):
    # get query params form url if they are set
    # https: // civitai.com / models / 25694?modelVersionId = 127742
    query_params = {}
    if '?' in url:
        query_string = url.split('?')[1]
        query_params = dict(qc.split("=") for qc in query_string.split("&"))

    # get model id from url
    model_id = url.split('/')[-1]
    # remove query params from model id
    if '?' in model_id:
        model_id = model_id.split('?')[0]
    if model_id.isdigit():
        model_id = int(model_id)
    else:
        raise ValueError(f"Invalid model id: {model_id}")

    model_cache = ModelCache()
    model_path = model_cache.get_model_path(model_id, query_params.get('modelVersionId', None))
    if model_path is not None:
        return model_path
    else:
        # download model
        file_info, model_version_id = get_model_download_info(model_id, query_params.get('modelVersionId', None))

        download_url = file_info['downloadUrl']  # url does not work directly
        size_kb = file_info['sizeKB']
        filename = file_info['name']
        model_path = os.path.join(MODELS_PATH, filename)

        # download model
        print(f"Did not find model locally, downloading from model from: {download_url}")

        # use tqdm to show status of downlod
        response = requests.get(download_url, stream=True)
        response.raise_for_status()
        total_size_in_bytes = int(response.headers.get('content-length', 0))
        block_size = 1024  # 1 Kibibyte
        progress_bar = tqdm.tqdm(total=total_size_in_bytes, unit='iB', unit_scale=True)
        tmp_path = os.path.join(MODELS_PATH, f".download_tmp_{filename}")
        os.makedirs(os.path.dirname(model_path), exist_ok=True)
        # remove tmp file if it exists
        if os.path.exists(tmp_path):
            os.remove(tmp_path)

        try:

            with open(tmp_path, 'wb') as f:
                for data in response.iter_content(block_size):
                    progress_bar.update(len(data))
                    f.write(data)
                progress_bar.close()
                # move to final path
                os.rename(tmp_path, model_path)
                model_cache.update_cache(model_id, model_version_id, model_path)

                return model_path
        except Exception as e:
            # remove tmp file
            os.remove(tmp_path)
            raise e


# if is main
if __name__ == '__main__':
    model_path = get_model_path_from_url("https://civitai.com/models/25694?modelVersionId=127742")
    print(model_path)