mirror of
https://github.com/invoke-ai/InvokeAI.git
synced 2026-02-02 21:45:19 -05:00
3d6d89feb4
* pass model config to _load_model * make conversion work again * do not write diffusers to disk when convert_cache set to 0 * adding same model to cache twice is a no-op, not an assertion error * fix issues identified by psychedelicious during pr review * following conversion, avoid redundant read of cached submodels * fix error introduced while merging --------- Co-authored-by: Lincoln Stein <lstein@gmail.com>
59 lines
2.2 KiB
Python
59 lines
2.2 KiB
Python
# Copyright (c) 2024, Lincoln D. Stein and the InvokeAI Development Team
|
|
"""Class for ControlNet model loading in InvokeAI."""
|
|
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
|
|
from invokeai.backend.model_manager import (
|
|
AnyModel,
|
|
AnyModelConfig,
|
|
BaseModelType,
|
|
ModelFormat,
|
|
ModelType,
|
|
)
|
|
from invokeai.backend.model_manager.config import CheckpointConfigBase
|
|
from invokeai.backend.model_manager.convert_ckpt_to_diffusers import convert_controlnet_to_diffusers
|
|
|
|
from .. import ModelLoaderRegistry
|
|
from .generic_diffusers import GenericDiffusersLoader
|
|
|
|
|
|
@ModelLoaderRegistry.register(base=BaseModelType.Any, type=ModelType.ControlNet, format=ModelFormat.Diffusers)
|
|
@ModelLoaderRegistry.register(base=BaseModelType.Any, type=ModelType.ControlNet, format=ModelFormat.Checkpoint)
|
|
class ControlNetLoader(GenericDiffusersLoader):
|
|
"""Class to load ControlNet models."""
|
|
|
|
def _needs_conversion(self, config: AnyModelConfig, model_path: Path, dest_path: Path) -> bool:
|
|
if not isinstance(config, CheckpointConfigBase):
|
|
return False
|
|
elif (
|
|
dest_path.exists()
|
|
and (dest_path / "config.json").stat().st_mtime >= (config.converted_at or 0.0)
|
|
and (dest_path / "config.json").stat().st_mtime >= model_path.stat().st_mtime
|
|
):
|
|
return False
|
|
else:
|
|
return True
|
|
|
|
def _convert_model(self, config: AnyModelConfig, model_path: Path, output_path: Optional[Path] = None) -> AnyModel:
|
|
assert isinstance(config, CheckpointConfigBase)
|
|
image_size = (
|
|
512
|
|
if config.base == BaseModelType.StableDiffusion1
|
|
else 768
|
|
if config.base == BaseModelType.StableDiffusion2
|
|
else 1024
|
|
)
|
|
|
|
self._logger.info(f"Converting {model_path} to diffusers format")
|
|
with open(self._app_config.legacy_conf_path / config.config_path, "r") as config_stream:
|
|
result = convert_controlnet_to_diffusers(
|
|
model_path,
|
|
output_path,
|
|
original_config_file=config_stream,
|
|
image_size=image_size,
|
|
precision=self._torch_dtype,
|
|
from_safetensors=model_path.suffix == ".safetensors",
|
|
)
|
|
return result
|