stable-diffusion-webui

1
import base64
2
import io
3
import time
4

5
import gradio as gr
6
from pydantic import BaseModel, Field
7

8
from modules.shared import opts
9

10
import modules.shared as shared
11
from collections import OrderedDict
12
import string
13
import random
14
from typing import List
15

16
current_task = None
17
pending_tasks = OrderedDict()
18
finished_tasks = []
19
recorded_results = []
20
recorded_results_limit = 2
21

22

23
def start_task(id_task):
24
    global current_task
25

26
    current_task = id_task
27
    pending_tasks.pop(id_task, None)
28

29

30
def finish_task(id_task):
31
    global current_task
32

33
    if current_task == id_task:
34
        current_task = None
35

36
    finished_tasks.append(id_task)
37
    if len(finished_tasks) > 16:
38
        finished_tasks.pop(0)
39

40
def create_task_id(task_type):
41
    N = 7
42
    res = ''.join(random.choices(string.ascii_uppercase +
43
    string.digits, k=N))
44
    return f"task({task_type}-{res})"
45

46
def record_results(id_task, res):
47
    recorded_results.append((id_task, res))
48
    if len(recorded_results) > recorded_results_limit:
49
        recorded_results.pop(0)
50

51

52
def add_task_to_queue(id_job):
53
    pending_tasks[id_job] = time.time()
54

55
class PendingTasksResponse(BaseModel):
56
    size: int = Field(title="Pending task size")
57
    tasks: List[str] = Field(title="Pending task ids")
58

59
class ProgressRequest(BaseModel):
60
    id_task: str = Field(default=None, title="Task ID", description="id of the task to get progress for")
61
    id_live_preview: int = Field(default=-1, title="Live preview image ID", description="id of last received last preview image")
62
    live_preview: bool = Field(default=True, title="Include live preview", description="boolean flag indicating whether to include the live preview image")
63

64

65
class ProgressResponse(BaseModel):
66
    active: bool = Field(title="Whether the task is being worked on right now")
67
    queued: bool = Field(title="Whether the task is in queue")
68
    completed: bool = Field(title="Whether the task has already finished")
69
    progress: float = Field(default=None, title="Progress", description="The progress with a range of 0 to 1")
70
    eta: float = Field(default=None, title="ETA in secs")
71
    live_preview: str = Field(default=None, title="Live preview image", description="Current live preview; a data: uri")
72
    id_live_preview: int = Field(default=None, title="Live preview image ID", description="Send this together with next request to prevent receiving same image")
73
    textinfo: str = Field(default=None, title="Info text", description="Info text used by WebUI.")
74

75

76
def setup_progress_api(app):
77
    app.add_api_route("/internal/pending-tasks", get_pending_tasks, methods=["GET"])
78
    return app.add_api_route("/internal/progress", progressapi, methods=["POST"], response_model=ProgressResponse)
79

80

81
def get_pending_tasks():
82
    pending_tasks_ids = list(pending_tasks)
83
    pending_len = len(pending_tasks_ids)
84
    return PendingTasksResponse(size=pending_len, tasks=pending_tasks_ids)
85

86

87
def progressapi(req: ProgressRequest):
88
    active = req.id_task == current_task
89
    queued = req.id_task in pending_tasks
90
    completed = req.id_task in finished_tasks
91

92
    if not active:
93
        textinfo = "Waiting..."
94
        if queued:
95
            sorted_queued = sorted(pending_tasks.keys(), key=lambda x: pending_tasks[x])
96
            queue_index = sorted_queued.index(req.id_task)
97
            textinfo = "In queue: {}/{}".format(queue_index + 1, len(sorted_queued))
98
        return ProgressResponse(active=active, queued=queued, completed=completed, id_live_preview=-1, textinfo=textinfo)
99

100
    progress = 0
101

102
    job_count, job_no = shared.state.job_count, shared.state.job_no
103
    sampling_steps, sampling_step = shared.state.sampling_steps, shared.state.sampling_step
104

105
    if job_count > 0:
106
        progress += job_no / job_count
107
    if sampling_steps > 0 and job_count > 0:
108
        progress += 1 / job_count * sampling_step / sampling_steps
109

110
    progress = min(progress, 1)
111

112
    elapsed_since_start = time.time() - shared.state.time_start
113
    predicted_duration = elapsed_since_start / progress if progress > 0 else None
114
    eta = predicted_duration - elapsed_since_start if predicted_duration is not None else None
115

116
    live_preview = None
117
    id_live_preview = req.id_live_preview
118

119
    if opts.live_previews_enable and req.live_preview:
120
        shared.state.set_current_image()
121
        if shared.state.id_live_preview != req.id_live_preview:
122
            image = shared.state.current_image
123
            if image is not None:
124
                buffered = io.BytesIO()
125

126
                if opts.live_previews_image_format == "png":
127
                    # using optimize for large images takes an enormous amount of time
128
                    if max(*image.size) <= 256:
129
                        save_kwargs = {"optimize": True}
130
                    else:
131
                        save_kwargs = {"optimize": False, "compress_level": 1}
132

133
                else:
134
                    save_kwargs = {}
135

136
                image.save(buffered, format=opts.live_previews_image_format, **save_kwargs)
137
                base64_image = base64.b64encode(buffered.getvalue()).decode('ascii')
138
                live_preview = f"data:image/{opts.live_previews_image_format};base64,{base64_image}"
139
                id_live_preview = shared.state.id_live_preview
140

141
    return ProgressResponse(active=active, queued=queued, completed=completed, progress=progress, eta=eta, live_preview=live_preview, id_live_preview=id_live_preview, textinfo=shared.state.textinfo)
142

143

144
def restore_progress(id_task):
145
    while id_task == current_task or id_task in pending_tasks:
146
        time.sleep(0.1)
147

148
    res = next(iter([x[1] for x in recorded_results if id_task == x[0]]), None)
149
    if res is not None:
150
        return res
151

152
    return gr.update(), gr.update(), gr.update(), f"Couldn't restore progress for {id_task}: results either have been discarded or never were obtained"
153