hololens/stable-diffusion-webui-depthmap-script

1# DepthMap can be run inside stable-diffusion-webui, but also separately.
2# All the stable-diffusion-webui stuff that the DepthMap relies on
3# must be resided in this file (or in the scripts folder).
4import pathlib
5from datetime import datetime
6import enum
7import sys
8
9
10class BackboneType(enum.Enum):
11    WEBUI = 1
12    STANDALONE = 2
13
14
15try:
16    # stable-diffusion-webui backbone
17    from modules.images import save_image  # Should fail if not on stable-diffusion-webui
18    from modules.devices import torch_gc  # TODO: is this really sufficient?
19    from modules.images import get_next_sequence_number
20    from modules.call_queue import wrap_gradio_gpu_call
21    from modules.shared import listfiles
22
23    def get_opt(name, default):
24        from modules.shared import opts
25        if hasattr(opts, name):
26            return opts.__getattr__(name)
27        return default
28
29    def get_cmd_opt(name, default):
30        """Get command line argument"""
31        from modules.shared import cmd_opts
32        if hasattr(cmd_opts, name):
33            return cmd_opts.__getattribute__(name)
34        return default
35
36    def gather_ops():
37        """Parameters for depthmap generation"""
38        ops = {}
39        for s in ['boost_rmax', 'precision', 'no_half', 'marigold_ensembles', 'marigold_steps']:
40            c = get_opt('depthmap_script_' + s, None)
41            if c is None:
42                c = get_cmd_opt(s, None)
43            if c is not None:
44                ops[s] = c
45        # sanitize for integers.
46        for s in ['marigold_ensembles', 'marigold_steps']:
47            if s in ops:
48                ops[s] = int(ops[s])
49        return ops
50
51
52    def get_outpath():
53        """Get path where results are saved by default"""
54        path = get_opt('outdir_samples', None)
55        if path is None or len(path) == 0:
56            path = get_opt('outdir_extras_samples', None)
57        assert path is not None and len(path) > 0
58        return path
59
60
61    def unload_sd_model():
62        from modules import shared, devices
63        if shared.sd_model is not None:
64            if shared.sd_model.cond_stage_model is not None:
65                shared.sd_model.cond_stage_model.to(devices.cpu)
66            if shared.sd_model.first_stage_model is not None:
67                shared.sd_model.first_stage_model.to(devices.cpu)
68        # Maybe something else???
69
70
71    def reload_sd_model():
72        from modules import shared, devices
73        if shared.sd_model is not None:
74            if shared.sd_model.cond_stage_model is not None:
75                shared.sd_model.cond_stage_model.to(devices.device)
76            if shared.sd_model.first_stage_model:
77                shared.sd_model.first_stage_model.to(devices.device)
78        # Maybe something else???
79
80    def get_hide_dirs():
81        import modules.shared
82        return modules.shared.hide_dirs
83
84    USED_BACKBONE = BackboneType.WEBUI
85except:
86    # Standalone backbone
87    print(  # "  DepthMap did not detect stable-diffusion-webui; launching with the standalone backbone.\n"
88          "  The standalone mode is not on par with the stable-diffusion-webui mode.\n"
89          "  Some features may be missing or work differently. Please report bugs.\n")
90
91    def save_image(image, path, basename, **kwargs):
92        import os
93        os.makedirs(path, exist_ok=True)
94        if 'suffix' not in kwargs or len(kwargs['suffix']) == 0:
95            kwargs['suffix'] = ''
96        else:
97            kwargs['suffix'] = f"-{kwargs['suffix']}"
98        format = get_opt('samples_format', kwargs['extension'])
99        fullfn = os.path.join(
100            path, f"{basename}-{get_next_sequence_number(path, basename)}{kwargs['suffix']}.{format}")
101        image.save(fullfn, format=format)
102
103    def torch_gc():
104        # TODO: is this really sufficient?
105        import torch
106        if torch.cuda.is_available():
107            with torch.cuda.device('cuda'):
108                torch.cuda.empty_cache()
109                torch.cuda.ipc_collect()
110
111    launched_at = int(datetime.now().timestamp())
112    backbone_current_seq_number = 0
113
114    # Make sure to preserve the function signature when calling!
115    def get_next_sequence_number(outpath, basename):
116        global backbone_current_seq_number
117        backbone_current_seq_number += 1
118        return int(f"{launched_at}{backbone_current_seq_number:04}")
119
120    def wrap_gradio_gpu_call(f): return f  # Displaying various stats is not supported
121
122    def listfiles(dirname):
123        import os
124        filenames = [os.path.join(dirname, x) for x in sorted(os.listdir(dirname)) if not x.startswith(".")]
125        return [file for file in filenames if os.path.isfile(file)]
126
127    def get_opt(name, default): return default  # Configuring is not supported
128
129
130    def get_cmd_opt(name, default): return default  # Configuring is not supported
131
132    def gather_ops():  # Configuring is not supported
133        return {'boost_rmax': 1600,
134                'precision': 'autocast',
135                'no_half': False,
136                'marigold_ensembles': 5,
137                'marigold_steps': 12}
138
139    def get_outpath(): return str(pathlib.Path('.', 'outputs'))
140
141    def unload_sd_model(): pass  # Not needed
142
143    def reload_sd_model(): pass  # Not needed
144
145    def get_hide_dirs(): return {}  # Directories will not be hidden from traversal (except when starts with the dot)
146
147
148    USED_BACKBONE = BackboneType.STANDALONE
149