chore: import upstream snapshot with attribution

invokeai/backend/model_manager/metadata/__init__.py

@@ -0,0 +1,40 @@
+"""
+Initialization file for invokeai.backend.model_manager.metadata
+
+Usage:
+
+from invokeai.backend.model_manager.metadata import(
+   AnyModelRepoMetadata,
+   CommercialUsage,
+   LicenseRestrictions,
+   HuggingFaceMetadata,
+)
+
+from invokeai.backend.model_manager.metadata.fetch import HuggingFaceMetadataFetch
+
+data = HuggingFaceMetadataFetch().from_id("<REPO_ID>")
+assert isinstance(data, HuggingFaceMetadata)
+"""
+
+from invokeai.backend.model_manager.metadata.fetch import HuggingFaceMetadataFetch, ModelMetadataFetchBase
+from invokeai.backend.model_manager.metadata.metadata_base import (
+    AnyModelRepoMetadata,
+    AnyModelRepoMetadataValidator,
+    BaseMetadata,
+    HuggingFaceMetadata,
+    ModelMetadataWithFiles,
+    RemoteModelFile,
+    UnknownMetadataException,
+)
+
+__all__ = [
+    "AnyModelRepoMetadata",
+    "AnyModelRepoMetadataValidator",
+    "HuggingFaceMetadata",
+    "HuggingFaceMetadataFetch",
+    "ModelMetadataFetchBase",
+    "BaseMetadata",
+    "ModelMetadataWithFiles",
+    "RemoteModelFile",
+    "UnknownMetadataException",
+]

invokeai/backend/model_manager/metadata/fetch/__init__.py

@@ -0,0 +1,16 @@
+"""
+Initialization file for invokeai.backend.model_manager.metadata.fetch
+
+Usage:
+from invokeai.backend.model_manager.metadata.fetch import (
+    HuggingFaceMetadataFetch,
+)
+
+data = HuggingFaceMetadataFetch().from_id("<repo_id>")
+assert isinstance(data, HuggingFaceMetadata)
+"""
+
+from invokeai.backend.model_manager.metadata.fetch.fetch_base import ModelMetadataFetchBase
+from invokeai.backend.model_manager.metadata.fetch.huggingface import HuggingFaceMetadataFetch
+
+__all__ = ["ModelMetadataFetchBase", "HuggingFaceMetadataFetch"]

invokeai/backend/model_manager/metadata/fetch/fetch_base.py

@@ -0,0 +1,68 @@
+# Copyright (c) 2023 Lincoln D. Stein and the InvokeAI Development Team
+
+"""
+This module is the base class for subclasses that fetch metadata from model repositories
+
+Usage:
+
+from invokeai.backend.model_manager.metadata.fetch import HuggingFaceMetadataFetch
+
+data = HuggingFaceMetadataFetch().from_id("<REPO_ID>")
+assert isinstance(data, HuggingFaceMetadata)
+"""
+
+from abc import ABC, abstractmethod
+from typing import Optional
+
+from pydantic.networks import AnyHttpUrl
+from requests.sessions import Session
+
+from invokeai.backend.model_manager.metadata.metadata_base import (
+    AnyModelRepoMetadata,
+    AnyModelRepoMetadataValidator,
+    BaseMetadata,
+)
+from invokeai.backend.model_manager.taxonomy import ModelRepoVariant
+
+
+class ModelMetadataFetchBase(ABC):
+    """Fetch metadata from remote generative model repositories."""
+
+    @abstractmethod
+    def __init__(self, session: Optional[Session] = None):
+        """
+        Initialize the fetcher with an optional requests.sessions.Session object.
+
+        By providing a configurable Session object, we can support unit tests on
+        this module without an internet connection.
+        """
+        pass
+
+    @abstractmethod
+    def from_url(self, url: AnyHttpUrl) -> AnyModelRepoMetadata:
+        """
+        Given a URL to a model repository, return a ModelMetadata object.
+
+        This method will raise a `UnknownMetadataException`
+        in the event that the requested model metadata is not found at the provided location.
+        """
+        pass
+
+    @abstractmethod
+    def from_id(self, id: str, variant: Optional[ModelRepoVariant] = None) -> AnyModelRepoMetadata:
+        """
+        Given an ID for a model, return a ModelMetadata object.
+
+        :param id: An ID.
+        :param variant: A model variant from the ModelRepoVariant enum.
+
+        This method will raise a `UnknownMetadataException`
+        in the event that the requested model's metadata is not found at the provided id.
+        """
+        pass
+
+    @classmethod
+    def from_json(cls, json: str) -> AnyModelRepoMetadata:
+        """Given the JSON representation of the metadata, return the corresponding Pydantic object."""
+        metadata: BaseMetadata = AnyModelRepoMetadataValidator.validate_json(json)  # type: ignore
+        return metadata

