minimax unpack latents

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

pipelines/minimax/minimax_latents.py

@@ -0,0 +1,37 @@
+import diffusers
+
+
+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 = 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
+    latent_width = width // components.vae_spatial_compression_ratio
+    patch_t, patch_h, patch_w = components.patch_size
+    channels = components.vae_latent_channels
+    rows = state.latents[state.num_condition_video_rows :]
+    rows = rows.reshape(
+        -1,
+        num_latent_frames // patch_t,
+        latent_height // patch_h,
+        latent_width // patch_w,
+        channels,
+        patch_t,
+        patch_h,
+        patch_w,
+    )
+    rows = rows.permute(0, 4, 1, 5, 2, 6, 3, 7)
+    latents = rows.reshape(
+        -1,
+        channels,
+        num_latent_frames,
+        latent_height,
+        latent_width,
+    ).contiguous()
+    return latents

pipelines/model_minimax.py

@@ -3,42 +3,6 @@ from modules import shared, devices, sd_models
 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 = 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
-    latent_width = width // components.vae_spatial_compression_ratio
-    patch_t, patch_h, patch_w = components.patch_size
-    channels = components.vae_latent_channels
-    rows = state.latents[state.num_condition_video_rows :]
-    rows = rows.reshape(
-        -1,
-        num_latent_frames // patch_t,
-        latent_height // patch_h,
-        latent_width // patch_w,
-        channels,
-        patch_t,
-        patch_h,
-        patch_w,
-    )
-    rows = rows.permute(0, 4, 1, 5, 2, 6, 3, 7)
-    latents = rows.reshape(
-        -1,
-        channels,
-        num_latent_frames,
-        latent_height,
-        latent_width,
-    ).contiguous()
-    return latents
-
-
 def load_minimax(checkpoint_info, diffusers_load_config = None, workflow: str | None = None):
     from modules.video_models import video_load
     from modules.modular_load import load_modular_pipe
@@ -75,6 +39,7 @@ def load_minimax(checkpoint_info, diffusers_load_config = None, workflow: str |
     if hasattr(pipe, 'vae') and hasattr(pipe.vae, 'enable_tiling'):
         pipe.vae.enable_tiling()
 
+    from pipelines.minimax.minimax_latents import unpack_latents
     pipe.custom_unpack_latents = unpack_latents # add a helper to unpack the video latents from the block state
     devices.torch_gc()
     return pipe