refactor(mm): wip schema changes

invokeai/backend/model_manager/config.py

@@ -26,7 +26,7 @@ from typing import Literal, Optional, Type, Union
 
 import torch
 from diffusers.models.modeling_utils import ModelMixin
-from pydantic import BaseModel, ConfigDict, Discriminator, Field, Tag, TypeAdapter
+from pydantic import BaseModel, ConfigDict, Discriminator, Field, JsonValue, Tag, TypeAdapter
 from typing_extensions import Annotated, Any, Dict
 
 from ..raw_model import RawModel
@@ -142,6 +142,10 @@ class ModelConfigBase(BaseModel):
     description: Optional[str] = Field(description="Model description", default=None)
     source: str = Field(description="The original source of the model (path, URL or repo_id).")
     source_type: ModelSourceType = Field(description="The type of source")
+    source_api_response: Optional[JsonValue] = Field(
+        description="The original API response from the source", default=None
+    )
+    trigger_words: Optional[set[str]] = Field(description="Set of trigger words for this model", default=None)
 
     model_config = ConfigDict(use_enum_values=False, validate_assignment=True)
 

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

@@ -19,7 +19,6 @@ from typing import Optional
 
 import requests
 from huggingface_hub import HfApi, configure_http_backend, hf_hub_url
-from huggingface_hub.hf_api import RepoSibling
 from huggingface_hub.utils._errors import RepositoryNotFoundError, RevisionNotFoundError
 from pydantic.networks import AnyHttpUrl
 from requests.sessions import Session
@@ -61,6 +60,7 @@ class HuggingFaceMetadataFetch(ModelMetadataFetchBase):
         # 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
         while not model_info:
             try:
@@ -73,12 +73,23 @@ class HuggingFaceMetadataFetch(ModelMetadataFetchBase):
                 else:
                     variant = None
 
+        files: list[RemoteModelFile] = []
+
         _, name = id.split("/")
-        return HuggingFaceMetadata(
-            id=model_info.id,
-            name=name,
-            files=parse_siblings(id, model_info.siblings, variant),
-        )
+
+        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(id, s.rfilename, revision=variant),
+                    path=Path(name, s.rfilename),
+                    size=s.size,
+                    sha256=s.lfs.get("sha256") if s.lfs else None,
+                )
+            )
+
+        return HuggingFaceMetadata(id=model_info.id, name=name, files=files)
 
     def from_url(self, url: AnyHttpUrl) -> AnyModelRepoMetadata:
         """
@@ -91,27 +102,3 @@ class HuggingFaceMetadataFetch(ModelMetadataFetchBase):
             return self.from_id(repo_id)
         else:
             raise UnknownMetadataException(f"'{url}' does not look like a HuggingFace model page")
-
-
-def parse_siblings(
-    repo_id: str, siblings: Optional[list[RepoSibling]] = None, variant: Optional[ModelRepoVariant] = None
-) -> list[RemoteModelFile]:
-    """Parse the siblings list from the HuggingFace API into a list of RemoteModelFile objects."""
-    if not siblings:
-        return []
-
-    files: list[RemoteModelFile] = []
-
-    for s in siblings:
-        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.value if variant else None),
-                path=Path(s.rfilename),
-                size=s.size,
-                sha256=s.lfs.get("sha256") if s.lfs else None,
-            )
-        )
-
-    return files

invokeai/backend/model_manager/probe.py

@@ -8,6 +8,7 @@ import torch
 from picklescan.scanner import scan_file_path
 
 import invokeai.backend.util.logging as logger
+from invokeai.app.util.misc import uuid_string
 from invokeai.backend.util.util import SilenceWarnings
 
 from .config import (
@@ -149,6 +150,9 @@ class ModelProbe(object):
 
         probe = probe_class(model_path)
 
+        fields["source_type"] = fields.get("source_type")
+        fields["source"] = fields.get("source") or model_path.as_posix()
+        fields["key"] = fields.get("key", uuid_string())
         fields["path"] = model_path.as_posix()
         fields["type"] = fields.get("type") or model_type
         fields["base"] = fields.get("base") or probe.get_base_type()
@@ -162,7 +166,7 @@ class ModelProbe(object):
         fields["format"] = fields.get("format") or probe.get_format()
         fields["hash"] = fields.get("hash") or ModelHash().hash(model_path)
 
-        if format_type == ModelFormat.Diffusers and hasattr(probe, "get_repo_variant"):
+        if format_type == ModelFormat.Diffusers and isinstance(probe, FolderProbeBase):
             fields["repo_variant"] = fields.get("repo_variant") or probe.get_repo_variant()
 
         # additional fields needed for main and controlnet models