From d8b90ac121cbf0c18b1dc9d56a5e1d14ca51e74e Mon Sep 17 00:00:00 2001 From: AUTOMATIC <16777216c@gmail.com> Date: Sun, 15 Jan 2023 18:50:56 +0300 Subject: big rework of progressbar/preview system to allow multiple users to prompts at the same time and do not get previews of each other --- modules/progress.py | 96 +++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 96 insertions(+) create mode 100644 modules/progress.py (limited to 'modules/progress.py') diff --git a/modules/progress.py b/modules/progress.py new file mode 100644 index 00000000..3327b883 --- /dev/null +++ b/modules/progress.py @@ -0,0 +1,96 @@ +import base64 +import io +import time + +import gradio as gr +from pydantic import BaseModel, Field + +from modules.shared import opts + +import modules.shared as shared + + +current_task = None +pending_tasks = {} +finished_tasks = [] + + +def start_task(id_task): + global current_task + + current_task = id_task + pending_tasks.pop(id_task, None) + + +def finish_task(id_task): + global current_task + + if current_task == id_task: + current_task = None + + finished_tasks.append(id_task) + if len(finished_tasks) > 16: + finished_tasks.pop(0) + + +def add_task_to_queue(id_job): + pending_tasks[id_job] = time.time() + + +class ProgressRequest(BaseModel): + id_task: str = Field(default=None, title="Task ID", description="id of the task to get progress for") + id_live_preview: int = Field(default=-1, title="Live preview image ID", description="id of last received last preview image") + + +class ProgressResponse(BaseModel): + active: bool = Field(title="Whether the task is being worked on right now") + queued: bool = Field(title="Whether the task is in queue") + completed: bool = Field(title="Whether the task has already finished") + progress: float = Field(default=None, title="Progress", description="The progress with a range of 0 to 1") + eta: float = Field(default=None, title="ETA in secs") + live_preview: str = Field(default=None, title="Live preview image", description="Current live preview; a data: uri") + id_live_preview: int = Field(default=None, title="Live preview image ID", description="Send this together with next request to prevent receiving same image") + textinfo: str = Field(default=None, title="Info text", description="Info text used by WebUI.") + + +def setup_progress_api(app): + return app.add_api_route("/internal/progress", progressapi, methods=["POST"], response_model=ProgressResponse) + + +def progressapi(req: ProgressRequest): + active = req.id_task == current_task + queued = req.id_task in pending_tasks + completed = req.id_task in finished_tasks + + if not active: + return ProgressResponse(active=active, queued=queued, completed=completed, id_live_preview=-1, textinfo="In queue..." if queued else "Waiting...") + + progress = 0 + + if shared.state.job_count > 0: + progress += shared.state.job_no / shared.state.job_count + if shared.state.sampling_steps > 0: + progress += 1 / shared.state.job_count * shared.state.sampling_step / shared.state.sampling_steps + + progress = min(progress, 1) + + elapsed_since_start = time.time() - shared.state.time_start + predicted_duration = elapsed_since_start / progress if progress > 0 else None + eta = predicted_duration - elapsed_since_start if predicted_duration is not None else None + + id_live_preview = req.id_live_preview + shared.state.set_current_image() + if opts.live_previews_enable and shared.state.id_live_preview != req.id_live_preview: + image = shared.state.current_image + if image is not None: + buffered = io.BytesIO() + image.save(buffered, format="png") + live_preview = 'data:image/png;base64,' + base64.b64encode(buffered.getvalue()).decode("ascii") + id_live_preview = shared.state.id_live_preview + else: + live_preview = None + else: + live_preview = None + + 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) + -- cgit v1.2.3 From 308b51012a5def38edb1c2e127e736c43aa6e1a3 Mon Sep 17 00:00:00 2001 From: AUTOMATIC <16777216c@gmail.com> Date: Thu, 19 Jan 2023 08:41:37 +0300 Subject: fix an unlikely division by 0 error --- modules/progress.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) (limited to 'modules/progress.py') diff --git a/modules/progress.py b/modules/progress.py index 3327b883..f9e005d3 100644 --- a/modules/progress.py +++ b/modules/progress.py @@ -67,10 +67,13 @@ def progressapi(req: ProgressRequest): progress = 0 - if shared.state.job_count > 0: - progress += shared.state.job_no / shared.state.job_count - if shared.state.sampling_steps > 0: - progress += 1 / shared.state.job_count * shared.state.sampling_step / shared.state.sampling_steps + job_count, job_no = shared.state.job_count, shared.state.job_no + sampling_steps, sampling_step = shared.state.sampling_steps, shared.state.sampling_step + + if job_count > 0: + progress += job_no / job_count + if sampling_steps > 0: + progress += 1 / job_count * sampling_step / sampling_steps progress = min(progress, 1) -- cgit v1.2.3 From c7e50425f63c07242068f8dcccce70a4ef28a17f Mon Sep 17 00:00:00 2001 From: AUTOMATIC <16777216c@gmail.com> Date: Thu, 19 Jan 2023 09:25:37 +0300 Subject: add progress bar to modelmerger --- javascript/ui.js | 11 +++++++++++ modules/extras.py | 18 +++++++++++++++--- modules/progress.py | 2 +- modules/ui.py | 13 ++++++++----- style.css | 5 +++++ 5 files changed, 40 insertions(+), 9 deletions(-) (limited to 'modules/progress.py') diff --git a/javascript/ui.js b/javascript/ui.js index 7d3d57a3..428375d4 100644 --- a/javascript/ui.js +++ b/javascript/ui.js @@ -172,6 +172,17 @@ function submit_img2img(){ return res } +function modelmerger(){ + var id = randomId() + requestProgress(id, gradioApp().getElementById('modelmerger_results_panel'), null, function(){}) + + gradioApp().getElementById('modelmerger_result').innerHTML = '' + + var res = create_submit_args(arguments) + res[0] = id + return res +} + function ask_for_style_name(_, prompt_text, negative_prompt_text) { name_ = prompt('Style name:') diff --git a/modules/extras.py b/modules/extras.py index 367c15cc..034f28e4 100644 --- a/modules/extras.py +++ b/modules/extras.py @@ -274,14 +274,15 @@ def create_config(ckpt_result, config_source, a, b, c): shutil.copyfile(cfg, checkpoint_filename) -def run_modelmerger(primary_model_name, secondary_model_name, tertiary_model_name, interp_method, multiplier, save_as_half, custom_name, checkpoint_format, config_source): +def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_model_name, interp_method, multiplier, save_as_half, custom_name, checkpoint_format, config_source): shared.state.begin() shared.state.job = 'model-merge' + shared.state.job_count = 1 def fail(message): shared.state.textinfo = message shared.state.end() - return [message, *[gr.update() for _ in range(4)]] + return [*[gr.update() for _ in range(4)], message] def weighted_sum(theta0, theta1, alpha): return ((1 - alpha) * theta0) + (alpha * theta1) @@ -320,9 +321,12 @@ def run_modelmerger(primary_model_name, secondary_model_name, tertiary_model_nam theta_1 = sd_models.read_state_dict(secondary_model_info.filename, map_location='cpu') if theta_func1: + shared.state.job_count += 1 + print(f"Loading {tertiary_model_info.filename}...") theta_2 = sd_models.read_state_dict(tertiary_model_info.filename, map_location='cpu') + shared.state.sampling_steps = len(theta_1.keys()) for key in tqdm.tqdm(theta_1.keys()): if 'model' in key: if key in theta_2: @@ -330,8 +334,12 @@ def run_modelmerger(primary_model_name, secondary_model_name, tertiary_model_nam theta_1[key] = theta_func1(theta_1[key], t2) else: theta_1[key] = torch.zeros_like(theta_1[key]) + + shared.state.sampling_step += 1 del theta_2 + shared.state.nextjob() + shared.state.textinfo = f"Loading {primary_model_info.filename}..." print(f"Loading {primary_model_info.filename}...") theta_0 = sd_models.read_state_dict(primary_model_info.filename, map_location='cpu') @@ -340,6 +348,7 @@ def run_modelmerger(primary_model_name, secondary_model_name, tertiary_model_nam chckpoint_dict_skip_on_merge = ["cond_stage_model.transformer.text_model.embeddings.position_ids"] + shared.state.sampling_steps = len(theta_0.keys()) for key in tqdm.tqdm(theta_0.keys()): if 'model' in key and key in theta_1: @@ -367,6 +376,8 @@ def run_modelmerger(primary_model_name, secondary_model_name, tertiary_model_nam if save_as_half: theta_0[key] = theta_0[key].half() + shared.state.sampling_step += 1 + # I believe this part should be discarded, but I'll leave it for now until I am sure for key in theta_1.keys(): if 'model' in key and key not in theta_0: @@ -393,6 +404,7 @@ def run_modelmerger(primary_model_name, secondary_model_name, tertiary_model_nam output_modelname = os.path.join(ckpt_dir, filename) + shared.state.nextjob() shared.state.textinfo = f"Saving to {output_modelname}..." print(f"Saving to {output_modelname}...") @@ -410,4 +422,4 @@ def run_modelmerger(primary_model_name, secondary_model_name, tertiary_model_nam shared.state.textinfo = "Checkpoint saved to " + output_modelname shared.state.end() - return ["Checkpoint saved to " + output_modelname] + [gr.Dropdown.update(choices=sd_models.checkpoint_tiles()) for _ in range(4)] + return [*[gr.Dropdown.update(choices=sd_models.checkpoint_tiles()) for _ in range(4)], "Checkpoint saved to " + output_modelname] diff --git a/modules/progress.py b/modules/progress.py index f9e005d3..c69ecf3d 100644 --- a/modules/progress.py +++ b/modules/progress.py @@ -72,7 +72,7 @@ def progressapi(req: ProgressRequest): if job_count > 0: progress += job_no / job_count - if sampling_steps > 0: + if sampling_steps > 0 and job_count > 0: progress += 1 / job_count * sampling_step / sampling_steps progress = min(progress, 1) diff --git a/modules/ui.py b/modules/ui.py index 09a3c92e..aeee7853 100644 --- a/modules/ui.py +++ b/modules/ui.py @@ -1208,8 +1208,9 @@ def create_ui(): with gr.Row(): modelmerger_merge = gr.Button(elem_id="modelmerger_merge", value="Merge", variant='primary') - with gr.Column(variant='panel'): - submit_result = gr.Textbox(elem_id="modelmerger_result", show_label=False) + with gr.Column(variant='compact', elem_id="modelmerger_results_container"): + with gr.Group(elem_id="modelmerger_results_panel"): + modelmerger_result = gr.HTML(elem_id="modelmerger_result", show_label=False) with gr.Blocks(analytics_enabled=False) as train_interface: with gr.Row().style(equal_height=False): @@ -1753,12 +1754,14 @@ def create_ui(): print("Error loading/saving model file:", file=sys.stderr) print(traceback.format_exc(), file=sys.stderr) modules.sd_models.list_models() # to remove the potentially missing models from the list - return [f"Error merging checkpoints: {e}"] + [gr.Dropdown.update(choices=modules.sd_models.checkpoint_tiles()) for _ in range(4)] + return [*[gr.Dropdown.update(choices=modules.sd_models.checkpoint_tiles()) for _ in range(4)], f"Error merging checkpoints: {e}"] return results modelmerger_merge.click( - fn=modelmerger, + fn=wrap_gradio_gpu_call(modelmerger, extra_outputs=lambda: [gr.update() for _ in range(4)]), + _js='modelmerger', inputs=[ + dummy_component, primary_model_name, secondary_model_name, tertiary_model_name, @@ -1770,11 +1773,11 @@ def create_ui(): config_source, ], outputs=[ - submit_result, primary_model_name, secondary_model_name, tertiary_model_name, component_dict['sd_model_checkpoint'], + modelmerger_result, ] ) diff --git a/style.css b/style.css index a6abd93d..32ba4753 100644 --- a/style.css +++ b/style.css @@ -737,6 +737,11 @@ footer { line-height: 2.4em; } +#modelmerger_results_container{ + margin-top: 1em; + overflow: visible; +} + /* The following handles localization for right-to-left (RTL) languages like Arabic. The rtl media type will only be activated by the logic in javascript/localization.js. If you change anything above, you need to make sure it is RTL compliant by just running -- cgit v1.2.3