Class WeightModule

lightx2v/common/modules/weight_module.py:4–188

class WeightModule:
    def __init__(self):
        self._modules = {}
        self._parameters = {}

    def is_empty(self):
        return len(self._modules) == 0 and len(self._parameters) == 0

    def add_module(self, name, module):
        self._modules[name] = module
        setattr(self, name, module)

    def register_parameter(self, name, param):
        self._parameters[name] = param
        setattr(self, name, param)

    def load(self, weight_dict):
        for _, module in self._modules.items():
            if hasattr(module, "load"):
                module.load(weight_dict)

        for _, parameter in self._parameters.items():
            if hasattr(parameter, "load"):
                parameter.load(weight_dict)

    def register_diff(self, weight_dict):
        for _, module in self._modules.items():
            if hasattr(module, "register_diff"):
                module.register_diff(weight_dict)

        for _, parameter in self._parameters.items():
            if hasattr(parameter, "register_diff"):
                parameter.register_diff(weight_dict)

    def register_lora(self, weight_dict, strength):
        for _, module in self._modules.items():
            if hasattr(module, "register_lora"):
                module.register_lora(weight_dict, strength)

        for _, parameter in self._parameters.items():
            if hasattr(parameter, "register_lora"):
                parameter.register_lora(weight_dict, strength)

    def update_lora(self, weight_dict, strength):
        for _, module in self._modules.items():
            if hasattr(module, "update_lora"):
                module.update_lora(weight_dict, strength)

        for _, parameter in self._parameters.items():
            if hasattr(parameter, "update_lora"):
                parameter.update_lora(weight_dict, strength)

    def remove_lora(self):
        for _, module in self._modules.items():
            if hasattr(module, "remove_lora"):
                module.remove_lora()

        for _, parameter in self._parameters.items():