handle change to Civitai metadata schema for commercial usage

invoke-ai · Feb 28, 2024 · fce2b15 · fce2b15
1 parent 4418c11
commit fce2b15
Show file tree

Hide file tree

Showing 4 changed files with 9 additions and 6 deletions.
diff --git a/invokeai/backend/model_manager/metadata/fetch/civitai.py b/invokeai/backend/model_manager/metadata/fetch/civitai.py
@@ -160,7 +160,7 @@ def _from_model_json(self, model_json: Dict[str, Any], version_id: Optional[int]
             nsfw=model_json["nsfw"],
             restrictions=LicenseRestrictions(
                 AllowNoCredit=model_json["allowNoCredit"],
-                AllowCommercialUse=CommercialUsage(model_json["allowCommercialUse"]),
+                AllowCommercialUse={CommercialUsage(x) for x in model_json["allowCommercialUse"]},
                 AllowDerivatives=model_json["allowDerivatives"],
                 AllowDifferentLicense=model_json["allowDifferentLicense"],
             ),

diff --git a/invokeai/backend/model_manager/metadata/metadata_base.py b/invokeai/backend/model_manager/metadata/metadata_base.py
@@ -54,8 +54,8 @@ class LicenseRestrictions(BaseModel):
     AllowDifferentLicense: bool = Field(
         description="if true, derivatives of this model be redistributed under a different license", default=False
     )
-    AllowCommercialUse: Optional[CommercialUsage] = Field(
-        description="Type of commercial use allowed or 'No' if no commercial use is allowed.", default=None
+    AllowCommercialUse: Optional[Set[CommercialUsage] | CommercialUsage] = Field(
+        description="Type of commercial use allowed if no commercial use is allowed.", default=None
     )
 
 
@@ -142,7 +142,10 @@ def allow_commercial_use(self) -> bool:
         if self.restrictions.AllowCommercialUse is None:
             return False
         else:
-            return self.restrictions.AllowCommercialUse != CommercialUsage("None")
+            # accommodate schema change
+            acu = self.restrictions.AllowCommercialUse
+            commercial_usage = acu if isinstance(acu, set) else {acu}
+            return CommercialUsage.No not in commercial_usage
 
     @property
     def allow_derivatives(self) -> bool: