minimax preview

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

CHANGELOG.md

@@ -34,6 +34,7 @@
   - intercept and profiling hooks
   - on-demand convert standard model on-demand 
 - **Other**
+  - Video Preview: TAESD support for MiniMax
   - new optional transformer hooks: *settings -> compute add-ons*  
     *PAG: Perturbed attention guidance, PAB: Pyramid attention broadcast, FBC: First Block Cache, FC: Faster Cache, LS: Layer Skip, MC: Mag Cache, TS: TaylorSeer*  
     *note*: compatibility of different methods varies across different models  

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:

pipelines/model_minimax.py

@@ -6,9 +6,11 @@ from modules.logger import log
 def unpack_latents(latents, components: diffusers.modular_pipelines.ModularPipeline, state: diffusers.modular_pipelines.BlockState):
     from diffusers.modular_pipelines.minimax_h3.modular_pipeline import align_num_frames, video_latent_num_frames
     from modules import processing_callbacks
-    frames = processing_callbacks.p.num_frames
-    width = processing_callbacks.p.width
-    height = processing_callbacks.p.height
+    frames = getattr(processing_callbacks.p, 'frames', 1)
+    width = getattr(processing_callbacks.p, 'width', 1024)
+    height = getattr(processing_callbacks.p, 'height', 1024)
+    if frames <= 0 or width <= 0 or height <= 0:
+        return latents
     num_frames = align_num_frames(frames, components.vae_frames_per_chunk, components.vae_latents_per_chunk)
     num_latent_frames = video_latent_num_frames(num_frames, components.vae_frames_per_chunk, components.vae_latents_per_chunk)
     latent_height = height // components.vae_spatial_compression_ratio

ui/dist/sdnext.mjs

@@ -11102,6 +11102,7 @@ function requestProgress(id_task = "undefined", progressEl = null, galleryEl = n
   let img;
   const initLivePreview = () => {
     if (!parentGallery) return;
+    debug("initLivePreview", { el: galleryEl, parent: parentGallery });
     const footers = Array.from(gradioApp().querySelectorAll(".gallery_footer"));
     for (const footer of footers) {
       if (footer.id !== "gallery_footer") footer.style.display = "none";
@@ -11152,7 +11153,7 @@ function requestProgress(id_task = "undefined", progressEl = null, galleryEl = n
   };
   const previewVisible = () => {
     try {
-      return !galleryEl?.closest(".section")?.classList.contains("minimize");
+      return galleryEl ? !galleryEl.closest(".section")?.classList.contains("minimize") : true;
     } catch {
       return true;
     }
@@ -11528,7 +11529,7 @@ function submit_framepack(...args) {
 function submit_ltx(...args) {
   const id = randomId();
   log("submitFramepack", id);
-  requestProgress(id, null, null);
+  requestProgress(id, null, gradioApp().getElementById("ltx_output_video"));
   window.submit_state = "";
   args[0] = id;
   return args;
@@ -11536,7 +11537,7 @@ function submit_ltx(...args) {
 function submit_minimax(...args) {
   const id = randomId();
   log("submitMiniMax", id);
-  requestProgress(id, null, null);
+  requestProgress(id, null, gradioApp().getElementById("minimax_output_video"));
   window.submit_state = "";
   args[0] = id;
   return args;

ui/progressBar.ts

@@ -126,6 +126,7 @@ export function requestProgress(id_task = 'undefined', progressEl = null, galler
 
   const initLivePreview = () => {
     if (!parentGallery) return;
+    debug('initLivePreview', { el: galleryEl, parent: parentGallery });
     const footers = Array.from<any>(gradioApp().querySelectorAll('.gallery_footer'));
     for (const footer of footers) {
       if (footer.id !== 'gallery_footer') footer.style.display = 'none'; // remove all footers
@@ -182,7 +183,7 @@ export function requestProgress(id_task = 'undefined', progressEl = null, galler
 
   const previewVisible = () => {
     try {
-      return !galleryEl?.closest('.section')?.classList.contains('minimize');
+      return galleryEl ? !galleryEl.closest('.section')?.classList.contains('minimize') : true;
     } catch {
       return true;
     }

ui/ui.ts

@@ -341,7 +341,7 @@ function submit_framepack(...args) {
 function submit_ltx(...args) {
   const id = randomId();
   log('submitFramepack', id);
-  requestProgress(id, null, null);
+  requestProgress(id, null, gradioApp().getElementById('ltx_output_video'));
   window.submit_state = '';
   args[0] = id;
   return args;
@@ -350,7 +350,7 @@ function submit_ltx(...args) {
 function submit_minimax(...args) {
   const id = randomId();
   log('submitMiniMax', id);
-  requestProgress(id, null, null);
+  requestProgress(id, null, gradioApp().getElementById('minimax_output_video'));
   window.submit_state = '';
   args[0] = id;
   return args;