# Documentation - Class name: SUPIR_model_loader_v2 - Category: SUPIR - Output node: False - Repo Ref: https://github.com/comfyanonymous/ComfyUI SUPIR_model_loader_v2节点用于初始化和准备SUPIR模型以供使用，主要通过从指定路径加载模型状态来实现，并可选择性地调整其精度以提高内存效率。该节点在为后续图像处理任务设置SUPIR模型中扮演着至关重要的角色，确保模型被正确加载并准备就绪。 # Input types ## Required - model - 指定SUPIR模型加载器的主要模型配置，集成了CLIP和VAE等各种组件，以增强图像处理能力。 - Comfy dtype: MODEL - Python dtype: str - clip - 指定与SUPIR模型配合使用的CLIP模型，用于增强其功能，特别是在涉及图像和文本处理的任务中。 - Comfy dtype: CLIP - Python dtype: str - vae - 指定与SUPIR模型一起使用的VAE模型，为整体图像处理和生成能力做出贡献。 - Comfy dtype: VAE - Python dtype: str - supir_model - 指定SUPIR模型检查点的文件路径，这对于加载模型状态以用于图像处理任务至关重要。 - Comfy dtype: COMBO[STRING] - Python dtype: List[str] - fp8_unet - 一个布尔标志，设置后将模型权重转换为较低精度格式（torch.float8_e4m3fn）以节省VRAM，但会略微影响质量。 - Comfy dtype: BOOLEAN - Python dtype: bool - diffusion_dtype - 确定扩散操作的数据类型，提供fp16、bf16、fp32或auto等选项，auto选项会根据执行环境自动选择。 - Comfy dtype: COMBO[STRING] - Python dtype: str ## Optional - high_vram - 一个布尔标志，启用时使用高VRAM模式加载模型权重，可能会提高模型加载速度。 - Comfy dtype: BOOLEAN - Python dtype: bool # Output types - SUPIR_model - 加载完毕的SUPIR模型，可以立即用于图像处理任务。 - Comfy dtype: SUPIRMODEL - Python dtype: torch.nn.Module - SUPIR_VAE - 加载完毕的SUPIR VAE组件，是模型图像处理能力的重要组成部分。 - Comfy dtype: SUPIRVAE - Python dtype: torch.nn.Module ## Usage tips - Infra type: `GPU` - Common nodes: unknown ## Source code ```python class SUPIR_model_loader_v2: @classmethod def INPUT_TYPES(s): return {"required": { "model" :("MODEL",), "clip": ("CLIP",), "vae": ("VAE",), "supir_model": (folder_paths.get_filename_list("checkpoints"),), "fp8_unet": ("BOOLEAN", {"default": False}), "diffusion_dtype": ( [ 'fp16', 'bf16', 'fp32', 'auto' ], { "default": 'auto' }), }, "optional": { "high_vram": ("BOOLEAN", {"default": False}), } } RETURN_TYPES = ("SUPIRMODEL", "SUPIRVAE") RETURN_NAMES = ("SUPIR_model","SUPIR_VAE",) FUNCTION = "process" CATEGORY = "SUPIR" DESCRIPTION = """ Loads the SUPIR model and merges it with the SDXL model. Diffusion type should be kept on auto, unless you have issues loading the model. fp8_unet casts the unet weights to torch.float8_e4m3fn, which saves a lot of VRAM but has slight quality impact. high_vram: uses Accelerate to load weights to GPU, slightly faster model loading. """ def process(self, supir_model, diffusion_dtype, fp8_unet, model, clip, vae, high_vram=False): if high_vram: device = mm.get_torch_device() else: device = mm.unet_offload_device() print("Loading weights to: ", device) mm.unload_all_models() SUPIR_MODEL_PATH = folder_paths.get_full_path("checkpoints", supir_model) config_path = os.path.join(script_directory, "options/SUPIR_v0.yaml") clip_config_path = os.path.join(script_directory, "configs/clip_vit_config.json") tokenizer_path = os.path.join(script_directory, "configs/tokenizer") custom_config = { 'diffusion_dtype': diffusion_dtype, 'supir_model': supir_model, 'fp8_unet': fp8_unet, 'model': model, "clip": clip, "vae": vae } if diffusion_dtype == 'auto': try: if mm.should_use_fp16(): print("Diffusion using fp16") dtype = torch.float16 elif mm.should_use_bf16(): print("Diffusion using bf16") dtype = torch.bfloat16 else: print("Diffusion using fp32") dtype = torch.float32 except: raise AttributeError("ComfyUI version too old, can't autodecet properly. Set your dtypes manually.") else: print(f"Diffusion using {diffusion_dtype}") dtype = convert_dtype(diffusion_dtype) if not hasattr(self, "model") or self.model is None or self.current_config != custom_config: self.current_config = custom_config self.model = None mm.soft_empty_cache() config = OmegaConf.load(config_path) if mm.XFORMERS_IS_AVAILABLE: print("Using XFORMERS") config.model.params.control_stage_config.params.spatial_transformer_attn_type = "softmax-xformers" config.model.params.network_config.params.spatial_transformer_attn_type = "softmax-xformers" config.model.params.first_stage_config.params.ddconfig.attn_type = "vanilla-xformers" config.model.target = ".SUPIR.models.SUPIR_model_v2.SUPIRModel" pbar = comfy.utils.ProgressBar(5) #with (init_empty_weights() if is_accelerate_available else nullcontext()): self.model = instantiate_from_config(config.model).cpu() self.model.model.dtype = dtype pbar.update(1) try: print(f"Attempting to load SDXL model from node inputs") mm.load_model_gpu(model) sdxl_state_dict = model.model.state_dict_for_saving(None, vae.get_sd(), None) if is_accelerate_available: for key in sdxl_state_dict: set_module_tensor_to_device(self.model, key, device=device, dtype=dtype, value=sdxl_state_dict[key]) else: self.model.load_state_dict(sdxl_state_dict, strict=False) if fp8_unet: self.model.model.to(torch.float8_e4m3fn) else: self.model.model.to(dtype) del sdxl_state_dict pbar.update(1) except: raise Exception("Failed to load SDXL model") gc.collect() mm.soft_empty_cache() #first clip model from SDXL checkpoint try: print("Loading first clip model from SDXL checkpoint") clip_sd = None clip_model = clip.load_model() mm.load_model_gpu(clip_model) clip_sd = clip.get_sd() clip_sd = model.model.model_config.process_clip_state_dict_for_saving(clip_sd) replace_prefix = {} replace_prefix["conditioner.embedders.0.transformer."] = "" clip_l_sd = comfy.utils.state_dict_prefix_replace(clip_sd, replace_prefix, filter_keys=True) clip_text_config = CLIPTextConfig.from_pretrained(clip_config_path) self.model.conditioner.embedders[0].tokenizer = CLIPTokenizer.from_pretrained(tokenizer_path) with (init_empty_weights() if is_accelerate_available else nullcontext()): self.model.conditioner.embedders[0].transformer = CLIPTextModel(clip_text_config) if is_accelerate_available: for key in clip_l_sd: set_module_tensor_to_device(self.model.conditioner.embedders[0].transformer, key, device=device, dtype=dtype, value=clip_l_sd[key]) else: self.model.conditioner.embedders[0].transformer.load_state_dict(clip_l_sd, strict=False) self.model.conditioner.embedders[0].eval() for param in self.model.conditioner.embedders[0].parameters(): param.requires_grad = False self.model.conditioner.embedders[0].to(dtype) del clip_l_sd pbar.update(1) except: raise Exception("Failed to load first clip model from SDXL checkpoint") gc.collect() mm.soft_empty_cache() #second clip model from SDXL checkpoint try: print("Loading second clip model from SDXL checkpoint") replace_prefix2 = {} replace_prefix2["conditioner.embedders.1.model."] = "" clip_g_sd = comfy.utils.state_dict_prefix_replace(clip_sd, replace_prefix2, filter_keys=True) clip_g = build_text_model_from_openai_state_dict(clip_g_sd, device, cast_dtype=dtype) self.model.conditioner.embedders[1].model = clip_g self.model.conditioner.embedders[1].model.to(dtype) del clip_g_sd pbar.update(1) except: raise Exception("Failed to load second clip model from SDXL checkpoint") try: print(f'Attempting to load SUPIR model: [{SUPIR_MODEL_PATH}]') supir_state_dict = load_state_dict(SUPIR_MODEL_PATH) if "Q" not in supir_model or not is_accelerate_available: #I don't know why this doesn't work with the Q model. for key in supir_state_dict: set_module_tensor_to_device(self.model, key, device=device, dtype=dtype, value=supir_state_dict[key]) else: self.model.load_state_dict(supir_state_dict, strict=False) if fp8_unet: self.model.model.to(torch.float8_e4m3fn) else: self.model.model.to(dtype) del supir_state_dict pbar.update(1) except: raise Exception("Failed to load SUPIR model") mm.soft_empty_cache() return (self.model, self.model.first_stage_model,) ```