Public/automatic
1
0
Fork 0
mirror of https://github.com/vladmandic/automatic synced 2026-09-19 01:04:32 +02:00
Code Issues Packages Projects Releases Wiki Activity

cleanup/refactor state history

Browse Source 
Signed-off-by: Vladimir Mandic <mandic00@live.com>
This commit is contained in:
Vladimir Mandic
2025-09-12 16:12:43 -04:00
parent a8b850adf4
commit 175e9cbe29
41 changed files with 172 additions and 171 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/api/control.py
+3 -4
View File
@@ -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 []
modules/api/generate.py
+4 -4
View File
@@ -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:
modules/api/models.py
+2 -3
View File
@@ -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):
modules/api/process.py
+6 -6
View File
@@ -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):
modules/call_queue.py
+2 -2
View File
@@ -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"<div class='error'>{html.escape(type(e).__name__+': '+str(e))}</div>"]
shared.state.end()
shared.state.end(jobid)
if not add_stats:
return tuple(res)
elapsed = time.perf_counter() - t
modules/civitai/download_civitai.py
+5 -4
View File
@@ -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:
modules/extras.py
+5 -5
View File
@@ -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)
modules/hashes.py
+2 -5
View File
@@ -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
modules/images.py
+3 -1
View File
@@ -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'):
modules/interrogate/openclip.py
+4 -4
View File
@@ -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)
modules/interrogate/vqa.py
+4 -5
View File
@@ -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)
modules/lora/extra_networks_lora.py
+4 -6
View File
@@ -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)
modules/lora/lora_extract.py
+2 -2
View File
@@ -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:
modules/ltx/ltx_util.py
-2
View File
@@ -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
modules/modelloader.py
+4 -2
View File
@@ -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
modules/postprocess/yolo.py
+2 -3
View File
@@ -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)
modules/processing.py
+5 -5
View File
@@ -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
modules/processing_args.py
+2
View File
@@ -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
modules/processing_diffusers.py
+10 -7
View File
@@ -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')
modules/processing_helpers.py
+5 -3
View File
@@ -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
modules/processing_vae.py
+5 -6
View File
@@ -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
modules/scripts_postprocessing.py
+4 -2
View File
@@ -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)
modules/sd_hijack_te.py
+2 -2
View File
@@ -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
modules/sd_hijack_vae.py
+4 -4
View File
@@ -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
modules/sd_models.py
+4 -4
View File
@@ -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
modules/shared_state.py
+58 -54
View File
@@ -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}')
modules/ui_control.py
+2 -1
View File
@@ -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):
modules/ui_history.py
+3 -4
View File
@@ -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,
modules/upscaler.py
+2 -4
View File
@@ -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
modules/video_models/video_load.py
+2 -2
View File
@@ -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