InvokeAI/invokeai/app/services/generate_initializer.py

import os
import sys
import torch
import traceback
from argparse import Namespace
from omegaconf import OmegaConf

import invokeai.version
from ...backend import ModelManager, InvokeAIGeneratorBasicParams, InvokeAIGeneratorFactory
from ...backend.util import choose_precision, choose_torch_device
from ...backend import Globals

# TODO: most of this code should be split into individual services as the Generate.py code is deprecated
def get_generator_factory(args, config) -> InvokeAIGeneratorFactory:
    if not args.conf:
        config_file = os.path.join(Globals.root, "configs", "models.yaml")
        if not os.path.exists(config_file):
            report_model_error(
                args, FileNotFoundError(f"The file {config_file} could not be found.")
            )

    print(f">> {invokeai.version.__app_name__}, version {invokeai.version.__version__}")
    print(f'>> InvokeAI runtime directory is "{Globals.root}"')

    # these two lines prevent a horrible warning message from appearing
    # when the frozen CLIP tokenizer is imported
    import transformers  # type: ignore

    transformers.logging.set_verbosity_error()
    import diffusers

    diffusers.logging.set_verbosity_error()

    # Loading Face Restoration and ESRGAN Modules
    gfpgan, codeformer, esrgan = load_face_restoration(args)

    # normalize the config directory relative to root
    if not os.path.isabs(args.conf):
        args.conf = os.path.normpath(os.path.join(Globals.root, args.conf))

    if args.embeddings:
        if not os.path.isabs(args.embedding_path):
            embedding_path = os.path.normpath(
                os.path.join(Globals.root, args.embedding_path)
            )
        else:
            embedding_path = args.embedding_path
    else:
        embedding_path = None

    # migrate legacy models
    ModelManager.migrate_models()

    # load the infile as a list of lines
    if args.infile:
        try:
            if os.path.isfile(args.infile):
                infile = open(args.infile, "r", encoding="utf-8")
            elif args.infile == "-":  # stdin
                infile = sys.stdin
            else:
                raise FileNotFoundError(f"{args.infile} not found.")
        except (FileNotFoundError, IOError) as e:
            print(f"{e}. Aborting.")
            sys.exit(-1)

    # creating an InvokeAIGeneratorFactory object:
    try:
        device = torch.device(choose_torch_device())
        precision = 'float16' if args.precision=='float16' \
        else 'float32' if args.precision=='float32' \
        else choose_precision(device)
        
        model_manager = ModelManager(
            OmegaConf.load(args.conf),
            precision=precision,
            device_type=device,
            max_loaded_models=args.max_loaded_models,
        )
        # TO DO: initialize and pass safety checker!!!
        params = InvokeAIGeneratorBasicParams(
            precision=precision,
        )
        factory = InvokeAIGeneratorFactory(model_manager, params)
    except (FileNotFoundError, TypeError, AssertionError) as e:
        report_model_error(args, e)
    except (IOError, KeyError) as e:
        print(f"{e}. Aborting.")
        sys.exit(-1)

    if args.seamless:
        #TODO: do something here ?
        print(">> changed to seamless tiling mode")

    # try to autoconvert new models
    # autoimport new .ckpt files
    if path := args.autoconvert:
        model_manager.autoconvert_weights(
            conf_path=args.conf,
            weights_directory=path,
        )

    return factory

def load_face_restoration(opt):
    try:
        gfpgan, codeformer, esrgan = None, None, None
        if opt.restore or opt.esrgan:
            from invokeai.backend.restoration import Restoration

            restoration = Restoration()
            if opt.restore:
                gfpgan, codeformer = restoration.load_face_restore_models(
                    opt.gfpgan_model_path
                )
            else:
                print(">> Face restoration disabled")
            if opt.esrgan:
                esrgan = restoration.load_esrgan(opt.esrgan_bg_tile)
            else:
                print(">> Upscaling disabled")
        else:
            print(">> Face restoration and upscaling disabled")
    except (ModuleNotFoundError, ImportError):
        print(traceback.format_exc(), file=sys.stderr)
        print(">> You may need to install the ESRGAN and/or GFPGAN modules")
    return gfpgan, codeformer, esrgan


def report_model_error(opt: Namespace, e: Exception):
    print(f'** An error occurred while attempting to initialize the model: "{str(e)}"')
    print(
        "** This can be caused by a missing or corrupted models file, and can sometimes be fixed by (re)installing the models."
    )
    yes_to_all = os.environ.get("INVOKE_MODEL_RECONFIGURE")
    if yes_to_all:
        print(
            "** Reconfiguration is being forced by environment variable INVOKE_MODEL_RECONFIGURE"
        )
    else:
        response = input(
            "Do you want to run invokeai-configure script to select and/or reinstall models? [y] "
        )
        if response.startswith(("n", "N")):
            return

    print("invokeai-configure is launching....\n")

    # Match arguments that were set on the CLI
    # only the arguments accepted by the configuration script are parsed
    root_dir = ["--root", opt.root_dir] if opt.root_dir is not None else []
    config = ["--config", opt.conf] if opt.conf is not None else []
    previous_args = sys.argv
    sys.argv = ["invokeai-configure"]
    sys.argv.extend(root_dir)
    sys.argv.extend(config)
    if yes_to_all is not None:
        for arg in yes_to_all.split():
            sys.argv.append(arg)

    from invokeai.frontend.install import invokeai_configure

    invokeai_configure()
    # TODO: Figure out how to restart
    # print('** InvokeAI will now restart')
    # sys.argv = previous_args
    # main() # would rather do a os.exec(), but doesn't exist?
    # sys.exit(0)