feat: file size field

create_example_json.py

@@ -42,6 +42,7 @@ def main():
         trigger=["trigger1", "some other_trigger"],
         homepage="https://www.not.a.real_website.com",
         nsfw=False,
+        size_on_disk_bytes=123456789,
     )
 
     example_record_2_name = "Example anime model"
@@ -67,6 +68,7 @@ def main():
         trigger=["anime", "some other_anime_trigger"],
         homepage="https://www.another_fake_website.com",
         nsfw=True,
+        size_on_disk_bytes=123456789,
     )
 
     reference = StableDiffusion_ModelReference(

horde_model_reference/legacy/classes/raw_legacy_model_database_records.py

@@ -58,6 +58,7 @@ class RawLegacy_StableDiffusion_ModelRecord(BaseModel):
     config: Mapping[str, list[RawLegacy_FileRecord | RawLegacy_DownloadRecord]]
     available: bool | None = None
     features_not_supported: list[FEATURE_SUPPORTED] | None = None
+    size_on_disk_bytes: int | None = None
 
 
 class RawLegacy_StableDiffusion_ModelReference(RootModel[Mapping[str, RawLegacy_StableDiffusion_ModelRecord]]):

horde_model_reference/legacy/classes/staging_model_database_records.py

@@ -82,6 +82,7 @@ class Legacy_StableDiffusion_ModelRecord(StagingLegacy_Generic_ModelRecord):
     min_bridge_version: int | None = None
     trigger: list[str] | None = None
     homepage: str | None = None
+    size_on_disk_bytes: int | None = None
 
 
 class Legacy_Generic_ModelReference(BaseModel):

horde_model_reference/legacy/get_all_filesizes.py

@@ -0,0 +1,62 @@
+import json
+from pathlib import Path
+
+from loguru import logger
+
+from horde_model_reference.legacy.classes.raw_legacy_model_database_records import (
+    RawLegacy_StableDiffusion_ModelRecord,
+)
+from horde_model_reference.path_consts import AIWORKER_CACHE_HOME
+
+
+def get_all_file_sizes(sd_db: Path, write_to_path: Path | None = None) -> bool:
+    raw_json_sd_db: str
+    with open(sd_db) as sd_db_file:
+        raw_json_sd_db = sd_db_file.read()
+    try:
+        loaded_json_sd_db = json.loads(raw_json_sd_db)
+    except Exception as e:
+        logger.exception(e)
+        logger.exception(f"ERROR: The stable diffusion database specified ({sd_db}) is not a valid json file.")
+        if __name__ == "__main__":
+            exit(1)
+        else:
+            return False
+
+    parsed_db_records: dict[str, RawLegacy_StableDiffusion_ModelRecord] = {
+        k: RawLegacy_StableDiffusion_ModelRecord.model_validate(v) for k, v in loaded_json_sd_db.items()
+    }
+
+    for _, model_details in parsed_db_records.items():
+        filename = model_details.config["files"][0].path
+
+        full_file_path = Path(AIWORKER_CACHE_HOME) / "compvis" / filename
+        if not full_file_path.exists():
+            logger.error(f"File {full_file_path} does not exist.")
+            continue
+
+        model_details.size_on_disk_bytes = full_file_path.stat().st_size
+
+    correct_json_layout = json.dumps(
+        {
+            k: v.model_dump(
+                exclude_none=True,
+                exclude_defaults=False,
+                by_alias=True,
+            )
+            for k, v in parsed_db_records.items()
+        },
+        indent=4,
+    )
+    correct_json_layout += "\n"  # Add a newline to the end of the file, for consistency with formatters.
+
+    with open(write_to_path, "w") as corrected_sd_db_file:
+        corrected_sd_db_file.write(correct_json_layout)
+    return True
+
+
+if __name__ == "__main__":
+    get_all_file_sizes(
+        Path("t:/_NATAILI_CACHE_HOME_/horde_model_reference/legacy/stable_diffusion.json"),
+        "with_sizes.json",
+    )

horde_model_reference/model_reference_records.py

@@ -77,6 +77,8 @@ class StableDiffusion_ModelRecord(Generic_ModelRecord):
     style: MODEL_STYLE | None = None
     """The style of the model."""
 
+    size_on_disk_bytes: int | None = None
+
     @model_validator(mode="after")  # type: ignore # FIXME
     def validator_set_arrays_to_empty_if_none(self) -> StableDiffusion_ModelRecord:
         """Set any `None` values to empty lists."""

legacy_stable_diffusion.schema.json

@@ -11,8 +11,8 @@
             "title": "FEATURE_SUPPORTED",
             "type": "string"
         },
