Maryhipp/trigger phrases

invokeai/app/api/routers/model_manager.py

@@ -1,10 +1,11 @@
 # Copyright (c) 2023 Lincoln D. Stein
 """FastAPI route for model configuration records."""
 
 import pathlib
 import shutil
 from hashlib import sha1
 from random import randbytes
+import traceback
 from typing import Any, Dict, List, Optional, Set
 
 from fastapi import Body, Path, Query, Response
@@ -14,6 +15,7 @@
 from typing_extensions import Annotated
 
 from invokeai.app.services.model_install import ModelInstallJob
+from invokeai.app.services.model_metadata.metadata_store_base import ModelMetadataChanges
 from invokeai.app.services.model_records import (
     DuplicateModelException,
     InvalidModelException,
@@ -32,6 +34,7 @@
 )
 from invokeai.backend.model_manager.merge import MergeInterpolationMethod, ModelMerger
 from invokeai.backend.model_manager.metadata import AnyModelRepoMetadata
+from invokeai.backend.model_manager.metadata.metadata_base import BaseMetadata
 from invokeai.backend.model_manager.search import ModelSearch
 
 from ..dependencies import ApiDependencies
@@ -242,6 +245,43 @@
 
     return result
 
+@model_manager_router.patch(
+    "/i/{key}/metadata",
+    operation_id="update_model_metadata",
+    responses={
+        201: {
+            "description": "The model metadata was updated successfully",
+            "content": {"application/json": {"example": example_model_metadata}},
+        },
+        400: {"description": "Bad request"},
+    },
+)
+async def update_model_metadata(
+    key: str = Path(description="Key of the model repo metadata to fetch."),
+    changes: ModelMetadataChanges = Body(description="The changes")
+) -> Optional[AnyModelRepoMetadata]:
+    """Updates or creates a model metadata object."""
+    record_store = ApiDependencies.invoker.services.model_manager.store
+    metadata_store = ApiDependencies.invoker.services.model_manager.store.metadata_store
+
+    try:
+        original_metadata = record_store.get_metadata(key)
+        if original_metadata:
+            original_metadata.trigger_phrases = changes.trigger_phrases
+
+            metadata_store.update_metadata(key, original_metadata)
+        else:
+            metadata_store.add_metadata(key, BaseMetadata(name="", author="",trigger_phrases=changes.trigger_phrases))
+    except Exception as e:
+        raise HTTPException(
+            status_code=500,
+            detail=f"An error occurred while updating the model metadata: {e}",
+        )
+
+    result: Optional[AnyModelRepoMetadata] = record_store.get_metadata(key)
+
+    return result
+
 
 @model_manager_router.get(
     "/tags",

invokeai/app/services/model_metadata/metadata_store_base.py

@@ -3,11 +3,24 @@
 Storage for Model Metadata
 """
 
 from abc import ABC, abstractmethod
-from typing import List, Set, Tuple
+from typing import List, Optional, Set, Tuple
+
+from pydantic import Field
+from invokeai.app.util.model_exclude_null import BaseModelExcludeNull
 
 from invokeai.backend.model_manager.metadata import AnyModelRepoMetadata
 
+class ModelMetadataChanges(BaseModelExcludeNull, extra="allow"):
+    """A set of changes to apply to model metadata.
+
+    Only limited changes are valid:
+      - `trigger_phrases`: the list of trigger phrases for this model
+    """
+
+    trigger_phrases: Optional[List[str]] = Field(default=None, description="The model's list of trigger phrases")
+    """The model's list of trigger phrases"""
+
 
 class ModelMetadataStoreBase(ABC):
     """Store, search and fetch model metadata retrieved from remote repositories."""

invokeai/app/services/model_metadata/metadata_store_sql.py

@@ -115,6 +115,8 @@ def update_metadata(self, model_key: str, metadata: AnyModelRepoMetadata) -> Any
             except sqlite3.Error as e:
                 self._db.conn.rollback()
                 raise e
+            except Exception as e:
+                raise e
 
         return self.get_metadata(model_key)
 
@@ -179,44 +181,45 @@ def search_by_name(self, name: str) -> Set[str]:
         )
         return {x[0] for x in self._cursor.fetchall()}
 
-    def _update_tags(self, model_key: str, tags: Set[str]) -> None:
+    def _update_tags(self, model_key: str, tags: Optional[Set[str]]) -> None:
         """Update tags for the model referenced by model_key."""
+        if tags:
         # remove previous tags from this model
-        self._cursor.execute(
-            """--sql
-            DELETE FROM model_tags
-            WHERE model_id=?;
-            """,
-            (model_key,),
-        )
-
-        for tag in tags:
-            self._cursor.execute(
-                """--sql
-                INSERT OR IGNORE INTO tags (
-                  tag_text
-                  )
-                VALUES (?);
-                """,
-                (tag,),
-            )
-            self._cursor.execute(
-                """--sql
-                SELECT tag_id
-                FROM tags
-                WHERE tag_text = ?
-                LIMIT 1;
-                """,
-                (tag,),
-            )
-            tag_id = self._cursor.fetchone()[0]
             self._cursor.execute(
                 """--sql
-                INSERT OR IGNORE INTO model_tags (
-                   model_id,
-                   tag_id
-                  )
-                VALUES (?,?);
+                DELETE FROM model_tags
+                WHERE model_id=?;
                 """,
-                (model_key, tag_id),
+                (model_key,),
             )
