"""Per-keyframe injection frame and strength for the `fl2va` branch of the staged-denoising Space. Two independent knobs per keyframe, both optional (defaults reproduce the stock pipeline exactly): * **Injection frame.** A keyframe's conditioning rows carry one rotary time, `anchor_time`. Stock supports exactly two values: "first" = `T` and "last" = `T + 5/3 * (N - 1)`, where `T` is the text-token count and `N` the number of pixel frames the latent grid spans. The rotary clock runs at 5/3 units per *pixel* frame, so injecting at pixel frame `p` is just `anchor_time = T + 5/3 * p`. Nothing else about the keyframe changes. NOTE: the released model was only trained with the two stock positions. Other positions are untested. * **Strength.** A keyframe is held at `t = 0.999` ("just short of clean": `x_t = t*x0 + (1-t)*noise`). Strength `s` in `[0, 1]` maps to that `t`: `s = 1` is stock (0.999), `s = 0` is pure noise (`t = 0`, the keyframe carries no information). The same `t` is what the rows are presented to the transformer at, so the label stays honest. Blocks: MiniMaxH3KeyframePlacementStep - rewrites the keyframe rows' rotary time after `prepare_layout` MiniMaxH3KeyframeConditionStep - noises each keyframe to its own level instead of the shared 0.999 (build_row_timesteps patch) - presents each keyframe's rows at its own level, in every sampler The patch of `MiniMaxH3SetTimestepsStep.build_row_timesteps` is deliberate: the dpmpp_2s / SDE / seeds_2 samplers build their own second-evaluation row timesteps through that same static method, so patching it keeps all of them consistent with the plan without touching their files. It is a no-op for anything with no conditioning rows (`t2va`, the momentum pipeline) and whenever no per-keyframe levels were set for the request. """ import torch from diffusers.modular_pipelines.minimax_h3.before_denoise import ( _ROPE_FRAME_RESCALE, _ROPE_FRAMES_PER_LATENT, MiniMaxH3FL2VAPrepareLatentsStep, MiniMaxH3PrepareConditionLatentsStep, MiniMaxH3PrepareLayoutStep, MiniMaxH3SetTimestepsStep, patchify_video_latents, ) from diffusers.modular_pipelines.minimax_h3.before_denoise import MiniMaxH3PrepareLatentsStep from diffusers.modular_pipelines.minimax_h3.decoders import MiniMaxH3AfterDenoiseStep from diffusers.modular_pipelines.minimax_h3.denoise import MiniMaxH3DenoiseStep from diffusers.modular_pipelines.minimax_h3.modular_blocks_minimax_h3 import ( MiniMaxH3CoreDenoiseStep, MiniMaxH3DecodeStep, ) from diffusers.modular_pipelines.minimax_h3.modular_pipeline import MiniMaxH3ModularPipeline from diffusers.modular_pipelines.modular_pipeline import ( ConditionalPipelineBlocks, ModularPipelineBlocks, PipelineState, SequentialPipelineBlocks, ) from diffusers.modular_pipelines.modular_pipeline_utils import InputParam, OutputParam from diffusers.utils.torch_utils import randn_tensor from h3_split_blocks import MiniMaxH3SplitBeforeEncodeStep, MiniMaxH3SplitVaeEncoderStep # --------------------------------------------------------------------------------------------------------------- # Per-request state shared with build_row_timesteps. Set by the condition step, read by the patched static method. # --------------------------------------------------------------------------------------------------------------- _LEVELS: tuple[float, ...] | None = None _AUG: float = 0.999 def strength_to_level(strength: float, aug: float) -> float: """Slider strength in [0, 1] -> the keyframe's `t`. 1 (or anything above `aug`) is stock.""" return max(0.0, min(float(strength), float(aug))) _original_build_row_timesteps = MiniMaxH3SetTimestepsStep.build_row_timesteps def _build_row_timesteps( video_indices, audio_indices, num_condition_video_rows, num_condition_audio_rows, num_text_tokens, video_timestep, audio_timestep, condition_video_timestep, condition_audio_timestep, ): levels = _LEVELS n = len(levels) if levels else 0 if n == 0 or num_condition_video_rows <= 0 or num_condition_video_rows % n != 0: return _original_build_row_timesteps( video_indices, audio_indices, num_condition_video_rows, num_condition_audio_rows, num_text_tokens, video_timestep, audio_timestep, condition_video_timestep, condition_audio_timestep, ) sequence_length = int(video_indices.numel() + audio_indices.numel() + num_text_tokens) row_timesteps = torch.full((sequence_length,), video_timestep, dtype=torch.float32) rows_per_keyframe = num_condition_video_rows // n for i, level in enumerate(levels): block = video_indices[i * rows_per_keyframe : (i + 1) * rows_per_keyframe] # At or above the stock level the keyframe is presented exactly as stock does (max with the current step's # timestep). Below it, the keyframe really is that noisy and stays labelled that way for every step. row_timesteps[block] = condition_video_timestep if level >= _AUG else level row_timesteps[audio_indices[num_condition_audio_rows:]] = audio_timestep row_timesteps[audio_indices[:num_condition_audio_rows]] = condition_audio_timestep return torch.unique(row_timesteps, sorted=True, return_inverse=True) MiniMaxH3SetTimestepsStep.build_row_timesteps = staticmethod(_build_row_timesteps) # --------------------------------------------------------------------------------------------------------------- # Blocks # --------------------------------------------------------------------------------------------------------------- class MiniMaxH3KeyframePlacementStep(ModularPipelineBlocks): model_name = "minimax-h3" @property def description(self) -> str: return ( "Moves each keyframe's conditioning rows to a chosen pixel frame on the rotary time axis. The stock " "layout only knows 'first' (frame 0) and 'last' (the final frame); `anchor_time = T + 5/3 * frame`." ) @property def inputs(self) -> list[InputParam]: return [ InputParam(name="position_ids", type_hint=torch.Tensor, required=True), InputParam(name="text_indices", type_hint=torch.Tensor, required=True), InputParam(name="num_latent_frames", type_hint=int, required=True), InputParam(name="latent_height", type_hint=int, required=True), InputParam(name="latent_width", type_hint=int, required=True), InputParam(name="keyframe_anchors", type_hint=tuple, default=()), InputParam( name="keyframe_frames", type_hint=tuple, default=None, description="Requested pixel frame per keyframe, in packed order. None / missing entries keep stock.", ), ] @property def intermediate_outputs(self) -> list[OutputParam]: return [OutputParam("position_ids", type_hint=torch.Tensor)] @torch.no_grad() def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState): block_state = self.get_block_state(state) anchors = tuple(block_state.keyframe_anchors or ()) frames = block_state.keyframe_frames if anchors and frames is not None: _, patch_h, patch_w = components.patch_size rows_per_frame = (block_state.latent_height // patch_h) * (block_state.latent_width // patch_w) num_text = int(block_state.text_indices.numel()) pattern = _ROPE_FRAMES_PER_LATENT total_pixel_frames = sum(pattern[i % len(pattern)] for i in range(block_state.num_latent_frames)) last_frame = total_pixel_frames - 1 position_ids = block_state.position_ids for index, (anchor, requested) in enumerate(zip(anchors, frames)): if requested is None: continue frame = max(0, min(int(requested), last_frame)) # Defaults reproduce the stock anchors bit for bit: leave those rows untouched. if (anchor == "first" and frame <= 0) or (anchor == "last" and frame >= last_frame): continue rows = slice(num_text + index * rows_per_frame, num_text + (index + 1) * rows_per_frame) position_ids[rows, 0] = float(num_text) + _ROPE_FRAME_RESCALE * frame block_state.position_ids = position_ids self.set_block_state(state, block_state) return components, state class MiniMaxH3KeyframeConditionStep(MiniMaxH3PrepareConditionLatentsStep): model_name = "minimax-h3" @property def description(self) -> str: return ( "`MiniMaxH3PrepareConditionLatentsStep` with a noise level per keyframe: `x_t = t*x0 + (1-t)*noise`, " "`t = strength` capped at the stock 0.999. Draw order is unchanged (one draw per condition)." ) @property def inputs(self) -> list[InputParam]: return super().inputs + [ InputParam( name="keyframe_strengths", type_hint=tuple, default=None, description="Strength in [0, 1] per keyframe, in packed order. None keeps the stock 0.999.", ) ] @property def intermediate_outputs(self) -> list[OutputParam]: return super().intermediate_outputs @torch.no_grad() def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState): global _LEVELS, _AUG block_state = self.get_block_state(state) device = components._execution_device patch_size = components.patch_size aug = float(components.keyframe_noise_aug) strengths = block_state.keyframe_strengths levels = [] packed = [] for index, condition in enumerate(block_state.condition_latents): noise = randn_tensor(condition.shape, generator=block_state.generator, device=device, dtype=torch.float32) strength = 1.0 if strengths is None or index >= len(strengths) or strengths[index] is None else strengths[index] level = strength_to_level(strength, aug) levels.append(level) noised = components.scheduler.scale_noise(condition.to(device), level, noise) packed.append(patchify_video_latents(noised, patch_size)) block_state.condition_rows = torch.cat(packed) if block_state.condition_rows.shape[0] != block_state.num_condition_video_rows: raise ValueError( f"The layout reserved {block_state.num_condition_video_rows} conditioning rows but the encoded " f"conditioning latents pack into {block_state.condition_rows.shape[0]}." ) _AUG = aug _LEVELS = tuple(levels) if any(level < aug for level in levels) else None self.set_block_state(state, block_state) return components, state class MiniMaxH3KeyframeControlCoreDenoiseStep(SequentialPipelineBlocks): """`MiniMaxH3FL2VACoreDenoiseStep` with the placement step after the layout and the per-keyframe condition step.""" model_name = "minimax-h3" block_classes = [ MiniMaxH3PrepareLayoutStep, MiniMaxH3KeyframePlacementStep, MiniMaxH3KeyframeConditionStep, MiniMaxH3PrepareLatentsStep, MiniMaxH3FL2VAPrepareLatentsStep, MiniMaxH3SetTimestepsStep, MiniMaxH3DenoiseStep, MiniMaxH3AfterDenoiseStep, ] block_names = [ "prepare_layout", "keyframe_placement", "prepare_condition_latents", "prepare_latents", "prepare_latents_fl2va", "set_timesteps", "denoise", "after_denoise", ] @property def description(self) -> str: return "Core denoising workflow for `fl2va` with per-keyframe injection frame and strength." @property def outputs(self): return [ OutputParam("latents", type_hint=torch.Tensor, description="The generated video latents."), OutputParam("audio_latents", type_hint=torch.Tensor, description="The generated audio latents."), ] class MiniMaxH3KeyframeControlDenoiseStep(ConditionalPipelineBlocks): """Split-deployment denoise selector with the keyframe-control `fl2va` core in place of the stock one.""" model_name = "minimax-h3" block_classes = [MiniMaxH3KeyframeControlCoreDenoiseStep, MiniMaxH3CoreDenoiseStep] block_names = ["fl2va", "t2va"] block_trigger_inputs = ["image", "last_image"] default_block_name = "t2va" def select_block(self, **kwargs) -> str | None: if kwargs.get("image") is not None or kwargs.get("last_image") is not None: return "fl2va" return None @property def description(self): return "`fl2va` (keyframe control) when a keyframe is provided, `t2va` otherwise." class MiniMaxH3KeyframeGeneratorBlocks(SequentialPipelineBlocks): """Drop-in for `h3_split_blocks.MiniMaxH3GeneratorBlocks`, with keyframe injection frame / strength.""" model_name = "minimax-h3" block_classes = [ MiniMaxH3SplitBeforeEncodeStep, MiniMaxH3SplitVaeEncoderStep, MiniMaxH3KeyframeControlDenoiseStep, MiniMaxH3DecodeStep, ] block_names = ["resize", "vae_encoder", "denoise", "decode"] @property def description(self): return "Denoising half of the split deployment with per-keyframe injection frame and strength." @property def outputs(self): return [ OutputParam.template("videos", description="The generated video."), OutputParam("audio", type_hint=torch.Tensor, description="The soundtrack, `(1, 2, num_samples)`."), OutputParam("sampling_rate", type_hint=int, description="Sample rate of the soundtrack in Hz."), ]