From 6532d9ffa1f9bc3b47c47af6727b8f02e3b2d6bd Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Thu, 12 Oct 2023 22:04:54 -0400 Subject: [PATCH] closes #4768 --- invokeai/frontend/install/model_install.py | 43 +++++++++++----------- invokeai/frontend/install/widgets.py | 12 +++--- 2 files changed, 27 insertions(+), 28 deletions(-) diff --git a/invokeai/frontend/install/model_install.py b/invokeai/frontend/install/model_install.py index 1fb3b61891..0ea2570b2b 100644 --- a/invokeai/frontend/install/model_install.py +++ b/invokeai/frontend/install/model_install.py @@ -20,6 +20,7 @@ from multiprocessing import Process from multiprocessing.connection import Connection, Pipe from pathlib import Path from shutil import get_terminal_size +from typing import Optional import npyscreen import torch @@ -630,21 +631,23 @@ def ask_user_for_prediction_type(model_path: Path, tui_conn: Connection = None) return _ask_user_for_pt_cmdline(model_path) -def _ask_user_for_pt_cmdline(model_path: Path) -> SchedulerPredictionType: +def _ask_user_for_pt_cmdline(model_path: Path) -> Optional[SchedulerPredictionType]: choices = [SchedulerPredictionType.Epsilon, SchedulerPredictionType.VPrediction, None] print( f""" -Please select the type of the V2 checkpoint named {model_path.name}: -[1] A model based on Stable Diffusion v2 trained on 512 pixel images (SD-2-base) -[2] A model based on Stable Diffusion v2 trained on 768 pixel images (SD-2-768) -[3] Skip this model and come back later. +Please select the scheduler prediction type of the checkpoint named {model_path.name}: +[1] "epsilon" - most v1.5 models and v2 models trained on 512 pixel images +[2] "vprediction" - v2 models trained on 768 pixel images and a few v1.5 models +[3] Accept the best guess; you can fix it in the Web UI later """ ) choice = None ok = False while not ok: try: - choice = input("select> ").strip() + choice = input("select [3]> ").strip() + if not choice: + return None choice = choices[int(choice) - 1] ok = True except (ValueError, IndexError): @@ -655,22 +658,18 @@ Please select the type of the V2 checkpoint named {model_path.name}: def _ask_user_for_pt_tui(model_path: Path, tui_conn: Connection) -> SchedulerPredictionType: - try: - tui_conn.send_bytes(f"*need v2 config for:{model_path}".encode("utf-8")) - # note that we don't do any status checking here - response = tui_conn.recv_bytes().decode("utf-8") - if response is None: - return None - elif response == "epsilon": - return SchedulerPredictionType.epsilon - elif response == "v": - return SchedulerPredictionType.VPrediction - elif response == "abort": - logger.info("Conversion aborted") - return None - else: - return response - except Exception: + tui_conn.send_bytes(f"*need v2 config for:{model_path}".encode("utf-8")) + # note that we don't do any status checking here + response = tui_conn.recv_bytes().decode("utf-8") + if response is None: + return None + elif response == "epsilon": + return SchedulerPredictionType.epsilon + elif response == "v": + return SchedulerPredictionType.VPrediction + elif response == "guess": + return None + else: return None diff --git a/invokeai/frontend/install/widgets.py b/invokeai/frontend/install/widgets.py index 06d5473fa3..4a37aba9b8 100644 --- a/invokeai/frontend/install/widgets.py +++ b/invokeai/frontend/install/widgets.py @@ -381,12 +381,12 @@ def select_stable_diffusion_config_file( wrap: bool = True, model_name: str = "Unknown", ): - message = f"Please select the correct base model for the V2 checkpoint named '{model_name}'. Press to skip installation." + message = f"Please select the correct prediction type for the checkpoint named '{model_name}'. Press to skip installation." title = "CONFIG FILE SELECTION" options = [ - "An SD v2.x base model (512 pixels; no 'parameterization:' line in its yaml file)", - "An SD v2.x v-predictive model (768 pixels; 'parameterization: \"v\"' line in its yaml file)", - "Skip installation for now and come back later", + "'epsilon' - most v1.5 models and v2 models trained on 512 pixel images", + "'vprediction' - v2 models trained on 768 pixel images and a few v1.5 models)", + "Accept the best guess; you can fix it in the Web UI later", ] F = ConfirmCancelPopup( @@ -410,7 +410,7 @@ def select_stable_diffusion_config_file( choice = F.add( npyscreen.SelectOne, values=options, - value=[0], + value=[2], max_height=len(options) + 1, scroll_exit=True, ) @@ -420,5 +420,5 @@ def select_stable_diffusion_config_file( if not F.value: return None assert choice.value[0] in range(0, 3), "invalid choice" - choices = ["epsilon", "v", "abort"] + choices = ["epsilon", "v", "guess"] return choices[choice.value[0]]