mirror of
https://github.com/invoke-ai/InvokeAI.git
synced 2025-04-04 22:43:40 +08:00
a44bfb4658
Before FLUX Fill was merged, we didn't do any checks for the model variant. We always returned "normal". To determine if a model is a FLUX Fill model, we need to check the state dict for a specific key. Initially, this logic was too strict and rejected quantized FLUX models. This issue was resolved, but it turns out there is another failure mode - some fine-tunes use a different key. This change further reduces the strictness, handling the alternate key and also falling back to "normal" if we don't see either key. This effectively restores the previous probing behaviour for all FLUX models. Closes #7856 Closes #7859
24 lines
840 B
Python
24 lines
840 B
Python
from typing import TYPE_CHECKING
|
|
|
|
if TYPE_CHECKING:
|
|
from invokeai.backend.model_manager.legacy_probe import CkptType
|
|
|
|
|
|
def get_flux_in_channels_from_state_dict(state_dict: "CkptType") -> int | None:
|
|
"""Gets the in channels from the state dict."""
|
|
|
|
# "Standard" FLUX models use "img_in.weight", but some community fine tunes use
|
|
# "model.diffusion_model.img_in.weight". Known models that use the latter key:
|
|
# - https://civitai.com/models/885098?modelVersionId=990775
|
|
# - https://civitai.com/models/1018060?modelVersionId=1596255
|
|
# - https://civitai.com/models/978314/ultrareal-fine-tune?modelVersionId=1413133
|
|
|
|
keys = {"img_in.weight", "model.diffusion_model.img_in.weight"}
|
|
|
|
for key in keys:
|
|
val = state_dict.get(key)
|
|
if val is not None:
|
|
return val.shape[1]
|
|
|
|
return None
|