Download ComfyUI-MiniMaxH3-Context-Loop/masked_context.py from dzmdd/Fast_deploy: direct link, hf CLI and curl.
- Browser
- Download file 31.8 kB
-
https://huggingface.co/dzmdd/Fast_deploy/resolve/main/ComfyUI-MiniMaxH3-Context-Loop/masked_context.py
- Command line
-
hf download hf://dzmdd/Fast_deploy/ComfyUI-MiniMaxH3-Context-Loop/masked_context.py
-
curl -L -o masked_context.py https://huggingface.co/dzmdd/Fast_deploy/resolve/main/ComfyUI-MiniMaxH3-Context-Loop/masked_context.py
31.8 kB
| """Masked target-prefix continuation for recursive MiniMax H3 chains. | |
| Generated scenes continue directly from the previous sampled H3 video/audio | |
| latent tails, avoiding a lossy video decode/re-encode round trip. Imported | |
| scene-1 video and audio still use their respective VAEs because no sampled H3 | |
| latent exists yet. A nested AV denoise mask protects the copied prefix while | |
| the sampler generates only the future portion. | |
| The mask design follows ComfyUI PR #15375 and seitanism's GPL-3.0 | |
| ComfyUI-H3-Motion-Context-MultiRef masked-extension work. Runtime support is | |
| enabled lazily so ordinary guide-mode chains do not patch H3 mask handling. | |
| """ | |
| from __future__ import annotations | |
| import logging | |
| import math | |
| import torch | |
| from .contracts_v05 import CONTEXT_SPATIAL_PROXY_RECIPE, DETAIL_AV_RECIPE | |
| from .nodes import ( | |
| AV_RUN_GRID, | |
| AUDIO_HZ, | |
| FPS, | |
| VIDEO_RUN_GRID, | |
| _audio_tail_from_latent, | |
| _pixel_frames, | |
| _resize, | |
| _streams_from_latent, | |
| ) | |
| _LOG = logging.getLogger("minimax_h3_context_loop.masked_prefix") | |
| def _latent_context_spatial_proxy(video_prefix): | |
| """Low-pass a copied AV prefix through the fixed 5/6 latent canvas. | |
| Time, channels, dtype, device, and final target geometry are unchanged. | |
| Only this disposable copy is filtered; the predecessor checkpoint remains | |
| byte-for-byte untouched. | |
| """ | |
| if not torch.is_tensor(video_prefix) or video_prefix.ndim != 5: | |
| raise ValueError( | |
| "h3_context_spatial_proxy: video prefix must be " | |
| "[B,C,T,H,W], got %s." % | |
| (tuple(getattr(video_prefix, "shape", ())),)) | |
| if not video_prefix.is_floating_point(): | |
| raise ValueError( | |
| "h3_context_spatial_proxy: video prefix must be floating point.") | |
| batch, channels, steps, latent_height, latent_width = ( | |
| int(value) for value in video_prefix.shape) | |
| alignment = int(CONTEXT_SPATIAL_PROXY_RECIPE["pixel_alignment"]) | |
| numerator = int(CONTEXT_SPATIAL_PROXY_RECIPE["scale_numerator"]) | |
| denominator = int(CONTEXT_SPATIAL_PROXY_RECIPE["scale_denominator"]) | |
| def proxy_latent_size(value): | |
| pixels = int(value) * 16 | |
| units = max(1, int(round( | |
| pixels * numerator / float(denominator * alignment)))) | |
| return max(1, units * alignment // 16) | |
| proxy_height = proxy_latent_size(latent_height) | |
| proxy_width = proxy_latent_size(latent_width) | |
| if (proxy_height, proxy_width) == (latent_height, latent_width): | |
| return video_prefix.detach().contiguous().clone(), ( | |
| proxy_height, proxy_width) | |
| source = video_prefix.permute(0, 2, 1, 3, 4).reshape( | |
| batch * steps, channels, latent_height, latent_width) | |
| work = source.float() | |
| reduced = torch.nn.functional.interpolate( | |
| work, size=(proxy_height, proxy_width), mode="area") | |
| restored = torch.nn.functional.interpolate( | |
| reduced, size=(latent_height, latent_width), mode="bilinear", | |
| align_corners=False) | |
| restored = restored.to( | |
| device=video_prefix.device, dtype=video_prefix.dtype) | |
| restored = restored.reshape( | |
| batch, steps, channels, latent_height, latent_width).permute( | |
| 0, 2, 1, 3, 4).contiguous() | |
| return restored, (proxy_height, proxy_width) | |
| def _detail_av_alpha(position, step_count): | |
| """Return the latent-noise amount for one carried video step.""" | |
| position = int(position) | |
| step_count = int(step_count) | |
| if step_count < 1 or position < 0 or position >= step_count: | |
| raise ValueError("h3_detail_av: position is outside its context span.") | |
| ramp = min(int(DETAIL_AV_RECIPE["ramp_steps"]), step_count) | |
| from_end = step_count - 1 - position | |
| alpha = float(DETAIL_AV_RECIPE["alpha"]) | |
| if from_end >= ramp: | |
| return alpha | |
| alpha_end = float(DETAIL_AV_RECIPE["alpha_end"]) | |
| if from_end == 0: | |
| return alpha_end | |
| return alpha + (alpha_end - alpha) * (ramp - from_end) / ramp | |
| def _detail_av_video_prefix(video_prefix, seed): | |
| """Build a deterministic, disposable noisy copy of an AV video prefix. | |
| Only the video tensor is accepted here. The caller retains the clean | |
| predecessor latent and the audio prefix independently, preventing this | |
| one-shot context treatment from contaminating recursive state. | |
| """ | |
| if not torch.is_tensor(video_prefix) or video_prefix.ndim != 5: | |
| raise ValueError( | |
| "h3_detail_av: video prefix must be [B,C,T,H,W], got %s." % | |
| (tuple(getattr(video_prefix, "shape", ())),) | |
| ) | |
| if not video_prefix.is_floating_point(): | |
| raise ValueError("h3_detail_av: video prefix must be floating point.") | |
| steps = int(video_prefix.shape[2]) | |
| expected_steps = int(DETAIL_AV_RECIPE["video_steps"]) | |
| if steps != expected_steps: | |
| raise ValueError( | |
| "h3_detail_av: the v2 recipe requires %d video steps, got %d." % | |
| (expected_steps, steps) | |
| ) | |
| out = video_prefix.detach().contiguous().clone() | |
| scale = float(out.detach().float().std().item()) | |
| if not scale or not math.isfinite(scale): | |
| _LOG.warning( | |
| "h3_detail_av: degenerate carried-latent standard deviation; " | |
| "using unit Gaussian scale") | |
| scale = 1.0 | |
| generator = torch.Generator(device="cpu") | |
| generator.manual_seed(int(seed) & 0xFFFFFFFFFFFFFFFF) | |
| alphas = [] | |
| for position in range(steps): | |
| amount = _detail_av_alpha(position, steps) | |
| alphas.append(amount) | |
| block = out[:, :, position] | |
| noise = torch.randn( | |
| tuple(block.shape), generator=generator, dtype=torch.float32, | |
| device="cpu") | |
| noise = noise.mul_(scale).to( | |
| device=block.device, dtype=block.dtype) | |
| out[:, :, position] = block * (1.0 - amount) + noise * amount | |
| return out, scale, tuple(alphas) | |
| def _require_h3_mask_support(): | |
| """Compatibility alias retained for focused tests and chain callers.""" | |
| from .masking_support import require_h3_mask_support | |
| return require_h3_mask_support("masked AV continuation") | |
| def _snap_prefix_length( | |
| requested, available, target_frames, preserve_audio_prefix=True): | |
| """Resolve an AV prefix on the grid required by its active streams.""" | |
| cap = min(int(requested), int(available), int(target_frames) - 1) | |
| preserve_audio_prefix = bool(preserve_audio_prefix) | |
| grid = AV_RUN_GRID if preserve_audio_prefix else VIDEO_RUN_GRID | |
| minimum = 39 if preserve_audio_prefix else 5 | |
| run = next((value for value in grid if value <= cap), 0) | |
| if run < minimum: | |
| requirement = ( | |
| "at least 39 previous frames and an exact shared video/audio " | |
| "boundary" if preserve_audio_prefix else | |
| "at least 5 previous frames on H3's video latent grid" | |
| ) | |
| raise ValueError( | |
| "h3_masked_prefix: masked continuation needs %s and a target " | |
| "longer than the prefix." % requirement | |
| ) | |
| if run != int(requested): | |
| if preserve_audio_prefix: | |
| detail = "exact shared AV runs are 39, 90, 141, 192, and 243" | |
| else: | |
| detail = "video-only prefix snapped to H3's video latent grid" | |
| _LOG.warning( | |
| "h3_masked_prefix: context_length %d -> exact H3 prefix %d (%s)", | |
| int(requested), run, detail) | |
| return run | |
| def _validate_target_streams(latent, strict_audio_grid=True): | |
| video, audio = _streams_from_latent(latent)[:2] | |
| if video.ndim == 4: | |
| video = video.unsqueeze(0) | |
| if audio.ndim == 3: | |
| audio = audio.unsqueeze(0) | |
| if video.ndim != 5: | |
| raise ValueError( | |
| "h3_masked_prefix: video latent must be [B,C,T,H,W], got %s" % | |
| (tuple(video.shape),) | |
| ) | |
| if audio.ndim != 4: | |
| raise ValueError( | |
| "h3_masked_prefix: audio latent must be [B,C,2,T], got %s" % | |
| (tuple(audio.shape),) | |
| ) | |
| if int(video.shape[0]) != 1 or int(audio.shape[0]) != 1: | |
| raise ValueError( | |
| "h3_masked_prefix: masked continuation supports H3 batch size 1." | |
| ) | |
| target_frames = _pixel_frames(int(video.shape[2])) | |
| expected_audio = int(round(target_frames / float(FPS) * AUDIO_HZ)) | |
| if int(audio.shape[-1]) != expected_audio: | |
| message = ( | |
| "h3_masked_prefix: target latent has %d audio steps for %d video " | |
| "frames; expected %d on H3's nominal 40 Hz grid." % | |
| (int(audio.shape[-1]), target_frames, expected_audio) | |
| ) | |
| if strict_audio_grid: | |
| raise RuntimeError(message) | |
| _LOG.warning("%s The target audio length is authoritative.", message) | |
| return video, audio, target_frames | |
| def _generated_video_tail(previous_latent, frames, target_video): | |
| """Copy a phase-aligned video prefix from a generated H3 AV latent.""" | |
| parts = _streams_from_latent(previous_latent) | |
| if len(parts) < 2: | |
| raise ValueError( | |
| "h3_masked_prefix: previous sampled latent has no audio stream. " | |
| "Wire the joint H3 AV sampler output, not a video-only latent." | |
| ) | |
| source_video, source_audio = parts[:2] | |
| if source_video.ndim == 4: | |
| source_video = source_video.unsqueeze(0) | |
| if source_audio.ndim == 3: | |
| source_audio = source_audio.unsqueeze(0) | |
| if source_video.ndim != 5: | |
| raise ValueError( | |
| "h3_masked_prefix: previous video latent must be [B,C,T,H,W], " | |
| "got %s." % (tuple(source_video.shape),) | |
| ) | |
| if source_audio.ndim != 4: | |
| raise ValueError( | |
| "h3_masked_prefix: previous audio latent must be [B,C,2,T], got " | |
| "%s." % (tuple(source_audio.shape),) | |
| ) | |
| if int(source_video.shape[0]) != 1 or int(source_audio.shape[0]) != 1: | |
| raise ValueError( | |
| "h3_masked_prefix: generated-latent continuation supports H3 " | |
| "batch size 1." | |
| ) | |
| # Valid H3 prefix runs 5/22/39/... map to 2/7/12/... latent steps. Full | |
| # H3 clips and these prefixes are both 2 mod 5 latent steps, so slicing the | |
| # source tail and placing it at target step zero preserves temporal phase. | |
| video_steps = 2 + 5 * ((int(frames) - 5) // 17) | |
| if _pixel_frames(video_steps) != int(frames): | |
| raise RuntimeError( | |
| "h3_masked_prefix: internal H3 context mapping failed for %d " | |
| "frames." % int(frames) | |
| ) | |
| if int(source_video.shape[2]) < video_steps: | |
| raise ValueError( | |
| "h3_masked_prefix: previous sampled latent has too few video " | |
| "steps for the %d-frame context." % int(frames) | |
| ) | |
| if tuple(source_video.shape[1:2] + source_video.shape[3:]) != tuple( | |
| target_video.shape[1:2] + target_video.shape[3:]): | |
| raise ValueError( | |
| "h3_masked_prefix: previous/target video latent geometry differs: " | |
| "%s vs %s. Keep chained clips at the same H3 resolution." % | |
| (tuple(source_video.shape), tuple(target_video.shape)) | |
| ) | |
| return source_video[:1, :, -video_steps:].clone(), video_steps | |
| def _encoded_video_tail(vae, previous_frames, frames, target_video, crop): | |
| """Encode an imported decoded-video tail when no H3 source latent exists.""" | |
| if getattr(previous_frames, "ndim", 0) != 4: | |
| raise ValueError( | |
| "h3_masked_prefix: imported previous frames must be IMAGE " | |
| "[N,H,W,C]." | |
| ) | |
| available = int(previous_frames.shape[0]) | |
| if available < int(frames): | |
| raise ValueError( | |
| "h3_masked_prefix: imported video has %d frames but the resolved " | |
| "prefix needs %d." % (available, int(frames)) | |
| ) | |
| width = int(target_video.shape[4]) * 16 | |
| height = int(target_video.shape[3]) * 16 | |
| video_tail = _resize( | |
| previous_frames[available - int(frames):], width, height, crop) | |
| video_prefix = vae.encode(video_tail) | |
| if getattr(video_prefix, "ndim", 0) != 5: | |
| raise ValueError( | |
| "h3_masked_prefix: video VAE returned %s; expected " | |
| "[B,C,T,H,W]." % | |
| (tuple(getattr(video_prefix, "shape", ())),) | |
| ) | |
| video_steps = int(video_prefix.shape[2]) | |
| covered = _pixel_frames(video_steps) | |
| if covered != int(frames): | |
| raise RuntimeError( | |
| "h3_masked_prefix: %d imported context frames encoded to %d " | |
| "video steps covering %d frames; refusing a phase-shifted seam." % | |
| (int(frames), video_steps, covered) | |
| ) | |
| return video_prefix, video_steps | |
| def _encode_imported_audio(audio_vae, audio, frames): | |
| if audio_vae is None: | |
| raise ValueError( | |
| "h3_masked_prefix: imported-video scene 1 needs the H3 audio VAE " | |
| "connected to Chain Context." | |
| ) | |
| if audio is None: | |
| raise ValueError( | |
| "h3_masked_prefix: imported-video scene 1 has no context audio. " | |
| "Reconnect source audio to Existing Video Context." | |
| ) | |
| waveform = audio["waveform"] | |
| sample_rate = int(audio["sample_rate"]) | |
| vae_rate = int(getattr(audio_vae, "audio_sample_rate", 32000)) | |
| if sample_rate != vae_rate: | |
| try: | |
| import torchaudio | |
| except ImportError as exc: | |
| raise RuntimeError( | |
| "h3_masked_prefix: imported audio is %d Hz but the H3 audio " | |
| "VAE wants %d Hz and torchaudio is unavailable." % | |
| (sample_rate, vae_rate) | |
| ) from exc | |
| waveform = torchaudio.functional.resample( | |
| waveform, sample_rate, vae_rate) | |
| wanted = int(round(int(frames) / float(FPS) * vae_rate)) | |
| if int(waveform.shape[-1]) < wanted: | |
| raise ValueError( | |
| "h3_masked_prefix: imported context audio is shorter than the " | |
| "%d-frame masked prefix." % int(frames) | |
| ) | |
| encoded = audio_vae.encode(waveform[:1, ..., -wanted:].movedim(1, -1)) | |
| if getattr(encoded, "ndim", 0) != 4: | |
| raise ValueError( | |
| "h3_masked_prefix: audio VAE returned %s; expected [B,C,2,T]." % | |
| (tuple(getattr(encoded, "shape", ())),) | |
| ) | |
| steps = int(round(int(frames) / float(FPS) * AUDIO_HZ)) | |
| if int(encoded.shape[-1]) < steps: | |
| raise RuntimeError( | |
| "h3_masked_prefix: %d frames need %d audio steps but the audio " | |
| "VAE produced %d." % (int(frames), steps, int(encoded.shape[-1])) | |
| ) | |
| return encoded[:1, ..., -steps:], steps, "imported decoded audio" | |
| def _existing_mask_streams(latent, video, audio): | |
| mask = latent.get("noise_mask") | |
| if mask is None: | |
| return ( | |
| torch.ones( | |
| (1, 1, int(video.shape[2]), int(video.shape[3]), | |
| int(video.shape[4])), | |
| device=video.device, dtype=torch.float32, | |
| ), | |
| torch.ones( | |
| (1, 1, int(audio.shape[2]), int(audio.shape[3])), | |
| device=audio.device, dtype=torch.float32, | |
| ), | |
| ) | |
| if hasattr(mask, "unbind"): | |
| parts = list(mask.unbind()) | |
| elif isinstance(mask, (tuple, list)): | |
| parts = list(mask) | |
| else: | |
| raise ValueError( | |
| "h3_masked_prefix: an existing target noise_mask is not a nested " | |
| "H3 video/audio mask and cannot be composed safely." | |
| ) | |
| if len(parts) < 2: | |
| raise ValueError( | |
| "h3_masked_prefix: existing target noise_mask has no audio stream." | |
| ) | |
| video_mask, audio_mask = parts[:2] | |
| if video_mask.ndim == 4: | |
| video_mask = video_mask.unsqueeze(0) | |
| if audio_mask.ndim == 3: | |
| audio_mask = audio_mask.unsqueeze(0) | |
| expected_video = (1, 1, int(video.shape[2]), int(video.shape[3]), | |
| int(video.shape[4])) | |
| expected_audio = (1, 1, int(audio.shape[2]), int(audio.shape[3])) | |
| try: | |
| video_mask = torch.broadcast_to(video_mask, expected_video).clone() | |
| audio_mask = torch.broadcast_to(audio_mask, expected_audio).clone() | |
| except RuntimeError as exc: | |
| raise ValueError( | |
| "h3_masked_prefix: existing AV noise-mask shapes %s / %s cannot " | |
| "broadcast to target %s / %s." % | |
| (tuple(video_mask.shape), tuple(audio_mask.shape), expected_video, | |
| expected_audio) | |
| ) from exc | |
| return video_mask.float(), audio_mask.float() | |
| def apply_locked_source_audio_target( | |
| latent, audio_vae, source_audio, *, lip_sync_options=None, | |
| voice=None, clip_start_seconds=0.0, | |
| voice_clip_start_seconds=None, | |
| force_active_voice_prefix_seconds=0.0): | |
| """Place exact scene-local source audio in H3's protected target stream. | |
| Reuse the standalone master-audio implementation so Chain Policy receives | |
| the same authoritative target-grid sizing, short encoder-lookahead retry, | |
| channel normalization, and mask composition as the public masking node. | |
| The Current Shot window is already scene-local, so its clock begins at 0. | |
| """ | |
| if audio_vae is None: | |
| raise ValueError( | |
| "h3_source_audio_target: Lock source audio requires the H3 audio " | |
| "VAE connected to Chain Context.") | |
| if not isinstance(source_audio, dict): | |
| raise ValueError( | |
| "h3_source_audio_target: Current Shot state has no source-audio " | |
| "target window. Keep Current Shot state connected to Chain Context.") | |
| from .master_audio_context import MiniMaxH3ContexMasterAudioMaskedAV | |
| option_kwargs = {} | |
| if lip_sync_options is not None: | |
| option_kwargs = { | |
| key: lip_sync_options[key] | |
| for key in ( | |
| "preroll_seconds", "lookahead_seconds", "audio_denoise", | |
| "gap_denoise", "gate_hold_seconds") | |
| } | |
| out, prefix_frames, _clip_audio = ( | |
| MiniMaxH3ContexMasterAudioMaskedAV().prepare( | |
| latent=latent, | |
| audio_vae=audio_vae, | |
| master_audio=source_audio, | |
| clip_start_seconds=clip_start_seconds, | |
| context_length=0, | |
| voice=voice, | |
| voice_clip_start_seconds=voice_clip_start_seconds, | |
| force_active_voice_prefix_seconds=( | |
| force_active_voice_prefix_seconds), | |
| **option_kwargs, | |
| )) | |
| if int(prefix_frames) != 0: | |
| raise RuntimeError( | |
| "h3_source_audio_target: audio-only target unexpectedly changed " | |
| "the video prefix.") | |
| return out | |
| def _feather_preserved_prefix(video_mask, audio_mask, video_steps, audio_steps): | |
| """Apply a narrow, high-denoise handoff at a protected AV tail.""" | |
| video_steps = int(video_steps) | |
| audio_steps = int(audio_steps) | |
| video_feather = min(4, max(0, video_steps - 1)) | |
| if video_feather < 1: | |
| video_mask[:, :, :video_steps] = 0.0 | |
| audio_mask[..., :audio_steps] = 0.0 | |
| return 0, 0 | |
| # Fractional H3 masks are most useful close to full denoise. Keep the | |
| # accepted prefix exact until the final four video-latent steps, then | |
| # give the model a deliberately narrow 0.85..0.95 reconstruction band. | |
| # Audio uses its own shorter 200 ms ramp instead of inheriting the much | |
| # wider video-to-audio grid conversion, which used to start the audible | |
| # handoff roughly 575 ms before the prefix boundary. | |
| audio_feather = min(max(0, audio_steps - 1), 8) | |
| video_ramp = torch.linspace( | |
| 0.85, 0.95, video_feather, | |
| device=video_mask.device, dtype=video_mask.dtype, | |
| ) | |
| video_hard_steps = video_steps - video_feather | |
| video_mask[:, :, :video_hard_steps] = 0.0 | |
| video_mask[:, :, video_hard_steps:video_steps] = torch.minimum( | |
| video_mask[:, :, video_hard_steps:video_steps], | |
| video_ramp.view(1, 1, video_feather, 1, 1), | |
| ) | |
| if audio_feather: | |
| audio_ramp = torch.linspace( | |
| 0.85, 0.95, audio_feather, | |
| device=audio_mask.device, dtype=audio_mask.dtype, | |
| ) | |
| audio_hard_steps = audio_steps - audio_feather | |
| audio_mask[..., :audio_hard_steps] = 0.0 | |
| audio_mask[..., audio_hard_steps:audio_steps] = torch.minimum( | |
| audio_mask[..., audio_hard_steps:audio_steps], | |
| audio_ramp.view(1, 1, 1, audio_feather), | |
| ) | |
| return video_feather, audio_feather | |
| def _audio_feather_preserved_prefix( | |
| video_mask, audio_mask, video_steps, audio_steps, | |
| ): | |
| """Keep picture exact and release only the final audio-prefix ticks.""" | |
| video_steps = int(video_steps) | |
| audio_steps = int(audio_steps) | |
| video_mask[:, :, :video_steps] = 0.0 | |
| audio_feather = min(8, max(0, audio_steps)) | |
| audio_hard_steps = audio_steps - audio_feather | |
| audio_mask[..., :audio_hard_steps] = 0.0 | |
| if audio_feather: | |
| indices = torch.arange( | |
| 1, audio_feather + 1, | |
| device=audio_mask.device, dtype=audio_mask.dtype, | |
| ) | |
| audio_ramp = 0.5 - 0.5 * torch.cos( | |
| torch.pi * indices / float(audio_feather)) | |
| audio_mask[..., audio_hard_steps:audio_steps] = torch.minimum( | |
| audio_mask[..., audio_hard_steps:audio_steps], | |
| audio_ramp.view(1, 1, 1, audio_feather), | |
| ) | |
| return audio_feather | |
| def _drop_prefix_guides(conditioning, prefix_frames): | |
| """Remove target guides that conflict with the preserved latent prefix.""" | |
| out = [] | |
| dropped = [] | |
| for embedding, extra in conditioning: | |
| metadata = extra.copy() | |
| kept = [] | |
| for guide in metadata.get("minimax_keyframes") or []: | |
| position = float(guide.get( | |
| "resolved_frame_index", guide.get("frame_index", 0))) | |
| if 0 <= position < int(prefix_frames): | |
| dropped.append(position) | |
| else: | |
| kept.append(guide) | |
| if "minimax_keyframes" in metadata: | |
| metadata["minimax_keyframes"] = kept | |
| out.append([embedding, metadata]) | |
| if dropped: | |
| _LOG.warning( | |
| "h3_masked_prefix: dropped %d target guide(s) inside preserved " | |
| "frames 0..%d; the clean target latent already owns that prefix.", | |
| len(dropped), int(prefix_frames) - 1, | |
| ) | |
| return out | |
| def apply_masked_prefix( | |
| conditioning, | |
| vae, | |
| latent, | |
| previous_frames, | |
| context_length, | |
| crop, | |
| previous_latent=None, | |
| audio_vae=None, | |
| previous_audio=None, | |
| temporal_feather=False, | |
| audio_only_feather=False, | |
| preserve_audio_prefix=True, | |
| detail_video_taper=False, | |
| detail_video_seed=0, | |
| context_spatial_proxy="off", | |
| latent_color_carry=None, | |
| ): | |
| """Return conditioning, masked target latent, and repeated trim length.""" | |
| _require_h3_mask_support() | |
| target_video, target_audio, target_frames = _validate_target_streams(latent) | |
| width = int(target_video.shape[4]) * 16 | |
| height = int(target_video.shape[3]) * 16 | |
| if previous_latent is not None: | |
| source_video = _streams_from_latent(previous_latent)[0] | |
| if source_video.ndim == 4: | |
| source_video = source_video.unsqueeze(0) | |
| if source_video.ndim != 5: | |
| raise ValueError( | |
| "h3_masked_prefix: previous video latent must be " | |
| "[B,C,T,H,W], got %s." % (tuple(source_video.shape),) | |
| ) | |
| available = _pixel_frames(int(source_video.shape[2])) | |
| frames = _snap_prefix_length( | |
| context_length, available, target_frames, | |
| preserve_audio_prefix=preserve_audio_prefix) | |
| video_prefix, video_steps = _generated_video_tail( | |
| previous_latent, frames, target_video) | |
| video_source = "previous sampled latent" | |
| else: | |
| if getattr(previous_frames, "ndim", 0) != 4: | |
| raise ValueError( | |
| "h3_masked_prefix: imported previous frames must be IMAGE " | |
| "[N,H,W,C]." | |
| ) | |
| available = int(previous_frames.shape[0]) | |
| frames = _snap_prefix_length( | |
| context_length, available, target_frames, | |
| preserve_audio_prefix=preserve_audio_prefix) | |
| video_prefix, video_steps = _encoded_video_tail( | |
| vae, previous_frames, frames, target_video, crop) | |
| video_source = "imported decoded frames via video VAE" | |
| if video_steps >= int(target_video.shape[2]): | |
| raise ValueError( | |
| "h3_masked_prefix: video prefix consumes the whole target latent." | |
| ) | |
| context_spatial_proxy = str(context_spatial_proxy or "off").strip().lower() | |
| if context_spatial_proxy == "latent_5_6": | |
| if previous_latent is None: | |
| raise ValueError( | |
| "h3_context_spatial_proxy: latent 5/6 requires a sampled " | |
| "predecessor latent; imported RGB context is not eligible.") | |
| original_height = int(video_prefix.shape[3]) | |
| original_width = int(video_prefix.shape[4]) | |
| video_prefix, proxy_size = _latent_context_spatial_proxy(video_prefix) | |
| video_source += " + 5/6 latent spatial proxy" | |
| _LOG.info( | |
| "h3_context_spatial_proxy: AV boundary latent %dx%d -> %dx%d " | |
| "-> %dx%d; audio and saved predecessor are unchanged.", | |
| original_width, original_height, int(proxy_size[1]), | |
| int(proxy_size[0]), original_width, original_height) | |
| elif context_spatial_proxy != "off": | |
| raise ValueError( | |
| "h3_context_spatial_proxy: masked AV accepts only off or " | |
| "latent_5_6, got %r." % context_spatial_proxy) | |
| latent_color_summary = None | |
| if latent_color_carry is not None: | |
| if previous_latent is None: | |
| raise ValueError( | |
| "h3_latent_color_carry: Color-Stable Drift AV requires a " | |
| "sampled generated predecessor latent; imported scene-1 " | |
| "context is not eligible.") | |
| if not isinstance(latent_color_carry, dict): | |
| raise ValueError( | |
| "h3_latent_color_carry: scene anchor data is malformed.") | |
| from .latent_color_carry import apply_delta_vae_color_carry | |
| video_prefix, latent_color_summary = apply_delta_vae_color_carry( | |
| video_prefix, | |
| vae, | |
| latent_color_carry.get("anchor_stats"), | |
| latent_color_carry.get("current_stats"), | |
| ) | |
| if bool(latent_color_summary.get("applied", False)): | |
| video_source += " + tapered scene-one VAE-delta color carry" | |
| else: | |
| video_source += " + scene-one color carry (neutral)" | |
| detail_scale = None | |
| detail_alphas = () | |
| if bool(detail_video_taper): | |
| if int(frames) != int(DETAIL_AV_RECIPE["context_frames"]): | |
| raise ValueError( | |
| "h3_detail_av: the experimental v2 recipe requires exactly " | |
| "%d context frames, got %d." % | |
| (int(DETAIL_AV_RECIPE["context_frames"]), int(frames)) | |
| ) | |
| video_prefix, detail_scale, detail_alphas = _detail_av_video_prefix( | |
| video_prefix, detail_video_seed) | |
| video_source += " + disposable Detail AV latent taper" | |
| if not bool(preserve_audio_prefix): | |
| audio_prefix = None | |
| audio_steps = 0 | |
| audio_source = "open target (generated continuity off)" | |
| elif previous_latent is not None: | |
| audio_prefix, audio_steps, overhang = _audio_tail_from_latent( | |
| previous_latent, frames) | |
| audio_source = "previous sampled latent" | |
| if abs(overhang) > 1e-9: | |
| _LOG.warning( | |
| "h3_masked_prefix: predecessor audio grid ends %.3f latent " | |
| "steps from its last video frame; the copied %d-step prefix " | |
| "is end-aligned. Use 39/90/141/... context frames for an exact " | |
| "prefix duration.", | |
| overhang, audio_steps, | |
| ) | |
| else: | |
| audio_prefix, audio_steps, audio_source = _encode_imported_audio( | |
| audio_vae, previous_audio, frames) | |
| if bool(preserve_audio_prefix): | |
| expected_audio_steps = int(round(frames / float(FPS) * AUDIO_HZ)) | |
| if audio_steps != expected_audio_steps: | |
| raise RuntimeError( | |
| "h3_masked_prefix: %d video frames require %d audio steps, " | |
| "got %d." % (frames, expected_audio_steps, audio_steps) | |
| ) | |
| if audio_steps >= int(target_audio.shape[-1]): | |
| raise ValueError( | |
| "h3_masked_prefix: audio prefix consumes the whole target " | |
| "latent.") | |
| out_video = target_video.clone() | |
| out_audio = target_audio.clone() | |
| vp = video_prefix[:1].to(out_video.device, out_video.dtype) | |
| ap = (audio_prefix[:1].to(out_audio.device, out_audio.dtype) | |
| if audio_prefix is not None else None) | |
| if (int(vp.shape[1]) != int(out_video.shape[1]) | |
| or tuple(vp.shape[3:]) != tuple(out_video.shape[3:])): | |
| raise ValueError( | |
| "h3_masked_prefix: encoded video prefix shape %s does not match " | |
| "target %s." % (tuple(vp.shape), tuple(out_video.shape)) | |
| ) | |
| if ap is not None and tuple(ap.shape[1:3]) != tuple(out_audio.shape[1:3]): | |
| raise ValueError( | |
| "h3_masked_prefix: audio prefix shape %s does not match target %s." | |
| % (tuple(ap.shape), tuple(out_audio.shape)) | |
| ) | |
| out_video[:, :, :video_steps] = vp | |
| if ap is not None: | |
| out_audio[..., :audio_steps] = ap | |
| video_mask, audio_mask = _existing_mask_streams( | |
| latent, out_video, out_audio) | |
| video_feather_steps = audio_feather_steps = 0 | |
| if bool(audio_only_feather): | |
| audio_feather_steps = _audio_feather_preserved_prefix( | |
| video_mask, audio_mask, video_steps, audio_steps) | |
| elif bool(temporal_feather): | |
| video_feather_steps, audio_feather_steps = _feather_preserved_prefix( | |
| video_mask, audio_mask, video_steps, audio_steps) | |
| else: | |
| video_mask[:, :, :video_steps] = 0.0 | |
| audio_mask[..., :audio_steps] = 0.0 | |
| import comfy.nested_tensor | |
| out_latent = latent.copy() | |
| out_latent["samples"] = comfy.nested_tensor.NestedTensor( | |
| (out_video, out_audio)) | |
| out_latent["noise_mask"] = comfy.nested_tensor.NestedTensor( | |
| (video_mask, audio_mask)) | |
| out_conditioning = _drop_prefix_guides(conditioning, frames) | |
| if not bool(preserve_audio_prefix): | |
| mask_summary = "audio fully denoisable (generated continuity off)" | |
| elif audio_only_feather: | |
| mask_summary = "audio-only half-cosine feather %d audio steps" % ( | |
| audio_feather_steps) | |
| elif temporal_feather: | |
| mask_summary = "temporal feather %d video / %d audio steps" % ( | |
| video_feather_steps, audio_feather_steps) | |
| else: | |
| mask_summary = "hard prefix mask" | |
| if detail_scale is not None: | |
| mask_summary += ( | |
| "; Detail AV video-only Gaussian taper %.3f -> %.3f over " | |
| "%d steps (sigma %.4f, seed %d)" % | |
| (float(detail_alphas[0]), float(detail_alphas[-1]), | |
| int(DETAIL_AV_RECIPE["ramp_steps"]), float(detail_scale), | |
| int(detail_video_seed))) | |
| if latent_color_summary is not None: | |
| mask_summary += ( | |
| "; scene-one latent color delta %s, brightness %+.4f, " | |
| "saturation %.4f, audio unchanged" % | |
| ("applied" if latent_color_summary.get("applied") else "neutral", | |
| float(latent_color_summary.get("brightness", 0.0)), | |
| float(latent_color_summary.get("saturation", 1.0)))) | |
| _LOG.info( | |
| "h3_masked_prefix: preserved %d target frames = %d video steps / %d " | |
| "audio steps (%.3fs, video from %s, audio from %s); %s; target %d " | |
| "frames at %dx%d; trim %d", | |
| frames, video_steps, audio_steps, frames / float(FPS), video_source, | |
| audio_source, | |
| mask_summary, | |
| target_frames, width, height, frames, | |
| ) | |
| return out_conditioning, out_latent, frames | |