fred-dev/comfy_ui_ali

1"""2    This file is part of ComfyUI.3    Copyright (C) 2024 Comfy4 5    This program is free software: you can redistribute it and/or modify6    it under the terms of the GNU General Public License as published by7    the Free Software Foundation, either version 3 of the License, or8    (at your option) any later version.9 10    This program is distributed in the hope that it will be useful,11    but WITHOUT ANY WARRANTY; without even the implied warranty of12    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the13    GNU General Public License for more details.14 15    You should have received a copy of the GNU General Public License16    along with this program.  If not, see <https://www.gnu.org/licenses/>.17"""18 19import torch20from . import model_base21from . import utils22from . import latent_formats23 24class ClipTarget:25    def __init__(self, tokenizer, clip):26        self.clip = clip27        self.tokenizer = tokenizer28        self.params = {}29 30class BASE:31    unet_config = {}32    unet_extra_config = {33        "num_heads": -1,34        "num_head_channels": 64,35    }36 37    required_keys = {}38 39    clip_prefix = []40    clip_vision_prefix = None41    noise_aug_config = None42    sampling_settings = {}43    latent_format = latent_formats.LatentFormat44    vae_key_prefix = ["first_stage_model."]45    text_encoder_key_prefix = ["cond_stage_model."]46    supported_inference_dtypes = [torch.float16, torch.bfloat16, torch.float32]47 48    memory_usage_factor = 2.049 50    manual_cast_dtype = None51    custom_operations = None52    scaled_fp8 = None53    optimizations = {"fp8": False}54 55    @classmethod56    def matches(s, unet_config, state_dict=None):57        for k in s.unet_config:58            if k not in unet_config or s.unet_config[k] != unet_config[k]:59                return False60        if state_dict is not None:61            for k in s.required_keys:62                if k not in state_dict:63                    return False64        return True65 66    def model_type(self, state_dict, prefix=""):67        return model_base.ModelType.EPS68 69    def inpaint_model(self):70        return self.unet_config["in_channels"] > 471 72    def __init__(self, unet_config):73        self.unet_config = unet_config.copy()74        self.sampling_settings = self.sampling_settings.copy()75        self.latent_format = self.latent_format()76        self.optimizations = self.optimizations.copy()77        for x in self.unet_extra_config:78            self.unet_config[x] = self.unet_extra_config[x]79 80    def get_model(self, state_dict, prefix="", device=None):81        if self.noise_aug_config is not None:82            out = model_base.SD21UNCLIP(self, self.noise_aug_config, model_type=self.model_type(state_dict, prefix), device=device)83        else:84            out = model_base.BaseModel(self, model_type=self.model_type(state_dict, prefix), device=device)85        if self.inpaint_model():86            out.set_inpaint()87        return out88 89    def process_clip_state_dict(self, state_dict):90        state_dict = utils.state_dict_prefix_replace(state_dict, {k: "" for k in self.text_encoder_key_prefix}, filter_keys=True)91        return state_dict92 93    def process_unet_state_dict(self, state_dict):94        return state_dict95 96    def process_vae_state_dict(self, state_dict):97        return state_dict98 99    def process_clip_state_dict_for_saving(self, state_dict):100        replace_prefix = {"": self.text_encoder_key_prefix[0]}101        return utils.state_dict_prefix_replace(state_dict, replace_prefix)102 103    def process_clip_vision_state_dict_for_saving(self, state_dict):104        replace_prefix = {}105        if self.clip_vision_prefix is not None:106            replace_prefix[""] = self.clip_vision_prefix107        return utils.state_dict_prefix_replace(state_dict, replace_prefix)108 109    def process_unet_state_dict_for_saving(self, state_dict):110        replace_prefix = {"": "model.diffusion_model."}111        return utils.state_dict_prefix_replace(state_dict, replace_prefix)112 113    def process_vae_state_dict_for_saving(self, state_dict):114        replace_prefix = {"": self.vae_key_prefix[0]}115        return utils.state_dict_prefix_replace(state_dict, replace_prefix)116 117    def set_inference_dtype(self, dtype, manual_cast_dtype):118        self.unet_config['dtype'] = dtype119        self.manual_cast_dtype = manual_cast_dtype120