File size: 5,509 Bytes
e04dce3
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
# DepthMap can be run inside stable-diffusion-webui, but also separately.
# All the stable-diffusion-webui stuff that the DepthMap relies on
# must be resided in this file (or in the scripts folder).
import pathlib
from datetime import datetime
import enum
import sys


class BackboneType(enum.Enum):
    WEBUI = 1
    STANDALONE = 2


try:
    # stable-diffusion-webui backbone
    from modules.images import save_image  # Should fail if not on stable-diffusion-webui
    from modules.devices import torch_gc  # TODO: is this really sufficient?
    from modules.images import get_next_sequence_number
    from modules.call_queue import wrap_gradio_gpu_call
    from modules.shared import listfiles

    def get_opt(name, default):
        from modules.shared import opts
        if hasattr(opts, name):
            return opts.__getattr__(name)
        return default

    def get_cmd_opt(name, default):
        """Get command line argument"""
        from modules.shared import cmd_opts
        if hasattr(cmd_opts, name):
            return cmd_opts.__getattribute__(name)
        return default

    def gather_ops():
        """Parameters for depthmap generation"""
        ops = {}
        for s in ['boost_rmax', 'precision', 'no_half', 'marigold_ensembles', 'marigold_steps']:
            c = get_opt('depthmap_script_' + s, None)
            if c is None:
                c = get_cmd_opt(s, None)
            if c is not None:
                ops[s] = c
        # sanitize for integers.
        for s in ['marigold_ensembles', 'marigold_steps']:
            if s in ops:
                ops[s] = int(ops[s])
        return ops


    def get_outpath():
        """Get path where results are saved by default"""
        path = get_opt('outdir_samples', None)
        if path is None or len(path) == 0:
            path = get_opt('outdir_extras_samples', None)
        assert path is not None and len(path) > 0
        return path


    def unload_sd_model():
        from modules import shared, devices
        if shared.sd_model is not None:
            if shared.sd_model.cond_stage_model is not None:
                shared.sd_model.cond_stage_model.to(devices.cpu)
            if shared.sd_model.first_stage_model is not None:
                shared.sd_model.first_stage_model.to(devices.cpu)
        # Maybe something else???


    def reload_sd_model():
        from modules import shared, devices
        if shared.sd_model is not None:
            if shared.sd_model.cond_stage_model is not None:
                shared.sd_model.cond_stage_model.to(devices.device)
            if shared.sd_model.first_stage_model:
                shared.sd_model.first_stage_model.to(devices.device)
        # Maybe something else???

    def get_hide_dirs():
        import modules.shared
        return modules.shared.hide_dirs

    USED_BACKBONE = BackboneType.WEBUI
except:
    # Standalone backbone
    print(  # "  DepthMap did not detect stable-diffusion-webui; launching with the standalone backbone.\n"
          "  The standalone mode is not on par with the stable-diffusion-webui mode.\n"
          "  Some features may be missing or work differently. Please report bugs.\n")

    def save_image(image, path, basename, **kwargs):
        import os
        os.makedirs(path, exist_ok=True)
        if 'suffix' not in kwargs or len(kwargs['suffix']) == 0:
            kwargs['suffix'] = ''
        else:
            kwargs['suffix'] = f"-{kwargs['suffix']}"
        format = get_opt('samples_format', kwargs['extension'])
        fullfn = os.path.join(
            path, f"{basename}-{get_next_sequence_number(path, basename)}{kwargs['suffix']}.{format}")
        image.save(fullfn, format=format)

    def torch_gc():
        # TODO: is this really sufficient?
        import torch
        if torch.cuda.is_available():
            with torch.cuda.device('cuda'):
                torch.cuda.empty_cache()
                torch.cuda.ipc_collect()

    launched_at = int(datetime.now().timestamp())
    backbone_current_seq_number = 0

    # Make sure to preserve the function signature when calling!
    def get_next_sequence_number(outpath, basename):
        global backbone_current_seq_number
        backbone_current_seq_number += 1
        return int(f"{launched_at}{backbone_current_seq_number:04}")

    def wrap_gradio_gpu_call(f): return f  # Displaying various stats is not supported

    def listfiles(dirname):
        import os
        filenames = [os.path.join(dirname, x) for x in sorted(os.listdir(dirname)) if not x.startswith(".")]
        return [file for file in filenames if os.path.isfile(file)]

    def get_opt(name, default): return default  # Configuring is not supported


    def get_cmd_opt(name, default): return default  # Configuring is not supported

    def gather_ops():  # Configuring is not supported
        return {'boost_rmax': 1600,
                'precision': 'autocast',
                'no_half': False,
                'marigold_ensembles': 5,
                'marigold_steps': 12}

    def get_outpath(): return str(pathlib.Path('.', 'outputs'))

    def unload_sd_model(): pass  # Not needed

    def reload_sd_model(): pass  # Not needed

    def get_hide_dirs(): return {}  # Directories will not be hidden from traversal (except when starts with the dot)


    USED_BACKBONE = BackboneType.STANDALONE