nodes.py

# mengqin@gmail.com || Apache-2.0 (apache.org/licenses/LICENSE-2.0)

import logging
import os
import re
import uuid
import copy
import weakref
import torch

import comfy.sd
import comfy.model_patcher
import folder_paths

from .ops import LazyOps
from .loader import safetensors_sd_loader

class UnetBnbModelPatcher(comfy.model_patcher.ModelPatcher):    
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        # module_key list of (strength_patch, patch_obj, strength_model, None, None)
        self.bnb_lora_patches = {}
        # optional backups for module-local temp data when partially_unload
        self._bnb_lora_module_backups = {}

    def add_patches(self, patches, strength_patch=1.0, strength_model=1.0):        
        added = []
        for key in patches:
            if not isinstance(key, str):
                continue

            module_key = key.rsplit('.', 1)[0]
            try:
                module = comfy.utils.get_attr(self.model, module_key)
                is_bnb = hasattr(module, 'is_bnb_quantized') and module.is_bnb_quantized()
            except Exception:                
                is_bnb = False

            if is_bnb:
                self.bnb_lora_patches.setdefault(module_key, []).append(
                    (strength_patch, patches[key], strength_model, None, None)
                )
            else:
                current = self.patches.get(key, [])
                current.append((strength_patch, patches[key], strength_model, None, None))
                self.patches[key] = current
            added.append(key)

        self.patches_uuid = uuid.uuid4()
        return added

    def clone(self):
        cloned = super().clone()
        if not isinstance(cloned, UnetBnbModelPatcher):
            new_cloned = UnetBnbModelPatcher(cloned.model, cloned.load_device, cloned.offload_device, cloned.size)
            new_cloned.patches = cloned.patches
            new_cloned.object_patches = cloned.object_patches
            cloned = new_cloned
        
        cloned.bnb_lora_patches = copy.deepcopy(self.bnb_lora_patches)
        cloned._bnb_lora_module_backups = {}
        return cloned

    def pre_run(self, *args, **kwargs):        
        super().pre_run(*args, **kwargs)

        for name, module in self.model.named_modules():            
            try:
                if isinstance(module, LazyOps.Linear):                    
                    try:
                        module.patcher = weakref.proxy(self)
                    except Exception:                        
                        module.patcher = self
                    
                    module.module_key_name = name
                    module.weight_key_name = f"{name}.weight"
            except Exception:                
                logging.debug(f"pre_run: skip module {name} assignment due to exception", exc_info=True)
        
        self.apply_bnb_patches()

    def apply_bnb_patches(self):        
        for module_key, p_list in list(self.bnb_lora_patches.items()):
            try:
                module = comfy.utils.get_attr(self.model, module_key)
            except Exception:
                continue
            if getattr(module, "_bnb_lora_attached", False):
                continue            
            module._bnb_lora_attached = True
            module._bnb_lora_patch_count = len(p_list)

    def get_patches_for_module(self, module_key_name, *, is_bnb=False):        
        if is_bnb:
            return self.bnb_lora_patches.get(module_key_name, None)
        else:
            return self.patches.get(f"{module_key_name}.weight", None)

    def remove_bnb_patches(self):        
        for module_key in list(self.bnb_lora_patches.keys()):
            try:
                module = comfy.utils.get_attr(self.model, module_key)
            except Exception:
                continue

            for attr in ("_bnb_lora_attached", "_bnb_lora_patch_count",):
                if hasattr(module, attr):
                    try:
                        delattr(module, attr)
                    except Exception:
                        logging.debug(f"remove_bnb_patches: could not del {attr} on {module_key}", exc_info=True)
            
            self._bnb_lora_module_backups.pop(module_key, None)

    def clear_bnb_patches(self):        
        self.bnb_lora_patches.clear()
        self._bnb_lora_module_backups.clear()

    def unpatch_model(self, device_to=None, unpatch_weights=True):        
        super().unpatch_model(device_to=device_to, unpatch_weights=unpatch_weights)

        self.remove_bnb_patches()
        
        for name, module in self.model.named_modules():            
            if hasattr(module, "patcher"):
                try:
                    p = getattr(module, "patcher")                    
                    delattr(module, "patcher")
                except Exception:
                    pass

            for attr in ("module_key_name", "weight_key_name"):
                if hasattr(module, attr):
                    try:
                        delattr(module, attr)
                    except Exception:
                        pass
        
        return

    def partially_unload(self, device_to, memory_to_free=0):

        memory_freed = super().partially_unload(device_to, memory_to_free=memory_to_free)

        for name, module in self.model.named_modules():
            if getattr(module, "_bnb_lora_attached", False):
                module_key = getattr(module, "module_key_name", name)                
                if hasattr(module, "_bnb_lora_patch_count"):
                    self._bnb_lora_module_backups[module_key] = getattr(module, "_bnb_lora_patch_count", None)                
                for attr in ("_bnb_lora_attached", "_bnb_lora_patch_count",):
                    if hasattr(module, attr):
                        try:
                            delattr(module, attr)
                        except Exception:
                            pass
        return memory_freed
    
    def calculate_weight_with_patches(self, module_key_name, base_weight_fp32, is_bnb=False):                
        if is_bnb:
            patches = self.bnb_lora_patches.get(module_key_name, None)
        else:
            patches = self.patches.get(f"{module_key_name}.weight", None)

        if not patches:
            return None
        
        try:            
            base = base_weight_fp32.to(torch.float32)
            weight_final_fp32 = comfy.lora.calculate_weight(patches, base, f"{module_key_name}.weight")            
            return weight_final_fp32.to(torch.float32)
        except Exception:            
            logging.exception(f"calculate_weight_with_patches failed for {module_key_name}")
            return None
        
