Merge pull request #84 from Haidra-Org/main

feat: support `optimization` and `requirements` fields
Haidra-Org · Mar 3, 2024 · 507bea9 · 507bea9
2 parents 9ef7cd7 + 79cc22b
commit 507bea9
Show file tree

Hide file tree

Showing 6 changed files with 110 additions and 0 deletions.
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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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": [
  {