separate progress monitoring from live preview, live preview improvements, progress details

Signed-off-by: Vladimir Mandic <mandic00@live.com>

modules/processing_callbacks.py

@@ -61,6 +61,7 @@ def diffusers_callback(pipe, step: int = 0, timestep: int = 0, kwargs: dict | No
             torch.xpu.synchronize(devices.device)
         elif devices.backend in {"cuda", "zluda", "rocm"}:
             torch.cuda.synchronize(devices.device)
+        time.sleep(0.001) # 1ms yield frees GIL for the preview thread
 
     t1 = time.time()
 
@@ -82,7 +83,7 @@ def diffusers_callback(pipe, step: int = 0, timestep: int = 0, kwargs: dict | No
 
     latents = kwargs.get('latents', None)
     if debug:
-        debug_callback(f'Callback: step={step} timestep={timestep} latents={latents.shape if latents is not None else None} kwargs={list(kwargs)}')
+        debug_callback(f'Callback: step={step} timestep={timestep} latents={latents.shape if latents is not None else None} sync={shared.opts.torch_sync} kwargs={list(kwargs)}')
     if shared.state.sampling_steps == 0 and getattr(pipe, 'num_timesteps', 0) > 0:
         shared.state.sampling_steps = pipe.num_timesteps
     shared.state.step()
@@ -91,14 +92,6 @@ def diffusers_callback(pipe, step: int = 0, timestep: int = 0, kwargs: dict | No
     if latents is None or p is None:
         return kwargs
 
-    """
-    if torch.isnan(latents).any().item():
-        log.error(f'Callback: step={step} timestep={timestep} latents={latents.shape}:{latents.device}:{latents.dtype} error="contains NaN values"')
-        if (shared.state.current_latent is not None) and (shared.state.current_latent.shape == latents.shape):
-            log.error(f'Callback: step={step} timestep={timestep} latents={latents.shape}:{latents.device}:{latents.dtype} error="replacing with previous latent"')
-            latents = shared.state.current_latent
-    """
-
     if len(getattr(p, 'ip_adapter_names', [])) > 0 and p.ip_adapter_names[0] != 'None':
         ip_adapter_scales = list(p.ip_adapter_scales)
         ip_adapter_starts = list(p.ip_adapter_starts)
@@ -122,6 +115,7 @@ def diffusers_callback(pipe, step: int = 0, timestep: int = 0, kwargs: dict | No
     cfg_end = getattr(p, "cfg_end", 1.0) or 1.0
     total_steps = getattr(pipe, "num_timesteps", 0)
     target_step = int(total_steps * cfg_end) if total_steps else 0
+
     if (cfg_end < 1.0) and not getattr(pipe, "_cfg_end_applied", False) and (step >= target_step):
         pipe._cfg_end_applied = True # pylint: disable=protected-access
         if "PAG" in shared.sd_model.__class__.__name__:
@@ -145,7 +139,7 @@ def diffusers_callback(pipe, step: int = 0, timestep: int = 0, kwargs: dict | No
             else:
                 width = getattr(p, 'width', 1024)
                 height = getattr(p, 'height', 1024)
-            shared.state.current_latent = pipe._unpack_latents(kwargs['latents'], height, width, pipe.vae_scale_factor) # pylint: disable=protected-access
+            shared.state.current_latent = pipe._unpack_latents(latents, height, width, pipe.vae_scale_factor) # pylint: disable=protected-access
             if current_noise_pred is not None:
                 shared.state.current_noise_pred = pipe._unpack_latents(current_noise_pred, height, width, pipe.vae_scale_factor) # pylint: disable=protected-access
             else:
@@ -158,7 +152,6 @@ def diffusers_callback(pipe, step: int = 0, timestep: int = 0, kwargs: dict | No
             else:
                 width = getattr(p, 'width', 1024)
                 height = getattr(p, 'height', 1024)
-            latents = kwargs['latents']
             if len(latents.shape) == 4:
                 latents = pipe._unpatchify_latents(latents) # [B, C*4, h/2, w/2] -> [B, C, h, w] # pylint: disable=protected-access
             elif len(latents.shape) == 3:  # packed format [B, seq_len, patch_channels]
@@ -183,7 +176,6 @@ def diffusers_callback(pipe, step: int = 0, timestep: int = 0, kwargs: dict | No
                 current_noise_pred = current_noise_pred.permute(0, 3, 1, 4, 2, 5).reshape(b, channels, h_patches * 2, w_patches * 2)
             shared.state.current_noise_pred = current_noise_pred
         elif 'Ideogram4' in pipe.__class__.__name__:  # packed normalized [B, seq, 128] -> Flux.2 latent space for TAE FLUX.2
-            latents = kwargs['latents']
             if latents.ndim == 3:
                 b, seq_len, packed_ch = latents.shape
                 vae_scale = getattr(pipe, 'vae_scale_factor', 8)
@@ -203,7 +195,7 @@ def diffusers_callback(pipe, step: int = 0, timestep: int = 0, kwargs: dict | No
                 shared.state.current_latent = latents
             shared.state.current_noise_pred = current_noise_pred
         else:
-            shared.state.current_latent = kwargs['latents']
+            shared.state.current_latent = latents
             shared.state.current_noise_pred = current_noise_pred
 
         # Video latent preview: extract middle frame from 5D [B,C,T,H,W] to 4D [B,C,H,W]

modules/progress.py

@@ -48,6 +48,8 @@ class ProgressRequest(BaseModel):
 
 class InternalProgressResponse(BaseModel):
     job: str = Field(default=None, title="Job name", description="Internal job name")
+    job_timestamp: str|None = Field(default=None, title="Job timestamp", description="Timestamp of the job start")
+    job_time: float|None = Field(default=None, title="Job start time", description="Time of the job start")
     textinfo: str|None = Field(default=None, title="Info text", description="Info text used by WebUI.")
     # status fields
     active: bool = Field(title="Whether the task is being worked on right now")
@@ -122,7 +124,8 @@ def api_progress(req: ProgressRequest):
         steps=steps,
         batch_no=batch_no,
         batch_count=batch_count,
-        job_timestamp=shared.state.time_start,
+        job_timestamp=shared.state.job_timestamp,
+        job_time=shared.state.time_start,
         eta=eta,
         live_preview=live_preview,
         id_live_preview=id_live_preview,

modules/sd_samplers_common.py

@@ -1,3 +1,4 @@
+import os
 import time
 import threading
 from collections import namedtuple
@@ -14,6 +15,7 @@ approximation_indexes = { "Simple": 0, "Approximate": 1, "TAESD": 2, "Full VAE":
 flow_models = ['f1', 'f2', 'sd3', 'lumina', 'auraflow', 'sana', 'zimage', 'lumina2', 'cogview4', 'h1', 'cosmos', 'anima', 'chroma', 'omnigen', 'omnigen2', 'longcat', 'ideogram4', 'krea2']
 warned = False
 queue_lock = threading.Lock()
+debug = os.environ.get('SD_PREVIEW_DEBUG', None) is not None
 
 
 def warn_once(message):
@@ -35,10 +37,12 @@ def setup_img2img_steps(p, steps=None):
     return steps, t_enc
 
 
-def single_sample_to_image(sample, approximation=None):
+def single_sample_to_image(sample, approximation=None, fast=False):
     with queue_lock:
         t0 = time.time()
         approximation = approximation or shared.opts.show_progress_type
+        if debug:
+            log.debug(f'Preview sample: shape={list(sample.shape)} dtype={sample.dtype} method={approximation}')
         try:
             if (sample.dtype == torch.bfloat16) and (approximation in ["Simple", "Approximate"]):
                 sample = sample.to(torch.float16)
@@ -59,7 +63,7 @@ def single_sample_to_image(sample, approximation=None):
                     sample = torch.nn.functional.interpolate(sample.unsqueeze(0), scale_factor=[scale, scale], mode='bilinear', align_corners=False)[0]
                 except Exception:
                     pass
-            x_sample = sd_vae_taesd.decode(sample)
+            x_sample = sd_vae_taesd.decode(sample, fast=fast)
             # x_sample = (1.0 + x_sample) / 2.0 # preview requires smaller range
         elif shared.sd_model_type == 'sc' and approximation != "Full":
             x_sample = sd_vae_stablecascade.decode(sample)
@@ -98,8 +102,8 @@ def sample_to_image(samples, index=0, approximation=None):
     return single_sample_to_image(samples[index], approximation)
 
 
-def samples_to_image_grid(samples, approximation=None):
-    return images.image_grid([single_sample_to_image(sample, approximation) for sample in samples])
+def samples_to_image_grid(samples, approximation=None, fast=False):
+    return images.image_grid([single_sample_to_image(sample, approximation, fast=fast) for sample in samples])
 
 
 def store_latent(decoded):

modules/shared_state.py

@@ -293,8 +293,9 @@ class State:
                         elif self.prediction_type == "v_prediction":
                             sample = self.current_noise_pred * (-self.current_sigma / (self.current_sigma**2 + 1) ** 0.5) + (original_sample / (self.current_sigma**2 + 1)) # pylint: disable=invalid-unary-operand-type
                 except Exception:
+                    # log.error(f'State image sigma: last={self.id_live_preview} step={self.sampling_step} {e}')
                     pass # ignore sigma errors
-                image = sd_samplers_common.samples_to_image_grid(sample)
+                image = sd_samplers_common.samples_to_image_grid(sample, fast=self.sampling_step > 1)
                 self.assign_current_image(image)
                 self.preview_job = -1
                 return True

modules/ui_common.py

@@ -296,7 +296,7 @@ def open_folder(result_gallery, gallery_index = 0):
         subprocess.Popen([opener, path])  # pylint: disable=consider-using-with
 
 
-def create_output_panel(tabname, preview=True, prompt=None, height=None, transfer=True, scale=1, result_info=None):
+def create_output_panel(tabname, preview=True, prompt=None, height=None, transfer=True, scale=1, result_info=None, html_log_val=''):
     with gr.Column(variant='panel', elem_id=f"{tabname}_results", scale=scale):
         with gr.Group(elem_id=f"{tabname}_gallery_container"):
             if tabname == "txt2img":
@@ -348,7 +348,7 @@ def create_output_panel(tabname, preview=True, prompt=None, height=None, transfe
                 html_info = gr.HTML(elem_id=f'html_info_{tabname}', elem_classes="infotext", visible=False) # contains raw infotext as returned by wrapped call
                 html_info_formatted = gr.HTML(elem_id=f'html_info_formatted_{tabname}', elem_classes="infotext", visible=True) # contains html formatted infotext
                 html_info.change(fn=infotext_to_html, inputs=[html_info], outputs=[html_info_formatted], show_progress='hidden')
-                html_log = gr.HTML(elem_id=f'html_log_{tabname}', elem_classes=["hint"])
+                html_log = gr.HTML(elem_id=f'html_log_{tabname}', elem_classes=["hint"], value=html_log_val)
                 generation_info = gr.Textbox(visible=False, elem_id=f'generation_info_{tabname}')
                 generation_info_button = gr.Button(visible=False, elem_id=f"{tabname}_generation_info_button")
 

modules/ui_control.py

@@ -21,7 +21,7 @@ use_generator = os.environ.get('SD_USE_GENERATOR', None) is not None
 
 def return_stats(t: float | None = None):
     if t is None:
-        elapsed_text = ''
+        elapsed_text = '⏱ Idle'
     else:
         elapsed = time.perf_counter() - t
         elapsed_m = int(elapsed // 60)
@@ -40,10 +40,11 @@ def return_stats(t: float | None = None):
         if peak > 0:
             gpu += f"| 🕮 GPU {peak} MB"
             gpu += f" {used}%" if used > 0 else ''
-            gpu += f" | retries {retries} oom {ooms}" if retries > 0 or ooms > 0 else ''
+            gpu += f" | Retries {retries} OOM {ooms}" if retries > 0 or ooms > 0 else ''
     ram = ram_stats()
     if ram['used'] > 0:
-        cpu += f" RAM {ram['used']} GB"
+        # change emoji/symbol for ram to something better
+        cpu += f"| 🗒 RAM {ram['used']} GB"
         cpu += f" {round(100.0 * ram['used'] / ram['total'])}%" if ram['total'] > 0 else ''
     return f"<div class='performance hint' id='control-performance'><p>{elapsed_text} {summary} {gpu} {cpu}</p></div>"
 
@@ -247,7 +248,7 @@ def create_ui(_blocks: gr.Blocks=None):
                     gr.HTML('<span id="control-output-button">Output</p>')
                     with gr.Tabs(elem_classes=['control-tabs'], elem_id='control-tab-output') as output_tabs:
                         with gr.Tab('Gallery', id='out-gallery'):
-                            output_gallery, _output_gen_info, _output_html_info, _output_html_info_formatted, output_html_log = ui_common.create_output_panel("control", preview=False, prompt=prompt, height=gr_height, result_info=result_txt)
+                            output_gallery, _output_gen_info, _output_html_info, _output_html_info_formatted, output_html_log = ui_common.create_output_panel("control", preview=False, prompt=prompt, height=gr_height, result_info=result_txt, html_log_val=return_stats())
                         with gr.Tab('Image', id='out-image'):
                             output_image = gr.Image(label="Output", show_label=False, type="pil", interactive=False, tool="editor", height=gr_height, elem_id='control_output_image', elem_classes=['control-image'])
                         with gr.Tab('Video', id='out-video'):

modules/ui_extra_networks.py

@@ -57,9 +57,8 @@ preview_map = None
 def init_api():
 
     def get_thumb(filename: str = ""):
-        global allowed_dirs # pylint: disable=global-statement
-        if len(allowed_dirs) == 0:
-            allowed_dirs = shared.demo.allowed_paths
+        if os.path.join('ui', 'assets') not in allowed_dirs:
+            allowed_dirs.append(os.path.join('ui', 'assets'))
         if filename is None or len(filename) == 0:
             return JSONResponse({ "error": "no filename" }, status_code=400)
         if not any(Path(folder).absolute() in Path(filename).absolute().parents for folder in allowed_dirs):

modules/vae/sd_vae_taesd.py

@@ -43,7 +43,8 @@ prev_warnings = False
 first_run = True
 prev_cls = ''
 prev_type = ''
-prev_model = ''
+prev_variant = ''
+prev_model = None
 lock = threading.Lock()
 
 
@@ -82,12 +83,12 @@ def get_model(model_cls, variant=None):
         warn_once(f'cls={shared.sd_model.__class__.__name__} type={shared.sd_model_type} unsuppported', variant=variant)
         return model_cls, None
     if debug:
-        log.debug(f'TAESD detect: cls={model_cls} variant={variant}')
+        log.debug(f'TAESD detect: cls={model_cls} variant="{variant}"')
     return model_cls, variant
 
 
 def load_model(model_type = 'decoder', variant = None, vae_file: str | None = None):
-    global prev_cls, prev_type, prev_model, prev_warnings # pylint: disable=global-statement
+    global prev_cls, prev_type, prev_variant, prev_warnings # pylint: disable=global-statement
     model_cls = shared.sd_model_type if shared.sd_loaded else None
     if vae_file is not None and os.path.exists(vae_file):
         model_cls = 'sdxl'
@@ -101,7 +102,7 @@ def load_model(model_type = 'decoder', variant = None, vae_file: str | None = No
     os.makedirs(folder, exist_ok=True)
     if variant.startswith('TAE'):
         cfg = TAESD_MODELS[variant]
-        if (model_cls == prev_cls) and (model_type == prev_type) and (variant == prev_model) and (cfg['model'] is not None):
+        if (model_cls == prev_cls) and (model_type == prev_type) and (variant == prev_variant) and (cfg['model'] is not None):
             return cfg['model'], variant
         fn = os.path.join(folder, cfg['fn'] + model_type + '_' + model_cls + '.pth')
         if not os.path.exists(fn):
@@ -117,7 +118,7 @@ def load_model(model_type = 'decoder', variant = None, vae_file: str | None = No
         if os.path.exists(fn):
             prev_cls = model_cls
             prev_type = model_type
-            prev_model = variant
+            prev_variant = variant
             log.print() # new line
             log.debug(f'Decode: type="taesd" variant="{variant}" fn="{fn}" layers={shared.opts.taesd_layers} load')
             vae = None
@@ -133,12 +134,6 @@ def load_model(model_type = 'decoder', variant = None, vae_file: str | None = No
             else:
                 from modules.taesd.taesd import TAESD
                 vae = TAESD(decoder_path=fn if model_type=='decoder' else None, encoder_path=fn if model_type=='encoder' else None)
-                """
-                _vae = diffusers.AutoencoderKL()
-                from installer import Dot
-                _config = diffusers.AutoencoderKL().config.copy()
-                vae.config = Dot(_config) # set config for compatibility with standard vae
-                """
             if vae is not None:
                 prev_warnings = False # reset warnings for new model
                 vae = vae.to(devices.device, dtype=dtype)
@@ -147,7 +142,7 @@ def load_model(model_type = 'decoder', variant = None, vae_file: str | None = No
             return vae, variant
     elif variant.startswith('Hybrid'):
         cfg = CQYAN_MODELS[variant].get(model_cls, None)
-        if (model_cls == prev_cls) and (model_type == prev_type) and (variant == prev_model) and (cfg['model'] is not None):
+        if (model_cls == prev_cls) and (model_type == prev_type) and (variant == prev_variant) and (cfg['model'] is not None):
             return cfg['model'], variant
         if cfg is None:
             warn_once(f'cls={shared.sd_model.__class__.__name__} type={model_cls} unsuppported', variant=variant)
@@ -155,7 +150,7 @@ def load_model(model_type = 'decoder', variant = None, vae_file: str | None = No
         repo = cfg['repo']
         prev_cls = model_cls
         prev_type = model_type
-        prev_model = variant
+        prev_variant = variant
         log.debug(f'Decode: type="taesd" variant="{variant}" id="{repo}" load')
         if 'tiny' in repo:
             from diffusers.models import AutoencoderTiny
@@ -190,13 +185,20 @@ def restore_preview_size(image, vae):
     return image
 
 
-def decode(latents):
-    global first_run # pylint: disable=global-statement
+def decode(latents, fast=False):
+    global first_run, prev_model, prev_variant # pylint: disable=global-statement
     with lock:
         try:
-            vae, variant = load_model(model_type='decoder')
-            if vae is None or max(latents.shape) > 256: # safetey check of large tensors
-                return latents
+            if fast and prev_model is not None:
+                vae = prev_model
+                variant = prev_variant
+            else:
+                vae, variant = load_model(model_type='decoder')
+                if vae is None or max(latents.shape) > 256: # safety check of large tensors
+                    return latents
+                prev_model = vae
+                prev_variant = variant
+                fast = False
         except Exception as e:
             # from modules import errors
             # errors.display(e, 'taesd"')
@@ -208,7 +210,7 @@ def decode(latents):
                 tensor = latents.unsqueeze(0) if len(latents.shape) == 3 else latents
                 tensor = tensor.detach().clone().to(devices.device, dtype=dtype)
                 if debug:
-                    log.debug(f'Decode: type="taesd" variant="{variant}" input={latents.shape} tensor={tensor.shape}')
+                    log.debug(f'Decode: type="taesd" variant="{variant}" input={latents.shape} fast={fast} tensor={tensor.shape}')
                 # Fallback: reshape packed 128-channel latents to 32 channels if not already unpacked
                 if (variant == 'TAE FLUX.2') and (len(tensor.shape) == 4) and (tensor.shape[1] == 128):
                     b, _c, h, w = tensor.shape
@@ -221,7 +223,7 @@ def decode(latents):
                     image = (image / 2.0 + 0.5).clamp(0, 1).detach()
                 image = restore_preview_size(image, vae)
                 t1 = time.time()
-                if (t1 - t0) > 3.0 and not first_run:
+                if (t1 - t0) > 5.0 and not first_run:
                     log.warning(f'Decode: type="taesd" variant="{variant}" long decode time={t1 - t0:.2f}')
                 first_run = False
                 return image