patch 'state_dict' to retrieve the original state_dict

[wkpark]

Description

• after partial update/modify block level UNet, shared.sd_model.state_dict() works nicely, but after using any lora in the prompt, sd_mode.state_dict() does not work as expected.

this fix adds more patch to LoraPatches to hook state_dict() to retrieve original weights/bias.

without this fix we can get the original state_dict() by the following patch/hook

import torch

from modules import patches, script_callbacks

class StateDictPatches:
    def __init__(self):
        import networks

        def network_Linear_state_dict(module, *args, **kwargs):
            with torch.no_grad():
                networks.network_restore_weights_from_backup(module)

            return self.Linear_state_dict(module, *args, **kwargs)


        def network_Conv2d_state_dict(module, *args, **kwargs):
            with torch.no_grad():
                networks.network_restore_weights_from_backup(module)

            return self.Conv2d_state_dict(module, *args, **kwargs)


        def network_GroupNorm_state_dict(module, *args, **kwargs):
            with torch.no_grad():
                networks.network_restore_weights_from_backup(module)

            return self.GroupNorm_state_dict(module, *args, **kwargs)


        def network_LayerNorm_state_dict(module, *args, **kwargs):
            with torch.no_grad():
                networks.network_restore_weights_from_backup(module)

            return self.LayerNorm_state_dict(module, *args, **kwargs)


        def network_MultiheadAttention_state_dict(module, *args, **kwargs):
            with torch.no_grad():
                networks.network_restore_weights_from_backup(module)

            return self.MultiheadAttention_state_dict(module, *args, **kwargs)


        self.Linear_state_dict = patches.patch(__name__, torch.nn.Linear, 'state_dict', network_Linear_state_dict)
        self.Conv2d_state_dict = patches.patch(__name__, torch.nn.Conv2d, 'state_dict', network_Conv2d_state_dict)
        self.GroupNorm_state_dict = patches.patch(__name__, torch.nn.GroupNorm, 'state_dict', network_GroupNorm_state_dict)
        self.LayerNorm_state_dict = patches.patch(__name__, torch.nn.LayerNorm, 'state_dict', network_LayerNorm_state_dict)
        self.MultiheadAttention_state_dict = patches.patch(__name__, torch.nn.MultiheadAttention, 'state_dict', network_MultiheadAttention_state_dict)


    def undo(self):
        self.Linear_state_dict = patches.undo(__name__, torch.nn.Linear, 'state_dict')
        self.Conv2d_state_dict = patches.undo(__name__, torch.nn.Conv2d, 'state_dict')
        self.GroupNorm_state_dict = patches.undo(__name__, torch.nn.GroupNorm, 'state_dict')
        self.LayerNorm_state_dict = patches.undo(__name__, torch.nn.LayerNorm, 'state_dict')
        self.MultiheadAttention_state_dict = patches.undo(__name__, torch.nn.MultiheadAttention, 'state_dict')

Checklist:

• I have read contributing wiki page
• I have performed a self-review of my own code
• My code follows the style guidelines
• My code passes tests

[AUTOMATIC1111]

The problem is that you're expecting state_dict of the original checkpoint and are getting state_dict of checkpoint with loras applied, is that correct? I kind of thought that getting weights with loras was the desired outcome.

[wkpark]

The problem is that you're expecting state_dict of the original checkpoint and are getting state_dict of checkpoint with loras applied, is that correct? I kind of thought that getting weights with loras was the desired outcome.

I don't think this is the intended result, but the current behavior can be useful for merging LoRAs into a checkpoint

anyway, without this PR, we can get the original state_dict by partial hook/patches onto state_dict like this PR.