handle change to Civitai metadata schema for commercial usage

This commit is contained in:
Lincoln Stein 2024-02-27 23:54:28 -05:00 committed by psychedelicious
parent 4418c118db
commit 98a13aa7dc
4 changed files with 9 additions and 6 deletions

View File

@ -160,7 +160,7 @@ class CivitaiMetadataFetch(ModelMetadataFetchBase):
nsfw=model_json["nsfw"], nsfw=model_json["nsfw"],
restrictions=LicenseRestrictions( restrictions=LicenseRestrictions(
AllowNoCredit=model_json["allowNoCredit"], AllowNoCredit=model_json["allowNoCredit"],
AllowCommercialUse=CommercialUsage(model_json["allowCommercialUse"]), AllowCommercialUse={CommercialUsage(x) for x in model_json["allowCommercialUse"]},
AllowDerivatives=model_json["allowDerivatives"], AllowDerivatives=model_json["allowDerivatives"],
AllowDifferentLicense=model_json["allowDifferentLicense"], AllowDifferentLicense=model_json["allowDifferentLicense"],
), ),

View File

@ -54,8 +54,8 @@ class LicenseRestrictions(BaseModel):
AllowDifferentLicense: bool = Field( AllowDifferentLicense: bool = Field(
description="if true, derivatives of this model be redistributed under a different license", default=False description="if true, derivatives of this model be redistributed under a different license", default=False
) )
AllowCommercialUse: Optional[CommercialUsage] = Field( AllowCommercialUse: Optional[Set[CommercialUsage] | CommercialUsage] = Field(
description="Type of commercial use allowed or 'No' if no commercial use is allowed.", default=None description="Type of commercial use allowed if no commercial use is allowed.", default=None
) )
@ -142,7 +142,10 @@ class CivitaiMetadata(ModelMetadataWithFiles):
if self.restrictions.AllowCommercialUse is None: if self.restrictions.AllowCommercialUse is None:
return False return False
else: 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 @property
def allow_derivatives(self) -> bool: def allow_derivatives(self) -> bool:

File diff suppressed because one or more lines are too long

View File

@ -133,7 +133,7 @@ def test_metadata_civitai_fetch(mm2_session: Session) -> None:
assert metadata.id == 215485 assert metadata.id == 215485
assert metadata.author == "test_author" # note that this is not the same as the original from Civitai assert metadata.author == "test_author" # note that this is not the same as the original from Civitai
assert metadata.allow_commercial_use # changed to make sure we are reading locally not remotely assert metadata.allow_commercial_use # changed to make sure we are reading locally not remotely
assert metadata.restrictions.AllowCommercialUse == CommercialUsage("RentCivit") assert CommercialUsage("RentCivit") in metadata.restrictions.AllowCommercialUse
assert metadata.version_id == 242807 assert metadata.version_id == 242807
assert metadata.tags == {"tool", "turbo", "sdxl turbo"} assert metadata.tags == {"tool", "turbo", "sdxl turbo"}