baulab/Erasing-Concepts-In-Diffusion

1import copy2import re3import torch4import util5 6class FineTunedModel(torch.nn.Module):7 8    def __init__(self,9                 model,10                 modules,11                 frozen_modules=[]12                 ):13 14        super().__init__()15 16        if isinstance(modules, str):17            modules = [modules]18 19        self.model = model20        self.ft_modules = {}21        self.orig_modules = {}22 23        util.freeze(self.model)24 25        for module_name, module in model.named_modules():26            for ft_module_regex in modules:27 28                match = re.search(ft_module_regex, module_name)29                30                if match is not None:31 32                    ft_module = copy.deepcopy(module)33                    34                    self.orig_modules[module_name] = module35                    self.ft_modules[module_name] = ft_module36 37                    util.unfreeze(ft_module)38 39                    print(f"=> Finetuning {module_name}")40       41                    for ft_module_name, module in ft_module.named_modules():42 43                        ft_module_name = f"{module_name}.{ft_module_name}"44 45                        for freeze_module_name in frozen_modules:46 47                            match = re.search(freeze_module_name, ft_module_name)48 49                            if match:50                                print(f"=> Freezing {ft_module_name}")51                                util.freeze(module)52 53        self.ft_modules_list = torch.nn.ModuleList(self.ft_modules.values())54        self.orig_modules_list = torch.nn.ModuleList(self.orig_modules.values())55 56 57    @classmethod58    def from_checkpoint(cls, model, checkpoint, frozen_modules=[]):59 60        if isinstance(checkpoint, str):61            checkpoint = torch.load(checkpoint)62 63        modules = [f"{key}$" for key in list(checkpoint.keys())]64 65        ftm = FineTunedModel(model, modules, frozen_modules=frozen_modules)66        ftm.load_state_dict(checkpoint)67 68        return ftm69 70        71    def __enter__(self):72 73        for key, ft_module in self.ft_modules.items():74            util.set_module(self.model, key, ft_module)75 76    def __exit__(self, exc_type, exc_value, tb):77 78        for key, module in self.orig_modules.items():79            util.set_module(self.model, key, module)80 81    def parameters(self):82 83        parameters = []84 85        for ft_module in self.ft_modules.values():86 87            parameters.extend(list(ft_module.parameters()))88 89        return parameters90 91    def state_dict(self):92 93        state_dict = {key: module.state_dict() for key, module in self.ft_modules.items()}94 95        return state_dict96 97    def load_state_dict(self, state_dict):98 99        for key, sd in state_dict.items():100            101            self.ft_modules[key].load_state_dict(sd)