def get_safetensors_model_list(folder_path_key):
    shard_pattern = re.compile(r'.*-(\d{5})-of-(\d{5})\.safetensors$')
    try:
        initial_list = folder_paths.get_filename_list(folder_path_key)
    except KeyError:
        logging.error(f"Path type '{folder_path_key}' is not registered.")
        return []

    sharded_files_to_remove = set()
    parent_dirs_to_add = set()

    for item in initial_list:        
        is_dir = False
        for basedir in folder_paths.get_folder_paths(folder_path_key):
             if os.path.isdir(os.path.join(basedir, item)):
                 is_dir = True
                 break
        if is_dir:
            continue

        filename = os.path.basename(item)
        if shard_pattern.match(filename):
            sharded_files_to_remove.add(item)
            parent_dir = os.path.dirname(item)
            if parent_dir and parent_dir != ".":                
                parent_dirs_to_add.add(parent_dir.replace(os.sep, '/'))

    final_list = [item for item in initial_list if item not in sharded_files_to_remove]
    final_list.extend(list(parent_dirs_to_add))
    
    return sorted(list(set(final_list)))

def is_bnb_4bit(sd: dict) -> bool:
    if not isinstance(sd, dict):
        return False
    for k in sd.keys():
        if k.endswith(".quant_state.bitsandbytes__nf4") or k.endswith(".quant_state.bitsandbytes__fp4"):
            return True
    return False

class UnetBnbModelLoader:
    FOLDER_PATH_KEY = "unet"

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "model_name": (get_safetensors_model_list(s.FOLDER_PATH_KEY),),
            }
        }

    RETURN_TYPES = ("MODEL",)
    FUNCTION = "load_model"
    CATEGORY = "loaders"
    TITLE = "Unet Bnb Model Loader"

    def load_model(self, model_name):        
        model_path = folder_paths.get_full_path(self.FOLDER_PATH_KEY, model_name)
        
        if model_path is None:
            for basedir in folder_paths.get_folder_paths(self.FOLDER_PATH_KEY):
                candidate_path = os.path.join(basedir, model_name)
                if os.path.isdir(candidate_path):
                    model_path = candidate_path
                    break
        
        if model_path is None:
            raise FileNotFoundError(f"Model not found in the directory configured by '{self.FOLDER_PATH_KEY}' class: {model_name}")        
        
        state_dict = safetensors_sd_loader(model_path)        
        
        if is_bnb_4bit(state_dict):
            model_patcher  = comfy.sd.load_diffusion_model_state_dict(state_dict, {"custom_operations": LazyOps()})
            custom_patcher = UnetBnbModelPatcher(model_patcher.model, model_patcher.load_device, model_patcher.offload_device, model_patcher.size)
        else:
            model_patcher  = comfy.sd.load_diffusion_model_state_dict(state_dict)
            custom_patcher = comfy.model_patcher.ModelPatcher(model_patcher.model, model_patcher.load_device, model_patcher.offload_device, model_patcher.size)
        
        if model_patcher is None:
            raise RuntimeError(f"Unable to detect or load UNet model: {model_path}")        
        
        return (custom_patcher,)

NODE_CLASS_MAPPINGS = {
    "UnetBnbModelLoader": UnetBnbModelLoader,
}
NODE_DISPLAY_NAME_MAPPINGS = {
    "UnetBnbModelLoader": "Unet Bnb Model Loader",
}