mirror of
https://github.com/invoke-ai/InvokeAI
synced 2024-08-30 20:32:17 +00:00
fix(mm): handle depth and inpainting models when converting to diffusers
"Normal" models have 4 in-channels, while "Depth" models have 5 and "Inpaint" models have 9. We need to explicitly tell diffusers the channel count when converting models. Closes #6058
This commit is contained in:
parent
536bb4f053
commit
eb33303e79
@ -14,12 +14,18 @@ from invokeai.backend.model_manager import (
|
||||
SchedulerPredictionType,
|
||||
SubModelType,
|
||||
)
|
||||
from invokeai.backend.model_manager.config import CheckpointConfigBase, MainCheckpointConfig
|
||||
from invokeai.backend.model_manager.config import CheckpointConfigBase, MainCheckpointConfig, ModelVariantType
|
||||
from invokeai.backend.model_manager.convert_ckpt_to_diffusers import convert_ckpt_to_diffusers
|
||||
|
||||
from .. import ModelLoaderRegistry
|
||||
from .generic_diffusers import GenericDiffusersLoader
|
||||
|
||||
VARIANT_TO_IN_CHANNEL_MAP = {
|
||||
ModelVariantType.Normal: 4,
|
||||
ModelVariantType.Depth: 5,
|
||||
ModelVariantType.Inpaint: 9,
|
||||
}
|
||||
|
||||
|
||||
@ModelLoaderRegistry.register(base=BaseModelType.Any, type=ModelType.Main, format=ModelFormat.Diffusers)
|
||||
@ModelLoaderRegistry.register(base=BaseModelType.Any, type=ModelType.Main, format=ModelFormat.Checkpoint)
|
||||
@ -87,6 +93,7 @@ class StableDiffusionDiffusersModel(GenericDiffusersLoader):
|
||||
)
|
||||
|
||||
self._logger.info(f"Converting {model_path} to diffusers format")
|
||||
|
||||
convert_ckpt_to_diffusers(
|
||||
model_path,
|
||||
output_path,
|
||||
@ -99,5 +106,6 @@ class StableDiffusionDiffusersModel(GenericDiffusersLoader):
|
||||
image_size=image_size,
|
||||
upcast_attention=upcast_attention,
|
||||
load_safety_checker=False,
|
||||
num_in_channels=VARIANT_TO_IN_CHANNEL_MAP[config.variant],
|
||||
)
|
||||
return output_path
|
||||
|
Loading…
Reference in New Issue
Block a user