diff --git a/pipelines/minimax/__init__.py b/pipelines/minimax/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/pipelines/minimax/minimax_latents.py b/pipelines/minimax/minimax_latents.py new file mode 100644 index 000000000..576befc37 --- /dev/null +++ b/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 diff --git a/pipelines/model_minimax.py b/pipelines/model_minimax.py index 53ae78930..964f8629a 100644 --- a/pipelines/model_minimax.py +++ b/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