From 79cc22b4335156a9a96d86172f12d92f9a30f90d Mon Sep 17 00:00:00 2001 From: tazlin Date: Sun, 3 Mar 2024 13:40:12 -0500 Subject: [PATCH] feat: support `optimization` and `requirements` fields --- .../raw_legacy_model_database_records.py | 2 + .../classes/staging_model_database_records.py | 2 + .../model_reference_records.py | 4 ++ legacy_stable_diffusion.schema.json | 49 +++++++++++++++++++ stable_diffusion.example.json | 4 ++ stable_diffusion.schema.json | 49 +++++++++++++++++++ 6 files changed, 110 insertions(+) diff --git a/horde_model_reference/legacy/classes/raw_legacy_model_database_records.py b/horde_model_reference/legacy/classes/raw_legacy_model_database_records.py index c2b15e1..67e0f1b 100644 --- a/horde_model_reference/legacy/classes/raw_legacy_model_database_records.py +++ b/horde_model_reference/legacy/classes/raw_legacy_model_database_records.py @@ -46,6 +46,7 @@ class RawLegacy_StableDiffusion_ModelRecord(BaseModel): name: str baseline: str + optimization: str | None = None type: str inpainting: bool description: str | None = None @@ -58,6 +59,7 @@ class RawLegacy_StableDiffusion_ModelRecord(BaseModel): homepage: str | None = None nsfw: bool download_all: bool + requirements: dict[str, int | str | list[int] | list[str] | bool] | None = None config: Mapping[str, list[RawLegacy_FileRecord | RawLegacy_DownloadRecord]] available: bool | None = None features_not_supported: list[FEATURE_SUPPORTED] | None = None diff --git a/horde_model_reference/legacy/classes/staging_model_database_records.py b/horde_model_reference/legacy/classes/staging_model_database_records.py index 3142411..06892c1 100644 --- a/horde_model_reference/legacy/classes/staging_model_database_records.py +++ b/horde_model_reference/legacy/classes/staging_model_database_records.py @@ -95,6 +95,8 @@ class Legacy_StableDiffusion_ModelRecord(StagingLegacy_Generic_ModelRecord): trigger: list[str] | None = None homepage: str | None = None size_on_disk_bytes: int | None = None + optimization: str | None = None + requirements: dict[str, int | str | list[int] | list[str] | bool] | None = None class Legacy_Generic_ModelReference(BaseModel): diff --git a/horde_model_reference/model_reference_records.py b/horde_model_reference/model_reference_records.py index 021fbb8..b8db3a7 100644 --- a/horde_model_reference/model_reference_records.py +++ b/horde_model_reference/model_reference_records.py @@ -72,6 +72,8 @@ class StableDiffusion_ModelRecord(Generic_ModelRecord): """If this is an inpainting model or not.""" baseline: STABLE_DIFFUSION_BASELINE_CATEGORY | str """The model on which this model is based.""" + optimization: str | None = None + """The optimization type of the model.""" tags: list[str] | None = [] """Any tags associated with the model which may be useful for searching.""" showcases: list[str] | None = [] @@ -88,6 +90,8 @@ class StableDiffusion_ModelRecord(Generic_ModelRecord): style: MODEL_STYLE | str | None = None """The style of the model.""" + requirements: dict[str, int | str | list[int] | list[str] | bool] | None = None + size_on_disk_bytes: int | None = None @model_validator(mode="after") diff --git a/legacy_stable_diffusion.schema.json b/legacy_stable_diffusion.schema.json index 0d82588..217f305 100644 --- a/legacy_stable_diffusion.schema.json +++ b/legacy_stable_diffusion.schema.json @@ -109,6 +109,18 @@ "title": "Baseline", "type": "string" }, + "optimization": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Optimization" + }, "type": { "title": "Type", "type": "string" @@ -222,6 +234,43 @@ "title": "Download All", "type": "boolean" }, + "requirements": { + "anyOf": [ + { + "additionalProperties": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "string" + }, + { + "items": { + "type": "integer" + }, + "type": "array" + }, + { + "items": { + "type": "string" + }, + "type": "array" + }, + { + "type": "boolean" + } + ] + }, + "type": "object" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Requirements" + }, "config": { "additionalProperties": { "items": { diff --git a/stable_diffusion.example.json b/stable_diffusion.example.json index 870b459..ffab846 100644 --- a/stable_diffusion.example.json +++ b/stable_diffusion.example.json @@ -18,6 +18,7 @@ "features_not_supported": null, "inpainting": false, "baseline": "stable_diffusion_1", + "optimization": null, "tags": [ "anime", "faces" @@ -33,6 +34,7 @@ "homepage": "https://www.not.a.real_website.com", "nsfw": false, "style": "generalist", + "requirements": null, "size_on_disk_bytes": 123456789 }, "example model 2": { @@ -54,6 +56,7 @@ "features_not_supported": null, "inpainting": false, "baseline": "stable_diffusion_1", + "optimization": null, "tags": [ "anime", "faces" @@ -69,6 +72,7 @@ "homepage": "https://www.another_fake_website.com", "nsfw": true, "style": "anime", + "requirements": null, "size_on_disk_bytes": 123456789 } } diff --git a/stable_diffusion.schema.json b/stable_diffusion.schema.json index f69df4e..3475fb6 100644 --- a/stable_diffusion.schema.json +++ b/stable_diffusion.schema.json @@ -174,6 +174,18 @@ ], "title": "Baseline" }, + "optimization": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Optimization" + }, "tags": { "anyOf": [ { @@ -262,6 +274,43 @@ "default": null, "title": "Style" }, + "requirements": { + "anyOf": [ + { + "additionalProperties": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "string" + }, + { + "items": { + "type": "integer" + }, + "type": "array" + }, + { + "items": { + "type": "string" + }, + "type": "array" + }, + { + "type": "boolean" + } + ] + }, + "type": "object" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Requirements" + }, "size_on_disk_bytes": { "anyOf": [ {