mirror of
https://github.com/invoke-ai/InvokeAI.git
synced 2026-02-14 16:25:18 -05:00
814406d98a
In the previous commit, the LLaVA model was updated to support partial loading. In this commit, the SigLIP model is updated in the same way. This model is used for FLUX Redux. It's <4GB and only ever run in isolation, so it won't benefit from partial loading for the vast majority of users. Regardless, I think it is best if we make _all_ models work with partial loading. PS: I also fixed the initial load dtype issue, described in the prev commit. It's probably a non-issue for this model, but we may as well fix it.
29 lines
1.1 KiB
Python
29 lines
1.1 KiB
Python
from pathlib import Path
|
|
from typing import Optional
|
|
|
|
from transformers import SiglipVisionModel
|
|
|
|
from invokeai.backend.model_manager.config import (
|
|
AnyModelConfig,
|
|
)
|
|
from invokeai.backend.model_manager.load.load_default import ModelLoader
|
|
from invokeai.backend.model_manager.load.model_loader_registry import ModelLoaderRegistry
|
|
from invokeai.backend.model_manager.taxonomy import AnyModel, BaseModelType, ModelFormat, ModelType, SubModelType
|
|
|
|
|
|
@ModelLoaderRegistry.register(base=BaseModelType.Any, type=ModelType.SigLIP, format=ModelFormat.Diffusers)
|
|
class SigLIPModelLoader(ModelLoader):
|
|
"""Class for loading SigLIP models."""
|
|
|
|
def _load_model(
|
|
self,
|
|
config: AnyModelConfig,
|
|
submodel_type: Optional[SubModelType] = None,
|
|
) -> AnyModel:
|
|
if submodel_type is not None:
|
|
raise ValueError("Unexpected submodel requested for LLaVA OneVision model.")
|
|
|
|
model_path = Path(config.path)
|
|
model = SiglipVisionModel.from_pretrained(model_path, local_files_only=True, torch_dtype=self._torch_dtype)
|
|
return model
|