ollama 없이 모델 파일을 ollama에서 다운 받기

import json
import os
import sys
from urllib.parse import urljoin, urlunparse

import requests
import urllib3

urllib3.disable_warnings()

DEFAULT_REGISTRY = "registry.ollama.ai"
DEFAULT_NAMESPACE = "library"
DEFAULT_TAG = "latest"
DEFAULT_PROTOCOL_SCHEME = "https"


class Manifest:
    def __init__(
        self,
        schema_version,
        media_type,
        config,
        layers,
        filepath=None,
        file_info=None,
        digest=None,
    ):
        self.schema_version = schema_version
        self.media_type = media_type
        self.config = config
        self.layers = layers
        self.filepath = filepath
        self.file_info = file_info
        self.digest = digest


class Layer:
    def __init__(self, mediaType, digest, size, source=None, status=None):
        self.media_type = mediaType
        self.digest = digest
        self.size = size
        self.source = source
        self.status = status


class ModelPath:
    def __init__(
        self,
        protocol_scheme=DEFAULT_PROTOCOL_SCHEME,
        registry=DEFAULT_REGISTRY,
        namespace=DEFAULT_NAMESPACE,
        repository="",
        tag=DEFAULT_TAG,
    ):
        self.protocol_scheme = protocol_scheme
        self.registry = registry
        self.namespace = namespace
        self.repository = repository
        self.tag = tag

    @classmethod
    def parse_model_path(cls, name):
        mp = cls()

        if "://" in name:
            mp.protocol_scheme, name = name.split("://", 1)

        name = name.replace(os.path.sep, "/")
        parts = name.split("/")

        if len(parts) == 3:
            mp.registry, mp.namespace, mp.repository = parts
        elif len(parts) == 2:
            mp.namespace, mp.repository = parts
        elif len(parts) == 1:
            mp.repository = parts[0]

        if ":" in mp.repository:
            mp.repository, mp.tag = mp.repository.split(":", 1)

        return mp

    def get_namespace_repository(self):
        return f"{self.namespace}/{self.repository}"

    def base_url(self):
        return urlunparse((self.protocol_scheme, self.registry, "", "", "", ""))


def get_manifest(model_path, reg_opts=None):
    request_url = urljoin(
        model_path.base_url(),
        f"v2/{model_path.get_namespace_repository()}/manifests/{model_path.tag}",
    )
    headers = {"Accept": "application/vnd.docker.distribution.manifest.v2+json"}

    try:
        response = requests.get(
            request_url,
            headers=headers,
            auth=reg_opts,
            verify=False,
            proxies={"https": "http://12.26.204.100:8080"},
        )
        response.raise_for_status()
    except requests.RequestException as e:
        print(f"매니페스트를 가져오는 중 오류 발생: {e}")
        return None, request_url

    try:
        manifest_data = response.json()
        config = Layer(**manifest_data.get("config", {}))
        layers = [Layer(**layer) for layer in manifest_data.get("layers", [])]
        manifest = Manifest(
            schema_version=manifest_data.get("schemaVersion"),
            media_type=config.media_type,
            config=config,
            layers=layers,
        )
        return manifest, request_url
    except json.JSONDecodeError as e:
        print(f"JSON 디코딩 오류: {e}")
        return None, request_url


def download_blob(blob_url, output_path, use_tqdm=True):

    try:
        if use_tqdm is None:
            from tqdm import tqdm

            use_tqdm = True
        else:
            use_tqdm = False
    except ImportError:
        use_tqdm = False

    try:
        with requests.get(
            blob_url,
            stream=True,
            verify=False,
            proxies={"https": "http://12.26.204.100:8080"},
        ) as response:
            response.raise_for_status()
            total_size = int(response.headers.get("content-length", 0))
            print(f"Downloading: {blob_url}..., {output_path}")

            with open(output_path, "wb") as file:
                if use_tqdm:
                    with tqdm(
                        desc=os.path.basename(output_path),
                        total=total_size,
                        unit="B",
                        unit_scale=True,
                        unit_divisor=1024,
                    ) as bar:
                        for chunk in response.iter_content(chunk_size=8192):
                            file.write(chunk)
                            bar.update(len(chunk))
                else:
                    print(f"다운로드 중: {os.path.basename(output_path)}")
                    for chunk in response.iter_content(chunk_size=8192):
                        file.write(chunk)
        print(f"Downloaded: {output_path}")
    except requests.RequestException as e:
        print(f"Error downloading blob from {blob_url}: {e}")


def download_model(model_name, output_dir):
    os.makedirs(output_dir, exist_ok=True)

    parsed_model_path = ModelPath.parse_model_path(model_name)
    manifest, manifest_link = get_manifest(parsed_model_path)

    if manifest is None:
        print(f"Failed to fetch manifest for model: {model_name}")
        sys.exit(1)

    layers = manifest.layers + [manifest.config]
    for layer in layers:
        if layer.media_type != "application/vnd.ollama.image.model":
            continue

        blob_url = urljoin(
            parsed_model_path.base_url(),
            f"v2/{parsed_model_path.get_namespace_repository()}/blobs/{layer.digest}",
        )

        output_path = os.path.join(output_dir, model_name.replace(":", "-") + ".gguf")
        if os.path.exists(output_path):
            print("Skipped " + output_path)
        else:
            download_blob(blob_url, output_path)

    print(f"All blobs for model {model_name} downloaded successfully.")


if __name__ == "__main__":
    model_name = "llama3.2:1b"
    output_folder = "temp"
    download_model(model_name, output_folder)