From 175e9cbe29e993846676afc3e0cf7061d65a58ce Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Fri, 12 Sep 2025 16:12:43 -0400 Subject: [PATCH] cleanup/refactor state history Signed-off-by: Vladimir Mandic --- modules/api/control.py | 7 +- modules/api/generate.py | 8 +- modules/api/models.py | 5 +- modules/api/process.py | 12 +-- modules/call_queue.py | 4 +- modules/civitai/download_civitai.py | 9 ++- modules/extras.py | 10 +-- modules/hashes.py | 7 +- modules/images.py | 4 +- modules/interrogate/openclip.py | 8 +- modules/interrogate/vqa.py | 9 +-- modules/lora/extra_networks_lora.py | 10 +-- modules/lora/lora_extract.py | 4 +- modules/ltx/ltx_util.py | 2 - modules/modelloader.py | 6 +- modules/postprocess/yolo.py | 5 +- modules/processing.py | 10 +-- modules/processing_args.py | 2 + modules/processing_diffusers.py | 17 +++-- modules/processing_helpers.py | 8 +- modules/processing_vae.py | 11 ++- modules/scripts_postprocessing.py | 6 +- modules/sd_hijack_te.py | 4 +- modules/sd_hijack_vae.py | 8 +- modules/sd_models.py | 8 +- modules/shared_state.py | 112 ++++++++++++++-------------- modules/ui_control.py | 3 +- modules/ui_history.py | 7 +- modules/upscaler.py | 6 +- modules/video_models/video_load.py | 4 +- scripts/cogvideo.py | 2 - scripts/loopback.py | 1 - scripts/outpainting_mk_2.py | 1 - scripts/poor_mans_outpainting.py | 1 - scripts/prompt_enhance.py | 4 +- scripts/prompts_from_file.py | 1 - scripts/pulid_ext.py | 4 +- scripts/sd_upscale.py | 1 - scripts/xyz_grid.py | 4 +- scripts/xyz_grid_on.py | 2 + webui.py | 6 +- 41 files changed, 172 insertions(+), 171 deletions(-) diff --git a/modules/api/control.py b/modules/api/control.py index 9d6130d93..1f74c9ba0 100644 --- a/modules/api/control.py +++ b/modules/api/control.py @@ -1,8 +1,7 @@ -from typing import Optional, List, Union +from typing import Optional, List from threading import Lock from pydantic import BaseModel, Field # pylint: disable=no-name-in-module from modules import errors, shared, processing_helpers -from modules.processing import StableDiffusionProcessingControl from modules.api import models, helpers from modules.control import run @@ -180,7 +179,7 @@ class APIControl(): # run with self.queue_lock: - shared.state.begin('API-CTL', api=True) + jobid = shared.state.begin('API-CTL', api=True) output_images = [] output_processed = [] output_info = '' @@ -198,7 +197,7 @@ class APIControl(): output_info += item else: pass - shared.state.end(api=False) + shared.state.end(jobid) # return b64images = list(map(helpers.encode_pil_to_base64, output_images)) if send_images else [] diff --git a/modules/api/generate.py b/modules/api/generate.py index 1d589eb52..f03a3a360 100644 --- a/modules/api/generate.py +++ b/modules/api/generate.py @@ -109,7 +109,7 @@ class APIGenerate(): p.outpath_samples = shared.opts.outdir_samples or shared.opts.outdir_txt2img_samples for key, value in getattr(txt2imgreq, "extra", {}).items(): setattr(p, key, value) - shared.state.begin('API TXT', api=True) + jobid = shared.state.begin('API-TXT', api=True) script_args = script.init_script_args(p, txt2imgreq, self.default_script_arg_txt2img, selectable_scripts, selectable_script_idx, script_runner) p.script_args = tuple(script_args) # Need to pass args as tuple here if selectable_scripts is not None: @@ -118,7 +118,7 @@ class APIGenerate(): processed = process_images(p) processed = scripts_manager.scripts_txt2img.after(p, processed, *script_args) p.close() - shared.state.end(api=False) + shared.state.end(jobid) if processed is None or processed.images is None or len(processed.images) == 0: b64images = [] else: @@ -161,7 +161,7 @@ class APIGenerate(): p.outpath_samples = shared.opts.outdir_img2img_samples for key, value in getattr(img2imgreq, "extra", {}).items(): setattr(p, key, value) - shared.state.begin('API-IMG', api=True) + jobid = shared.state.begin('API-IMG', api=True) script_args = script.init_script_args(p, img2imgreq, self.default_script_arg_img2img, selectable_scripts, selectable_script_idx, script_runner) p.script_args = tuple(script_args) # Need to pass args as tuple here if selectable_scripts is not None: @@ -170,7 +170,7 @@ class APIGenerate(): processed = process_images(p) processed = scripts_manager.scripts_img2img.after(p, processed, *script_args) p.close() - shared.state.end(api=False) + shared.state.end(jobid) if processed is None or processed.images is None or len(processed.images) == 0: b64images = [] else: diff --git a/modules/api/models.py b/modules/api/models.py index 8245eca98..94dd9e3a5 100644 --- a/modules/api/models.py +++ b/modules/api/models.py @@ -339,9 +339,8 @@ class ResProgress(BaseModel): class ResHistory(BaseModel): id: Union[int, str, None] = Field(title="ID", description="Task ID") job: str = Field(title="Job", description="Job name") - op: str = Field(title="Operation", description="Operation name") - start: Union[float, None] = Field(title="Start", description="Start time") - end: Union[float, None] = Field(title="End", description="End time") + op: str = Field(title="Operation", description="Job state") + timestamp: Union[float, None] = Field(title="Timestamp", description="Job timestamp") outputs: List[str] = Field(title="Outputs", description="List of filenames") class ResStatus(BaseModel): diff --git a/modules/api/process.py b/modules/api/process.py index f1b3cd97b..3151907a5 100644 --- a/modules/api/process.py +++ b/modules/api/process.py @@ -77,10 +77,10 @@ class APIProcess(): for k, v in req.params.items(): if k not in processors.config[processor.processor_id]['params']: return JSONResponse(status_code=400, content={"error": f"Processor invalid parameter: id={req.model} {k}={v}"}) - shared.state.begin('API-PRE', api=True) + jobid = shared.state.begin('API-PRE', api=True) processed = processor(image, local_config=req.params) image = encode_pil_to_base64(processed) - shared.state.end(api=False) + shared.state.end(jobid) return ResPreprocess(model=processor.processor_id, image=image) def get_mask(self): @@ -103,10 +103,10 @@ class APIProcess(): return JSONResponse(status_code=400, content={"error": f"Mask invalid parameter: {k}={v}"}) else: setattr(masking.opts, k, v) - shared.state.begin('API-MASK', api=True) + jobid = shared.state.begin('API-MASK', api=True) with self.queue_lock: processed = masking.run_mask(input_image=image, input_mask=mask, return_type=req.type) - shared.state.end(api=False) + shared.state.end(jobid) if processed is None: return JSONResponse(status_code=400, content={"error": "Mask is none"}) image = encode_pil_to_base64(processed) @@ -115,7 +115,7 @@ class APIProcess(): def post_detect(self, req: ReqFace): from modules.shared import yolo # pylint: disable=no-name-in-module image = decode_base64_to_image(req.image) - shared.state.begin('API-FACE', api=True) + jobid = shared.state.begin('API-FACE', api=True) images = [] scores = [] classes = [] @@ -129,7 +129,7 @@ class APIProcess(): classes.append(item.cls) labels.append(item.label) boxes.append(item.box) - shared.state.end(api=False) + shared.state.end(jobid) return ResFace(classes=classes, labels=labels, scores=scores, boxes=boxes, images=images) def post_prompt_enhance(self, req: models.ReqPromptEnhance): diff --git a/modules/call_queue.py b/modules/call_queue.py index 694d0cd5b..38b372ca2 100644 --- a/modules/call_queue.py +++ b/modules/call_queue.py @@ -50,7 +50,7 @@ def wrap_gradio_call(func, extra_outputs=None, add_stats=False, name=None): task_id = args[0] else: task_id = 0 - shared.state.begin(job_name, task_id=task_id) + jobid = shared.state.begin(job_name, task_id=task_id) try: if shared.cmd_opts.profile: pr = cProfile.Profile() @@ -70,7 +70,7 @@ def wrap_gradio_call(func, extra_outputs=None, add_stats=False, name=None): if extra_outputs_array is None: extra_outputs_array = [None, ''] res = extra_outputs_array + [f"
{html.escape(type(e).__name__+': '+str(e))}
"] - shared.state.end() + shared.state.end(jobid) if not add_stats: return tuple(res) elapsed = time.perf_counter() - t diff --git a/modules/civitai/download_civitai.py b/modules/civitai/download_civitai.py index 932f8ce7e..a767d9c9f 100644 --- a/modules/civitai/download_civitai.py +++ b/modules/civitai/download_civitai.py @@ -62,7 +62,7 @@ def download_civit_preview(model_path: str, preview_url: str): block_size = 16384 # 16KB blocks written = 0 img = None - shared.state.begin('CivitAI') + jobid = shared.state.begin('Download') if pbar is None: pbar = p.Progress(p.TextColumn('[cyan]Download'), p.DownloadColumn(), p.BarColumn(), p.TaskProgressColumn(), p.TimeRemainingColumn(), p.TimeElapsedColumn(), p.TransferSpeedColumn(), p.TextColumn('[yellow]{task.description}'), console=shared.console) try: @@ -82,8 +82,9 @@ def download_civit_preview(model_path: str, preview_url: str): img = Image.open(preview_file) except Exception as e: shared.log.error(f'CivitAI download error: url={preview_url} file="{preview_file}" written={written} {e}') + shared.state.end(jobid) return 500, '', str(e) - shared.state.end() + shared.state.end(jobid) if img is None: return 500, '', 'image is none' shared.log.info(f'CivitAI download: url={preview_url} file="{preview_file}" size={total_size} image={img.size}') @@ -138,7 +139,7 @@ def download_civit_model_thread(model_name: str, model_url: str, model_path: str res += f' size={round((starting_pos + total_size)/1024/1024, 2)}Mb' shared.log.info(res) - shared.state.begin('CivitAI') + jobid = shared.state.begin('Download') block_size = 16384 # 16KB blocks written = starting_pos global pbar # pylint: disable=global-statement @@ -171,7 +172,7 @@ def download_civit_model_thread(model_name: str, model_url: str, model_path: str elif os.path.exists(temp_file): shared.log.debug(f'Model download complete: temp="{temp_file}" path="{model_file}"') os.rename(temp_file, model_file) - shared.state.end() + shared.state.end(jobid) if os.path.exists(model_file): return model_file else: diff --git a/modules/extras.py b/modules/extras.py index 9d5c706d2..2914d6c1e 100644 --- a/modules/extras.py +++ b/modules/extras.py @@ -30,12 +30,12 @@ def to_half(tensor, enable): def run_modelmerger(id_task, **kwargs): # pylint: disable=unused-argument - shared.state.begin('Merge') + jobid = shared.state.begin('Merge') t0 = time.time() def fail(message): shared.state.textinfo = message - shared.state.end() + shared.state.end(jobid) return [*[gr.update() for _ in range(4)], message] kwargs["models"] = { @@ -177,7 +177,7 @@ def run_modelmerger(id_task, **kwargs): # pylint: disable=unused-argument if created_model: created_model.calculate_shorthash() devices.torch_gc(force=True, reason='merge') - shared.state.end() + shared.state.end(jobid) return [*[gr.Dropdown.update(choices=sd_models.checkpoint_titles()) for _ in range(4)], f"Model saved to {output_modelname}"] @@ -209,7 +209,7 @@ def run_model_modules(model_type:str, model_name:str, custom_name:str, yield msg("input model not found", err=True) return fn = checkpoint_info.filename - shared.state.begin('Merge') + jobid = shared.state.begin('Merge') yield msg("modules merge starting") yield msg("unload current model") sd_models.unload_model_weights(op='model') @@ -257,4 +257,4 @@ def run_model_modules(model_type:str, model_name:str, custom_name:str, sd_models.set_diffuser_options(shared.sd_model, offload=False) sd_models.set_diffuser_offload(shared.sd_model) yield msg("pipeline loaded") - shared.state.end() + shared.state.end(jobid) diff --git a/modules/hashes.py b/modules/hashes.py index a003f4840..2be5eff0a 100644 --- a/modules/hashes.py +++ b/modules/hashes.py @@ -1,4 +1,3 @@ -import copy import hashlib import os.path from rich import progress, errors @@ -75,8 +74,7 @@ def sha256(filename, title, use_addnet_hash=False): return None if not os.path.isfile(filename): return None - orig_state = copy.deepcopy(shared.state) - shared.state.begin("Hash") + jobid = shared.state.begin("Hash") if use_addnet_hash: if progress_ok: try: @@ -94,8 +92,7 @@ def sha256(filename, title, use_addnet_hash=False): "mtime": os.path.getmtime(filename), "sha256": sha256_value } - shared.state.end() - shared.state = orig_state + shared.state.end(jobid) dump_cache() return sha256_value diff --git a/modules/images.py b/modules/images.py index 1b8eb6cce..8fd0e04f8 100644 --- a/modules/images.py +++ b/modules/images.py @@ -46,6 +46,7 @@ def atomically_save_image(): Image.MAX_IMAGE_PIXELS = None # disable check in Pillow and rely on check below to allow large custom image sizes while True: image, filename, extension, params, exifinfo, filename_txt, is_grid = save_queue.get() + jobid = shared.state.begin('Save') shared.state.image_history += 1 if len(exifinfo) > 2: with open(paths.params_path, "w", encoding="utf8") as file: @@ -126,6 +127,8 @@ def atomically_save_image(): entries.append(entry) shared.writefile(entries, fn, mode='w', silent=True) shared.log.info(f'Save: json="{fn}" records={len(entries)}') + shared.state.outputs(filename) + shared.state.end(jobid) save_queue.task_done() @@ -206,7 +209,6 @@ def save_image(image, exifinfo += params.pnginfo.get(pnginfo_section_name, '') filename, extension = os.path.splitext(params.filename) filename_txt = f"{filename}.txt" if shared.opts.save_txt and len(exifinfo) > 0 else None - shared.state.outputs(params.filename) save_queue.put((params.image, filename, extension, params, exifinfo, filename_txt, grid)) # actual save is executed in a thread that polls data from queue save_queue.join() if not hasattr(params.image, 'already_saved_as'): diff --git a/modules/interrogate/openclip.py b/modules/interrogate/openclip.py index 0995ad036..aa8c907ce 100644 --- a/modules/interrogate/openclip.py +++ b/modules/interrogate/openclip.py @@ -135,7 +135,7 @@ def interrogate(image, mode, caption=None): def interrogate_image(image, clip_model, blip_model, mode): - shared.state.begin('Interrogate') + jobid = shared.state.begin('Interrogate') try: if shared.sd_loaded: from modules.sd_models import apply_balanced_offload # prevent circular import @@ -148,7 +148,7 @@ def interrogate_image(image, clip_model, blip_model, mode): prompt = f"Exception {type(e)}" shared.log.error(f'Interrogate: {e}') errors.display(e, 'Interrogate') - shared.state.end() + shared.state.end(jobid) return prompt @@ -164,7 +164,7 @@ def interrogate_batch(batch_files, batch_folder, batch_str, clip_model, blip_mod if len(files) == 0: shared.log.warning('Interrogate batch: type=clip no images') return '' - shared.state.begin('Interrogate batch') + jobid = shared.state.begin('Interrogate batch') prompts = [] load_interrogator(clip_model, blip_model) @@ -191,7 +191,7 @@ def interrogate_batch(batch_files, batch_folder, batch_str, clip_model, blip_mod writer.close() ci.config.quiet = False unload_clip_model() - shared.state.end() + shared.state.end(jobid) return '\n\n'.join(prompts) diff --git a/modules/interrogate/vqa.py b/modules/interrogate/vqa.py index 62bef069b..958eebdd5 100644 --- a/modules/interrogate/vqa.py +++ b/modules/interrogate/vqa.py @@ -605,8 +605,7 @@ def sa2(question: str, image: Image.Image, repo: str = None): def interrogate(question:str='', system_prompt:str=None, prompt:str=None, image:Image.Image=None, model_name:str=None, quiet:bool=False): global quant_args # pylint: disable=global-statement - if not quiet: - shared.state.begin('Interrogate') + jobid = shared.state.begin('Interrogate') t0 = time.time() quant_args = model_quant.create_config(module='LLM') model_name = model_name or shared.opts.interrogate_vlm_model @@ -691,7 +690,7 @@ def interrogate(question:str='', system_prompt:str=None, prompt:str=None, image: t1 = time.time() if not quiet: shared.log.debug(f'Interrogate: type=vlm model="{model_name}" repo="{vqa_model}" args={get_kwargs()} time={t1-t0:.2f}') - shared.state.end() + shared.state.end(jobid) return answer @@ -725,7 +724,7 @@ def batch(model_name, system_prompt, batch_files, batch_folder, batch_str, quest if len(files) == 0: shared.log.warning('Interrogate batch: type=vlm no images') return '' - shared.state.begin('Interrogate batch') + jobid = shared.state.begin('Interrogate batch') prompts = [] if write: mode = 'w' if not append else 'a' @@ -751,5 +750,5 @@ def batch(model_name, system_prompt, batch_files, batch_folder, batch_str, quest if write: writer.close() shared.opts.interrogate_offload = orig_offload - shared.state.end() + shared.state.end(jobid) return '\n\n'.join(prompts) diff --git a/modules/lora/extra_networks_lora.py b/modules/lora/extra_networks_lora.py index ee60f30de..d18444ea0 100644 --- a/modules/lora/extra_networks_lora.py +++ b/modules/lora/extra_networks_lora.py @@ -180,11 +180,10 @@ class ExtraNetworkLora(extra_networks.ExtraNetwork): if load_method == 'diffusers': has_changed = False # diffusers handles its own loading if len(exclude) == 0: - job = shared.state.job - shared.state.job = 'LoRA' + jobid = shared.state.begin('LoRA') lora_load.network_load(names, te_multipliers, unet_multipliers, dyn_dims) # load only on first call sd_models.set_diffuser_offload(shared.sd_model, op="model") - shared.state.job = job + shared.state.end(jobid) elif load_method == 'nunchaku': from modules.lora import lora_nunchaku has_changed = lora_nunchaku.load_nunchaku(names, unet_multipliers) @@ -192,16 +191,15 @@ class ExtraNetworkLora(extra_networks.ExtraNetwork): lora_load.network_load(names, te_multipliers, unet_multipliers, dyn_dims) # load has_changed = self.changed(requested, include, exclude) if has_changed: - job = shared.state.job - shared.state.job = 'LoRA' + jobid = shared.state.begin('LoRA') if len(l.previously_loaded_networks) > 0: shared.log.info(f'Network unload: type=LoRA apply={[n.name for n in l.previously_loaded_networks]} mode={"fuse" if shared.opts.lora_fuse_diffusers else "backup"}') networks.network_deactivate(include, exclude) networks.network_activate(include, exclude) if len(exclude) > 0: # only update on last activation l.previously_loaded_networks = l.loaded_networks.copy() - shared.state.job = job debug_log(f'Network load: type=LoRA previous={[n.name for n in l.previously_loaded_networks]} current={[n.name for n in l.loaded_networks]} changed') + shared.state.end(jobid) if len(l.loaded_networks) > 0 and (len(networks.applied_layers) > 0 or load_method=='diffusers' or load_method=='nunchaku') and step == 0: infotext(p) diff --git a/modules/lora/lora_extract.py b/modules/lora/lora_extract.py index 18682804b..25cf4a438 100644 --- a/modules/lora/lora_extract.py +++ b/modules/lora/lora_extract.py @@ -135,7 +135,7 @@ def make_lora(fn, maxrank, auto_rank, rank_ratio, modules, overwrite): maxrank = int(maxrank) rank_ratio = 1 if not auto_rank else rank_ratio shared.log.debug(f'LoRA extract: modules={modules} maxrank={maxrank} auto={auto_rank} ratio={rank_ratio} fn="{fn}"') - shared.state.begin('LoRA extract') + jobid = shared.state.begin('LoRA extract') with rp.Progress(rp.TextColumn('[cyan]LoRA extract'), rp.BarColumn(), rp.TaskProgressColumn(), rp.TimeRemainingColumn(), rp.TimeElapsedColumn(), rp.TextColumn('[cyan]{task.description}'), console=shared.console) as progress: @@ -226,7 +226,7 @@ def make_lora(fn, maxrank, auto_rank, rank_ratio, modules, overwrite): yield msg return - shared.state.end() + shared.state.end(jobid) meta = make_meta(fn, maxrank, rank_ratio) shared.log.debug(f'LoRA metadata: {meta}') try: diff --git a/modules/ltx/ltx_util.py b/modules/ltx/ltx_util.py index 07708b25b..a07434734 100644 --- a/modules/ltx/ltx_util.py +++ b/modules/ltx/ltx_util.py @@ -28,7 +28,6 @@ def load_model(engine: str, model: str): def load_upsample(upsample_pipe, upsample_repo_id): if upsample_pipe is None: t0 = time.time() - shared.state.begin('Load') from diffusers.pipelines.ltx.pipeline_ltx_latent_upsample import LTXLatentUpsamplePipeline shared.log.info(f'Video load: cls={LTXLatentUpsamplePipeline.__class__.__name__} repo="{upsample_repo_id}"') upsample_pipe = LTXLatentUpsamplePipeline.from_pretrained( @@ -37,7 +36,6 @@ def load_upsample(upsample_pipe, upsample_repo_id): cache_dir=shared.opts.hfcache_dir, torch_dtype=devices.dtype, ) - shared.state.end() t1 = time.time() timer.process.add('load', t1 - t0) return upsample_pipe diff --git a/modules/modelloader.py b/modules/modelloader.py index ae3abc7f3..0eaccf1bd 100644 --- a/modules/modelloader.py +++ b/modules/modelloader.py @@ -53,7 +53,7 @@ def download_diffusers_model(hub_id: str, cache_dir: str = None, download_config if hub_id is None or len(hub_id) == 0: return None from diffusers import DiffusionPipeline - shared.state.begin('HuggingFace') + jobid = shared.state.begin('Download') if hub_id.startswith('huggingface/'): hub_id = hub_id.replace('huggingface/', '') if download_config is None: @@ -98,9 +98,11 @@ def download_diffusers_model(hub_id: str, cache_dir: str = None, download_config debug(f'Diffusers download error: id="{hub_id}" {e}') if 'gated' in str(e): shared.log.error(f'Diffusers download error: id="{hub_id}" model access requires login') + shared.state.end(jobid) return None if pipeline_dir is None: shared.log.error(f'Diffusers download error: id="{hub_id}" {err}') + shared.state.end(jobid) return None try: model_info_dict = hf.model_info(hub_id).cardData if pipeline_dir is not None else None @@ -113,7 +115,7 @@ def download_diffusers_model(hub_id: str, cache_dir: str = None, download_config f.write("True") if pipeline_dir is not None: shared.writefile(model_info_dict, os.path.join(pipeline_dir, "model_info.json")) - shared.state.end() + shared.state.end(jobid) return pipeline_dir diff --git a/modules/postprocess/yolo.py b/modules/postprocess/yolo.py index 300f9d957..74c67cf78 100644 --- a/modules/postprocess/yolo.py +++ b/modules/postprocess/yolo.py @@ -332,7 +332,6 @@ class YoloRestorer(Detailer): mask_all = [] p.state = '' - prev_state = shared.state.job pc = copy(p) pc.ops.append('detailer') @@ -348,8 +347,9 @@ class YoloRestorer(Detailer): pc.image_mask = [item.mask] pc.overlay_images = [] pc.recursion = True - shared.state.job = 'Detailer' + jobid = shared.state.begin('Detailer') pp = processing.process_images_inner(pc) + shared.state.end(jobid) del pc.recursion if pp is not None and pp.images is not None and len(pp.images) > 0: image = pp.images[0] # update image to be reused for next item @@ -369,7 +369,6 @@ class YoloRestorer(Detailer): p.image_mask = orig_p.get('image_mask', None) p.state = orig_p.get('state', None) p.ops = orig_p.get('ops', []) - shared.state.job = prev_state shared.opts.data['mask_apply_overlay'] = orig_apply_overlay np_image = np.array(image) diff --git a/modules/processing.py b/modules/processing.py index 28185d6cc..3dea7829c 100644 --- a/modules/processing.py +++ b/modules/processing.py @@ -369,7 +369,6 @@ def process_samples(p: StableDiffusionProcessing, samples): def process_images_inner(p: StableDiffusionProcessing) -> Processed: - """this is the main loop that both txt2img and img2img use; it calls func_init once inside all the scopes and func_sample once per batch""" if type(p.prompt) == list: assert len(p.prompt) > 0 else: @@ -383,7 +382,7 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed: if p.scripts is not None and isinstance(p.scripts, scripts_manager.ScriptRunner): p.scripts.process(p) - shared.state.begin('Process') + jobid = shared.state.begin('Process') shared.state.batch_count = p.n_iter with devices.inference_context(): t0 = time.time() @@ -445,8 +444,9 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed: batch_images, batch_infotexts = process_samples(p, samples) for batch_image, batch_infotext in zip(batch_images, batch_infotexts): - output_images.append(batch_image) - infotexts.append(batch_infotext) + if batch_image is not None and batch_image not in output_images: + output_images.append(batch_image) + infotexts.append(batch_infotext) if shared.cmd_opts.lowvram: devices.torch_gc(force=True, reason='lowvram') @@ -495,5 +495,5 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed: if shared.cmd_opts.lowvram or shared.cmd_opts.medvram: devices.torch_gc(force=True, reason='final') - shared.state.end() + shared.state.end(jobid) return results diff --git a/modules/processing_args.py b/modules/processing_args.py index 91e41003b..a82bf11f3 100644 --- a/modules/processing_args.py +++ b/modules/processing_args.py @@ -168,6 +168,7 @@ def set_pipeline_args(p, model, prompts:list, negative_prompts:list, prompts_2:t 'Chroma' in model.__class__.__name__ or 'HiDreamImagePipeline' in model.__class__.__name__ ): + jobid = shared.state.begin('TE Encode') try: prompt_parser_diffusers.embedder = prompt_parser_diffusers.PromptEmbedder(prompts, negative_prompts, steps, clip_skip, p) parser = shared.opts.prompt_attention @@ -176,6 +177,7 @@ def set_pipeline_args(p, model, prompts:list, negative_prompts:list, prompts_2:t if os.environ.get('SD_PROMPT_DEBUG', None) is not None: errors.display(e, 'Prompt parser encode') timer.process.record('prompt', reset=False) + shared.state.end(jobid) else: prompt_parser_diffusers.embedder = None diff --git a/modules/processing_diffusers.py b/modules/processing_diffusers.py index 642755e1a..3b9099621 100644 --- a/modules/processing_diffusers.py +++ b/modules/processing_diffusers.py @@ -118,7 +118,7 @@ def process_post(p: processing.StableDiffusionProcessing): def process_base(p: processing.StableDiffusionProcessing): - shared.state.begin('Base') + jobid = shared.state.begin('Base') txt2img = is_txt2img() use_refiner_start = is_refiner_enabled(p) and (not p.is_hr_pass) use_denoise_start = not txt2img and p.refiner_start > 0 and p.refiner_start < 1 @@ -164,7 +164,9 @@ def process_base(p: processing.StableDiffusionProcessing): base_args['gate_step'] = p.gate_step output = shared.sd_model.tgate(**base_args) # pylint: disable=not-callable else: + taskid = shared.state.begin('Model') output = shared.sd_model(**base_args) + shared.state.end(taskid) if isinstance(output, dict): output = SimpleNamespace(**output) if isinstance(output, list): @@ -207,8 +209,8 @@ def process_base(p: processing.StableDiffusionProcessing): finally: process_post(p) + shared.state.end(jobid) shared.state.nextjob() - shared.state.end() return output @@ -217,6 +219,7 @@ def process_hires(p: processing.StableDiffusionProcessing, output): if (output is None) or (output.images is None): return output if p.enable_hr: + jobid = shared.state.begin('Hires') p.is_hr_pass = True if hasattr(p, 'init_hr'): p.init_hr(p.hr_scale, p.hr_upscaler, force=p.hr_force) @@ -228,7 +231,6 @@ def process_hires(p: processing.StableDiffusionProcessing, output): p.hr_resize_context = p.resize_context p.hr_upscale_to_x = p.width * p.hr_scale if p.hr_resize_x == 0 else p.hr_resize_x p.hr_upscale_to_y = p.height * p.hr_scale if p.hr_resize_y == 0 else p.hr_resize_y - prev_job = shared.state.job # hires runs on original pipeline if hasattr(shared.sd_model, 'restore_pipeline') and (shared.sd_model.restore_pipeline is not None) and (not shared.opts.control_hires): @@ -240,7 +242,6 @@ def process_hires(p: processing.StableDiffusionProcessing, output): p.ops.append('upscale') if shared.opts.samples_save and not p.do_not_save_samples and shared.opts.save_images_before_highres_fix and hasattr(shared.sd_model, 'vae'): save_intermediate(p, latents=output.images, suffix="-before-hires") - shared.state.update('Upscale', 0, 1) output.images = resize_hires(p, latents=output.images) sd_hijack_hypertile.hypertile_set(p, hr=True) elif torch.is_tensor(output.images) and output.images.shape[-1] == 3: # nhwc @@ -295,7 +296,9 @@ def process_hires(p: processing.StableDiffusionProcessing, output): try: if 'base' in p.skip: extra_networks.activate(p) + taskid = shared.state.begin('Model') output = shared.sd_model(**hires_args) # pylint: disable=not-callable + shared.state.end(taskid) if isinstance(output, dict): output = SimpleNamespace(**output) if hasattr(output, 'images'): @@ -314,7 +317,7 @@ def process_hires(p: processing.StableDiffusionProcessing, output): if orig_image is not None: p.task_args['image'] = orig_image p.denoising_strength = orig_denoise - shared.state.job = prev_job + shared.state.end(jobid) shared.state.nextjob() p.is_hr_pass = False timer.process.record('hires') @@ -326,7 +329,6 @@ def process_refine(p: processing.StableDiffusionProcessing, output): if (output is None) or (output.images is None): return output if is_refiner_enabled(p): - prev_job = shared.state.job if shared.opts.samples_save and not p.do_not_save_samples and shared.opts.save_images_before_refiner and hasattr(shared.sd_model, 'vae'): save_intermediate(p, latents=output.images, suffix="-before-refiner") if shared.opts.diffusers_move_base: @@ -335,6 +337,7 @@ def process_refine(p: processing.StableDiffusionProcessing, output): if shared.state.interrupted or shared.state.skipped: shared.sd_model = orig_pipeline return output + jobid = shared.state.begin('Refine') shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model) if shared.opts.diffusers_move_refiner: sd_models.move_model(shared.sd_refiner, devices.device) @@ -400,7 +403,7 @@ def process_refine(p: processing.StableDiffusionProcessing, output): elif shared.opts.diffusers_move_refiner: shared.log.debug('Moving to CPU: model=refiner') sd_models.move_model(shared.sd_refiner, devices.cpu) - shared.state.job = prev_job + shared.state.end(jobid) shared.state.nextjob() p.is_refiner_pass = False timer.process.record('refine') diff --git a/modules/processing_helpers.py b/modules/processing_helpers.py index e3210ba66..63b38683b 100644 --- a/modules/processing_helpers.py +++ b/modules/processing_helpers.py @@ -205,8 +205,6 @@ def decode_first_stage(model, x): shared.log.debug(f'Decode VAE: skipped={shared.state.skipped} interrupted={shared.state.interrupted}') x_sample = torch.zeros((len(x), 3, x.shape[2] * 8, x.shape[3] * 8), dtype=devices.dtype_vae, device=devices.device) return x_sample - prev_job = shared.state.job - shared.state.job = 'VAE' with devices.autocast(disable = x.dtype==devices.dtype_vae): try: if hasattr(model, 'decode_first_stage'): @@ -220,7 +218,6 @@ def decode_first_stage(model, x): except Exception as e: x_sample = x shared.log.error(f'Decode VAE: {e}') - shared.state.job = prev_job return x_sample @@ -301,17 +298,21 @@ def resize_init_images(p): def resize_hires(p, latents): # input=latents output=pil if not latent_upscaler else latent + jobid = shared.state.begin('Resize') if not torch.is_tensor(latents): shared.log.warning('Hires: input is not tensor') decoded = processing_vae.vae_decode(latents=latents, model=shared.sd_model, vae_type=p.vae_type, output_type='pil', width=p.width, height=p.height) + shared.state.end(jobid) return decoded if (p.hr_upscale_to_x == 0 or p.hr_upscale_to_y == 0) and hasattr(p, 'init_hr'): shared.log.error('Hires: missing upscaling dimensions') + shared.state.end(jobid) return decoded if p.hr_upscaler.lower().startswith('latent'): resized = images.resize_image(p.hr_resize_mode, latents, p.hr_upscale_to_x, p.hr_upscale_to_y, upscaler_name=p.hr_upscaler, context=p.hr_resize_context) + shared.state.end(jobid) return resized decoded = processing_vae.vae_decode(latents=latents, model=shared.sd_model, vae_type=p.vae_type, output_type='pil', width=p.width, height=p.height) @@ -320,6 +321,7 @@ def resize_hires(p, latents): # input=latents output=pil if not latent_upscaler resize = images.resize_image(p.hr_resize_mode, image, p.hr_upscale_to_x, p.hr_upscale_to_y, upscaler_name=p.hr_upscaler, context=p.hr_resize_context) resized.append(resize) devices.torch_gc() + shared.state.end(jobid) return resized diff --git a/modules/processing_vae.py b/modules/processing_vae.py index f4e4f0e61..c97ffa12b 100644 --- a/modules/processing_vae.py +++ b/modules/processing_vae.py @@ -268,17 +268,16 @@ def vae_decode(latents, model, output_type='np', vae_type='Full', width=None, he model = model.pipe if latents is None or not torch.is_tensor(latents): # already decoded return latents - prev_job = shared.state.job if vae_type == 'Remote': - shared.state.job = 'Remote VAE' + jobid = shared.state.begin('Remote VAE') from modules.sd_vae_remote import remote_decode tensors = remote_decode(latents=latents, width=width, height=height) - shared.state.job = prev_job + shared.state.end(jobid) if tensors is not None and len(tensors) > 0: return vae_postprocess(tensors, model, output_type) - - shared.state.job = 'VAE' + + jobid = shared.state.begin('VAE Decode') if latents.shape[0] == 0: shared.log.error(f'VAE nothing to decode: {latents.shape}') return [] @@ -308,11 +307,11 @@ def vae_decode(latents, model, output_type='np', vae_type='Full', width=None, he decoded = 2.0 * decoded - 1.0 # typical normalized range images = vae_postprocess(decoded, model, output_type) - shared.state.job = prev_job if shared.cmd_opts.profile or debug: t1 = time.time() shared.log.debug(f'Profile: VAE decode: {t1-t0:.2f}') devices.torch_gc() + shared.state.end(jobid) return images diff --git a/modules/scripts_postprocessing.py b/modules/scripts_postprocessing.py index 6c9983801..95a370522 100644 --- a/modules/scripts_postprocessing.py +++ b/modules/scripts_postprocessing.py @@ -96,13 +96,14 @@ class ScriptPostprocessingRunner: def run(self, pp: PostprocessedImage, args): for script in self.scripts_in_preferred_order(): - shared.state.job = script.name + jobid = shared.state.begin(script.name) script_args = args[script.args_from:script.args_to] process_args = {} for (name, _component), value in zip(script.controls.items(), script_args): process_args[name] = value shared.log.debug(f'Process: script={script.name} args={process_args}') script.process(pp, **process_args) + shared.state.end(jobid) def create_args_for_run(self, scripts_args): if not self.ui_created: @@ -125,10 +126,11 @@ class ScriptPostprocessingRunner: for script in self.scripts_in_preferred_order(): if not hasattr(script, 'postprocess'): continue - shared.state.job = script.name + jobid = shared.state.begin(script.name) script_args = args[script.args_from:script.args_to] process_args = {} for (name, _component), value in zip(script.controls.items(), script_args): process_args[name] = value shared.log.debug(f'Postprocess: script={script.name} args={process_args}') script.postprocess(filenames, **process_args) + shared.state.end(jobid) diff --git a/modules/sd_hijack_te.py b/modules/sd_hijack_te.py index 82ca4ec5d..d753eaa9b 100644 --- a/modules/sd_hijack_te.py +++ b/modules/sd_hijack_te.py @@ -4,7 +4,7 @@ from modules import shared, errors, timer, sd_models def hijack_encode_prompt(*args, **kwargs): - shared.state.begin('TE') + jobid = shared.state.begin('TE Encode') t0 = time.time() if 'max_sequence_length' in kwargs and kwargs['max_sequence_length'] is not None: kwargs['max_sequence_length'] = max(kwargs['max_sequence_length'], os.environ.get('HIDREAM_MAX_SEQUENCE_LENGTH', 256)) @@ -22,7 +22,7 @@ def hijack_encode_prompt(*args, **kwargs): # if hasattr(shared.sd_model, "maybe_free_model_hooks"): # shared.sd_model.maybe_free_model_hooks() shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model) - shared.state.end() + shared.state.end(jobid) return res diff --git a/modules/sd_hijack_vae.py b/modules/sd_hijack_vae.py index 26c36b347..edd8427e6 100644 --- a/modules/sd_hijack_vae.py +++ b/modules/sd_hijack_vae.py @@ -8,7 +8,7 @@ debug = shared.log.trace if os.environ.get('SD_VIDEO_DEBUG', None) is not None e def hijack_vae_decode(*args, **kwargs): - shared.state.begin('VAE') + jobid = shared.state.begin('VAE Decode') t0 = time.time() res = None shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model, exclude=['vae']) @@ -27,12 +27,12 @@ def hijack_vae_decode(*args, **kwargs): res = None t1 = time.time() timer.process.add('vae', t1-t0) - shared.state.end() + shared.state.end(jobid) return res def hijack_vae_encode(*args, **kwargs): - shared.state.begin('VAE') + jobid = shared.state.begin('VAE Encode') t0 = time.time() res = None shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model, exclude=['vae']) @@ -51,7 +51,7 @@ def hijack_vae_encode(*args, **kwargs): res = None t1 = time.time() timer.process.add('vae', t1-t0) - shared.state.end() + shared.state.end(jobid) return res diff --git a/modules/sd_models.py b/modules/sd_models.py index fd7457e45..9a73eb787 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -1122,8 +1122,7 @@ def reload_model_weights(sd_model=None, info=None, op='model', force=False, revi if checkpoint_info is None: unload_model_weights(op=op) return None - orig_state = copy.deepcopy(shared.state) - shared.state.begin('Load') + jobid = shared.state.begin('Load') if sd_model is None: sd_model = model_data.sd_model if op == 'model' or op == 'dict' else model_data.sd_refiner if sd_model is None: # previous model load failed @@ -1131,6 +1130,7 @@ def reload_model_weights(sd_model=None, info=None, op='model', force=False, revi else: current_checkpoint_info = getattr(sd_model, 'sd_checkpoint_info', None) if current_checkpoint_info is not None and checkpoint_info is not None and current_checkpoint_info.filename == checkpoint_info.filename and not force: + shared.state.end(jobid) return None else: move_model(sd_model, devices.cpu) @@ -1142,14 +1142,14 @@ def reload_model_weights(sd_model=None, info=None, op='model', force=False, revi if sd_model is None or force: sd_model = None load_diffuser(checkpoint_info, op=op, revision=revision) - shared.state.end() - shared.state = orig_state + shared.state.end(jobid) if op == 'model': shared.opts.data["sd_model_checkpoint"] = checkpoint_info.title return model_data.sd_model else: shared.opts.data["sd_model_refiner"] = checkpoint_info.title return model_data.sd_refiner + shared.state.end(jobid) return None # should not be here diff --git a/modules/shared_state.py b/modules/shared_state.py index 616e92cce..53e044727 100644 --- a/modules/shared_state.py +++ b/modules/shared_state.py @@ -12,9 +12,9 @@ debug_history = debug_output or os.environ.get('SD_STATE_HISTORY', None) class State: - job_history = [] - task_history = [] state_history = [] + job_history = 0 + task_history = 0 image_history = 0 latent_history = 0 id = 0 @@ -45,7 +45,6 @@ class State: disable_preview = False preview_job = -1 time_start = None - time_end = None need_restart = False server_start = time.time() oom = False @@ -142,8 +141,14 @@ class State: res.status = 'running' if self.job != '' else 'idle' return res - def history(self, op:str): - job = { 'id': self.id, 'job': self.job.lower(), 'op': op.lower(), 'start': self.time_start, 'end': self.time_end, 'outputs': self.results } + def find(self, task_id:str): + for job in reversed(self.state_history): + if job['id'] == task_id: + return job + return None + + def history(self, op:str, task_id:str=None, results:list=[]): + job = { 'id': task_id or self.id, 'job': self.job.lower(), 'op': op.lower(), 'timestamp': self.time_start, 'outputs': results } self.state_history.append(job) l = len(self.state_history) if l > 10000: @@ -156,6 +161,8 @@ class State: self.results += results else: self.results.append(results) + if len(self.results) > 0: + self.history('output', self.id, results=self.results) def get_id(self, task_id:str=None): if task_id is None or task_id == 0: @@ -165,52 +172,7 @@ class State: match = re.search(r'\((.*?)\)', task_id) return match.group(1) if match else task_id - def begin(self, title="", task_id=0, api=None): - import modules.devices - self.job_history.append(title) - self.total_jobs += 1 - self.current_image = None - self.current_image_sampling_step = 0 - self.current_latent = None - self.current_noise_pred = None - self.current_sigma = None - self.current_sigma_next = None - self.id_live_preview = 0 - self.interrupted = False - self.preview_job = -1 - self.results = [] - self.id = self.get_id(task_id) - self.job = title - self.job_count = 1 # cannot be less than 1 on new job - self.frame_count = 0 - self.batch_no = 0 - self.batch_count = 0 - self.job_no = 0 - self.job_timestamp = datetime.datetime.now().strftime("%Y%m%d%H%M%S") - self.paused = False - self._sampling_step = 0 - self.sampling_steps = 0 - self.skipped = False - self.textinfo = None - self.prediction_type = "epsilon" - self.api = api or self.api - self.time_start = time.time() - self.time_end = None - self.history('begin') - if debug_output: - log.trace(f'State begin: {self}') - modules.devices.torch_gc() - - def end(self, api=None): - import modules.devices - if self.time_start is None: # someone called end before being - # fn = f'{sys._getframe(2).f_code.co_name}:{sys._getframe(1).f_code.co_name}' # pylint: disable=protected-access - # log.debug(f'Access state.end: {fn}') # pylint: disable=protected-access - self.time_start = time.time() - if debug_output: - log.trace(f'State end: {self}') - self.time_end = time.time() - self.history('end') + def clear(self): self.id = '' self.job = '' self.job_count = 0 @@ -220,14 +182,57 @@ class State: self.paused = False self.interrupted = False self.skipped = False + self.results = [] + + def begin(self, title="", task_id=0, api=None): + import modules.devices + self.clear() + self.job_history += 1 + self.total_jobs += 1 + self.current_image = None + self.current_image_sampling_step = 0 + self.current_latent = None + self.current_noise_pred = None + self.current_sigma = None + self.current_sigma_next = None + self.id_live_preview = 0 + self.id = self.get_id(task_id) + self.job = title + self.job_count = 1 # cannot be less than 1 on new job + self.batch_no = 0 + self.batch_count = 0 + self.job_timestamp = datetime.datetime.now().strftime("%Y%m%d%H%M%S") + self._sampling_step = 0 + self.sampling_steps = 0 + self.textinfo = None + self.prediction_type = "epsilon" self.api = api or self.api + self.time_start = time.time() + self.history('begin', self.id) + if debug_output: + log.trace(f'State begin: {self}') + modules.devices.torch_gc() + return self.id + + def end(self, task_id=None): + import modules.devices + if debug_output: + log.trace(f'State end: {self}') + if task_id is not None: + prev_job = self.find(task_id) + if prev_job is not None: + self.id = prev_job['id'] + self.job = prev_job['job'] + self.time_start = time.time() + self.history('end', task_id or self.id) + self.clear() modules.devices.torch_gc() def step(self, step:int=1): self.sampling_step += step def update(self, job:str, steps:int=0, jobs:int=0): - self.task_history.append(job) + self.task_history += 1 # self._sampling_step = 0 if job == 'Ignore': return @@ -237,8 +242,7 @@ class State: else: self.sampling_steps += (steps * jobs) self.job_count += jobs - self.job = job - self.history('update') + # self.job = job if debug_output: log.trace(f'State update: {self} steps={steps} jobs={jobs}') diff --git a/modules/ui_control.py b/modules/ui_control.py index 04c5470a6..b75842aee 100644 --- a/modules/ui_control.py +++ b/modules/ui_control.py @@ -101,6 +101,7 @@ def generate_click(job_id: str, state: str, active_tab: str, *args): with call_queue.queue_lock: yield [None, None, None, None, 'Control: starting', ''] shared.mem_mon.reset() + jobid = shared.state.begin('Control') progress.start_task(job_id) try: t = time.perf_counter() @@ -112,7 +113,7 @@ def generate_click(job_id: str, state: str, active_tab: str, *args): errors.display(e, 'Control') yield [None, None, None, None, f'Control: Exception: {e}', ''] progress.finish_task(job_id) - shared.state.end() + shared.state.end(jobid) def create_ui(_blocks: gr.Blocks=None): diff --git a/modules/ui_history.py b/modules/ui_history.py index 8c89b3b4d..9d642ec40 100644 --- a/modules/ui_history.py +++ b/modules/ui_history.py @@ -6,7 +6,7 @@ from modules import shared def refresh(): def ts(t): try: - return time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(t)) + return time.strftime('%Y-%m-%d %H:%M:%S.%f', time.localtime(t)) except Exception: return '' @@ -16,8 +16,7 @@ def refresh(): item['id'], item['job'], item['op'], - ts(item['start']), - ts(item['end']), + ts(item['timestamp']), len(item['outputs']), ]) shared.log.info(f"History: records={len(items)}") @@ -42,7 +41,7 @@ def create_ui(): with gr.Row(): history_table = gr.DataFrame( value=None, - headers=['ID', 'Job', 'Op', 'Start', 'End', 'Outputs'], + headers=['ID', 'Job', 'Op', 'Timestamp', 'Outputs'], label='History data', show_label=True, interactive=False, diff --git a/modules/upscaler.py b/modules/upscaler.py index 0bf94dea7..d9ae47828 100644 --- a/modules/upscaler.py +++ b/modules/upscaler.py @@ -90,8 +90,7 @@ class Upscaler: return img def upscale(self, img: Image, scale, selected_model: str = None): - orig_state = copy.deepcopy(shared.state) - shared.state.begin('Upscale') + jobid = shared.state.begin('Upscale') self.scale = scale if isinstance(img, Image.Image): dest_w = int(img.width * scale) @@ -111,8 +110,7 @@ class Upscaler: break if img.width != dest_w or img.height != dest_h: img = img.resize((int(dest_w), int(dest_h)), resample=Image.Resampling.LANCZOS) - shared.state.end() - shared.state = orig_state + shared.state.end(jobid) return img @abstractmethod diff --git a/modules/video_models/video_load.py b/modules/video_models/video_load.py index d3bfac729..7d0b9d15d 100644 --- a/modules/video_models/video_load.py +++ b/modules/video_models/video_load.py @@ -8,7 +8,6 @@ loaded_model = None def load_model(selected: models_def.Model): - shared.state.begin('Load') if selected is None: return '' global loaded_model # pylint: disable=global-statement @@ -16,6 +15,7 @@ def load_model(selected: models_def.Model): return '' sd_models.unload_model_weights() t0 = time.time() + jobid = shared.state.begin('Load') video_cache.apply_teacache_patch(selected.dit_cls) @@ -111,5 +111,5 @@ def load_model(selected: models_def.Model): loaded_model = selected.name msg = f'Video load: cls={shared.sd_model.__class__.__name__} model="{selected.name}" time={t1-t0:.2f}' shared.log.info(msg) - shared.state.end() + shared.state.end(jobid) return msg diff --git a/scripts/cogvideo.py b/scripts/cogvideo.py index cf72742df..1e7fbd076 100644 --- a/scripts/cogvideo.py +++ b/scripts/cogvideo.py @@ -184,7 +184,6 @@ class Script(scripts_manager.Script): # auto-executed by the script-callback def run(self, p: processing.StableDiffusionProcessing, model, sampler, frames, guidance, offload, override, video_type, duration, loop, pad, interpolate, image, video): # pylint: disable=arguments-differ, unused-argument - shared.state.begin('CogVideoX') processing.fix_seed(p) p.extra_generation_params['CogVideoX'] = model p.do_not_save_grid = True @@ -206,7 +205,6 @@ class Script(scripts_manager.Script): frames = self.generate(p, model) devices.torch_gc() processed = processing.get_processed(p, images_list=frames) - shared.state.end() return processed # auto-executed by the script-callback diff --git a/scripts/loopback.py b/scripts/loopback.py index ba50d6a43..a359dd83b 100644 --- a/scripts/loopback.py +++ b/scripts/loopback.py @@ -67,7 +67,6 @@ class Script(scripts_manager.Script): p.do_not_save_grid = True if opts.img2img_color_correction: p.color_corrections = initial_color_corrections - state.job = f"loopback iteration {i+1}/{loops} batch {n+1}/{initial_batch_count}" processed = processing.process_images(p) if processed is None: log.error("Loopback: processing output is none") diff --git a/scripts/outpainting_mk_2.py b/scripts/outpainting_mk_2.py index 28f71a221..0d587b2ca 100644 --- a/scripts/outpainting_mk_2.py +++ b/scripts/outpainting_mk_2.py @@ -208,7 +208,6 @@ class Script(scripts_manager.Script): all_processed_images = [] for i in range(batch_count): imgs = [init_img] * batch_size - state.job = f"outpainting batch {i+1}/{batch_count}" if left > 0: imgs = expand(imgs, batch_size, left, is_left=True) if right > 0: diff --git a/scripts/poor_mans_outpainting.py b/scripts/poor_mans_outpainting.py index 209811678..9d9d4fb22 100644 --- a/scripts/poor_mans_outpainting.py +++ b/scripts/poor_mans_outpainting.py @@ -90,7 +90,6 @@ class Script(scripts_manager.Script): p.init_images = [work[i]] p.image_mask = work_mask[i] p.latent_mask = work_latent_mask[i] - state.job = f"outpainting batch {i+1}/{batch_count}" processed = process_images(p) if initial_seed is None: initial_seed = processed.seed diff --git a/scripts/prompt_enhance.py b/scripts/prompt_enhance.py index c6ab184b1..aea0e5730 100644 --- a/scripts/prompt_enhance.py +++ b/scripts/prompt_enhance.py @@ -514,7 +514,7 @@ class Script(scripts_manager.Script): p.negative_prompt = shared.prompt_styles.apply_negative_styles_to_prompt(p.negative_prompt, p.styles) shared.prompt_styles.apply_styles_to_extra(p) p.styles = [] - shared.state.begin('LLM') + jobid = shared.state.begin('LLM') p.prompt = self.enhance( prompt=p.prompt, seed=p.seed, @@ -532,4 +532,4 @@ class Script(scripts_manager.Script): ) timer.process.record('prompt') p.extra_generation_params['LLM'] = llm_model - shared.state.end() + shared.state.end(jobid) diff --git a/scripts/prompts_from_file.py b/scripts/prompts_from_file.py index dea8ced6f..460b39ecf 100644 --- a/scripts/prompts_from_file.py +++ b/scripts/prompts_from_file.py @@ -136,7 +136,6 @@ class Script(scripts_manager.Script): all_negative = [] infotexts = [] for args in jobs: - state.job = f"{state.job_no + 1} out of {state.job_count}" copy_p = copy.copy(p) for k, v in args.items(): setattr(copy_p, k, v) diff --git a/scripts/pulid_ext.py b/scripts/pulid_ext.py index cd6c91005..a92480824 100644 --- a/scripts/pulid_ext.py +++ b/scripts/pulid_ext.py @@ -252,7 +252,7 @@ class Script(scripts_manager.Script): p.seed = processing_helpers.get_fixed_seed(p.seed) if direct: # run pipeline directly - shared.state.begin('PuLID') + jobid = shared.state.begin('PuLID') processing.fix_seed(p) p.prompt = shared.prompt_styles.apply_styles_to_prompt(p.prompt, p.styles) p.negative_prompt = shared.prompt_styles.apply_negative_styles_to_prompt(p.negative_prompt, p.styles) @@ -273,7 +273,7 @@ class Script(scripts_manager.Script): )[0] info = processing.create_infotext(p) processed = processing.get_processed(p, [output], info=info) - shared.state.end() + shared.state.end(jobid) else: # let processing run the pipeline p.task_args['id_embedding'] = id_embedding p.task_args['uncond_id_embedding'] = uncond_id_embedding diff --git a/scripts/sd_upscale.py b/scripts/sd_upscale.py index 0e71ff7d5..1e1e97229 100644 --- a/scripts/sd_upscale.py +++ b/scripts/sd_upscale.py @@ -70,7 +70,6 @@ class Script(scripts_manager.Script): for i in range(batch_count): p.batch_size = batch_size p.init_images = work[i * batch_size:(i + 1) * batch_size] - state.job = f"upscale batch {i+1+n*batch_count}/{state.job_count}" processed = processing.process_images(p) if initial_info is None: initial_info = processed.info diff --git a/scripts/xyz_grid.py b/scripts/xyz_grid.py index 9c7671353..2ec2452da 100644 --- a/scripts/xyz_grid.py +++ b/scripts/xyz_grid.py @@ -167,6 +167,7 @@ class Script(scripts_manager.Script): include_time, include_text, margin_size, create_video, video_type, video_duration, video_loop, video_pad, video_interpolate, ): # pylint: disable=W0221 + jobid = shared.state.begin('XYZ Grid') if not no_fixed_seeds: processing.fix_seed(p) if not shared.opts.return_grid: @@ -348,7 +349,7 @@ class Script(scripts_manager.Script): return processed, t1-t0 with SharedSettingsStackHelper(): - processed = draw_xyz_grid( + processed: processing.Processed = draw_xyz_grid( p, xs=xs, ys=ys, @@ -404,4 +405,5 @@ class Script(scripts_manager.Script): if create_video and video_type != 'None' and not shared.state.interrupted: images.save_video(p, filename=None, images=have_images, video_type=video_type, duration=video_duration, loop=video_loop, pad=video_pad, interpolate=video_interpolate) + shared.state.end(jobid) return processed diff --git a/scripts/xyz_grid_on.py b/scripts/xyz_grid_on.py index bfa0ca662..f5a910f93 100644 --- a/scripts/xyz_grid_on.py +++ b/scripts/xyz_grid_on.py @@ -184,6 +184,7 @@ class Script(scripts_manager.Script): processing.fix_seed(p) if not shared.opts.return_grid: p.batch_size = 1 + jobid = shared.state.begin('XYZ Grid') def process_axis(opt, vals, vals_dropdown): if opt.label == 'Nothing': @@ -430,6 +431,7 @@ class Script(scripts_manager.Script): p.disable_extra_networks = True active = False xyz_results_cache = processed + shared.state.end(jobid) return processed diff --git a/webui.py b/webui.py index d0e9b3d53..3fd7de7f0 100644 --- a/webui.py +++ b/webui.py @@ -138,7 +138,7 @@ def initialize(): # make the program just exit at ctrl+c without waiting for anything def sigint_handler(_sig, _frame): - log.trace(f'State history: uptime={round(time.time() - shared.state.server_start)} jobs={len(shared.state.job_history)} tasks={len(shared.state.task_history)} latents={shared.state.latent_history} images={shared.state.image_history}') + log.trace(f'State history: uptime={round(time.time() - shared.state.server_start)} jobs={shared.state.job_history} tasks={shared.state.task_history} latents={shared.state.latent_history} images={shared.state.image_history}') log.info('Exiting') try: for f in glob.glob("*.lock"): @@ -155,14 +155,14 @@ def load_model(): if not shared.opts.sd_checkpoint_autoload and shared.cmd_opts.ckpt is None: log.info('Model: autoload=False') else: - shared.state.begin('Load') + jobid = shared.state.begin('Load') thread_model = Thread(target=lambda: shared.sd_model) thread_model.start() thread_refiner = Thread(target=lambda: shared.sd_refiner) thread_refiner.start() thread_model.join() thread_refiner.join() - shared.state.end() + shared.state.end(jobid) timer.startup.record("checkpoint") shared.opts.onchange("sd_model_checkpoint", wrap_queued_call(lambda: modules.sd_models.reload_model_weights(op='model')), call=False) shared.opts.onchange("sd_model_refiner", wrap_queued_call(lambda: modules.sd_models.reload_model_weights(op='refiner')), call=False)