wehub-resource-sync
cddb07a176
docs / deploy (push) Has been cancelled
docs / changes (push) Has been cancelled
docs / check-and-build (push) Has been cancelled
build container image / cpu (push) Has been cancelled
build container image / cuda (push) Has been cancelled
build container image / rocm (push) Has been cancelled
frontend checks / frontend-checks (push) Has been cancelled
frontend tests / frontend-tests (push) Has been cancelled
lfs checks / lfs-check (push) Has been cancelled
python checks / python-checks (push) Has been cancelled
python tests / py3.12: macos-default (push) Has been cancelled
python tests / py3.11: windows-cpu (push) Has been cancelled
python tests / py3.12: windows-cpu (push) Has been cancelled
python tests / py3.11: linux-cpu (push) Has been cancelled
typegen checks / typegen-checks (push) Has been cancelled
uv lock checks / uv-lock-checks (push) Has been cancelled
openapi checks / openapi-checks (push) Has been cancelled
python tests / py3.11: macos-default (push) Has been cancelled
python tests / py3.12: linux-cpu (push) Has been cancelled
116 lines
4.9 KiB
Python
116 lines
4.9 KiB
Python
from pathlib import Path
|
|
|
|
import pytest
|
|
import torch
|
|
from diffusers import AutoencoderTiny
|
|
|
|
from invokeai.app.invocations.model import ModelIdentifierField
|
|
from invokeai.app.services.invocation_services import InvocationServices
|
|
from invokeai.app.services.model_manager import ModelManagerServiceBase
|
|
from invokeai.app.services.shared.invocation_context import (
|
|
InvocationContext,
|
|
InvocationContextData,
|
|
build_invocation_context,
|
|
)
|
|
from invokeai.backend.model_manager.configs.external_api import ExternalApiModelConfig, ExternalModelCapabilities
|
|
from invokeai.backend.model_manager.load.load_base import LoadedModelWithoutConfig
|
|
from tests.backend.model_manager.model_manager_fixtures import * # noqa F403
|
|
|
|
|
|
@pytest.fixture()
|
|
def mock_context(
|
|
mock_services: InvocationServices,
|
|
mm2_model_manager: ModelManagerServiceBase,
|
|
) -> InvocationContext:
|
|
mock_services.model_manager = mm2_model_manager
|
|
return build_invocation_context(
|
|
services=mock_services,
|
|
data=InvocationContextData(queue_item=None, invocation=None, source_invocation_id=None), # type: ignore
|
|
is_canceled=None, # type: ignore
|
|
)
|
|
|
|
|
|
def test_download_and_cache(mock_context: InvocationContext, mm2_root_dir: Path) -> None:
|
|
downloaded_path = mock_context.models.download_and_cache_model(
|
|
"https://www.test.foo/download/test_embedding.safetensors"
|
|
)
|
|
assert downloaded_path.is_file()
|
|
assert downloaded_path.exists()
|
|
assert downloaded_path.name == "test_embedding.safetensors"
|
|
assert downloaded_path.parent.parent == mm2_root_dir / "models/.download_cache"
|
|
|
|
downloaded_path_2 = mock_context.models.download_and_cache_model(
|
|
"https://www.test.foo/download/test_embedding.safetensors"
|
|
)
|
|
assert downloaded_path == downloaded_path_2
|
|
|
|
|
|
def test_load_from_path(mock_context: InvocationContext, embedding_file: Path) -> None:
|
|
downloaded_path = mock_context.models.download_and_cache_model(
|
|
"https://www.test.foo/download/test_embedding.safetensors"
|
|
)
|
|
loaded_model_1 = mock_context.models.load_local_model(downloaded_path)
|
|
assert isinstance(loaded_model_1, LoadedModelWithoutConfig)
|
|
|
|
loaded_model_2 = mock_context.models.load_local_model(downloaded_path)
|
|
assert isinstance(loaded_model_2, LoadedModelWithoutConfig)
|
|
assert loaded_model_1.model is loaded_model_2.model
|
|
|
|
loaded_model_3 = mock_context.models.load_local_model(embedding_file)
|
|
assert isinstance(loaded_model_3, LoadedModelWithoutConfig)
|
|
assert loaded_model_1.model is not loaded_model_3.model
|
|
assert isinstance(loaded_model_1.model, dict)
|
|
assert isinstance(loaded_model_3.model, dict)
|
|
assert torch.equal(loaded_model_1.model["emb_params"], loaded_model_3.model["emb_params"])
|
|
|
|
|
|
@pytest.mark.skip(reason="This requires a test model to load")
|
|
def test_load_from_dir(mock_context: InvocationContext, vae_directory: Path) -> None:
|
|
loaded_model = mock_context.models.load_local_model(vae_directory)
|
|
assert isinstance(loaded_model, LoadedModelWithoutConfig)
|
|
assert isinstance(loaded_model.model, AutoencoderTiny)
|
|
|
|
|
|
def test_download_and_load(mock_context: InvocationContext) -> None:
|
|
loaded_model_1 = mock_context.models.load_remote_model("https://www.test.foo/download/test_embedding.safetensors")
|
|
assert isinstance(loaded_model_1, LoadedModelWithoutConfig)
|
|
|
|
loaded_model_2 = mock_context.models.load_remote_model("https://www.test.foo/download/test_embedding.safetensors")
|
|
assert isinstance(loaded_model_2, LoadedModelWithoutConfig)
|
|
assert loaded_model_1.model is loaded_model_2.model # should be cached copy
|
|
|
|
|
|
def test_external_model_load_raises(
|
|
mock_context: InvocationContext, mm2_model_manager: ModelManagerServiceBase
|
|
) -> None:
|
|
config = ExternalApiModelConfig(
|
|
key="external_test",
|
|
name="External Test",
|
|
provider_id="openai",
|
|
provider_model_id="gpt-image-1",
|
|
capabilities=ExternalModelCapabilities(modes=["txt2img"]),
|
|
)
|
|
mm2_model_manager.store.add_model(config)
|
|
|
|
model_field = ModelIdentifierField.from_config(config)
|
|
|
|
with pytest.raises(ValueError, match="External API models"):
|
|
mock_context.models.load(model_field)
|
|
|
|
with pytest.raises(ValueError, match="External API models"):
|
|
mock_context.models.load_by_attrs(name=config.name, base=config.base, type=config.type)
|
|
|
|
|
|
def test_download_diffusers(mock_context: InvocationContext) -> None:
|
|
model_path = mock_context.models.download_and_cache_model("stabilityai/sdxl-turbo")
|
|
assert (model_path / "model_index.json").exists()
|
|
assert (model_path / "vae").is_dir()
|
|
|
|
|
|
def test_download_diffusers_subfolder(mock_context: InvocationContext) -> None:
|
|
model_path = mock_context.models.download_and_cache_model("stabilityai/sdxl-turbo::vae")
|
|
assert model_path.is_dir()
|
|
assert (model_path / "diffusion_pytorch_model.fp16.safetensors").exists() or (
|
|
model_path / "diffusion_pytorch_model.safetensors"
|
|
).exists()
|