Skip to content

Commit e380521

Browse files
fix(mm): fixes for migration 23
- Handle CLIP Embed and Main SD models missing variant field - Handle errors when calling the discriminator function, previously only handled ValidationError but it could be a ValueError or something else - Better logging for config migration
1 parent 644bb98 commit e380521

File tree

1 file changed

+43
-5
lines changed

1 file changed

+43
-5
lines changed

invokeai/app/services/shared/sqlite_migrator/migrations/migration_23.py

Lines changed: 43 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -12,9 +12,11 @@
1212
from invokeai.backend.model_manager.configs.unknown import Unknown_Config
1313
from invokeai.backend.model_manager.taxonomy import (
1414
BaseModelType,
15+
ClipVariantType,
1516
FluxVariantType,
1617
ModelFormat,
1718
ModelType,
19+
ModelVariantType,
1820
SchedulerPredictionType,
1921
)
2022

@@ -30,10 +32,20 @@ def __call__(self, cursor: sqlite3.Cursor) -> None:
3032
cursor.execute("SELECT id, config FROM models;")
3133
rows = cursor.fetchall()
3234

35+
migrated_count = 0
36+
fallback_count = 0
37+
3338
for model_id, config_json in rows:
3439
try:
3540
# Migrate the config JSON to the latest schema
36-
migrated_config = self._parse_and_migrate_config(config_json)
41+
config_dict: dict[str, Any] = json.loads(config_json)
42+
migrated_config = self._parse_and_migrate_config(config_dict)
43+
44+
if isinstance(migrated_config, Unknown_Config):
45+
fallback_count += 1
46+
else:
47+
migrated_count += 1
48+
3749
# Write the migrated config back to the database
3850
cursor.execute(
3951
"UPDATE models SET config = ? WHERE id = ?;",
@@ -46,9 +58,21 @@ def __call__(self, cursor: sqlite3.Cursor) -> None:
4658
self._logger.error("Invalid config JSON for model %s: %s", model_id, e)
4759
raise
4860

49-
def _parse_and_migrate_config(self, config_json: Any) -> AnyModelConfig:
50-
config_dict: dict[str, Any] = json.loads(config_json)
61+
if migrated_count > 0 and fallback_count == 0:
62+
self._logger.info(f"Migration complete: {migrated_count} model configs migrated")
63+
elif migrated_count > 0 and fallback_count > 0:
64+
self._logger.warning(
65+
f"Migration complete: {migrated_count} model configs migrated, "
66+
f"{fallback_count} model configs could not be migrated and were saved as unknown models",
67+
)
68+
elif migrated_count == 0 and fallback_count > 0:
69+
self._logger.warning(
70+
f"Migration complete: all {fallback_count} model configs could not be migrated and were saved as unknown models",
71+
)
72+
else:
73+
self._logger.info("Migration complete: no model configs needed migration")
5174

75+
def _parse_and_migrate_config(self, config_dict: dict[str, Any]) -> AnyModelConfig:
5276
# In v6.9.0 we made some improvements to the model taxonomy and the model config schemas. There are a changes
5377
# we need to make to old configs to bring them up to date.
5478

@@ -101,7 +125,12 @@ def _parse_and_migrate_config(self, config_json: Any) -> AnyModelConfig:
101125
# It's only on SD1.x, SD2.x, and SDXL main models.
102126
config_dict["prediction_type"] = config_dict.get("prediction_type", SchedulerPredictionType.Epsilon.value)
103127

104-
if base == BaseModelType.Flux and type == ModelType.LoRA.value and format == ModelFormat.Diffusers.value:
128+
# Prior to v6.9.0, the variant field was optional and would default to Normal if not present.
129+
# We now make it explicit and always present. Use the existing value if present, otherwise default to
130+
# Normal. It's only on SD main models.
131+
config_dict["variant"] = config_dict.get("variant", ModelVariantType.Normal.value)
132+
133+
if base == BaseModelType.Flux.value and type == ModelType.LoRA.value and format == ModelFormat.Diffusers.value:
105134
# Prior to v6.9.0, we used the Diffusers format for FLUX LoRA models that used the diffusers _key_
106135
# structure. This was misleading, as everywhere else in the application, we used the Diffusers format
107136
# to indicate that the model files were in the Diffusers _file_ format (i.e. a directory containing
@@ -124,9 +153,18 @@ def _parse_and_migrate_config(self, config_json: Any) -> AnyModelConfig:
124153
# as independent of any specific base model architecture.
125154
config_dict["base"] = BaseModelType.Any.value
126155

156+
if type == ModelType.CLIPEmbed.value:
157+
# Prior to v6.9.0, some CLIP Embed models did not have a variant set. The default was the L variant.
158+
# We now make it explicit and always present. Use the existing value if present, otherwise default to
159+
# L variant. Also, treat CLIP Embed models as independent of any specific base model architecture.
160+
config_dict["base"] = BaseModelType.Any.value
161+
config_dict["variant"] = config_dict.get("variant", ClipVariantType.L.value)
162+
127163
try:
128164
migrated_config = AnyModelConfigValidator.validate_python(config_dict)
129-
except ValidationError as e:
165+
# This could be a ValidationError or any other error that occurs during validation. A failure to generate a
166+
# union discriminator could raise a ValueError, for example. Who knows what else could fail - catch all.
167+
except Exception as e:
130168
self._logger.error("Failed to validate migrated config, attempting to save as unknown model: %s", e)
131169
cloned_config_dict = deepcopy(config_dict)
132170
cloned_config_dict.pop("base", None)

0 commit comments

Comments
 (0)