diff --git a/create_example_json.py b/create_example_json.py index 05ff519..41f9251 100644 --- a/create_example_json.py +++ b/create_example_json.py @@ -1,6 +1,6 @@ import json -from horde_model_reference import MODEL_PURPOSE, MODEL_STYLE, STABLE_DIFFUSION_BASELINE_CATEGORY +from horde_model_reference import IMAGE_GENERATION_BASELINE, MODEL_PURPOSE, MODEL_STYLE from horde_model_reference.legacy.classes.raw_legacy_model_database_records import ( RawLegacy_StableDiffusion_ModelReference, ) @@ -33,7 +33,7 @@ def main(): style=MODEL_STYLE.generalist, config={"download": [example_download_record]}, purpose=MODEL_PURPOSE.image_generation, - baseline=STABLE_DIFFUSION_BASELINE_CATEGORY.stable_diffusion_1, + baseline=IMAGE_GENERATION_BASELINE.stable_diffusion_1, tags=["anime", "faces"], showcases=[ "https://raw.githubusercontent.com/db0/AI-Horde-image-model-reference/main/showcase/test/test_general_01.png", @@ -59,7 +59,7 @@ def main(): style=MODEL_STYLE.anime, config={"download": [example_download_record_2]}, purpose=MODEL_PURPOSE.image_generation, - baseline=STABLE_DIFFUSION_BASELINE_CATEGORY.stable_diffusion_1, + baseline=IMAGE_GENERATION_BASELINE.stable_diffusion_1, tags=["anime", "faces"], showcases=[ "https://raw.githubusercontent.com/db0/AI-Horde-image-model-reference/main/showcase/test/anime_01.png", diff --git a/horde_model_reference/__init__.py b/horde_model_reference/__init__.py index e22c05e..8849d44 100644 --- a/horde_model_reference/__init__.py +++ b/horde_model_reference/__init__.py @@ -4,12 +4,12 @@ from .meta_consts import ( # noqa: E402 + IMAGE_GENERATION_BASELINE, KNOWN_TAGS, MODEL_PURPOSE, MODEL_PURPOSE_LOOKUP, MODEL_REFERENCE_CATEGORY, MODEL_STYLE, - STABLE_DIFFUSION_BASELINE_CATEGORY, ) from .path_consts import ( # noqa: E402 BASE_PATH, @@ -26,7 +26,7 @@ "MODEL_PURPOSE", "MODEL_PURPOSE_LOOKUP", "MODEL_STYLE", - "STABLE_DIFFUSION_BASELINE_CATEGORY", + "IMAGE_GENERATION_BASELINE", "BASE_PATH", "DEFAULT_SHOWCASE_FOLDER_NAME", "LEGACY_REFERENCE_FOLDER", diff --git a/horde_model_reference/legacy/convert_all_legacy_dbs.py b/horde_model_reference/legacy/convert_all_legacy_dbs.py index b83c8d7..d2d334b 100644 --- a/horde_model_reference/legacy/convert_all_legacy_dbs.py +++ b/horde_model_reference/legacy/convert_all_legacy_dbs.py @@ -46,6 +46,8 @@ def convert_all_legacy_model_references( non_generic_converter_categories = [ MODEL_REFERENCE_CATEGORY.stable_diffusion, MODEL_REFERENCE_CATEGORY.clip, + MODEL_REFERENCE_CATEGORY.lora, + MODEL_REFERENCE_CATEGORY.ti, ] generic_converted_categories = [x for x in MODEL_REFERENCE_CATEGORY if x not in non_generic_converter_categories] for model_category in generic_converted_categories: diff --git a/horde_model_reference/legacy/legacy_download_manager.py b/horde_model_reference/legacy/legacy_download_manager.py index 23e1031..8f4922d 100644 --- a/horde_model_reference/legacy/legacy_download_manager.py +++ b/horde_model_reference/legacy/legacy_download_manager.py @@ -6,7 +6,7 @@ from loguru import logger from horde_model_reference.legacy.convert_all_legacy_dbs import convert_all_legacy_model_references -from horde_model_reference.meta_consts import MODEL_REFERENCE_CATEGORY +from horde_model_reference.meta_consts import LOCAL_MODEL_REFERENCE_CATEGORIES, MODEL_REFERENCE_CATEGORY from horde_model_reference.path_consts import ( BASE_PATH, HORDE_PROXY_URL_BASE, @@ -43,6 +43,12 @@ def download_legacy_model_reference( model_category_name: MODEL_REFERENCE_CATEGORY, override_existing: bool = False, ) -> pathlib.Path | None: + target_file_path = get_model_reference_file_path(model_category_name, base_path=self.legacy_path) + + if model_category_name in LOCAL_MODEL_REFERENCE_CATEGORIES: + logger.debug(f"Skipping download of {model_category_name} reference file, as it is local.") + return target_file_path + response = requests.get(self.proxy_url + LEGACY_MODEL_GITHUB_URLS[model_category_name]) if response.status_code != 200: logger.error(f"Failed to download {model_category_name} reference file.") @@ -54,8 +60,6 @@ def download_legacy_model_reference( logger.error(f"Failed to parse {model_category_name} reference file as JSON.") return None - target_file_path = get_model_reference_file_path(model_category_name, base_path=self.legacy_path) - if target_file_path.exists() and not override_existing: logger.debug(f"File {target_file_path} already exists, skipping download.") return None diff --git a/horde_model_reference/meta_consts.py b/horde_model_reference/meta_consts.py index 8eab357..101c336 100644 --- a/horde_model_reference/meta_consts.py +++ b/horde_model_reference/meta_consts.py @@ -61,6 +61,14 @@ class MODEL_REFERENCE_CATEGORY(StrEnum): safety_checker = auto() stable_diffusion = auto() miscellaneous = auto() + ti = auto() + lora = auto() + + +LOCAL_MODEL_REFERENCE_CATEGORIES = [ + MODEL_REFERENCE_CATEGORY.ti, + MODEL_REFERENCE_CATEGORY.lora, +] class MODEL_PURPOSE(StrEnum): @@ -82,7 +90,7 @@ class MODEL_PURPOSE(StrEnum): miscellaneous = auto() -class STABLE_DIFFUSION_BASELINE_CATEGORY(StrEnum): +class IMAGE_GENERATION_BASELINE(StrEnum): """An enum of all the image generation baselines.""" stable_diffusion_1 = auto() @@ -90,7 +98,7 @@ class STABLE_DIFFUSION_BASELINE_CATEGORY(StrEnum): stable_diffusion_2_512 = auto() stable_diffusion_xl = auto() stable_cascade = auto() - flux_1 = auto() # TODO: Extract flux and create "IMAGE_GENERATION_BASELINE_CATEGORY" due to name inconsistency + flux_1 = auto() MODEL_PURPOSE_LOOKUP: dict[MODEL_REFERENCE_CATEGORY, MODEL_PURPOSE] = { @@ -105,18 +113,18 @@ class STABLE_DIFFUSION_BASELINE_CATEGORY(StrEnum): MODEL_REFERENCE_CATEGORY.miscellaneous: MODEL_PURPOSE.miscellaneous, } -STABLE_DIFFUSION_BASELINE_NATIVE_RESOLUTION_LOOKUP: dict[STABLE_DIFFUSION_BASELINE_CATEGORY, int] = { - STABLE_DIFFUSION_BASELINE_CATEGORY.stable_diffusion_1: 512, - STABLE_DIFFUSION_BASELINE_CATEGORY.stable_diffusion_2_768: 768, - STABLE_DIFFUSION_BASELINE_CATEGORY.stable_diffusion_2_512: 512, - STABLE_DIFFUSION_BASELINE_CATEGORY.stable_diffusion_xl: 1024, - STABLE_DIFFUSION_BASELINE_CATEGORY.stable_cascade: 1024, - STABLE_DIFFUSION_BASELINE_CATEGORY.flux_1: 1024, +STABLE_DIFFUSION_BASELINE_NATIVE_RESOLUTION_LOOKUP: dict[IMAGE_GENERATION_BASELINE, int] = { + IMAGE_GENERATION_BASELINE.stable_diffusion_1: 512, + IMAGE_GENERATION_BASELINE.stable_diffusion_2_768: 768, + IMAGE_GENERATION_BASELINE.stable_diffusion_2_512: 512, + IMAGE_GENERATION_BASELINE.stable_diffusion_xl: 1024, + IMAGE_GENERATION_BASELINE.stable_cascade: 1024, + IMAGE_GENERATION_BASELINE.flux_1: 1024, } """The single-side preferred resolution for each known stable diffusion baseline.""" -def get_baseline_native_resolution(baseline: STABLE_DIFFUSION_BASELINE_CATEGORY) -> int: +def get_baseline_native_resolution(baseline: IMAGE_GENERATION_BASELINE) -> int: """ Get the native resolution of a stable diffusion baseline. diff --git a/horde_model_reference/model_reference_records.py b/horde_model_reference/model_reference_records.py index ec31a4a..a257937 100644 --- a/horde_model_reference/model_reference_records.py +++ b/horde_model_reference/model_reference_records.py @@ -15,10 +15,10 @@ ) from horde_model_reference import ( + IMAGE_GENERATION_BASELINE, MODEL_PURPOSE, MODEL_REFERENCE_CATEGORY, MODEL_STYLE, - STABLE_DIFFUSION_BASELINE_CATEGORY, ) from horde_model_reference.meta_consts import CONTROLNET_STYLE @@ -70,7 +70,7 @@ class StableDiffusion_ModelRecord(Generic_ModelRecord): inpainting: bool | None = False """If this is an inpainting model or not.""" - baseline: STABLE_DIFFUSION_BASELINE_CATEGORY | str + baseline: IMAGE_GENERATION_BASELINE | str """The model on which this model is based.""" optimization: str | None = None """The optimization type of the model.""" @@ -108,7 +108,7 @@ def validator_set_arrays_to_empty_if_none(self) -> StableDiffusion_ModelRecord: @model_validator(mode="after") def validator_is_baseline_and_style_known(self) -> StableDiffusion_ModelRecord: """Check if the baseline is known.""" - if str(self.baseline) not in STABLE_DIFFUSION_BASELINE_CATEGORY.__members__: + if str(self.baseline) not in IMAGE_GENERATION_BASELINE.__members__: logger.warning(f"Unknown baseline {self.baseline} for model {self.name}") if self.style is not None and str(self.style) not in MODEL_STYLE.__members__: @@ -142,7 +142,7 @@ class Generic_ModelReference(RootModel[Mapping[str, Generic_ModelRecord]]): class StableDiffusion_ModelReference(Generic_ModelReference): """The combined metadata and model list.""" - _baseline: dict[STABLE_DIFFUSION_BASELINE_CATEGORY | str, int] = PrivateAttr(default_factory=dict) + _baseline: dict[IMAGE_GENERATION_BASELINE | str, int] = PrivateAttr(default_factory=dict) """A dictionary of all the baseline types and how many models use them.""" _styles: dict[MODEL_STYLE | str, int] = PrivateAttr(default_factory=dict) """A dictionary of all the styles and how many models use them.""" @@ -204,7 +204,7 @@ def rebuild_metadata(self) -> None: self._models_hosts[host] = self._models_hosts.get(host, 0) + 1 @property - def baseline(self) -> dict[STABLE_DIFFUSION_BASELINE_CATEGORY | str, int]: + def baseline(self) -> dict[IMAGE_GENERATION_BASELINE | str, int]: """Return a dictionary of all the baseline types and how many models use them.""" self.check_was_models_modified() return self._baseline @@ -232,7 +232,7 @@ def models_names(self) -> set[str]: """Return a list of all the model names.""" return set(self.root.keys()) - def get_model_baseline(self, model_name: str) -> STABLE_DIFFUSION_BASELINE_CATEGORY | str | None: + def get_model_baseline(self, model_name: str) -> IMAGE_GENERATION_BASELINE | str | None: """Return the baseline for a given model name.""" model: StableDiffusion_ModelRecord | None = self.root.get(model_name) diff --git a/stable_diffusion.schema.json b/stable_diffusion.schema.json index ad6edc9..f255a2e 100644 --- a/stable_diffusion.schema.json +++ b/stable_diffusion.schema.json @@ -48,6 +48,19 @@ "title": "DownloadRecord", "type": "object" }, + "IMAGE_GENERATION_BASELINE": { + "description": "An enum of all the image generation baselines.", + "enum": [ + "stable_diffusion_1", + "stable_diffusion_2_768", + "stable_diffusion_2_512", + "stable_diffusion_xl", + "stable_cascade", + "flux_1" + ], + "title": "IMAGE_GENERATION_BASELINE", + "type": "string" + }, "MODEL_PURPOSE": { "enum": [ "image_generation", @@ -73,19 +86,6 @@ "title": "MODEL_STYLE", "type": "string" }, - "STABLE_DIFFUSION_BASELINE_CATEGORY": { - "description": "An enum of all the image generation baselines.", - "enum": [ - "stable_diffusion_1", - "stable_diffusion_2_768", - "stable_diffusion_2_512", - "stable_diffusion_xl", - "stable_cascade", - "flux_1" - ], - "title": "STABLE_DIFFUSION_BASELINE_CATEGORY", - "type": "string" - }, "StableDiffusion_ModelRecord": { "description": "A model entry in the model reference.", "properties": { @@ -168,7 +168,7 @@ "baseline": { "anyOf": [ { - "$ref": "#/$defs/STABLE_DIFFUSION_BASELINE_CATEGORY" + "$ref": "#/$defs/IMAGE_GENERATION_BASELINE" }, { "type": "string"