stable-diffusion-webui

1
import torch.nn
2

3
from modules import script_callbacks, shared, devices
4

5
unet_options = []
6
current_unet_option = None
7
current_unet = None
8
original_forward = None  # not used, only left temporarily for compatibility
9

10
def list_unets():
11
    new_unets = script_callbacks.list_unets_callback()
12

13
    unet_options.clear()
14
    unet_options.extend(new_unets)
15

16

17
def get_unet_option(option=None):
18
    option = option or shared.opts.sd_unet
19

20
    if option == "None":
21
        return None
22

23
    if option == "Automatic":
24
        name = shared.sd_model.sd_checkpoint_info.model_name
25

26
        options = [x for x in unet_options if x.model_name == name]
27

28
        option = options[0].label if options else "None"
29

30
    return next(iter([x for x in unet_options if x.label == option]), None)
31

32

33
def apply_unet(option=None):
34
    global current_unet_option
35
    global current_unet
36

37
    new_option = get_unet_option(option)
38
    if new_option == current_unet_option:
39
        return
40

41
    if current_unet is not None:
42
        print(f"Dectivating unet: {current_unet.option.label}")
43
        current_unet.deactivate()
44

45
    current_unet_option = new_option
46
    if current_unet_option is None:
47
        current_unet = None
48

49
        if not shared.sd_model.lowvram:
50
            shared.sd_model.model.diffusion_model.to(devices.device)
51

52
        return
53

54
    shared.sd_model.model.diffusion_model.to(devices.cpu)
55
    devices.torch_gc()
56

57
    current_unet = current_unet_option.create_unet()
58
    current_unet.option = current_unet_option
59
    print(f"Activating unet: {current_unet.option.label}")
60
    current_unet.activate()
61

62

63
class SdUnetOption:
64
    model_name = None
65
    """name of related checkpoint - this option will be selected automatically for unet if the name of checkpoint matches this"""
66

67
    label = None
68
    """name of the unet in UI"""
69

70
    def create_unet(self):
71
        """returns SdUnet object to be used as a Unet instead of built-in unet when making pictures"""
72
        raise NotImplementedError()
73

74

75
class SdUnet(torch.nn.Module):
76
    def forward(self, x, timesteps, context, *args, **kwargs):
77
        raise NotImplementedError()
78

79
    def activate(self):
80
        pass
81

82
    def deactivate(self):
83
        pass
84

85

86
def create_unet_forward(original_forward):
87
    def UNetModel_forward(self, x, timesteps=None, context=None, *args, **kwargs):
88
        if current_unet is not None:
89
            return current_unet.forward(x, timesteps, context, *args, **kwargs)
90

91
        return original_forward(self, x, timesteps, context, *args, **kwargs)
92

93
    return UNetModel_forward
94

95