From 81f58d58d096920103b60a9de72715fe6b6ad9e3 Mon Sep 17 00:00:00 2001 From: cobryan05 <13701027+cobryan05@users.noreply.github.com> Date: Wed, 14 Sep 2022 15:48:56 -0500 Subject: [PATCH] Re-merge #611 - View/Cancel in-progress diffusions (#796) * JobManager: Re-merge #611 PR #611 seems to have got lost in the shuffle after the transition to 'dev'. This commit re-merges the feature branch. This adds support for viewing preview images as the image generates, as well as cancelling in-progress images and a couple fixes and clean-ups. * JobManager: Clear jobs that fail to start Sometimes if a job fails to start it will get stuck in the active job list. This commit ensures that jobs that raise exceptions are cleared, and also adds a start timer to clear out jobs that fail to start within a reasonable amount of time. --- frontend/job_manager.py | 205 +++++++++++++++++++++++++++++++--------- scripts/webui.py | 108 +++++++++++++++++++-- 2 files changed, 260 insertions(+), 53 deletions(-) diff --git a/frontend/job_manager.py b/frontend/job_manager.py index 0e0be64..026742f 100644 --- a/frontend/job_manager.py +++ b/frontend/job_manager.py @@ -1,7 +1,7 @@ ''' Provides simple job management for gradio, allowing viewing and stopping in-progress multi-batch generations ''' from __future__ import annotations import gradio as gr -from gradio.components import Component, Gallery +from gradio.components import Component, Gallery, Slider from threading import Event, Timer from typing import Callable, List, Dict, Tuple, Optional, Any from dataclasses import dataclass, field @@ -9,6 +9,7 @@ from functools import partial from PIL.Image import Image import uuid import traceback +import time @dataclass(eq=True, frozen=True) @@ -30,9 +31,21 @@ class JobInfo: session_key: str job_token: Optional[int] = None images: List[Image] = field(default_factory=list) + active_image: Image = None + rec_steps_enabled: bool = False + rec_steps_imgs: List[Image] = field(default_factory=list) + rec_steps_intrvl: int = None + rec_steps_to_gallery: bool = False + rec_steps_to_file: bool = False should_stop: Event = field(default_factory=Event) + refresh_active_image_requested: Event = field(default_factory=Event) + refresh_active_image_done: Event = field(default_factory=Event) + stop_cur_iter: Event = field(default_factory=Event) + active_iteration_cnt: int = field(default_factory=int) job_status: str = field(default_factory=str) finished: bool = False + started: bool = False + timestamp: float = None removed_output_idxs: List[int] = field(default_factory=list) @@ -76,7 +89,7 @@ class JobManagerUi: ''' return self._job_manager._wrap_func( func=func, inputs=inputs, outputs=outputs, - refresh_btn=self._refresh_btn, stop_btn=self._stop_btn, status_text=self._status_text + job_ui=self ) _refresh_btn: gr.Button @@ -84,10 +97,19 @@ class JobManagerUi: _status_text: gr.Textbox _stop_all_session_btn: gr.Button _free_done_sessions_btn: gr.Button + _active_image: gr.Image + _active_image_stop_btn: gr.Button + _active_image_refresh_btn: gr.Button + _rec_steps_intrvl_sldr: gr.Slider + _rec_steps_checkbox: gr.Checkbox + _save_rec_steps_to_gallery_chkbx: gr.Checkbox + _save_rec_steps_to_file_chkbx: gr.Checkbox _job_manager: JobManager class JobManager: + JOB_MAX_START_TIME = 5.0 # How long can a job be stuck 'starting' before assuming it isn't running + def __init__(self, max_jobs: int): self._max_jobs: int = max_jobs self._avail_job_tokens: List[Any] = list(range(max_jobs)) @@ -102,11 +124,23 @@ class JobManager: ''' assert gr.context.Context.block is not None, "draw_gradio_ui must be called within a 'gr.Blocks' 'with' context" with gr.Tabs(): - with gr.TabItem("Current Session"): + with gr.TabItem("Job Controls"): with gr.Row(): - stop_btn = gr.Button("Stop", elem_id="stop", variant="secondary") - refresh_btn = gr.Button("Refresh", elem_id="refresh", variant="secondary") + stop_btn = gr.Button("Stop All Batches", elem_id="stop", variant="secondary") + refresh_btn = gr.Button("Refresh Finished Batches", elem_id="refresh", variant="secondary") status_text = gr.Textbox(placeholder="Job Status", interactive=False, show_label=False) + with gr.Row(): + active_image_stop_btn = gr.Button("Skip Active Batch", variant="secondary") + active_image_refresh_btn = gr.Button("View Batch Progress", variant="secondary") + active_image = gr.Image(type="pil", interactive=False, visible=False, elem_id="active_iteration_image") + with gr.TabItem("Batch Progress Settings"): + with gr.Row(): + record_steps_checkbox = gr.Checkbox(value=False, label="Enable Batch Progress Grid") + record_steps_interval_slider = gr.Slider( + value=3, label="Record Interval (steps)", minimum=1, maximum=25, step=1) + with gr.Row() as record_steps_box: + steps_to_gallery_checkbox = gr.Checkbox(value=False, label="Save Progress Grid to Gallery") + steps_to_file_checkbox = gr.Checkbox(value=False, label="Save Progress Grid to File") with gr.TabItem("Maintenance"): with gr.Row(): gr.Markdown( @@ -118,9 +152,15 @@ class JobManager: free_done_sessions_btn = gr.Button( "Clear Finished Jobs", elem_id="clear_finished", variant="secondary" ) + return JobManagerUi(_refresh_btn=refresh_btn, _stop_btn=stop_btn, _status_text=status_text, _stop_all_session_btn=stop_all_sessions_btn, _free_done_sessions_btn=free_done_sessions_btn, - _job_manager=self) + _active_image=active_image, _active_image_stop_btn=active_image_stop_btn, + _active_image_refresh_btn=active_image_refresh_btn, + _rec_steps_checkbox=record_steps_checkbox, + _save_rec_steps_to_gallery_chkbx=steps_to_gallery_checkbox, + _save_rec_steps_to_file_chkbx=steps_to_file_checkbox, + _rec_steps_intrvl_sldr=record_steps_interval_slider, _job_manager=self) def clear_all_finished_jobs(self): ''' Removes all currently finished jobs, across all sessions. @@ -134,6 +174,7 @@ class JobManager: for session in self._sessions.values(): for job in session.jobs.values(): job.should_stop.set() + job.stop_cur_iter.set() def _get_job_token(self, block: bool = False) -> Optional[int]: ''' Attempts to acquire a job token, optionally blocking until available ''' @@ -175,6 +216,26 @@ class JobManager: job_info.should_stop.set() return "Stopping after current batch finishes" + def _refresh_cur_iter_func(self, func_key: FuncKey, session_key: str) -> List[Component]: + ''' Updates information from the active iteration ''' + session_info, job_info = self._get_call_info(func_key, session_key) + if job_info is None: + return [None, f"Session {session_key} was not running function {func_key}"] + + job_info.refresh_active_image_requested.set() + if job_info.refresh_active_image_done.wait(timeout=20.0): + job_info.refresh_active_image_done.clear() + return [gr.Image.update(value=job_info.active_image, visible=True), f"Sample iteration {job_info.active_iteration_cnt}"] + return [gr.Image.update(visible=False), "Timed out getting image"] + + def _stop_cur_iter_func(self, func_key: FuncKey, session_key: str) -> List[Component]: + ''' Marks that the active iteration should be stopped''' + session_info, job_info = self._get_call_info(func_key, session_key) + if job_info is None: + return [None, f"Session {session_key} was not running function {func_key}"] + job_info.stop_cur_iter.set() + return [gr.Image.update(visible=False), "Stopping current iteration"] + def _get_call_info(self, func_key: FuncKey, session_key: str) -> Tuple[SessionInfo, JobInfo]: ''' Helper to get the SessionInfo and JobInfo. ''' session_info = self._sessions.get(session_key, None) @@ -207,7 +268,8 @@ class JobManager: def _pre_call_func( self, func_key: FuncKey, output_dummy_obj: Component, refresh_btn: gr.Button, stop_btn: gr.Button, - status_text: gr.Textbox, session_key: str) -> List[Component]: + status_text: gr.Textbox, active_image: gr.Image, active_refresh_btn: gr.Button, active_stop_btn: gr.Button, + session_key: str) -> List[Component]: ''' Called when a job is about to start ''' session_info, job_info = self._get_call_info(func_key, session_key) @@ -219,7 +281,9 @@ class JobManager: return {output_dummy_obj: triggerChangeEvent(), refresh_btn: gr.Button.update(variant="primary", value=refresh_btn.value), stop_btn: gr.Button.update(variant="primary", value=stop_btn.value), - status_text: gr.Textbox.update(value="Generation has started. Click 'Refresh' for updates") + status_text: gr.Textbox.update(value="Generation has started. Click 'Refresh' to see finished images, 'View Batch Progress' for active images"), + active_refresh_btn: gr.Button.update(variant="primary", value=active_refresh_btn.value), + active_stop_btn: gr.Button.update(variant="primary", value=active_stop_btn.value), } def _call_func(self, func_key: FuncKey, session_key: str) -> List[Component]: @@ -228,12 +292,19 @@ class JobManager: if session_info is None or job_info is None: return [] + job_info.started = True try: + if job_info.should_stop.is_set(): + raise Exception(f"Job {job_info} requested a stop before execution began") outputs = job_info.func(*job_info.inputs, job_info=job_info) except Exception as e: job_info.job_status = f"Error: {e}" print(f"Exception processing job {job_info}: {e}\n{traceback.format_exc()}") - outputs = [] + raise + finally: + job_info.finished = True + session_info.finished_jobs[func_key] = session_info.jobs.pop(func_key) + self._release_job_token(job_info.job_token) # Filter the function output for any removed outputs filtered_output = [] @@ -241,11 +312,6 @@ class JobManager: if idx not in job_info.removed_output_idxs: filtered_output.append(output) - job_info.finished = True - session_info.finished_jobs[func_key] = session_info.jobs.pop(func_key) - - self._release_job_token(job_info.job_token) - # The wrapper added a dummy JSON output. Append a random text string # to fire the dummy objects 'change' event to notify that the job is done filtered_output.append(triggerChangeEvent()) @@ -254,12 +320,16 @@ class JobManager: def _post_call_func( self, func_key: FuncKey, output_dummy_obj: Component, refresh_btn: gr.Button, stop_btn: gr.Button, - status_text: gr.Textbox, session_key: str) -> List[Component]: + status_text: gr.Textbox, active_image: gr.Image, active_refresh_btn: gr.Button, active_stop_btn: gr.Button, + session_key: str) -> List[Component]: ''' Called when a job completes ''' return {output_dummy_obj: triggerChangeEvent(), refresh_btn: gr.Button.update(variant="secondary", value=refresh_btn.value), stop_btn: gr.Button.update(variant="secondary", value=stop_btn.value), - status_text: gr.Textbox.update(value="Generation has finished!") + status_text: gr.Textbox.update(value="Generation has finished!"), + active_refresh_btn: gr.Button.update(variant="secondary", value=active_refresh_btn.value), + active_stop_btn: gr.Button.update(variant="secondary", value=active_stop_btn.value), + active_image: gr.Image.update(visible=False) } def _update_gallery_event(self, func_key: FuncKey, session_key: str) -> List[Component]: @@ -270,21 +340,17 @@ class JobManager: if session_info is None or job_info is None: return [] - if job_info.finished: - session_info.finished_jobs.pop(func_key) - return job_info.images - def _wrap_func( - self, func: Callable, inputs: List[Component], outputs: List[Component], - refresh_btn: gr.Button = None, stop_btn: gr.Button = None, - status_text: Optional[gr.Textbox] = None) -> Tuple[Callable, List[Component]]: + def _wrap_func(self, func: Callable, inputs: List[Component], + outputs: List[Component], + job_ui: JobManagerUi) -> Tuple[Callable, List[Component]]: ''' handles JobManageUI's wrap_func''' assert gr.context.Context.block is not None, "wrap_func must be called within a 'gr.Blocks' 'with' context" # Create a unique key for this job - func_key = FuncKey(job_id=uuid.uuid4(), func=func) + func_key = FuncKey(job_id=uuid.uuid4().hex, func=func) # Create a unique session key (next gradio release can use gr.State, see https://gradio.app/state_in_blocks/) if self._session_key is None: @@ -302,9 +368,6 @@ class JobManager: del outputs[idx] break - # Add the session key to the inputs - inputs += [self._session_key] - # Create dummy objects update_gallery_obj = gr.JSON(visible=False, elem_id="JobManagerDummyObject") update_gallery_obj.change( @@ -314,21 +377,49 @@ class JobManager: queue=False ) - if refresh_btn: - refresh_btn.variant = 'secondary' - refresh_btn.click( + if job_ui._refresh_btn: + job_ui._refresh_btn.variant = 'secondary' + job_ui._refresh_btn.click( partial(self._refresh_func, func_key), [self._session_key], - [update_gallery_obj, status_text], + [update_gallery_obj, job_ui._status_text], queue=False ) - if stop_btn: - stop_btn.variant = 'secondary' - stop_btn.click( + if job_ui._stop_btn: + job_ui._stop_btn.variant = 'secondary' + job_ui._stop_btn.click( partial(self._stop_wrapped_func, func_key), [self._session_key], - [status_text], + [job_ui._status_text], + queue=False + ) + + if job_ui._active_image and job_ui._active_image_refresh_btn: + job_ui._active_image_refresh_btn.click( + partial(self._refresh_cur_iter_func, func_key), + [self._session_key], + [job_ui._active_image, job_ui._status_text], + queue=False + ) + + if job_ui._active_image_stop_btn: + job_ui._active_image_stop_btn.click( + partial(self._stop_cur_iter_func, func_key), + [self._session_key], + [job_ui._active_image, job_ui._status_text], + queue=False + ) + + if job_ui._stop_all_session_btn: + job_ui._stop_all_session_btn.click( + self.stop_all_jobs, [], [], + queue=False + ) + + if job_ui._free_done_sessions_btn: + job_ui._free_done_sessions_btn.click( + self.clear_all_finished_jobs, [], [], queue=False ) @@ -346,7 +437,8 @@ class JobManager: # Since some parameters are optional it makes sense to use the 'dict' return value type, which requires # the Component as a key... so group together the UI components that the event listeners are going to update # to make it easy to append to function calls and outputs - job_ui_params = [refresh_btn, stop_btn, status_text] + job_ui_params = [job_ui._refresh_btn, job_ui._stop_btn, job_ui._status_text, + job_ui._active_image, job_ui._active_image_refresh_btn, job_ui._active_image_stop_btn] job_ui_outputs = [comp for comp in job_ui_params if comp is not None] # Here a chain is constructed that will make a 'pre' call, a 'run' call, and a 'post' call, @@ -375,27 +467,50 @@ class JobManager: queue=False ) - # Now replace the original function with one that creates a JobInfo and triggers the dummy obj + # Add any components that we want the runtime values for + added_inputs = [self._session_key, job_ui._rec_steps_checkbox, job_ui._save_rec_steps_to_gallery_chkbx, + job_ui._save_rec_steps_to_file_chkbx, job_ui._rec_steps_intrvl_sldr] - def wrapped_func(*inputs): - session_key = inputs[-1] - inputs = inputs[:-1] + # Now replace the original function with one that creates a JobInfo and triggers the dummy obj + def wrapped_func(*wrapped_inputs): + # Remove the added_inputs (pop opposite order of list) + + wrapped_inputs = list(wrapped_inputs) + rec_steps_interval: int = wrapped_inputs.pop() + save_rec_steps_file: bool = wrapped_inputs.pop() + save_rec_steps_grid: bool = wrapped_inputs.pop() + record_steps_enabled: bool = wrapped_inputs.pop() + session_key: str = wrapped_inputs.pop() + job_inputs = tuple(wrapped_inputs) # Get or create a session for this key session_info = self._sessions.setdefault(session_key, SessionInfo()) # Is this session already running this job? if func_key in session_info.jobs: - return {status_text: "This session is already running that function!"} + job_info = session_info.jobs[func_key] + # If the job seems stuck in 'starting' then go ahead and toss it + if not job_info.started and time.time() > job_info.timestamp + JobManager.JOB_MAX_START_TIME: + job_info.should_stop.set() + job_info.stop_cur_iter.set() + session_info.jobs.pop(func_key) + return {job_ui._status_text: "Canceled possibly hung job. Try again"} + return {job_ui._status_text: "This session is already running that function!"} + + # Is this a new run of a previously finished job? Clear old info + if func_key in session_info.finished_jobs: + session_info.finished_jobs.pop(func_key) job_token = self._get_job_token(block=False) - job = JobInfo(inputs=inputs, func=func, removed_output_idxs=removed_idxs, session_key=session_key, - job_token=job_token) + job = JobInfo( + inputs=job_inputs, func=func, removed_output_idxs=removed_idxs, session_key=session_key, + job_token=job_token, rec_steps_enabled=record_steps_enabled, rec_steps_intrvl=rec_steps_interval, + rec_steps_to_gallery=save_rec_steps_grid, rec_steps_to_file=save_rec_steps_file, timestamp=time.time()) session_info.jobs[func_key] = job ret = {pre_call_dummyobj: triggerChangeEvent()} if job_token is None: - ret[status_text] = "Job is queued" + ret[job_ui._status_text] = "Job is queued" return ret - return wrapped_func, inputs, [pre_call_dummyobj, status_text] + return wrapped_func, inputs + added_inputs, [pre_call_dummyobj, job_ui._status_text] diff --git a/scripts/webui.py b/scripts/webui.py index eea8e1b..b8d9574 100644 --- a/scripts/webui.py +++ b/scripts/webui.py @@ -74,6 +74,7 @@ import copy from typing import List, Union, Dict, Callable, Any, Optional from pathlib import Path from collections import namedtuple +from functools import partial # tell the user which GPU the code is actually using if os.getenv("SD_WEBUI_DEBUG", 'False').lower() in ('true', '1', 'y'): @@ -307,15 +308,21 @@ class KDiffusionSampler: self.schedule = sampler def get_sampler_name(self): return self.schedule - def sample(self, S, conditioning, batch_size, shape, verbose, unconditional_guidance_scale, unconditional_conditioning, eta, x_T): + def sample(self, S, conditioning, batch_size, shape, verbose, unconditional_guidance_scale, unconditional_conditioning, eta, x_T, img_callback: Callable = None ): sigmas = self.model_wrap.get_sigmas(S) x = x_T * sigmas[0] model_wrap_cfg = CFGDenoiser(self.model_wrap) - samples_ddim = K.sampling.__dict__[f'sample_{self.schedule}'](model_wrap_cfg, x, sigmas, extra_args={'cond': conditioning, 'uncond': unconditional_conditioning, 'cond_scale': unconditional_guidance_scale}, disable=False) + samples_ddim = K.sampling.__dict__[f'sample_{self.schedule}'](model_wrap_cfg, x, sigmas, extra_args={'cond': conditioning, 'uncond': unconditional_conditioning, 'cond_scale': unconditional_guidance_scale}, disable=False, callback=partial(KDiffusionSampler.img_callback_wrapper, img_callback)) return samples_ddim, None + @classmethod + def img_callback_wrapper(cls, callback: Callable, *args): + ''' Converts a KDiffusion callback to the standard img_callback ''' + if callback: + arg_dict = args[0] + callback(image_sample=arg_dict['denoised'], iter_num=arg_dict['i']) def create_random_tensors(shape, seeds): xs = [] @@ -1003,6 +1010,7 @@ def process_images( if job_info: job_info.job_status = f"Processing Iteration {n+1}/{n_iter}. Batch size {batch_size}" + job_info.rec_steps_imgs.clear() for idx,(p,s) in enumerate(zip(prompts,seeds)): job_info.job_status += f"\nItem {idx}: Seed {s}\nPrompt: {p}" print(f"Current prompt: {p}") @@ -1057,7 +1065,78 @@ def process_images( # finally, slerp base_x noise to target_x noise for creating a variant x = slerp(device, max(0.0, min(1.0, cur_variant_amount)), base_x, target_x) - samples_ddim = func_sample(init_data=init_data, x=x, conditioning=c, unconditional_conditioning=uc, sampler_name=sampler_name) + # If optimized then use first stage for preview and store it on cpu until needed + if opt.optimized: + step_preview_model = modelFS + step_preview_model.cpu() + else: + step_preview_model = model + + def sample_iteration_callback(image_sample: torch.Tensor, iter_num: int): + ''' Called from the sampler every iteration ''' + if job_info: + job_info.active_iteration_cnt = iter_num + record_periodic_image = job_info.rec_steps_enabled and (0 == iter_num % job_info.rec_steps_intrvl) + if record_periodic_image or job_info.refresh_active_image_requested.is_set(): + preview_start_time = time.time() + if opt.optimized: + step_preview_model.to(device) + + decoded_batch: List[torch.Tensor] = [] + # Break up batch to save VRAM + for sample in image_sample: + sample = sample[None, :] # expands the tensor as if it still had a batch dimension + decoded_sample = step_preview_model.decode_first_stage(sample)[0] + decoded_sample = torch.clamp((decoded_sample + 1.0) / 2.0, min=0.0, max=1.0) + decoded_sample = decoded_sample.cpu() + decoded_batch.append(decoded_sample) + + batch_size = len(decoded_batch) + + if opt.optimized: + step_preview_model.cpu() + + images: List[Image.Image] = [] + # Convert tensor to image (copied from code below) + for ddim in decoded_batch: + x_sample = 255. * rearrange(ddim.numpy(), 'c h w -> h w c') + x_sample = x_sample.astype(np.uint8) + image = Image.fromarray(x_sample) + images.append(image) + + caption = f"Iter {iter_num}" + grid = image_grid(images, len(images), force_n_rows=1, captions=[caption]*len(images)) + + # Save the images if recording steps, and append existing saved steps + if job_info.rec_steps_enabled: + gallery_img_size = tuple(int(0.25*dim) for dim in images[0].size) + job_info.rec_steps_imgs.append(grid.resize(gallery_img_size)) + + # Notify the requester that the image is updated + if job_info.refresh_active_image_requested.is_set(): + if job_info.rec_steps_enabled: + grid_rows = None if batch_size == 1 else len(job_info.rec_steps_imgs) + grid = image_grid(imgs=job_info.rec_steps_imgs[::-1], batch_size=1, force_n_rows=grid_rows) + job_info.active_image = grid + job_info.refresh_active_image_done.set() + job_info.refresh_active_image_requested.clear() + + preview_elapsed_timed = time.time() - preview_start_time + if preview_elapsed_timed / job_info.rec_steps_intrvl > 1: + print( + f"Warning: Preview generation is slowing image generation. It took {preview_elapsed_timed:.2f}s to generate progress images for batch of {batch_size} images!") + + # Interrupt current iteration? + if job_info.stop_cur_iter.is_set(): + job_info.stop_cur_iter.clear() + raise StopIteration() + + try: + samples_ddim = func_sample(init_data=init_data, x=x, conditioning=c, unconditional_conditioning=uc, sampler_name=sampler_name, img_callback=sample_iteration_callback) + except StopIteration: + print("Skipping iteration") + job_info.job_status = "Skipping iteration" + continue if opt.optimized: modelFS.to(device) @@ -1196,6 +1275,19 @@ skip_grid, sort_samples, sampler_name, ddim_eta, n_iter, batch_size, i, denoisin if simple_templating: grid_captions.append( captions[i] ) + # Save the progress images? + if job_info: + if job_info.rec_steps_enabled and (job_info.rec_steps_to_file or job_info.rec_steps_to_gallery): + steps_grid = image_grid(job_info.rec_steps_imgs, 1) + if job_info.rec_steps_to_gallery: + gallery_img_size = tuple(2*dim for dim in image.size) + output_images.append( steps_grid.resize( gallery_img_size ) ) + if job_info.rec_steps_to_file: + steps_grid_filename = f"{original_filename}_step_grid" + save_sample(steps_grid, sample_path_i, steps_grid_filename, jpg_sample, prompts, seeds, width, height, steps, cfg_scale, + normalize_prompt_weights, use_GFPGAN, write_info_files, write_sample_info_to_log_file, prompt_matrix, init_img, uses_loopback, uses_random_seed_loopback, skip_save, + skip_grid, sort_samples, sampler_name, ddim_eta, n_iter, batch_size, i, denoising_strength, resize_mode, False) + if opt.optimized: mem = torch.cuda.memory_allocated()/1e6 modelFS.to("cpu") @@ -1215,7 +1307,7 @@ skip_grid, sort_samples, sampler_name, ddim_eta, n_iter, batch_size, i, denoisin import traceback print("Error creating prompt_matrix text:", file=sys.stderr) print(traceback.format_exc(), file=sys.stderr) - elif batch_size > 1 or n_iter > 1: + elif len(output_images) > 0 and (batch_size > 1 or n_iter > 1): grid = image_grid(output_images, batch_size) if grid is not None: grid_count = get_next_sequence_number(outpath, 'grid-') @@ -1308,8 +1400,8 @@ def txt2img(prompt: str, ddim_steps: int, sampler_name: str, toggles: List[int], def init(): pass - def sample(init_data, x, conditioning, unconditional_conditioning, sampler_name): - samples_ddim, _ = sampler.sample(S=ddim_steps, conditioning=conditioning, batch_size=int(x.shape[0]), shape=x[0].shape, verbose=False, unconditional_guidance_scale=cfg_scale, unconditional_conditioning=unconditional_conditioning, eta=ddim_eta, x_T=x) + def sample(init_data, x, conditioning, unconditional_conditioning, sampler_name, img_callback: Callable = None): + samples_ddim, _ = sampler.sample(S=ddim_steps, conditioning=conditioning, batch_size=int(x.shape[0]), shape=x[0].shape, verbose=False, unconditional_guidance_scale=cfg_scale, unconditional_conditioning=unconditional_conditioning, eta=ddim_eta, x_T=x, img_callback=img_callback) return samples_ddim try: @@ -1573,7 +1665,7 @@ def img2img(prompt: str, image_editor_mode: str, mask_mode: str, mask_blur_stren return init_latent, mask, - def sample(init_data, x, conditioning, unconditional_conditioning, sampler_name): + def sample(init_data, x, conditioning, unconditional_conditioning, sampler_name, img_callback: Callable = None): t_enc_steps = t_enc obliterate = False if ddim_steps == t_enc_steps: @@ -1595,7 +1687,7 @@ def img2img(prompt: str, image_editor_mode: str, mask_mode: str, mask_blur_stren sigma_sched = sigmas[ddim_steps - t_enc_steps - 1:] model_wrap_cfg = CFGMaskedDenoiser(sampler.model_wrap) - samples_ddim = K.sampling.__dict__[f'sample_{sampler.get_sampler_name()}'](model_wrap_cfg, xi, sigma_sched, extra_args={'cond': conditioning, 'uncond': unconditional_conditioning, 'cond_scale': cfg_scale, 'mask': z_mask, 'x0': x0, 'xi': xi}, disable=False) + samples_ddim = K.sampling.__dict__[f'sample_{sampler.get_sampler_name()}'](model_wrap_cfg, xi, sigma_sched, extra_args={'cond': conditioning, 'uncond': unconditional_conditioning, 'cond_scale': cfg_scale, 'mask': z_mask, 'x0': x0, 'xi': xi}, disable=False, callback=partial(KDiffusionSampler.img_callback_wrapper, img_callback)) else: x0, z_mask = init_data