-        "DownloadRecord": {
-            "description": "An entry in the `config` field of a `StableDiffusionModelRecord`.",
+        "RawLegacy_DownloadRecord": {
+            "description": "An entry in the `config` field of a `RawLegacy_StableDiffusion_ModelRecord`.",
             "properties": {
                 "file_name": {
                     "title": "File Name",
@@ -32,11 +32,11 @@
                 "file_path",
                 "file_url"
             ],
-            "title": "DownloadRecord",
+            "title": "RawLegacy_DownloadRecord",
             "type": "object"
         },
-        "FileRecord": {
-            "description": "An entry in the `config` field of a `StableDiffusionModelRecord`.",
+        "RawLegacy_FileRecord": {
+            "description": "An entry in the `config` field of a `RawLegacy_StableDiffusion_ModelRecord`.",
             "properties": {
                 "path": {
                     "title": "Path",
@@ -70,10 +70,10 @@
             "required": [
                 "path"
             ],
-            "title": "FileRecord",
+            "title": "RawLegacy_FileRecord",
             "type": "object"
         },
-        "StableDiffusionModelRecord": {
+        "RawLegacy_StableDiffusion_ModelRecord": {
             "additionalProperties": false,
             "description": "A model entry in the legacy model reference.",
             "properties": {
@@ -203,10 +203,10 @@
                         "items": {
                             "anyOf": [
                                 {
-                                    "$ref": "#/$defs/FileRecord"
+                                    "$ref": "#/$defs/RawLegacy_FileRecord"
                                 },
                                 {
-                                    "$ref": "#/$defs/DownloadRecord"
+                                    "$ref": "#/$defs/RawLegacy_DownloadRecord"
                                 }
                             ]
                         },
@@ -241,6 +241,18 @@
                     ],
                     "default": null,
                     "title": "Features Not Supported"
+                },
+                "size_on_disk_bytes": {
+                    "anyOf": [
+                        {
+                            "type": "integer"
+                        },
+                        {
+                            "type": "null"
+                        }
+                    ],
+                    "default": null,
+                    "title": "Size On Disk Bytes"
                 }
             },
             "required": [
@@ -253,13 +265,13 @@
                 "download_all",
                 "config"
             ],
-            "title": "StableDiffusionModelRecord",
+            "title": "RawLegacy_StableDiffusion_ModelRecord",
             "type": "object"
         }
     },
     "additionalProperties": {
-        "$ref": "#/$defs/StableDiffusionModelRecord"
+        "$ref": "#/$defs/RawLegacy_StableDiffusion_ModelRecord"
     },
-    "title": "StableDiffusionModelReference",
+    "title": "RawLegacy_StableDiffusion_ModelReference",
     "type": "object"
 }

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "horde_model_reference"
-version = "0.5.0"
+version = "0.5.1"
 description = "A helper library providing a way to work with the lists of diffusion models, utility models, and any other related files required for AI-Horde."
 authors = [
     {name = "tazlin", email = "tazlin.on.github@gmail.com"},

stable_diffusion.example.json

@@ -31,7 +31,8 @@
         ],
         "homepage": "https://www.not.a.real_website.com",
         "nsfw": false,
-        "style": "generalist"
+        "style": "generalist",
+        "size_on_disk_bytes": 123456789
     },
     "example model 2": {
         "name": "Example anime model",
@@ -65,6 +66,7 @@
         ],
         "homepage": "https://www.another_fake_website.com",
         "nsfw": true,
-        "style": "anime"
+        "style": "anime",
+        "size_on_disk_bytes": 123456789
     }
 }

stable_diffusion.schema.json

@@ -228,6 +228,18 @@
                         }
                     ],
                     "default": null
+                },
+                "size_on_disk_bytes": {
+                    "anyOf": [
+                        {
+                            "type": "integer"
+                        },
+                        {
+                            "type": "null"
+                        }
+                    ],
+                    "default": null,
+                    "title": "Size On Disk Bytes"
                 }
             },
             "required": [