invokeai/backend/model_manager/metadata/fetch/huggingface.py

@@ -0,0 +1,168 @@
+# Copyright (c) 2023 Lincoln D. Stein and the InvokeAI Development Team
+
+"""
+This module fetches model metadata objects from the HuggingFace model repository,
+using either a `repo_id` or the model page URL.
+
+Usage:
+
+from invokeai.backend.model_manager.metadata.fetch import HuggingFaceMetadataFetch
+
+fetcher = HuggingFaceMetadataFetch()
+metadata = fetcher.from_url("https://huggingface.co/stabilityai/sdxl-turbo")
+print(metadata.tags)
+"""
+
+import json
+import re
+from pathlib import Path
+from types import SimpleNamespace
+from typing import Optional
+
+import requests
+from huggingface_hub import HfApi, hf_hub_url
+from huggingface_hub.errors import RepositoryNotFoundError, RevisionNotFoundError
+from pydantic.networks import AnyHttpUrl
+from requests.sessions import Session
+
+from invokeai.backend.model_manager.metadata.fetch.fetch_base import ModelMetadataFetchBase
+from invokeai.backend.model_manager.metadata.metadata_base import (
+    AnyModelRepoMetadata,
+    HuggingFaceMetadata,
+    RemoteModelFile,
+    UnknownMetadataException,
+)
+from invokeai.backend.model_manager.taxonomy import ModelRepoVariant
+
+HF_MODEL_RE = r"https?://huggingface.co/([\w\-.]+/[\w\-.]+)"
+
+
+class HuggingFaceMetadataFetch(ModelMetadataFetchBase):
+    """Fetch model metadata from HuggingFace."""
+
+    def __init__(self, session: Optional[Session] = None):
+        """
+        Initialize the fetcher with an optional requests.sessions.Session object.
+
+        By providing a configurable Session object, we can support unit tests on
+        this module without an internet connection.
+        """
+        self._requests = session or requests.Session()
+        self._has_custom_session = session is not None
+
+    @classmethod
+    def from_json(cls, json: str) -> HuggingFaceMetadata:
+        """Given the JSON representation of the metadata, return the corresponding Pydantic object."""
+        metadata = HuggingFaceMetadata.model_validate_json(json)
+        return metadata
+
+    def _model_info_via_session(self, repo_id: str, variant: Optional[ModelRepoVariant] = None) -> SimpleNamespace:
+        """Fetch model info using the injected requests session (for testing/custom backends)."""
+        params = {"blobs": "true"}
+        url = f"https://huggingface.co/api/models/{repo_id}"
+        if variant is not None:
+            url += f"/revision/{variant}"
+        resp = self._requests.get(url, params=params)
+        if resp.status_code == 404:
+            error_code = resp.headers.get("X-Error-Code", "")
+            if error_code == "RevisionNotFound" or (variant is not None):
+                raise RevisionNotFoundError(f"Revision '{variant}' not found for repo '{repo_id}'.")
+            raise RepositoryNotFoundError(f"Repository '{repo_id}' not found.")
+        resp.raise_for_status()
+        data = resp.json()
+        # Convert siblings dicts to SimpleNamespace objects matching HfApi.model_info() shape
+        siblings = []
+        for s in data.get("siblings", []):
+            siblings.append(
+                SimpleNamespace(
+                    rfilename=s.get("rfilename"),
+                    size=s.get("size") or (s.get("lfs", {}) or {}).get("size"),
+                    lfs=s.get("lfs"),
+                )
+            )
+        return SimpleNamespace(id=data["id"], siblings=siblings)
+
+    def from_id(self, id: str, variant: Optional[ModelRepoVariant] = None) -> AnyModelRepoMetadata:
+        """Return a HuggingFaceMetadata object given the model's repo_id."""
+        # Little loop which tries fetching a revision corresponding to the selected variant.
+        # If not available, then set variant to None and get the default.
+        # If this too fails, raise exception.
+
+        model_info = None
+
+        # Handling for our special syntax - we only want the base HF `org/repo` here.
+        repo_id = id.split("::")[0] or id
+        while not model_info:
+            try:
+                # Use the injected session when provided (supports testing with mock adapters).
+                # Otherwise use HfApi which uses httpx internally.
+                if self._has_custom_session:
+                    model_info = self._model_info_via_session(repo_id, variant)
+                else:
+                    model_info = HfApi().model_info(repo_id=repo_id, files_metadata=True, revision=variant)
+            except RepositoryNotFoundError as excp:
+                raise UnknownMetadataException(f"'{repo_id}' not found. See trace for details.") from excp
+            except RevisionNotFoundError:
+                if variant is None:
+                    raise
+                else:
+                    variant = None
+
+        files: list[RemoteModelFile] = []
+
+        _, name = repo_id.split("/")
+
+        for s in model_info.siblings or []:
+            assert s.rfilename is not None
+            assert s.size is not None
+            files.append(
+                RemoteModelFile(
+                    url=hf_hub_url(repo_id, s.rfilename, revision=variant or "main"),
+                    path=Path(name, s.rfilename),
+                    size=s.size,
+                    sha256=s.lfs.get("sha256") if s.lfs else None,
+                )
+            )
+
+        # diffusers models have a `model_index.json` or `config.json` file
+        is_diffusers = any(str(f.url).endswith(("model_index.json", "config.json")) for f in files)
+
+        # These URLs will be exposed to the user - I think these are the only file types we fully support
+        ckpt_urls = (
+            None
+            if is_diffusers
+            else [
+                f.url
+                for f in files
+                if str(f.url).endswith(
+                    (
+                        ".safetensors",
+                        ".bin",
+                        ".pth",
+                        ".pt",
+                        ".ckpt",
+                    )
+                )
+            ]
+        )
+
+        return HuggingFaceMetadata(
+            id=model_info.id,
+            name=name,
+            files=files,
+            api_response=json.dumps(model_info.__dict__, default=str),
+            is_diffusers=is_diffusers,
+            ckpt_urls=ckpt_urls,
+        )
+
+    def from_url(self, url: AnyHttpUrl) -> AnyModelRepoMetadata:
+        """
+        Return a HuggingFaceMetadata object given the model's web page URL.
+
+        In the case of an invalid or missing URL, raises a ModelNotFound exception.
+        """
+        if match := re.match(HF_MODEL_RE, str(url), re.IGNORECASE):
+            repo_id = match.group(1)
+            return self.from_id(repo_id)
+        else:
+            raise UnknownMetadataException(f"'{url}' does not look like a HuggingFace model page")