+
+            for tag in tags:
+                self._cursor.execute(
+                    """--sql
+                    INSERT OR IGNORE INTO tags (
+                    tag_text
+                    )
+                    VALUES (?);
+                    """,
+                    (tag,),
+                )
+                self._cursor.execute(
+                    """--sql
+                    SELECT tag_id
+                    FROM tags
+                    WHERE tag_text = ?
+                    LIMIT 1;
+                    """,
+                    (tag,),
+                )
+                tag_id = self._cursor.fetchone()[0]
+                self._cursor.execute(
+                    """--sql
+                    INSERT OR IGNORE INTO model_tags (
+                    model_id,
+                    tag_id
+                    )
+                    VALUES (?,?);
+                    """,
+                    (model_key, tag_id),
+                )

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

@@ -164,6 +164,7 @@ def _from_model_json(self, model_json: Dict[str, Any], version_id: Optional[int]
                 AllowDerivatives=model_json["allowDerivatives"],
                 AllowDifferentLicense=model_json["allowDifferentLicense"],
             ),
+            trigger_phrases=version_json["trainedWords"],
         )
 
     def from_civitai_versionid(self, version_id: int, model_id: Optional[int] = None) -> CivitaiMetadata:

invokeai/backend/model_manager/metadata/metadata_base.py

@@ -73,7 +73,8 @@ class ModelMetadataBase(BaseModel):
 
     name: str = Field(description="model's name")
     author: str = Field(description="model's author")
-    tags: Set[str] = Field(description="tags provided by model source")
+    tags: Optional[Set[str]] = Field(description="tags provided by model source", default=None)
+    trigger_phrases: Optional[List[str]] = Field(description="trigger phrases for this model", default=None)
 
 
 class BaseMetadata(ModelMetadataBase):

invokeai/frontend/web/public/locales/en.json

@@ -78,6 +78,7 @@
         "aboutDesc": "Using Invoke for work? Check out:",
         "aboutHeading": "Own Your Creative Power",
         "accept": "Accept",
+        "add": "Add",
         "advanced": "Advanced",
         "advancedOptions": "Advanced Options",
         "ai": "ai",
@@ -303,6 +304,12 @@
             "method": "High Resolution Fix Method"
         }
     },
+    "prompt": {
+        "addPromptTrigger": "Add Prompt Trigger",
+        "compatibleEmbeddings": "Compatible Embeddings",
+        "noPromptTriggers": "No triggers available",
+        "noMatchingTriggers": "No matching triggers"
+    },
     "embedding": {
         "addEmbedding": "Add Embedding",
         "incompatibleModel": "Incompatible base model:",
@@ -768,6 +775,7 @@
         "mergedModelName": "Merged Model Name",
         "mergedModelSaveLocation": "Save Location",
         "mergeModels": "Merge Models",
+        "metadata": "Metadata",
         "model": "Model",
         "modelAdded": "Model Added",
         "modelConversionFailed": "Model Conversion Failed",
@@ -839,6 +847,8 @@
         "statusConverting": "Converting",
         "syncModels": "Sync Models",
         "syncModelsDesc": "If your models are out of sync with the backend, you can refresh them up using this option. This is generally handy in cases where you add models to the InvokeAI root folder or autoimport directory after the application has booted.",
+        "triggerPhrases": "Trigger Phrases",
+        "typePhraseHere": "Type phrase here",
         "upcastAttention": "Upcast Attention",
         "updateModel": "Update Model",
         "useCustomConfig": "Use Custom Config",

invokeai/frontend/web/src/features/modelManagerV2/subpanels/ModelPane.tsx

@@ -8,7 +8,7 @@ export const ModelPane = () => {
   const selectedModelKey = useAppSelector((s) => s.modelmanagerV2.selectedModelKey);
   return (
     <Box layerStyle="first" p={2} borderRadius="base" w="50%" h="full">
-      {selectedModelKey ? <Model /> : <ImportModels />}
+      {selectedModelKey ? <Model key={selectedModelKey} /> : <ImportModels />}
     </Box>
   );
 };

invokeai/frontend/web/src/features/modelManagerV2/subpanels/ModelPanel/Metadata/ModelMetadata.tsx

@@ -0,0 +1,21 @@
+import { Box, Flex } from '@invoke-ai/ui-library';
+import { skipToken } from '@reduxjs/toolkit/query';
+import { useAppSelector } from 'app/store/storeHooks';
+import DataViewer from 'features/gallery/components/ImageMetadataViewer/DataViewer';
+import { useGetModelMetadataQuery } from 'services/api/endpoints/models';
+
+import { TriggerPhrases } from './TriggerPhrases';
+
+export const ModelMetadata = () => {
+  const selectedModelKey = useAppSelector((s) => s.modelmanagerV2.selectedModelKey);
+  const { data: metadata } = useGetModelMetadataQuery(selectedModelKey ?? skipToken);
+
+  return (
+    <Flex flexDir="column" height="full" gap="3">
+      <Box layerStyle="second" borderRadius="base" p={3}>
+        <TriggerPhrases />
+      </Box>
+      <DataViewer label="metadata" data={metadata || {}} />
+    </Flex>
+  );
+};