Spaces:
Running on Zero
Running on Zero
| """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" | |
| 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`." | |
| ) | |
| 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.", | |
| ), | |
| ] | |
| def intermediate_outputs(self) -> list[OutputParam]: | |
| return [OutputParam("position_ids", type_hint=torch.Tensor)] | |
| 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" | |
| 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)." | |
| ) | |
| 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.", | |
| ) | |
| ] | |
| def intermediate_outputs(self) -> list[OutputParam]: | |
| return super().intermediate_outputs | |
| 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", | |
| ] | |
| def description(self) -> str: | |
| return "Core denoising workflow for `fl2va` with per-keyframe injection frame and strength." | |
| 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 | |
| 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"] | |
| def description(self): | |
| return "Denoising half of the split deployment with per-keyframe injection frame and strength." | |
| 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."), | |
| ] | |