invokeai/backend/model_manager/metadata/metadata_base.py

@@ -0,0 +1,136 @@
+# Copyright (c) 2023 Lincoln D. Stein and the InvokeAI Development Team
+
+"""This module defines core text-to-image model metadata fields.
+
+Metadata comprises any descriptive information that is not essential
+for getting the model to run. For example "author" is metadata, while
+"type", "base" and "format" are not. The latter fields are part of the
+model's config, as defined in invokeai.backend.model_manager.config.
+
+Note that the "name" and "description" are also present in `config`
+records. This is intentional. The config record fields are intended to
+be editable by the user as a form of customization. The metadata
+versions of these fields are intended to be kept in sync with the
+remote repo.
+"""
+
+from pathlib import Path
+from typing import List, Literal, Optional, Union
+
+from huggingface_hub import hf_hub_url
+from pydantic import BaseModel, Field, TypeAdapter
+from pydantic.networks import AnyHttpUrl
+from requests.sessions import Session
+from typing_extensions import Annotated
+
+from invokeai.backend.model_manager.taxonomy import ModelRepoVariant
+from invokeai.backend.model_manager.util.select_hf_files import filter_files
+
+
+class UnknownMetadataException(Exception):
+    """Raised when no metadata is available for a model."""
+
+
+class RemoteModelFile(BaseModel):
+    """Information about a downloadable file that forms part of a model."""
+
+    url: AnyHttpUrl = Field(description="The url to download this model file")
+    path: Path = Field(description="The path to the file, relative to the model root")
+    size: Optional[int] = Field(description="The size of this file, in bytes", default=0)
+    sha256: Optional[str] = Field(description="SHA256 hash of this model (not always available)", default=None)
+
+    def __hash__(self) -> int:
+        return hash(str(self))
+
+
+class ModelMetadataBase(BaseModel):
+    """Base class for model metadata information."""
+
+    name: str = Field(description="model's name")
+
+
+class BaseMetadata(ModelMetadataBase):
+    """Adds typing data for discriminated union."""
+
+    type: Literal["basemetadata"] = "basemetadata"
+
+
+class ModelMetadataWithFiles(ModelMetadataBase):
+    """Base class for metadata that contains a list of downloadable model file(s)."""
+
+    files: List[RemoteModelFile] = Field(description="model files and their sizes", default_factory=list)
+
+    def download_urls(
+        self,
+        variant: Optional[ModelRepoVariant] = None,
+        subfolder: Optional[Path] = None,
+        session: Optional[Session] = None,
+    ) -> List[RemoteModelFile]:
+        """
+        Return a list of URLs needed to download the model.
+
+        :param variant: Return files needed to reconstruct the indicated variant (e.g. ModelRepoVariant('fp16'))
+        :param subfolder: Return files in the designated subfolder only
+        :param session: A request.Session object for offline testing
+
+        Note that the "variant" and "subfolder" concepts currently only apply to HuggingFace.
+        However Civitai does have fields for the precision and format of its models, and may
+        provide variant selection criteria in the future.
+        """
+        return self.files
+
+
+class HuggingFaceMetadata(ModelMetadataWithFiles):
+    """Extended metadata fields provided by HuggingFace."""
+
+    type: Literal["huggingface"] = "huggingface"
+    id: str = Field(description="The HF model id")
+    api_response: Optional[str] = Field(description="Response from the HF API as stringified JSON", default=None)
+    is_diffusers: bool = Field(description="Whether the metadata is for a Diffusers format model", default=False)
+    ckpt_urls: Optional[List[AnyHttpUrl]] = Field(
+        description="URLs for all checkpoint format models in the metadata", default=None
+    )
+
+    def download_urls(
+        self,
+        variant: Optional[ModelRepoVariant] = None,
+        subfolder: Optional[Path] = None,
+        subfolders: Optional[List[Path]] = None,
+        session: Optional[Session] = None,
+    ) -> List[RemoteModelFile]:
+        """
+        Return list of downloadable files, filtering by variant and subfolder(s), if any.
+
+        :param variant: Return model files needed to reconstruct the indicated variant
+        :param subfolder: Return model files from the designated subfolder only (deprecated, use subfolders)
+        :param subfolders: Return model files from the designated subfolders
+        :param session: A request.Session object used for internet-free testing
+
+        Note that there is special variant-filtering behavior here:
+        When the fp16 variant is requested and not available, the
+        full-precision model is returned.
+        """
+        session = session or Session()
+
+        paths = filter_files([x.path for x in self.files], variant, subfolder, subfolders)  #  all files in the model
+
+        # Determine prefix for model_index.json check - only applies for single subfolder
+        prefix = ""
+        if subfolder and not subfolders:
+            prefix = f"{subfolder}/"
+
+        # the next step reads model_index.json to determine which subdirectories belong
+        # to the model (only for single subfolder case)
+        if Path(f"{prefix}model_index.json") in paths:
+            url = hf_hub_url(self.id, filename="model_index.json", subfolder=str(subfolder) if subfolder else None)
+            resp = session.get(url)
+            resp.raise_for_status()
+            submodels = resp.json()
+            paths = [Path(subfolder or "", x) for x in paths if Path(x).parent.as_posix() in submodels]
+            paths.insert(0, Path(f"{prefix}model_index.json"))
+
+        return [x for x in self.files if x.path in paths]
+
+
+AnyModelRepoMetadata = Annotated[Union[BaseMetadata, HuggingFaceMetadata], Field(discriminator="type")]
+AnyModelRepoMetadataValidator = TypeAdapter(AnyModelRepoMetadata)