minimax preview

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

modules/processing_callbacks.py

@@ -56,13 +56,7 @@ def diffusers_callback_legacy(step: int, timestep: int, latents: torch.FloatTens
             time.sleep(0.1)
 
 
-def diffusers_callback(pipe, step: int = 0, timestep: int = 0, kwargs: dict | None = None):
-    if kwargs is None:
-        kwargs = {}
-    t0 = time.time()
-    from modules.lora import lora_stack
-    lora_stack.on_step(step)
-
+def torch_sync():
     if shared.opts.torch_sync:
         if devices.backend == "ipex":
             torch.xpu.synchronize(devices.device)
@@ -70,6 +64,15 @@ def diffusers_callback(pipe, step: int = 0, timestep: int = 0, kwargs: dict | No
             torch.cuda.synchronize(devices.device)
         time.sleep(0.001) # 1ms yield frees GIL for the preview thread
 
+
+def diffusers_callback(pipe, step: int = 0, timestep: int = 0, kwargs: dict | None = None):
+    if kwargs is None:
+        kwargs = {}
+    t0 = time.time()
+    from modules.lora import lora_stack
+    lora_stack.on_step(step)
+    torch_sync()
+
     t1 = time.time()
 
     if shared.state.paused:

modules/progress.py

@@ -116,6 +116,8 @@ def api_progress(req: ProgressRequest):
             try:
                 buffered = io.BytesIO()
                 shared.state.current_image.save(buffered, format='jpeg', quality=60)
+                fn = f'/tmp/preview-{shared.state.preview_job}-{shared.state.current_image_sampling_step}.jpg'
+                shared.state.current_image.save(fn, quality=90)
                 b64 = base64.b64encode(buffered.getvalue())
                 live_preview = f'data:image/jpeg;base64,{b64.decode("ascii")}'
             except Exception:

modules/sd_hijack_modular.py

@@ -5,6 +5,7 @@ import torch
 import diffusers
 from modules.logger import log
 from modules import shared, sd_offload, timer
+from modules.processing_callbacks import torch_sync
 from modules.attention import context as attention_context
 from modules.lora import lora_stack
 
@@ -18,6 +19,7 @@ def modular_step(components: diffusers.modular_pipelines.ModularPipeline, state:
     if 'num_inference_steps' in keys:
         shared.state.sampling_steps = state.num_inference_steps
     if 'latents' in keys and state.latents.ndim > 1:
+        torch_sync()
         shared.state.step()
         if hasattr(components, 'custom_unpack_latents'):
             shared.state.current_latent = components.custom_unpack_latents(state.latents, components, state)

modules/shared_state.py

@@ -63,7 +63,7 @@ class State:
         status += 'oom ' if self.oom else ''
         status += 'api ' if self.api else ''
         fn = f'{sys._getframe(3).f_code.co_name}:{sys._getframe(2).f_code.co_name}' # pylint: disable=protected-access
-        return f'State: ts={self.job_timestamp} job={self.job} jobs={self.job_no+1}/{self.job_count}/{self.total_jobs} step={self.sampling_step}/{self.sampling_steps} preview={self.preview_job}/{self.id_live_preview}/{self.current_image_sampling_step} status="{status.strip()}" fn={fn}'
+        return f'State: ts={self.job_timestamp} job={self.job} jobs={self.job_no+1}/{self.job_count}/{self.total_jobs} step={self.sampling_step}/{self.sampling_steps} preview={self.preview_job}/{self.id_live_preview}/{self.current_image_sampling_step} status="{status.strip()}" image={self.current_image} latent={self.current_latent.shape if self.current_latent is not None else None} fn={fn}'
 
     @property
     def sampling_step(self):
@@ -293,7 +293,6 @@ 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, fast=self.sampling_step > 1)
                 self.assign_current_image(image)

modules/ui_definitions.py

@@ -643,6 +643,7 @@ def create_settings(cmd_opts):
         "live_preview_refresh_period": OptionInfo(500, "Progress update period", gr.Slider, {"minimum": 0, "maximum": 5000, "step": 25}),
         "taesd_variant": OptionInfo(shared_items.sd_taesd_items()[0], "TAESD variant", gr.Dropdown, {"choices": shared_items.sd_taesd_items()}),
         "taesd_layers": OptionInfo(3, "TAESD decode layers", gr.Slider, {"minimum": 1, "maximum": 3, "step": 1}),
+        "taesd_frames": OptionInfo(4, "TAESD video frames", gr.Slider, {"minimum": -1, "maximum": 32, "step": 1, "visible": False}),
         "live_preview_require_focus": OptionInfo(True, "Pause live previews when tab is not focused"),
         "live_preview_downscale": OptionInfo(True, "Downscale high resolution live previews"),
 

modules/vae/sd_vae_taesd.py

@@ -185,6 +185,23 @@ def restore_preview_size(image, vae):
     return image
 
 
+def tile_video_frames(tensor):
+    frame_count = tensor.shape[0]
+    requested = shared.opts.taesd_frames
+    if (frame_count <= 1) or (requested == 1):
+        return tensor[0]
+    if (requested == -1) or (requested >= frame_count):
+        indices = list(range(frame_count))
+    else:
+        indices = [round(i * (frame_count - 1) / (requested - 1)) for i in range(requested)]
+    selected = tensor[indices]
+    try:
+        tiled = torch.cat([selected[i] for i in range(selected.shape[0])], dim=-1)
+        return tiled
+    except Exception:
+        return tensor[0]
+
+
 def decode(latents, fast=False):
     global first_run, prev_model, prev_variant # pylint: disable=global-statement
     with lock:
@@ -221,11 +238,10 @@ def decode(latents, fast=False):
                 else:
                     image = vae.decode(tensor, return_dict=False)[0]
                     # image = (image / 2.0 + 0.5).clamp(0, 1).detach()
-                    if image.ndim == 4 and image.shape[0] > 1 and image.shape[1] == 3:
-                        # likely a video latent, just take the first frame
-                        # TODO video preview: tiled frames
-                        image = image[0]
                     image = image.clamp(0, 1).detach()
+                    if image.ndim == 4 and image.shape[0] > 1 and image.shape[1] == 3: # likely a video latent
+                        # image = tile_video_frames(image)
+                        image = image[0] # just take the first frame for now
                 image = restore_preview_size(image, vae)
                 t1 = time.time()
                 if (t1 - t0) > 5.0 and not first_run: