Remove check for models dir in model deletion, update tests to always assume the model path is an absolute path

This commit is contained in:
Brandon Rising 2024-03-08 11:55:28 -05:00 committed by Brandon
parent 6e2cef1db5
commit ee38fbe89c
2 changed files with 23 additions and 19 deletions

View File

@ -344,7 +344,7 @@ class ModelInstallService(ModelInstallServiceBase):
"""Unregister the model. Delete its files only if they are within our models directory."""
model = self.record_store.get_model(key)
models_dir = self.app_config.models_path
model_path = models_dir / model.path
model_path = Path(model.path)
if model_path.is_relative_to(models_dir):
self.unconditionally_delete(key)
else:
@ -352,11 +352,11 @@ class ModelInstallService(ModelInstallServiceBase):
def unconditionally_delete(self, key: str) -> None: # noqa D102
model = self.record_store.get_model(key)
path = self.app_config.models_path / model.path
if path.is_dir():
rmtree(path)
model_path = Path(model.path)
if model_path.is_dir():
rmtree(model_path)
else:
path.unlink()
model_path.unlink()
self.unregister(key)
def download_and_cache(

View File

@ -42,7 +42,9 @@ def test_registration_meta(mm2_installer: ModelInstallServiceBase, embedding_fil
assert model_record is not None
assert model_record.name == "test_embedding"
assert model_record.type == ModelType.TextualInversion
assert Path(model_record.path) == embedding_file
assert model_record.path.endswith(embedding_file.as_posix())
assert Path(model_record.path).is_absolute()
assert Path(model_record.path).exists()
assert model_record.base == BaseModelType("sd-1")
assert model_record.description is not None
assert model_record.source is not None
@ -73,7 +75,9 @@ def test_install(
store = mm2_installer.record_store
key = mm2_installer.install_path(embedding_file)
model_record = store.get_model(key)
assert model_record.path == "sd-1/embedding/test_embedding.safetensors"
assert model_record.path.endswith("sd-1/embedding/test_embedding.safetensors")
assert Path(model_record.path).is_absolute()
assert Path(model_record.path).exists()
assert model_record.source == embedding_file.as_posix()
@ -129,9 +133,11 @@ def test_background_install(
# see if the thing actually got installed at the expected location
model_record = mm2_installer.record_store.get_model(key)
assert model_record is not None
assert model_record.path == destination
assert model_record.path.endswith(destination)
assert Path(model_record.path).is_absolute()
assert Path(model_record.path).exists()
assert model_record.key != "<NOKEY>"
assert Path(mm2_app_config.models_dir / model_record.path).exists()
assert Path(model_record.path).exists()
# see if metadata was properly passed through
assert model_record.description == description
@ -153,7 +159,7 @@ def test_not_inplace_install(
assert job is not None
assert job.config_out is not None
assert Path(job.config_out.path) != embedding_file
assert Path(mm2_app_config.models_dir / job.config_out.path).exists()
assert Path(job.config_out.path).exists()
def test_inplace_install(
@ -168,17 +174,15 @@ def test_inplace_install(
def test_delete_install(
mm2_installer: ModelInstallServiceBase, embedding_file: Path, mm2_app_config: InvokeAIAppConfig
mm2_installer: ModelInstallServiceBase, embedding_file: Path
) -> None:
store = mm2_installer.record_store
key = mm2_installer.install_path(embedding_file)
model_record = store.get_model(key)
assert Path(mm2_app_config.models_dir / model_record.path).exists()
assert Path(model_record.path).exists()
assert embedding_file.exists() # original should still be there after installation
mm2_installer.delete(key)
assert not Path(
mm2_app_config.models_dir / model_record.path
).exists() # after deletion, installed copy should not exist
assert not Path(model_record.path).exists() # after deletion, installed copy should not exist
assert embedding_file.exists() # but original should still be there
with pytest.raises(UnknownModelException):
store.get_model(key)
@ -190,10 +194,10 @@ def test_delete_register(
store = mm2_installer.record_store
key = mm2_installer.register_path(embedding_file)
model_record = store.get_model(key)
assert Path(mm2_app_config.models_dir / model_record.path).exists()
assert Path(model_record.path).exists()
assert embedding_file.exists() # original should still be there after installation
mm2_installer.delete(key)
assert Path(mm2_app_config.models_dir / model_record.path).exists()
assert Path(model_record.path).exists()
with pytest.raises(UnknownModelException):
store.get_model(key)
@ -217,7 +221,7 @@ def test_simple_download(mm2_installer: ModelInstallServiceBase, mm2_app_config:
key = job.config_out.key
model_record = store.get_model(key)
assert Path(mm2_app_config.models_dir / model_record.path).exists()
assert Path(model_record.path).exists()
assert len(bus.events) == 3
event_names = [x.event_name for x in bus.events]
@ -241,7 +245,7 @@ def test_huggingface_download(mm2_installer: ModelInstallServiceBase, mm2_app_co
key = job.config_out.key
model_record = store.get_model(key)
assert Path(mm2_app_config.models_dir / model_record.path).exists()
assert Path(model_record.path).exists()
assert model_record.type == ModelType.Main
assert model_record.format == ModelFormat.Diffusers