modules/loader/module_loader.py

from diffusers import DDPMScheduler, DiffusionPipeline
from typing import List, Any, Union, Type
from utils.loader import get_class
from copy import deepcopy
from modules.loader.module_loader_config import ModuleLoaderConfig
import torch
import pytorch_lightning as pl
import jsonargparse


class bcolors:
    HEADER = '\033[95m'
    OKBLUE = '\033[94m'
    OKCYAN = '\033[96m'
    OKGREEN = '\033[92m'
    WARNING = '\033[93m'
    FAIL = '\033[91m'
    ENDC = '\033[0m'
    BOLD = '\033[1m'
    UNDERLINE = '\033[4m'


class GenericModuleLoader():

    def __init__(self,
                 pipeline_repo: str = None,
                 pipeline_obj: str = None,
                 set_prediction_type: str = "",
                 module_names: List[str] = [
                     "scheduler", "text_encoder", "tokenizer", "vae", "unet",],
                 module_config: dict[str,
                                     Union[ModuleLoaderConfig,  torch.nn.Module, Any]] = None,
                 fast_dev_run: Union[int, bool] = False,
                 root_cls: Type[Any] = None,
                 ) -> None:
        self.module_config = module_config
        self.pipeline_repo = pipeline_repo
        self.pipeline_obj = pipeline_obj
        self.set_prediction_type = set_prediction_type
        self.module_names = module_names
        self.fast_dev_run = fast_dev_run
        self.root_cls = root_cls

    def load_custom_scheduler(self):
        module_obj = DDPMScheduler.from_pretrained(
            self.pipeline_repo, subfolder="scheduler")

        if len(self.set_prediction_type) > 0:
            scheduler_config = module_obj.load_config(
                self.pipeline_repo, subfolder="scheduler")
            scheduler_config["prediction_type"] = self.set_prediction_type
            module_obj = module_obj.from_config(scheduler_config)
        return module_obj

    def load_pipeline(self):
        return DiffusionPipeline.from_pretrained(self.pipeline_repo) if self.pipeline_repo is not None else None

    def __call__(self, trainer: pl.LightningModule, diff_trainer_params):
        # load diffusers pipeline object if set
        if self.pipeline_obj is not None:
            pipe = self.load_pipeline()
        else:
            pipe = None

        if pipe is not None and self.pipeline_obj is not None:
            # store the entire diffusers pipeline object under the name given by pipeline_obj
            setattr(trainer, self.pipeline_obj, self.load_pipeline())

        for module_name in self.module_names:
            print(f" --- START:  Loading module: {module_name} ---")
            if module_name not in self.module_config.keys() and pipe is not None:
                # stores models from already loaded diffusers pipeline
                module_obj = getattr(pipe, module_name)
                if module_name == "scheduler":
                    module_obj = self.load_custom_scheduler()
                setattr(trainer, module_name, module_obj)
            else:
                if not isinstance(self.module_config[module_name], ModuleLoaderConfig):
                    # instantiate model by jsonargparse and store it
                    module = self.module_config[module_name]
                    # TODO we want to be able to load ckpt still.
                    config_obj = None
                else:
                    # instantiate object from class method (as used by Diffusers, e.g. DiffusionPipeline.load_from_pretrained)
                    config_obj = self.module_config[module_name]
                    # retrieve loader class
                    loader_cls = get_class(
                        config_obj.loader_cls_path)

                    # retrieve loader method
                    if config_obj.cls_func != "":
                        # we allow to specify a method for fast loading (e.g. in diffusers, from_config instead of from_pretrained)
                        # makes loading faster for quick testing
                        if not self.fast_dev_run or config_obj.cls_func_fast_dev_run == "":
                            cls_func = getattr(
                                loader_cls, config_obj.cls_func)
                        else:
                            print(
                                f"Model {module_name}: loading fast_dev_run class loader")
                            cls_func = getattr(
                                loader_cls, config_obj.cls_func_fast_dev_run)
                    else:
                        cls_func = loader_cls

                    # retrieve parameters
                    # load parameters specified in diff_trainer_params (so it links them)
                    kwargs_trainer_params = config_obj.kwargs_diff_trainer_params

                    kwargs_diffusers = config_obj.kwargs_diffusers

                    # names of dependent modules that we need as input
                    dependent_modules = config_obj.dependent_modules

                    # names of dependent modules that we need as input. Modules will be cloned
                    dependent_modules_cloned = config_obj.dependent_modules_cloned

                    # model kwargs. Can be just a dict, or a parameter class (derived from modules.params.params_mixin.AsDictMixin) so we have verification of inputs
                    model_params = config_obj.model_params

                    # kwargs used only if on fast_dev_run mode
                    model_params_fast_dev_run = config_obj.model_params_fast_dev_run

                    if model_params is not None:
                        if isinstance(model_params, dict):
                            model_dict = model_params
                        else:
                            model_dict = model_params.to_dict()
                    else:
                        model_dict = {}

                    if (model_params_fast_dev_run is None) or (not self.fast_dev_run):
                        model_params_fast_dev_run = {}
                    else:
                        print(
                            f"Module {module_name}: loading fast_dev_run params")

                    loaded_modules_dict = {}
                    if dependent_modules is not None:
                        for key, dependent_module in dependent_modules.items():
                            assert hasattr(
                                trainer, dependent_module), f"Module {dependent_module} not available. Set {dependent_module} before module {module_name} in module_loader.module_names. Current order: {self.module_names}"
                            loaded_modules_dict[key] = getattr(
                                trainer, dependent_module)

                    if dependent_modules_cloned is not None:
                        for key, dependent_module in dependent_modules_cloned.items():
                            assert hasattr(
                                trainer, dependent_module), f"Module {dependent_module} not available. Set {dependent_module} before module {module_name} in module_loader.module_names. Current order: {self.module_names}"
                            loaded_modules_dict[key] = getattr(
                                trainer, deepcopy(dependent_module))
                    if kwargs_trainer_params is not None:
                        for key, param in kwargs_trainer_params.items():
                            if param is not None:
                                kwargs_trainer_params[key] = getattr(
                                    diff_trainer_params, param)
                            else:
                                kwargs_trainer_params[key] = diff_trainer_params
                    else:
                        kwargs_trainer_params = {}

                    if kwargs_diffusers is None:
                        kwargs_diffusers = {}
                    else:
                        for key, value in kwargs_diffusers.items():
                            if key == "torch_dtype":
                                if value == "torch.float16":
                                    kwargs_diffusers[key] = torch.float16

                    kwargs = kwargs_diffusers | loaded_modules_dict | kwargs_trainer_params | model_dict | model_params_fast_dev_run
                    args = config_obj.args
                    # instantiate object
                    module = cls_func(*args, **kwargs)
                    module: torch.nn.Module
                    if self.root_cls is not None:
                        assert isinstance(module, self.root_cls)

                if config_obj is not None and config_obj.state_dict_path != "" and not self.fast_dev_run:
                    # TODO extend loading to hf spaces
                    print(
                        f"             * Loading checkpoint {config_obj.state_dict_path} - STARTED")
                    module_state_dict = torch.load(
                        config_obj.state_dict_path, map_location=torch.device("cpu"))
                    module_state_dict = module_state_dict["state_dict"]

                    if len(config_obj.state_dict_filters) > 0:
                        assert not config_obj.strict_loading
                        ckpt_params_dict = {}
                        for name, param in module.named_parameters(prefix=module_name):
                            for filter_str in config_obj.state_dict_filters:
                                filter_groups = filter_str.split("*")
                                has_all_parts = True
                                for filter_group in filter_groups:
                                    has_all_parts = has_all_parts and filter_group in name

                                if has_all_parts:
                                    validate_name = name
                                    for filter_group in filter_groups:
                                        if filter_group in validate_name:
                                            shift = validate_name.index(
                                                filter_group)
                                            validate_name = validate_name[shift+len(
                                                filter_group):]
                                        else:
                                            has_all_parts = False
                                            break
                                if has_all_parts:
                                    ckpt_params_dict[name[len(
                                        module_name+"."):]] = param
                    else:
                        ckpt_params_dict = dict(filter(lambda x: x[0].startswith(
                            module_name), module_state_dict.items()))
                        ckpt_params_dict = {
                            k.split(module_name+".")[1]: v for (k, v) in ckpt_params_dict.items()}

                    if len(ckpt_params_dict) > 0:
                        miss, unex = module.load_state_dict(
                            ckpt_params_dict, strict=config_obj.strict_loading)
                        ckpt_params_dict = {}
                        assert len(
                            unex) == 0, f"Unexpected parameters in checkpoint: {unex}"
                        if len(miss) > 0:
                            print(
                                f"Checkpoint {config_obj.state_dict_path} is missing parameters for module {module_name}.")
                            print(miss)
                    print(
                        f"             * Loading checkpoint {config_obj.state_dict_path} - FINISHED")
                if isinstance(module, jsonargparse.Namespace) or isinstance(module, dict):
                    print(bcolors.WARNING +
                          f"Warning: Seems object {module_name} was not build correct." + bcolors.ENDC)

                setattr(trainer, module_name, module)
            print(f" --- FINSHED:  Loading module: {module_name} ---")