model installer confirms deletion of models

This commit is contained in:
Lincoln Stein 2023-07-05 09:57:23 -04:00
parent 9edf78dd2e
commit bd82c4ace0
3 changed files with 20 additions and 7 deletions

View File

@ -284,16 +284,16 @@ class ModelInstall(object):
location = self._download_hf_model(repo_id, files, staging) location = self._download_hf_model(repo_id, files, staging)
break break
elif f'learned_embeds.{suffix}' in files: elif f'learned_embeds.{suffix}' in files:
location = self._download_hf_model(repo_id, ['learned_embeds.suffix'], staging) location = self._download_hf_model(repo_id, [f'learned_embeds.{suffix}'], staging)
break break
if not location: if not location:
logger.warning(f'Could not determine type of repo {repo_id}. Skipping install.') logger.warning(f'Could not determine type of repo {repo_id}. Skipping install.')
return return {}
info = ModelProbe().heuristic_probe(location, self.prediction_helper) info = ModelProbe().heuristic_probe(location, self.prediction_helper)
if not info: if not info:
logger.warning(f'Could not probe {location}. Skipping install.') logger.warning(f'Could not probe {location}. Skipping install.')
return return {}
dest = self.config.models_path / info.base_type.value / info.model_type.value / self._get_model_name(repo_id,location) dest = self.config.models_path / info.base_type.value / info.model_type.value / self._get_model_name(repo_id,location)
if dest.exists(): if dest.exists():
shutil.rmtree(dest) shutil.rmtree(dest)

View File

@ -64,8 +64,8 @@ class ModelProbe(object):
@classmethod @classmethod
def probe(cls, def probe(cls,
model_path: Path, model_path: Path,
model: Optional[Union[Dict, ModelMixin]], model: Optional[Union[Dict, ModelMixin]] = None,
prediction_type_helper: Callable[[Path],SchedulerPredictionType] = None)->ModelProbeInfo: prediction_type_helper: Optional[Callable[[Path],SchedulerPredictionType]] = None)->ModelProbeInfo:
''' '''
Probe the model at model_path and return sufficient information about it Probe the model at model_path and return sufficient information about it
to place it somewhere in the models directory hierarchy. If the model is to place it somewhere in the models directory hierarchy. If the model is

View File

@ -382,10 +382,21 @@ class addModelsForm(CyclingForm, npyscreen.FormMultiPage):
) )
return min(cols, len(self.installed_models)) return min(cols, len(self.installed_models))
def confirm_deletions(self, selections: InstallSelections)->bool:
remove_models = selections.remove_models
if len(remove_models) > 0:
mods = "\n".join([ModelManager.parse_key(x)[0] for x in remove_models])
return npyscreen.notify_ok_cancel(f"These unchecked models will be deleted from disk. Continue?\n---------\n{mods}")
else:
return True
def on_execute(self): def on_execute(self):
self.monitor.entry_widget.buffer(['Processing...'],scroll_end=True)
self.marshall_arguments() self.marshall_arguments()
app = self.parentApp app = self.parentApp
if not self.confirm_deletions(app.install_selections):
return
self.monitor.entry_widget.buffer(['Processing...'],scroll_end=True)
self.ok_button.hidden = True self.ok_button.hidden = True
self.display() self.display()
@ -417,6 +428,8 @@ class addModelsForm(CyclingForm, npyscreen.FormMultiPage):
def on_done(self): def on_done(self):
self.marshall_arguments() self.marshall_arguments()
if not self.confirm_deletions(self.parentApp.install_selections):
return
self.parentApp.setNextForm(None) self.parentApp.setNextForm(None)
self.parentApp.user_cancelled = False self.parentApp.user_cancelled = False
self.editing = False self.editing = False