diff --git a/src/h3_blackwell_runtime/packing.py b/src/h3_blackwell_runtime/packing.py index 511a7cd..e86e1de 100644 --- a/src/h3_blackwell_runtime/packing.py +++ b/src/h3_blackwell_runtime/packing.py @@ -96,6 +96,7 @@ class H3PromptPacker: *, text_token_tags: torch.Tensor | None = None, cond_latents: list[torch.Tensor] | None = None, + cond_frame_indices: list[int] | None = None, frame_count: int | None = None, seed: int = 0, ) -> tuple[torch.Tensor, torch.Tensor, list[tuple[int, int, int]], torch.Tensor, tuple[int, int, int], tuple[int, int, int]]: @@ -104,8 +105,9 @@ class H3PromptPacker: ``text`` is the refined text span (width 5376 when already refined, 5120 for raw Qwen states); ``text_token_tags`` is the per-token DiT modality tag (1=text, 0=video over vision pads). ``cond_latents`` are normalized keyframe - latents ``[1,24,1,H/16,W/16]`` spliced right after the text as non-denoised - cond rows with their own near-1 timestep. Returns + latents ``[1,24,1,H/16,W/16]`` and ``cond_frame_indices`` preserves each + keyframe's resolved first/last pixel index. They are spliced right after + the text as non-denoised cond rows with their own near-1 timestep. Returns ``(hidden, times, segments, positions, video_seg, audio_seg)`` where ``segments`` rows are ``t_row*3 + modality_tag``. """ @@ -122,7 +124,10 @@ class H3PromptPacker: video_rows = functional.linear(patchify_video(video.to(torch.bfloat16)).float(), self.video_weight, self.video_bias).to(torch.bfloat16) audio_rows = functional.linear(pack_audio(audio.to(torch.bfloat16)).float(), self.audio_weight, self.audio_bias).to(torch.bfloat16) cond_rows = None + cond_lengths = [] if cond_latents: + if cond_frame_indices is None or len(cond_frame_indices) != len(cond_latents): + raise ValueError("cond_frame_indices must match cond_latents") cond_patches = [] # every cond video restarts the same CPU RNG stream (Comfy _cond_video_rows) for idx, z in enumerate(cond_latents): @@ -132,6 +137,7 @@ class H3PromptPacker: noise = torch.randn(r.shape, generator=gen, dtype=torch.float32) r = self.VISUAL_COND_TIMESTEP * r + (1.0 - self.VISUAL_COND_TIMESTEP) * noise.to(r.device) cond_patches.append(r) + cond_lengths.append(r.shape[0]) cond_rows = functional.linear(torch.cat(cond_patches, dim=0), self.video_weight.to(torch.float32), self.video_bias.to(torch.float32)).to(torch.bfloat16) if model_timesteps is None: @@ -168,13 +174,17 @@ class H3PromptPacker: position_blocks = [text_positions] if cond_rows is not None and cond_latents: spans = _video_t_spans(latent_t) - cond_t_values = [ - float(text_length) if idx == 0 else (float(text_length) + sum(spans) - FRAME_RESCALE if frame_count is not None and idx == len(cond_latents) - 1 else float(text_length)) - for idx in range(len(cond_latents)) - ] + cond_t_values = [] + for pixel_index in cond_frame_indices: + if pixel_index == 0: + cond_t_values.append(float(text_length)) + elif frame_count is not None and pixel_index == frame_count - 1: + cond_t_values.append(float(text_length) + sum(spans) - FRAME_RESCALE) + else: + raise ValueError("only first/last keyframe anchors are supported") position_blocks.append(torch.cat([_cond_positions(frame_rows, cond_t, latent_h, latent_w) for cond_t in cond_t_values], dim=0)) - position_blocks.append(_audio_positions(audio.shape[-1], float(text_length + cond_length), latent_w, latent_h)) - position_blocks.append(_video_positions(latent_t, latent_h, latent_w, float(text_length + cond_length + audio_length))) + position_blocks.append(_audio_positions(audio.shape[-1], float(text_length), latent_w, latent_h)) + position_blocks.append(_video_positions(latent_t, latent_h, latent_w, float(text_length))) positions = torch.cat(position_blocks, dim=0) # mod_segments: (start, stop, t_row*3 + tag). @@ -192,8 +202,9 @@ class H3PromptPacker: cursor_start = text_length if cond_rows is not None: - segments.append((cursor_start, cursor_start + cond_rows.shape[0], t_row[cond_time] * 3 + 0)) - cursor_start += cond_rows.shape[0] + for length in cond_lengths: + segments.append((cursor_start, cursor_start + length, t_row[cond_time] * 3 + 0)) + cursor_start += length segments.append((cursor_start, cursor_start + audio_length, t_row[audio_time] * 3 + 2)) cursor_start += audio_length video_start = cursor_start diff --git a/src/h3_blackwell_runtime/qwen3vl_text.py b/src/h3_blackwell_runtime/qwen3vl_text.py index 3f641a0..522b769 100644 --- a/src/h3_blackwell_runtime/qwen3vl_text.py +++ b/src/h3_blackwell_runtime/qwen3vl_text.py @@ -194,7 +194,7 @@ class Qwen3VL32BTextEncoder(nn.Module): """Scaled token embeds ``[batch, tokens, 5120]`` in fp32 (pre-decoder).""" token_rows = F.embedding(input_ids.to(self.embed_tokens.device), self.embed_tokens).to(torch.float32) token_scales = F.embedding(input_ids.to(self.embed_scale.device), self.embed_scale) - return (token_rows * token_scales).to(torch.float32) + return (token_rows * token_scales).to(torch.bfloat16).to(torch.float32) def _run_layers( self, diff --git a/src/h3_blackwell_runtime/qwen3vl_vision.py b/src/h3_blackwell_runtime/qwen3vl_vision.py index 97c743f..d8fd7c1 100644 --- a/src/h3_blackwell_runtime/qwen3vl_vision.py +++ b/src/h3_blackwell_runtime/qwen3vl_vision.py @@ -25,6 +25,7 @@ from dataclasses import dataclass from pathlib import Path import torch +from torch.nn.attention import SDPBackend, sdpa_kernel from safetensors import safe_open from torch import nn from torch.nn import functional as F @@ -193,7 +194,7 @@ def mrope_freqs_cis(position_ids: torch.Tensor, *, theta: float = TEXT_ROPE_THET freqs = (inv_freq_expanded @ position_ids_expanded).transpose(1, 2) freqs_inter = freqs[0].clone() for axis_idx, offset in ((1, 1), (2, 2)): - length = rope_dims[axis_idx - 1] * 3 + length = rope_dims[axis_idx] * 3 idx = slice(offset, length, 3) freqs_inter[..., idx] = freqs[axis_idx, ..., idx] emb = torch.cat((freqs_inter, freqs_inter), dim=-1) @@ -214,7 +215,24 @@ class _VisionPatchEmbed(nn.Module): target = self.weight.dtype x = x.view(-1, 3, VISION_TEMPORAL, VISION_PATCH, VISION_PATCH) s = (VISION_TEMPORAL, VISION_PATCH, VISION_PATCH) - return F.conv3d(x.to(target), self.weight, self.bias, stride=s).view(-1, self.weight.shape[0]) + x = x.to(target) + if x.is_cuda and target in (torch.float16, torch.bfloat16): + # Match Comfy's NVIDIA Conv3d workaround dispatch exactly. + output = torch.cudnn_convolution( + x, + self.weight, + (0, 0, 0), + s, + (1, 1, 1), + 1, + benchmark=False, + deterministic=False, + allow_tf32=True, + ) + output += self.bias.view(1, -1, 1, 1, 1) + else: + output = F.conv3d(x, self.weight, self.bias, stride=s) + return output.view(-1, self.weight.shape[0]) class _VisionMLP(nn.Module): @@ -244,24 +262,27 @@ def _apply_rope_vision(q: torch.Tensor, k: torch.Tensor, freqs) -> tuple[torch.T their native dtype. """ cos, sin, neg_sin = freqs - q = (q * cos) + original_q = q + q = original_q * cos split = q.shape[-1] // 2 - q[..., :split] += q[..., split:] * neg_sin - q[..., split:] += q[..., :split] * sin - k = (k * cos) - k[..., :split] += k[..., split:] * neg_sin - k[..., split:] += k[..., :split] * sin + q[..., :split].addcmul_(original_q[..., split:], neg_sin) + q[..., split:].addcmul_(original_q[..., :split], sin) + original_k = k + k = original_k * cos + k[..., :split].addcmul_(original_k[..., split:], neg_sin) + k[..., split:].addcmul_(original_k[..., :split], sin) return q, k class _VisionAttention(nn.Module): - def __init__(self, qkv_w: torch.Tensor, qkv_b: torch.Tensor, proj_w: torch.Tensor, *, num_heads: int, head_dim: int): + def __init__(self, qkv_w: torch.Tensor, qkv_b: torch.Tensor, proj_w: torch.Tensor, proj_b: torch.Tensor, *, num_heads: int, head_dim: int): super().__init__() self.num_heads = num_heads self.head_dim = head_dim self.register_buffer("qkv_weight", qkv_w, persistent=False) self.register_buffer("qkv_bias", qkv_b, persistent=False) - self.proj_weight = proj_w # no bias (Qwen3.5 vision proj has none) + self.register_buffer("proj_weight", proj_w, persistent=False) + self.register_buffer("proj_bias", proj_b, persistent=False) def forward(self, x: torch.Tensor, cu_seqlens: torch.Tensor, position_embeddings: torch.Tensor) -> torch.Tensor: seq_length = x.shape[0] @@ -286,10 +307,24 @@ class _VisionAttention(nn.Module): torch.split(key_states, lengths, dim=0), torch.split(value_states, lengths, dim=0), ): - attn_outputs.append(F.scaled_dot_product_attention(q.transpose(0, 1).unsqueeze(0), k.transpose(0, 1).unsqueeze(0), v.transpose(0, 1).unsqueeze(0))) + with sdpa_kernel( + [ + SDPBackend.FLASH_ATTENTION, + SDPBackend.CUDNN_ATTENTION, + SDPBackend.EFFICIENT_ATTENTION, + SDPBackend.MATH, + ], + set_priority=True, + ): + output = F.scaled_dot_product_attention( + q.transpose(0, 1).unsqueeze(0), + k.transpose(0, 1).unsqueeze(0), + v.transpose(0, 1).unsqueeze(0), + ) + attn_outputs.append(output.transpose(1, 2).reshape(1, q.shape[0], -1)) attn_output = torch.cat(attn_outputs, dim=1) attn_output = attn_output.reshape(seq_length, -1) - return F.linear(attn_output, self.proj_weight) + return F.linear(attn_output, self.proj_weight, self.proj_bias) class _VisionBlock(nn.Module): @@ -329,16 +364,16 @@ class _VisionPatchMerger(nn.Module): self.out_hidden_size = out_hidden_size def forward(self, x: torch.Tensor) -> torch.Tensor: - # x: [t*h*w, hidden] (unmerged patches) for the main merger; - # [t*(h//2)*(w//2), merge_dim] (pre-merged 2x2) for the deepstack merger. - if x.shape[-1] == self.merge_dim: - # Deepstack merger: input is already 2x2-merged; norm over merge_dim. + # x is the existing 2x2-block-major [t*h*w, hidden] patch stream. + if self.norm_dim == self.merge_dim: + # DeepStack merger: group the existing block-major patch stream first. + x = x.view(-1, self.merge_dim) x = F.layer_norm(x, (self.merge_dim,), weight=self.norm_weight, bias=self.norm_bias, eps=1e-6) else: # Main merger: per-patch LayerNorm over hidden, then group 2x2 into merge_dim. x = F.layer_norm(x, (x.shape[-1],), weight=self.norm_weight, bias=self.norm_bias, eps=1e-6) x = x.view(-1, self.merge_dim) - return F.linear(F.gelu(F.linear(x, self.fc1_weight, self.fc1_bias), approximate="tanh"), self.fc2_weight, self.fc2_bias) + return F.linear(F.gelu(F.linear(x, self.fc1_weight, self.fc1_bias)), self.fc2_weight, self.fc2_bias) def resize_keyframe(image: torch.Tensor, width: int, height: int, *, crop: str = "disabled") -> torch.Tensor: @@ -377,15 +412,13 @@ def resize_keyframe(image: torch.Tensor, width: int, height: int, *, crop: str = return samples.clamp(0.0, 1.0).movedim(1, -1) # [1, H, W, 3] -def _text_run_ids(prompt: str) -> list[int]: +def _text_run_ids(tokenizer, prompt: str) -> list[int]: """Token ids for a raw text run (``add_special_tokens=False``, no template).""" - from .conditioning import H3PromptTokenizer - - tokenizer_dir = Path(__file__).with_name("qwen25_tokenizer") - if not tokenizer_dir.exists(): - raise FileNotFoundError(f"Qwen tokenizer directory missing: {tokenizer_dir}") - ids = H3PromptTokenizer(tokenizer_dir)(prompt or " ") - # input_ids is [1, seq]; flatten to a Python list of ints. + raw_tokenizer = getattr(tokenizer, "tokenizer", None) + if raw_tokenizer is not None: + ids = raw_tokenizer(prompt, add_special_tokens=False, return_tensors="pt").input_ids + else: + ids = tokenizer(prompt, device="cpu") return [int(t) for t in ids.reshape(-1).tolist()] @@ -441,11 +474,11 @@ def build_fl2va_presentation( # Build the entry list: (token_id/"text") runs and image placeholders. entries: list = [] for i in range(len(images)): - entries.extend((tid, "text") for tid in _text_run_ids(f": ")) entries.append((VISION_START, "text")) entries.append((i, "image")) entries.append((VISION_END, "text")) - entries.extend((tid, "text") for tid in _text_run_ids(prompt)) + entries.extend((tid, "text") for tid in _text_run_ids(tokenizer, prompt)) if not any(kind == "text" for _, kind in entries): entries = [(151643, "text")] @@ -547,6 +580,7 @@ class Qwen3VL32BVision(nn.Module): f"visual.blocks.{i}.attn.qkv.weight", f"visual.blocks.{i}.attn.qkv.bias", f"visual.blocks.{i}.attn.proj.weight", + f"visual.blocks.{i}.attn.proj.bias", f"visual.blocks.{i}.norm2.weight", f"visual.blocks.{i}.norm2.bias", f"visual.blocks.{i}.mlp.linear_fc1.weight", @@ -574,6 +608,8 @@ class Qwen3VL32BVision(nn.Module): dtype, lambda name: checkpoint.get_tensor(name).to(device=device, dtype=dtype), ) + # Comfy's Embedding is not dynamically cast to the FP32 vision stream. + self.pos_embed = checkpoint.get_tensor("visual.pos_embed.weight").to(device=device) def _init_modules(self, device, dtype, get) -> None: self.register_buffer("pos_embed", get("visual.pos_embed.weight"), persistent=False) @@ -600,7 +636,7 @@ class Qwen3VL32BVision(nn.Module): get(f"visual.blocks.{i}.norm1.weight"), get(f"visual.blocks.{i}.norm1.bias"), _VisionAttention( get(f"visual.blocks.{i}.attn.qkv.weight"), get(f"visual.blocks.{i}.attn.qkv.bias"), - get(f"visual.blocks.{i}.attn.proj.weight"), + get(f"visual.blocks.{i}.attn.proj.weight"), get(f"visual.blocks.{i}.attn.proj.bias"), num_heads=self.num_heads, head_dim=VISION_HEAD_DIM, ), get(f"visual.blocks.{i}.norm2.weight"), get(f"visual.blocks.{i}.norm2.bias"), @@ -631,7 +667,7 @@ class Qwen3VL32BVision(nn.Module): intra_row = torch.arange(merge_size, device=device) intra_col = torch.arange(merge_size, device=device) row_idx = (block_rows[:, None, None, None] * merge_size + intra_row[None, None, :, None]).expand(merged_h, merged_w, merge_size, merge_size).reshape(-1) - col_idx = (block_cols[None, :, None, None] * merge_size + intra_col[None, None, :, None]).expand(merged_h, merged_w, merge_size, merge_size).reshape(-1) + col_idx = (block_cols[None, :, None, None] * merge_size + intra_col[None, None, None, :]).expand(merged_h, merged_w, merge_size, merge_size).reshape(-1) coords = torch.stack((row_idx, col_idx), dim=-1) if num_frames > 1: coords = coords.repeat(num_frames, 1) @@ -688,38 +724,6 @@ class Qwen3VL32BVision(nn.Module): patch_pos_embeds_permute.append(pos_embed) return torch.cat(patch_pos_embeds_permute) - @staticmethod - def _merge_tokens(x: torch.Tensor, grid_thw: torch.Tensor) -> torch.Tensor: - """DeepStack layout: interleave 2x2 spatial neighbours and flatten. - - x: [t*h*w, C] -> [t*(h//2)*(w//2), 4*C]. - """ - first = grid_thw[0] - t = int(first[0].item()) - h = int(first[1].item()) - w = int(first[2].item()) - C = int(x.shape[-1]) - merge = 2 - x = x.view(t, h // merge, merge, w // merge, merge, C) - x = x.permute(0, 1, 3, 2, 4, 5).contiguous() - return x.reshape(-1, C * merge * merge) - - @staticmethod - def _interleave_2x2(x: torch.Tensor, grid_thw: torch.Tensor) -> torch.Tensor: - """Re-order [t*h*w, C] into 2x2-block-major order (no dim change). - - After this, x.view(-1, 4*C) groups spatially-adjacent 2x2 blocks. - """ - first = grid_thw[0] - t = int(first[0].item()) - h = int(first[1].item()) - w = int(first[2].item()) - C = int(x.shape[-1]) - merge = 2 - x = x.view(t, h // merge, merge, w // merge, merge, C) - x = x.permute(0, 1, 3, 2, 4, 5).contiguous() - return x.reshape(-1, C) # still [t*(h//2)*(w//2)*4, C], block-major - def forward(self, flatten_patches: torch.Tensor, grid_thw: torch.Tensor) -> tuple[torch.Tensor, list[torch.Tensor]]: """Run the visual tower -> (merged, deepstack).""" x = self.patch_embed(flatten_patches.to(self.dtype).to(self.device)) @@ -740,22 +744,18 @@ class Qwen3VL32BVision(nn.Module): x = block(x, cu_seqlens=cu_seqlens, position_embeddings=position_embeddings) # x: [t*h*w, hidden] (unmerged patches). if layer_num in self.deepstack_visual_indexes: - # DeepStack: merge 2x2 first, then project. deepstack_features.append( - self.deepstack_merger_list[self.deepstack_visual_indexes.index(layer_num)](self._merge_tokens(x, grid_thw)) + self.deepstack_merger_list[self.deepstack_visual_indexes.index(layer_num)](x) ) - # Main merger expects the 2x2-interleaved layout (per-patch norm is done - # over the UNMERGED hidden, but the grouping of 4 patches must be - # spatially-contiguous). Re-order x into 2x2-block-major order first. - x = self._interleave_2x2(x, grid_thw) return self.merger(x), deepstack_features class _VisionRotary(nn.Module): def __init__(self, dim: int, device, dtype): super().__init__() - inv_freq = 1.0 / (10000.0 ** (torch.arange(0, dim, 2, dtype=torch.float, device=device) / dim)) - self.register_buffer("inv_freq", inv_freq, persistent=False) + # Comfy constructs this buffer on CPU, then moves the model to CUDA. + inv_freq = 1.0 / (10000.0 ** (torch.arange(0, dim, 2, dtype=torch.float) / dim)) + self.register_buffer("inv_freq", inv_freq.to(device), persistent=False) def forward(self, seqlen: int) -> torch.Tensor: seq = torch.arange(seqlen, device=self.inv_freq.device, dtype=self.inv_freq.dtype) diff --git a/src/h3_blackwell_runtime/runtime.py b/src/h3_blackwell_runtime/runtime.py index 6a840f9..fa9423c 100644 --- a/src/h3_blackwell_runtime/runtime.py +++ b/src/h3_blackwell_runtime/runtime.py @@ -93,7 +93,7 @@ class H3HotRuntime: ) self.vision_tower = self._timed_load( "vision_tower_loaded", - lambda: Qwen3VL32BVision(config.text_encoder_path, device=config.device, dtype=torch.bfloat16), + lambda: Qwen3VL32BVision(config.text_encoder_path, device=config.device, dtype=torch.float32), ) def _timed_load(self, stage: str, fn): @@ -166,19 +166,18 @@ class H3HotRuntime: cond_latents = [] for kf in presentation.keyframes: resized = resize_keyframe(kf["image"].to(device), width, height, crop="disabled" if kf["resolved_frame_index"] == 0 else "center") - pix = resized.movedim(-1, 1).to(device, dtype=torch.float32) + pix = resized.movedim(-1, 1).to(device, dtype=torch.float32).mul(2.0).sub(1.0) cond_latents.append(self.vae_encoder.encode(pix)) return presentation, cond_latents, aligned_frames def _image_to_uint8_nhwc(self, img: torch.Tensor) -> torch.Tensor: - """Normalize a [1,3,H,W] image in either [-1,1] or [0,255] to [1,H,W,3] uint8.""" + """Normalize a [1,3,H,W] image in [0,1], [-1,1], or [0,255] to NHWC uint8.""" x = img.float() if x.numel() == 0: return x - mx = x.max() - if mx > 2.0: + if x.max() > 1.0: x = x / 255.0 # already 0..255 - else: + elif x.min() < 0.0: x = (x.clamp(-1, 1) + 1) * 0.5 # -1..1 -> 0..1 return (x.movedim(1, -1).clamp(0, 1) * 255).to(torch.uint8) @@ -237,6 +236,7 @@ class H3HotRuntime: pack_kwargs = { "text_token_tags": presentation.text_token_tags, "cond_latents": cond_latents, + "cond_frame_indices": [kf["resolved_frame_index"] for kf in presentation.keyframes], "frame_count": frame_count, "seed": seed, } diff --git a/src/h3_blackwell_runtime/sampler.py b/src/h3_blackwell_runtime/sampler.py index ee029e9..d6dbec1 100644 --- a/src/h3_blackwell_runtime/sampler.py +++ b/src/h3_blackwell_runtime/sampler.py @@ -14,13 +14,20 @@ def shifted_sigma(base: torch.Tensor, shift: float) -> torch.Tensor: def beta_sigmas(steps: int, *, device: torch.device | str, alpha: float = 0.6, beta: float = 0.6) -> torch.Tensor: """Comfy's discrete beta scheduler over H3's 1,000-entry shift-12 table.""" + import numpy as np from scipy.stats import beta as beta_distribution - table = shifted_sigma(torch.arange(1, 1001, device=device, dtype=torch.float32) / 1000, 12.0) - fractions = 1.0 - torch.arange(steps, device=device, dtype=torch.float64).cpu().numpy() / steps - indices = torch.from_numpy((999 * beta_distribution.ppf(fractions, alpha, beta)).round().astype("int64")).to(device) - indices = torch.unique_consecutive(indices) - return torch.cat((table[indices], table.new_zeros(1))) + timesteps = (torch.arange(1, 1001, 1) / 1000) * 1000 + table = shifted_sigma(timesteps / 1000, 12.0) + fractions = 1.0 - np.linspace(0, 1, steps, endpoint=False) + indices = np.rint(beta_distribution.ppf(fractions, alpha, beta) * 999) + sigmas = [] + last_index = -1 + for index in indices: + if index != last_index: + sigmas.append(float(table[int(index)])) + last_index = index + return torch.FloatTensor([*sigmas, 0.0]).to(device) def res_multistep_update(x: torch.Tensor, denoised: torch.Tensor, sigma: torch.Tensor, sigma_down: torch.Tensor, old_denoised: torch.Tensor | None, old_sigma_down: torch.Tensor | None, previous_sigma: torch.Tensor | None) -> torch.Tensor: @@ -80,6 +87,7 @@ def sample_video_res_multistep( seed: int = 0, text_token_tags: torch.Tensor | None = None, cond_latents: list[torch.Tensor] | None = None, + cond_frame_indices: list[int] | None = None, frame_count: int | None = None, cache_mode: str | None = None, cache_threshold: float = 0.0, @@ -148,6 +156,7 @@ def sample_video_res_multistep( step_timesteps, text_token_tags=text_token_tags, cond_latents=cond_latents, + cond_frame_indices=cond_frame_indices, frame_count=frame_count, seed=seed, ) diff --git a/src/h3_blackwell_runtime/vae_decoder.py b/src/h3_blackwell_runtime/vae_decoder.py index c2ded36..72fe038 100644 --- a/src/h3_blackwell_runtime/vae_decoder.py +++ b/src/h3_blackwell_runtime/vae_decoder.py @@ -32,12 +32,21 @@ def dtype_from_name(name: str) -> torch.dtype: def _rms_norm(x: torch.Tensor, weight: torch.Tensor | None, eps: float) -> torch.Tensor: - if os.getenv("H3_VAE_FAST_OPS", "").lower() in {"1", "true", "yes", "on"}: - if weight is None: - return F.rms_norm(x, (x.shape[-1],), eps=eps) - return F.rms_norm(x, weight.shape, weight=weight.to(device=x.device, dtype=x.dtype), eps=eps) - result = x * torch.rsqrt(x.float().square().mean(dim=-1, keepdim=True) + eps).to(x.dtype) - return result if weight is None else result * weight.to(dtype=x.dtype) + if weight is None: + return F.rms_norm(x, (x.shape[-1],), eps=eps) + return F.rms_norm(x, weight.shape, weight=weight.to(device=x.device, dtype=x.dtype), eps=eps) + + +def _conv3d(x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor | None) -> torch.Tensor: + if x.is_cuda and weight.dtype in (torch.float16, torch.bfloat16): + output = torch.cudnn_convolution( + x, weight, (0, 0, 0), (1, 1, 1), (1, 1, 1), 1, + benchmark=False, deterministic=False, allow_tf32=True, + ) + if bias is not None: + output += bias.reshape(1, -1, 1, 1, 1) + return output + return F.conv3d(x, weight, bias) def create_token_ids(patch_dims: tuple[int, int, int], device: torch.device, dtype: torch.dtype) -> torch.Tensor: @@ -50,16 +59,13 @@ class RotaryEmbeddingND(nn.Module): super().__init__() self.rotary_base = rotary_base self.step = 2 * n_dim / dim - inv_freq = 1 / rotary_base ** torch.arange(0, 1, self.step, dtype=torch.float32, device=device) + inv_freq = 1 / rotary_base ** torch.arange(0, 1, self.step, dtype=torch.float32) self.register_buffer("inv_freq", inv_freq, persistent=False) self.angle_scale = 2.0 * math.pi def forward(self, img_ids: torch.Tensor) -> torch.Tensor: inv_freq = self.inv_freq - if inv_freq.device.type == "meta": - inv_freq = 1 / self.rotary_base ** torch.arange(0, 1, self.step, dtype=torch.float32, device=img_ids.device) - else: - inv_freq = inv_freq.to(img_ids.device) + inv_freq = inv_freq.to(device=img_ids.device, dtype=img_ids.dtype) angles = self.angle_scale * img_ids[:, :, :, None].float() * inv_freq[None, None, None, :] angles = angles.flatten(2, 3) cos, sin = torch.cos(angles), torch.sin(angles) @@ -89,11 +95,10 @@ class FeedForward(nn.Module): def _apply_rope_split_half(x: torch.Tensor, table: torch.Tensor) -> torch.Tensor: """Apply the reference split-half RoPE layout to leading rotary channels.""" - if os.getenv("H3_VAE_FAST_OPS", "").lower() in {"1", "true", "yes", "on"}: - try: - return torch.ops.comfy_kitchen.apply_rope_split_half1(x, table) - except Exception: - pass + try: + return torch.ops.comfy_kitchen.apply_rope_split_half1(x, table) + except Exception: + pass pairs = table.shape[-3] rot = pairs * 2 first, second = x[..., :pairs], x[..., pairs:rot] @@ -117,15 +122,12 @@ class Attention(nn.Module): qkv = self.to_qkv(x).view(batch, sequence, self.heads, 3 * self.dim_head) query, key, value = qkv.chunk(3, dim=-1) query, key = self.norm_q(query), self.norm_k(key) - if os.getenv("H3_VAE_FAST_OPS", "").lower() in {"1", "true", "yes", "on"}: - try: - rot = rotary_pos_emb.shape[-3] * 2 - query_rot, key_rot = torch.ops.comfy_kitchen.apply_rope_split_half(query[..., :rot], key[..., :rot], rotary_pos_emb) - query = torch.cat((query_rot, query[..., rot:]), dim=-1) - key = torch.cat((key_rot, key[..., rot:]), dim=-1) - except Exception: - query, key = _apply_rope_split_half(query, rotary_pos_emb), _apply_rope_split_half(key, rotary_pos_emb) - else: + try: + rot = rotary_pos_emb.shape[-3] * 2 + query_rot, key_rot = torch.ops.comfy_kitchen.apply_rope_split_half(query[..., :rot], key[..., :rot], rotary_pos_emb) + query = torch.cat((query_rot, query[..., rot:]), dim=-1) + key = torch.cat((key_rot, key[..., rot:]), dim=-1) + except Exception: query, key = _apply_rope_split_half(query, rotary_pos_emb), _apply_rope_split_half(key, rotary_pos_emb) query, key, value = query.transpose(1, 2), key.transpose(1, 2), value.transpose(1, 2) try: @@ -250,7 +252,7 @@ class MiniMaxH3VideoVAE(nn.Module): return model def _decode_pixels(self, z: torch.Tensor) -> torch.Tensor: - return self.decoder(self.post_quant_conv(z)) + return self.decoder(_conv3d(z, self.post_quant_conv.weight, self.post_quant_conv.bias)) def split_tiles(self, length: int) -> tuple[list[int], list[int], list[int]]: if self.tile_size >= length: diff --git a/src/h3_blackwell_runtime/vae_encoder.py b/src/h3_blackwell_runtime/vae_encoder.py index 18711d1..0d64ced 100644 --- a/src/h3_blackwell_runtime/vae_encoder.py +++ b/src/h3_blackwell_runtime/vae_encoder.py @@ -1,8 +1,8 @@ """Direct, encoder-only MiniMax H3 video VAE implementation. Mirrors the encoder half of ``upstream_vae.py`` so keyframe/reference images can -be encoded without ComfyUI. The encoder runs in FP32 and the latent moments are -upcast to FP32 before mean/std normalization (the reference contract). +be encoded without ComfyUI. The encoder runs in FP16 like Comfy's keyframe VAE +path, and latent moments are upcast for mean/std normalization. Causal-conv semantics: spatial padding is reflect; temporal padding is causal (front-only zeros) with a stride grid that starts at the first input frame @@ -18,6 +18,7 @@ stateless kernels below. from __future__ import annotations import math +import os from pathlib import Path import torch @@ -41,6 +42,18 @@ NIN_LEVELS = frozenset({1, 3, 5}) DOWNSAMPLE_LEVELS = frozenset({0, 1, 2, 3}) +def _conv3d(x, weight, bias, stride=(1, 1, 1), padding=(0, 0, 0)): + if x.is_cuda and weight.dtype in (torch.float16, torch.bfloat16): + output = torch.cudnn_convolution( + x, weight, padding, stride, (1, 1, 1), 1, + benchmark=False, deterministic=False, allow_tf32=True, + ) + if bias is not None: + output += bias.reshape(1, -1, 1, 1, 1) + return output + return F.conv3d(x, weight, bias, stride, padding) + + def _causal_conv3d(x, weight, bias, *, kernel_size, stride, spatial_padding, temporal_causal): """Causal 3D conv (matches upstream_vae.CausalConv3d). @@ -48,23 +61,22 @@ def _causal_conv3d(x, weight, bias, *, kernel_size, stride, spatial_padding, tem - ``temporal_causal``: front-zero T by ``kernel_size - 1`` (the reference's ``causal_padding[0] * 2`` for ``causal_padding[0]=1``, which is every 3D-causal conv in the H3 VAE). A single-frame input truncates the - temporal taps to the center 1 instead of convolving zero rows. + temporal taps to the final input-aligned tap instead of convolving zero rows. - If neither applies: no padding at all (reference early-return). """ if x.shape[2] == 1: # Keyframe path (matches reference's `autopad="causal_zero"`): # apply spatial-reflect pad (if any), then run an effective 2D conv - # by slicing the kernel to its center temporal tap (T_out stays 1). + # by slicing the kernel to its final input-aligned temporal tap. if spatial_padding > 0: x = F.pad(x, (spatial_padding, spatial_padding, spatial_padding, spatial_padding, 0, 0), mode="reflect") - half = (kernel_size - 1) // 2 - kernel_5d = weight[:, :, half:half + 1, :, :] - return F.conv3d(x, kernel_5d, bias, (1, stride[1], stride[2]), (0, 0, 0)) + kernel_5d = weight[:, :, -1:, :, :] + return _conv3d(x, kernel_5d, bias, (1, stride[1], stride[2])) if spatial_padding > 0: x = F.pad(x, (spatial_padding, spatial_padding, spatial_padding, spatial_padding, 0, 0), mode="reflect") if temporal_causal: x = F.pad(x, (0, 0, 0, 0, kernel_size - 1, 0)) - return F.conv3d(x, weight, bias, stride, (0, 0, 0)) + return _conv3d(x, weight, bias, stride) def _group_norm_3d(x, weight, bias): @@ -76,7 +88,7 @@ def _group_norm_3d(x, weight, bias): def _resnet(x, p): # nin_shortcut uses CausalConv3d(k=1, padding=1) in the reference. - residual = x if p["nin"] is None else F.conv3d(x, p["nin"][0], p["nin"][1], (1, 1, 1)) + residual = x if p["nin"] is None else _conv3d(x, p["nin"][0], p["nin"][1]) h = _causal_conv3d(F.silu(_group_norm_3d(x, p["norm1_w"], p["norm1_b"])), p["conv1_w"], p["conv1_b"], kernel_size=3, stride=(1, 1, 1), spatial_padding=1, temporal_causal=True) h = _causal_conv3d(F.silu(_group_norm_3d(h, p["norm2_w"], p["norm2_b"])), p["conv2_w"], p["conv2_b"], kernel_size=3, stride=(1, 1, 1), spatial_padding=1, temporal_causal=True) return h.add_(residual) @@ -84,8 +96,8 @@ def _resnet(x, p): def _downsample(x, p): if p["space"] == 2: - # Reference Downsample3D pads H and W by +1 reflect before the conv. - x = F.pad(x, (1, 1, 1, 1, 0, 0), mode="reflect") + # Reference Downsample3D pads only the right and bottom edges. + x = F.pad(x, (0, 1, 0, 1, 0, 0), mode="reflect") # Conv uses padding=(1,0,0) -> causal_padding=(1,0,0), so spatial pad=0, # temporal front-zero is applied. return _causal_conv3d(x, p["w"], p["b"], kernel_size=3, stride=(p["time"], p["space"], p["space"]), spatial_padding=0, temporal_causal=True) @@ -109,8 +121,9 @@ class MiniMaxH3VideoVAEEncoder(nn.Module): ``self.W`` (a dict), so no ``nn.Module`` sub-hierarchy is needed. """ - def __init__(self, *, tiling: bool = True): + def __init__(self, *, tiling: bool = True, compute_dtype: torch.dtype = torch.float16): super().__init__() + self.compute_dtype = compute_dtype self.vae_ratio, self.vae_ratio_t = VAE_RATIO, 4 self.clip_length, self.token_drop = 17, 3 self.frame_pre_padding = (-self.clip_length) % self.vae_ratio_t @@ -122,7 +135,13 @@ class MiniMaxH3VideoVAEEncoder(nn.Module): self.register_buffer("pixel_std", torch.tensor(IMAGENET_STD).view(1, 3, 1, 1, 1), persistent=False) def _required_encoder_names(self) -> list[str]: - names = ["encoder.conv_in.weight", "encoder.conv_in.bias", "encoder.norm_out.weight", "encoder.norm_out.bias", "encoder.conv_out.weight", "encoder.conv_out.bias"] + names = [ + "encoder.conv_in.weight", "encoder.conv_in.bias", + "encoder.norm_out.weight", "encoder.norm_out.bias", + "encoder.conv_out.weight", "encoder.conv_out.bias", + "quant_conv.weight", "quant_conv.bias", + "latents_mean", "latents_std", + ] for i in range(len(CH_MULT)): for b in range(NUM_RES_BLOCKS): base = f"encoder.down.{i}.block.{b}." @@ -139,15 +158,37 @@ class MiniMaxH3VideoVAEEncoder(nn.Module): return names @classmethod - def from_safetensors(cls, path: str | Path, *, device: str | torch.device = "cuda", tiling: bool = True) -> "MiniMaxH3VideoVAEEncoder": - model = cls(tiling=tiling) + def from_safetensors(cls, path: str | Path, *, device: str | torch.device = "cuda", tiling: bool = True, dtype: torch.dtype = torch.float16) -> "MiniMaxH3VideoVAEEncoder": + model = cls(tiling=tiling, compute_dtype=dtype) names = model._required_encoder_names() - with safe_open(str(path), framework="pt", device=str(device)) as ck: - available = set(ck.keys()) + if os.getenv("H3_FAST_SAFETENSORS", "").lower() in {"1", "true", "yes", "on"}: + from fastsafetensors import fastsafe_open + + fast_device = "cuda:0" if str(device) == "cuda" else str(device) + with fastsafe_open(filenames=[str(path)], nogds=True, device=fast_device) as ck: + available = set(ck.keys()) + missing = [n for n in names if n not in available] + if missing: + raise ValueError(f"incompatible H3 VAE checkpoint; missing: {', '.join(sorted(missing)[:16])}") + W = {n: ck.get_tensor(n).clone().detach().to(dtype=dtype) for n in names} + elif os.getenv("H3_DISABLE_MMAP", "").lower() in {"1", "true", "yes", "on"}: + from safetensors.torch import load + + with open(path, "rb") as file: + available_weights = load(file.read()) + available = set(available_weights) missing = [n for n in names if n not in available] if missing: raise ValueError(f"incompatible H3 VAE checkpoint; missing: {', '.join(sorted(missing)[:16])}") - W = {n: ck.get_tensor(n).to(dtype=torch.float32).to(device) for n in names} + W = {n: available_weights[n].to(device=device, dtype=dtype) for n in names} + del available_weights + else: + with safe_open(str(path), framework="pt", device=str(device)) as ck: + available = set(ck.keys()) + missing = [n for n in names if n not in available] + if missing: + raise ValueError(f"incompatible H3 VAE checkpoint; missing: {', '.join(sorted(missing)[:16])}") + W = {n: ck.get_tensor(n).to(dtype=dtype).to(device) for n in names} # Build the structured params dict. down = [] @@ -179,14 +220,19 @@ class MiniMaxH3VideoVAEEncoder(nn.Module): "conv_out": (W["encoder.conv_out.weight"], W["encoder.conv_out.bias"]), } model.W = E - model.quant_conv.to(device, torch.float32) + model.quant_conv.to(device, dtype) + with torch.no_grad(): + model.quant_conv.weight.copy_(W["quant_conv.weight"]) + model.quant_conv.bias.copy_(W["quant_conv.bias"]) + model.latents_mean.copy_(W["latents_mean"].float().cpu()) + model.latents_std.copy_(W["latents_std"].float().cpu()) for b in ("latents_mean", "latents_std", "pixel_mean", "pixel_std"): getattr(model, b).to(device) return model @torch.inference_mode() def _encode_moments(self, x: torch.Tensor) -> torch.Tensor: - return F.conv3d(_encoder_run(x.to(torch.float32), self.W), self.quant_conv.weight, self.quant_conv.bias) + return _conv3d(_encoder_run(x.to(self.compute_dtype), self.W), self.quant_conv.weight, self.quant_conv.bias) def _adaptive_encode(self, x: torch.Tensor) -> torch.Tensor: if self.tiling: @@ -266,14 +312,13 @@ class MiniMaxH3VideoVAEEncoder(nn.Module): """``[B,3,H,W]`` or ``[B,3,T,H,W]`` pixels in ``[-1, 1]`` -> normalized latents ``[B,24,T_lat,H//16,W//16]``.""" if x.ndim == 4: x = x.unsqueeze(2) - x = (x + 1.0) * 0.5 - x = (x - self.pixel_mean.to(x)) / self.pixel_std.to(x) + # Comfy's VAE wrapper casts before entering the model, so image + # normalization rounds in the VAE compute dtype as well. + x = x.to(self.compute_dtype) + x = x.add(1.0).mul_(0.5).sub_(self.pixel_mean.to(x)).div_(self.pixel_std.to(x)) if x.shape[2] == 1: - # Pad with frame_pre_padding zero frames so the 17-tap stride-2 - # temporal downsampling produces exactly one latent frame (1+3-4). - pad = self.frame_pre_padding - x = torch.cat([torch.zeros_like(x[:, :, :pad]), x], dim=2) moments = self._adaptive_encode(x) + moments = moments[:, :, -1:, :, :] else: moments = self.encode_temporal(x) mean = torch.chunk(moments.float(), 2, dim=1)[0] diff --git a/tests/test_fl2va_contracts.py b/tests/test_fl2va_contracts.py new file mode 100644 index 0000000..6d74350 --- /dev/null +++ b/tests/test_fl2va_contracts.py @@ -0,0 +1,217 @@ +import math +from types import SimpleNamespace +import unittest +from unittest.mock import patch + +import torch +from torch import nn +from torch.nn import functional as F + +from h3_blackwell_runtime.packing import FRAME_RESCALE, H3PromptPacker, _video_t_spans +from h3_blackwell_runtime.qwen3vl_vision import ( + TEXT_HEAD_DIM, + TEXT_ROPE_DIMS, + TEXT_ROPE_THETA, + VISION_HIDDEN, + Qwen3VL32BVision, + _VisionAttention, + _VisionPatchMerger, + _apply_rope_vision, + _text_run_ids, + mrope_freqs_cis, +) +from h3_blackwell_runtime.runtime import H3HotRuntime +from h3_blackwell_runtime.vae_encoder import MiniMaxH3VideoVAEEncoder, _downsample + + +class Fl2vaVAEContracts(unittest.TestCase): + def test_quant_conv_is_a_required_checkpoint_weight(self): + names = MiniMaxH3VideoVAEEncoder()._required_encoder_names() + self.assertIn("quant_conv.weight", names) + self.assertIn("quant_conv.bias", names) + + def test_single_frame_is_encoded_without_temporal_prepad(self): + encoder = MiniMaxH3VideoVAEEncoder(tiling=False) + seen = [] + + def fake_encode(x): + seen.append(tuple(x.shape)) + return torch.zeros((x.shape[0], 48, x.shape[2], 1, 1), device=x.device) + + encoder._adaptive_encode = fake_encode + result = encoder.encode(torch.zeros(1, 3, 8, 8)) + self.assertEqual(seen, [(1, 3, 1, 8, 8)]) + self.assertEqual(tuple(result.shape), (1, 24, 1, 1, 1)) + + def test_downsample_pads_only_right_and_bottom(self): + x = torch.arange(16, dtype=torch.float32).reshape(1, 1, 1, 4, 4) + weight = torch.ones(1, 1, 3, 3, 3) + params = {"w": weight, "b": torch.zeros(1), "time": 1, "space": 2} + actual = _downsample(x, params) + padded = F.pad(x, (0, 1, 0, 1, 0, 0), mode="reflect") + expected = F.conv3d(padded, weight[:, :, -1:], params["b"], stride=(1, 2, 2)) + torch.testing.assert_close(actual, expected) + + def test_hot_runtime_preserves_zero_to_one_images(self): + runtime = H3HotRuntime.__new__(H3HotRuntime) + image = torch.tensor([[[[0.0, 0.5, 1.0]]]]) + converted = runtime._image_to_uint8_nhwc(image) + self.assertEqual(converted.flatten().tolist(), [0, 127, 255]) + + +class Fl2vaVisionContracts(unittest.TestCase): + def test_visual_rotary_coordinates_are_block_major(self): + class CoordinateTable(nn.Module): + def forward(self, length): + return torch.arange(length, dtype=torch.float32).unsqueeze(1) + + vision = Qwen3VL32BVision.__new__(Qwen3VL32BVision) + nn.Module.__init__(vision) + vision.spatial_merge_size = 2 + vision.rotary_pos_emb = CoordinateTable() + coordinates = vision.rot_pos_emb(torch.tensor([[1, 4, 4]])).tolist() + self.assertEqual(coordinates[:8], [ + [0.0, 0.0], [0.0, 1.0], [1.0, 0.0], [1.0, 1.0], + [0.0, 2.0], [0.0, 3.0], [1.0, 2.0], [1.0, 3.0], + ]) + + def test_sdpa_output_is_restored_to_token_major_layout(self): + torch.manual_seed(7) + sequence, heads, head_dim = 3, 2, 2 + hidden = heads * head_dim + qkv_weight = torch.randn(hidden * 3, hidden) + qkv_bias = torch.randn(hidden * 3) + proj_weight = torch.randn(hidden, hidden) + proj_bias = torch.randn(hidden) + module = _VisionAttention(qkv_weight, qkv_bias, proj_weight, proj_bias, num_heads=heads, head_dim=head_dim) + x = torch.randn(sequence, hidden) + cos = torch.ones(sequence, 1, head_dim) + sin = torch.zeros(sequence, 1, head_dim // 2) + actual = module(x, torch.tensor([0, sequence], dtype=torch.int32), (cos, sin, sin)) + + qkv = F.linear(x, qkv_weight, qkv_bias) + query, key, value = qkv.reshape(sequence, 3, heads, head_dim).permute(1, 0, 2, 3).unbind(0) + output = F.scaled_dot_product_attention( + query.transpose(0, 1).unsqueeze(0), + key.transpose(0, 1).unsqueeze(0), + value.transpose(0, 1).unsqueeze(0), + ) + expected = F.linear(output.transpose(1, 2).reshape(sequence, hidden), proj_weight, proj_bias) + torch.testing.assert_close(actual, expected) + + def test_vision_rope_uses_original_halves(self): + q = torch.tensor([[[1.0, 2.0, 3.0, 4.0]]]) + k = q + 4 + cos = torch.full_like(q, 0.5) + sin = torch.full_like(q[..., :2], 0.25) + neg_sin = -sin + actual_q, actual_k = _apply_rope_vision(q, k, (cos, sin, neg_sin)) + + def expected(x): + return torch.cat((x[..., :2] * 0.5 + x[..., 2:] * -0.25, + x[..., 2:] * 0.5 + x[..., :2] * 0.25), dim=-1) + + torch.testing.assert_close(actual_q, expected(q)) + torch.testing.assert_close(actual_k, expected(k)) + + def test_mrope_uses_reference_section_boundaries(self): + positions = torch.stack((torch.arange(8), torch.arange(8) + 10, torch.arange(8) + 20)) + actual = mrope_freqs_cis(positions)[0] + inv_freq = 1.0 / ( + TEXT_ROPE_THETA ** (torch.arange(0, TEXT_HEAD_DIM, 2, dtype=torch.float32) / TEXT_HEAD_DIM) + ) + freqs = (inv_freq[None, :, None].expand(3, -1, 1) @ positions[:, None, :].float()).transpose(1, 2) + interleaved = freqs[0].clone() + for axis, offset in ((1, 1), (2, 2)): + index = slice(offset, TEXT_ROPE_DIMS[axis] * 3, 3) + interleaved[..., index] = freqs[axis, ..., index] + expected = torch.cat((interleaved, interleaved), dim=-1).cos().unsqueeze(0) + torch.testing.assert_close(actual, expected) + + def test_mergers_preserve_existing_block_major_order(self): + x = torch.arange(4 * VISION_HIDDEN, dtype=torch.float32).reshape(4, VISION_HIDDEN) + passthrough = lambda value, *args, **kwargs: value + with ( + patch("h3_blackwell_runtime.qwen3vl_vision.F.layer_norm", side_effect=passthrough), + patch("h3_blackwell_runtime.qwen3vl_vision.F.linear", side_effect=passthrough), + patch("h3_blackwell_runtime.qwen3vl_vision.F.gelu", side_effect=passthrough), + ): + main = _VisionPatchMerger(*(torch.empty(1) for _ in range(6)), merge_size=2, + out_hidden_size=1, norm_dim=VISION_HIDDEN) + deepstack = _VisionPatchMerger(*(torch.empty(1) for _ in range(6)), merge_size=2, + out_hidden_size=1) + torch.testing.assert_close(main(x), x.reshape(1, -1)) + torch.testing.assert_close(deepstack(x), x.reshape(1, -1)) + + def test_configured_tokenizer_preserves_empty_prompt(self): + calls = [] + + class RawTokenizer: + def __call__(self, text, **kwargs): + calls.append((text, kwargs)) + return SimpleNamespace(input_ids=torch.empty((1, 0), dtype=torch.long)) + + self.assertEqual(_text_run_ids(SimpleNamespace(tokenizer=RawTokenizer()), ""), []) + self.assertEqual(calls[0][0], "") + self.assertFalse(calls[0][1]["add_special_tokens"]) + + +class _FakeCheckpoint: + def tensor(self, name, dtype=None): + if name == "video_patch_proj.weight": + value = torch.zeros(5376, 96) + elif name == "video_patch_proj.bias": + value = torch.zeros(5376) + elif name == "audio_patch_proj.weight": + value = torch.zeros(5376, 32) + elif name == "audio_patch_proj.bias": + value = torch.zeros(5376) + else: + value = torch.empty(0) + return value.to(dtype=dtype) if dtype is not None else value + + +class Fl2vaPackingContracts(unittest.TestCase): + def test_each_keyframe_keeps_its_own_condition_segment(self): + packer = H3PromptPacker(_FakeCheckpoint()) + text = torch.zeros(1, 3, 5376) + video = torch.zeros(1, 24, 2, 2, 2) + audio = torch.zeros(1, 32, 2, 2) + keyframes = [torch.zeros(1, 24, 1, 2, 2) for _ in range(2)] + _, _, segments, _, _, _ = packer( + text, + video, + audio, + 0.5, + cond_latents=keyframes, + cond_frame_indices=[0, 21], + frame_count=22, + ) + self.assertEqual(segments[1][:2], (3, 4)) + self.assertEqual(segments[2][:2], (4, 5)) + self.assertEqual(segments[1][2], segments[2][2]) + + def test_last_only_anchor_and_targets_share_reference_cursor(self): + packer = H3PromptPacker(_FakeCheckpoint()) + text = torch.zeros(1, 3, 5376) + video = torch.zeros(1, 24, 2, 2, 2) + audio = torch.zeros(1, 32, 2, 2) + last = torch.zeros(1, 24, 1, 2, 2) + _, _, _, positions, _, _ = packer( + text, + video, + audio, + 0.5, + cond_latents=[last], + cond_frame_indices=[21], + frame_count=22, + seed=1, + ) + expected_last_t = 3.0 + sum(_video_t_spans(2)) - FRAME_RESCALE + self.assertTrue(math.isclose(float(positions[3, 0]), expected_last_t)) + self.assertEqual(float(positions[4, 0]), 3.0) # target audio + self.assertEqual(float(positions[8, 0]), 3.0) # target video + + +if __name__ == "__main__": + unittest.main() diff --git a/tools/compare_fl2va_key_sampler.py b/tools/compare_fl2va_key_sampler.py new file mode 100644 index 0000000..b4e54b6 --- /dev/null +++ b/tools/compare_fl2va_key_sampler.py @@ -0,0 +1,128 @@ +"""Replay a matched Comfy keyframe sampler capture with exact static conditioning.""" + +import argparse +from pathlib import Path + +import torch + +from h3_blackwell_runtime.checkpoint import H3Checkpoint +from h3_blackwell_runtime.denoiser import H3PackedDenoiser +from h3_blackwell_runtime.packing import H3PromptPacker, unpatchify_video +from h3_blackwell_runtime.sampler import _audio_sigma, _model_sigma, _unpack_audio, beta_sigmas, res_multistep_update +from h3_blackwell_runtime.t2v import random_av_latents + + +parser = argparse.ArgumentParser() +parser.add_argument("--sampler", type=Path, required=True) +parser.add_argument("--dit", type=Path, required=True) +parser.add_argument("--model", default="/models/minimax_h3_fl2va_pruned_nvfp4.safetensors") +parser.add_argument("--result-latent", type=Path) +parser.add_argument("--production-trace", type=Path) +args = parser.parse_args() + +checkpoint = H3Checkpoint(args.model) +model = H3PackedDenoiser.from_checkpoint(checkpoint, attention_backend="sage2").eval() +packer = H3PromptPacker(checkpoint) +captured_input = torch.load(args.dit / "input.pt", map_location="cuda", weights_only=False) +text_length = next(start for start, _, code in captured_input["segments"] if code == 3) +prefix_stop = next(start for start, _, code in captured_input["segments"] if code % 3 == 2) +text = captured_input["hidden"][:text_length].unsqueeze(0).cuda() +prefix = captured_input["hidden"][:prefix_stop].cuda() +text_tags = torch.ones(text_length, dtype=torch.long, device="cuda") +for start, stop, code in captured_input["segments"]: + if stop <= text_length: + text_tags[start:stop] = code % 3 + +video, audio, frame_count = random_av_latents(384, 384, 22, 440207) +video_shape, audio_shape = video.shape, audio.shape +video_count, audio_count = video.numel(), audio.numel() +zero_cond = [torch.zeros(1, 24, 1, 24, 24, device="cuda") for _ in range(2)] +sigmas = torch.load(args.sampler / "initial.pt", map_location="cuda", weights_only=False)["sigmas"].cuda() +direct_sigmas = beta_sigmas(12, device="cuda") +sigma_delta = (direct_sigmas.float() - sigmas.float()).abs() +print({"stage": "sigmas", "direct": direct_sigmas.tolist(), "comfy": sigmas.tolist(), "mean_delta": float(sigma_delta.mean()), "max_delta": float(sigma_delta.max())}, flush=True) +video_history = audio_history = history_sigma = None + +for index, sigma in enumerate(sigmas[:-1]): + reference = torch.load(args.sampler / f"step_{index:02d}.pt", map_location="cuda", weights_only=False) + reference_x = reference["x"].cuda().reshape(-1) + reference_video = reference_x[:video_count].reshape(video_shape) + reference_audio = reference_x[video_count:video_count + audio_count].reshape(audio_shape) + pre_video = (video.float() - reference_video.float()).abs() + pre_audio = (audio.float() - reference_audio.float()).abs() + + sigma_audio = _audio_sigma(sigma) + carry = sigma_audio / sigma + hidden, times, segments, positions, video_segment, audio_segment = packer( + text, + video, + audio.to(torch.bfloat16) * carry, + _model_sigma(sigma), + text_token_tags=text_tags, + cond_latents=zero_cond, + cond_frame_indices=[0, frame_count - 1], + frame_count=frame_count, + seed=440207, + ) + hidden[:prefix_stop] = prefix.to(hidden) + input_hidden = hidden.detach().clone() + with torch.inference_mode(): + raw_video, raw_audio = model(hidden, times, positions, segments, video_segment, audio_segment) + raw_video = raw_video.to(torch.bfloat16).float() + raw_audio = raw_audio.to(torch.bfloat16) + video_denoised = video + sigma * unpatchify_video(raw_video, video.shape[2], video.shape[-2], video.shape[-1]) + audio_model_output = ( + (1.0 - 4.0) * (audio.to(torch.bfloat16) * carry.to(torch.bfloat16)) + + (1.0 + 3.0 * sigma_audio).to(torch.bfloat16) * (-_unpack_audio(raw_audio)) + ).float() + audio_denoised = audio - sigma * audio_model_output + + reference_denoised = reference["denoised"].cuda().reshape(-1) + reference_video_denoised = reference_denoised[:video_count].reshape(video_shape) + reference_audio_denoised = reference_denoised[video_count:video_count + audio_count].reshape(audio_shape) + denoised_video = (video_denoised.float() - reference_video_denoised.float()).abs() + denoised_audio = (audio_denoised.float() - reference_audio_denoised.float()).abs() + production = torch.load(args.production_trace / f"step_{index:02d}.pt", map_location="cuda", weights_only=False) if args.production_trace else None + production_video = (production["video"].float() - reference_video.float()).abs() if production else None + production_audio = (production["audio"].float() - reference_audio.float()).abs() if production else None + production_denoised = (production["video_denoised"].float() - reference_video_denoised.float()).abs() if production else None + production_hidden = (production["hidden"].float() - input_hidden.float()).abs() if production and "hidden" in production else None + production_raw = (production["raw_video"].float() - raw_video.float()).abs() if production and "raw_video" in production else None + production_segments = [ + (start, stop, code, float((production["hidden"][start:stop].float() - input_hidden[start:stop].float()).abs().mean())) + for start, stop, code in segments + ] if production and "hidden" in production else None + print({ + "step": index, + "pre_video_mean": float(pre_video.mean()), + "pre_video_max": float(pre_video.max()), + "pre_audio_mean": float(pre_audio.mean()), + "pre_audio_max": float(pre_audio.max()), + "denoised_video_mean": float(denoised_video.mean()), + "denoised_video_max": float(denoised_video.max()), + "denoised_audio_mean": float(denoised_audio.mean()), + "denoised_audio_max": float(denoised_audio.max()), + "production_video_mean": float(production_video.mean()) if production_video is not None else None, + "production_video_max": float(production_video.max()) if production_video is not None else None, + "production_audio_mean": float(production_audio.mean()) if production_audio is not None else None, + "production_audio_max": float(production_audio.max()) if production_audio is not None else None, + "production_denoised_mean": float(production_denoised.mean()) if production_denoised is not None else None, + "production_denoised_max": float(production_denoised.max()) if production_denoised is not None else None, + "production_hidden_mean": float(production_hidden.mean()) if production_hidden is not None else None, + "production_hidden_max": float(production_hidden.max()) if production_hidden is not None else None, + "production_raw_mean": float(production_raw.mean()) if production_raw is not None else None, + "production_raw_max": float(production_raw.max()) if production_raw is not None else None, + "production_segments": production_segments, + }, flush=True) + + previous_sigma = sigmas[index - 1] if index else None + sigma_down = sigmas[index + 1] + video = res_multistep_update(video, video_denoised, sigma, sigma_down, video_history, history_sigma, previous_sigma) + audio = res_multistep_update(audio, audio_denoised, sigma, sigma_down, audio_history, history_sigma, previous_sigma) + video_history, audio_history, history_sigma = video_denoised, audio_denoised, sigma_down + +if args.result_latent: + result = torch.load(args.result_latent, map_location="cuda", weights_only=False) + result_video = result["latent"] if isinstance(result, dict) else result + delta = (video.float() - result_video.cuda().float()).abs() + print({"stage": "final_video", "mean_delta": float(delta.mean()), "max_delta": float(delta.max())}, flush=True) diff --git a/tools/compare_fl2va_keyframe_input.py b/tools/compare_fl2va_keyframe_input.py new file mode 100644 index 0000000..f46fa8b --- /dev/null +++ b/tools/compare_fl2va_keyframe_input.py @@ -0,0 +1,193 @@ +"""Compare direct keyframe FL2VA packing with a matched Comfy DiT capture.""" + +import argparse +from pathlib import Path + +import numpy as np +from PIL import Image +import torch + +from h3_blackwell_runtime.checkpoint import H3Checkpoint +from h3_blackwell_runtime.denoiser import H3PackedDenoiser +from h3_blackwell_runtime.packing import H3PromptPacker +from h3_blackwell_runtime.qwen3vl_text import Qwen3VLPromptConditioner +from h3_blackwell_runtime.qwen3vl_vision import Qwen3VL32BVision, build_fl2va_presentation, resize_keyframe +from h3_blackwell_runtime.qwen3vl_vision import mrope_freqs_cis, mrope_position_ids +from h3_blackwell_runtime.sampler import _model_sigma, beta_sigmas +from h3_blackwell_runtime.t2v import random_av_latents +from h3_blackwell_runtime.token_refiner import H3TokenRefiner +from h3_blackwell_runtime.vae_encoder import MiniMaxH3VideoVAEEncoder + + +parser = argparse.ArgumentParser() +parser.add_argument("--capture", required=True) +parser.add_argument("--first", required=True) +parser.add_argument("--last", required=True) +parser.add_argument("--model", default="/models/minimax_h3_fl2va_pruned_nvfp4.safetensors") +parser.add_argument("--qwen", default="/text-encoders/qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors") +parser.add_argument("--vae", default="/vae/minimax_h3_video_vae_fp16.safetensors") +parser.add_argument("--tokenizer", default="/opt/h3-blackwell-runtime/src/h3_blackwell_runtime/qwen25_tokenizer") +parser.add_argument("--reference-vision-first") +parser.add_argument("--reference-vision-last") +parser.add_argument("--qwen-capture-dir") +parser.add_argument("--sampler-capture-dir") +parser.add_argument("--oracle-qwen-input", action="store_true") +parser.add_argument("--vae-dtype", choices=("float16", "bfloat16", "float32"), default="float32") +parser.add_argument("--vae-no-tiling", action="store_true") +parser.add_argument("--load-dit-first", action="store_true") +parser.add_argument("--preview-order", action="store_true") +args = parser.parse_args() + +prompt = "A studio time-lapse of the same pink peony bud opening into the same fully bloomed pink peony, fixed camera, cream background." +width = height = 384 +requested_frames = 22 +seed = 440207 + + +def load_image(path): + image = Image.open(path).convert("RGB") + return torch.from_numpy(np.asarray(image).copy()).unsqueeze(0).cuda().float().div(255.0) + + +def report(name, actual, expected): + expected = expected.to(actual.device) + delta = (actual.float() - expected.float()).abs() + print({ + "stage": name, + "shape": tuple(actual.shape), + "mean_delta": float(delta.mean()), + "max_delta": float(delta.max()), + }, flush=True) + + +capture = torch.load(args.capture, map_location="cuda", weights_only=False) +checkpoint = H3Checkpoint(args.model) +dit_probe = H3PackedDenoiser.from_checkpoint(checkpoint, attention_backend="sage2").eval() if args.load_dit_first else None +early_refiner = H3TokenRefiner(checkpoint, attention_backend="sage2") if args.preview_order else None +early_packer = H3PromptPacker(checkpoint) if args.preview_order else None +conditioner = Qwen3VLPromptConditioner(args.qwen, args.tokenizer) +vision = Qwen3VL32BVision(args.qwen, device="cuda", dtype=torch.float32) +if args.reference_vision_first and args.reference_vision_last: + class CapturedVision: + def __init__(self, paths): + self.outputs = [torch.load(path, map_location="cuda", weights_only=False) for path in paths] + + def __call__(self, flatten, grid): + output = self.outputs.pop(0) + return output["merged"].cuda(), [value.cuda() for value in output["deepstack"]] + + vision = CapturedVision([args.reference_vision_first, args.reference_vision_last]) +video, audio, frame_count = random_av_latents(width, height, requested_frames, seed) +if args.sampler_capture_dir: + initial = torch.load(Path(args.sampler_capture_dir) / "initial.pt", map_location="cuda", weights_only=False) + direct_initial = torch.cat((video.reshape(-1), audio.reshape(-1))) + report("sampler_initial", direct_initial, initial["initial_x"].reshape(-1)) +presentation = build_fl2va_presentation( + prompt, + load_image(args.first), + load_image(args.last), + width=width, + height=height, + frame_count=frame_count, + tokenizer=conditioner.tokenizer, + vision=vision, + text_encoder=conditioner.encoder, + device="cuda", +) +if args.qwen_capture_dir: + qwen_capture = Path(args.qwen_capture_dir) + expected_ids = torch.load(qwen_capture / "qwen_input_ids.pt", map_location="cuda", weights_only=False) + print({"stage": "qwen_input_ids", "equal": torch.equal(presentation.input_ids, expected_ids), "direct_shape": tuple(presentation.input_ids.shape), "comfy_shape": tuple(expected_ids.shape)}, flush=True) + direct_embeds = conditioner.encoder._embed_rows(presentation.input_ids) + visual_mask = torch.zeros((1, direct_embeds.shape[1]), dtype=torch.bool, device="cuda") + deepstack_by_index = {} + for embed in presentation.embeds_info: + start = embed["index"] + end = start + embed["size"] + direct_embeds[0, start:end] = embed["extra"]["merged"].to(direct_embeds) + visual_mask[0, start:end] = True + for index, value in enumerate(embed["extra"]["deepstack"]): + deepstack_by_index.setdefault(index, []).append(value) + compact_ids_path = qwen_capture / "qwen_compact_token_ids.pt" + if compact_ids_path.exists(): + compact_ids = torch.load(compact_ids_path, map_location="cuda", weights_only=False) + print({"stage": "qwen_compact_token_ids", "equal": torch.equal(presentation.input_ids[~visual_mask], compact_ids.reshape(-1)), "direct": presentation.input_ids[~visual_mask].tolist(), "comfy": compact_ids.reshape(-1).tolist()}, flush=True) + direct_deepstack = [torch.cat(values, dim=0) for _, values in sorted(deepstack_by_index.items())] + expected_embeds = torch.load(qwen_capture / "qwen_input_embeds.pt", map_location="cuda", weights_only=False) + raw_rows = torch.nn.functional.embedding(presentation.input_ids, conditioner.encoder.embed_tokens) + raw_scales = torch.nn.functional.embedding(presentation.input_ids, conditioner.encoder.embed_scale) + fp16_embeds = (raw_rows.to(torch.float16) * raw_scales.to(torch.float16)).float() + bf16_embeds = (raw_rows.to(torch.bfloat16) * raw_scales.to(torch.bfloat16)).float() + fp32_to_fp16_embeds = (raw_rows.float() * raw_scales.float()).half().float() + fp32_to_bf16_embeds = (raw_rows.float() * raw_scales.float()).bfloat16().float() + report("qwen_input_text_rows_fp16", fp16_embeds[~visual_mask], expected_embeds[~visual_mask]) + report("qwen_input_text_rows_bf16", bf16_embeds[~visual_mask], expected_embeds[~visual_mask]) + report("qwen_input_text_rows_fp32_to_fp16", fp32_to_fp16_embeds[~visual_mask], expected_embeds[~visual_mask]) + report("qwen_input_text_rows_fp32_to_bf16", fp32_to_bf16_embeds[~visual_mask], expected_embeds[~visual_mask]) + report("qwen_input_embeds", direct_embeds, expected_embeds) + report("qwen_input_text_rows", direct_embeds[~visual_mask], expected_embeds.to(direct_embeds.device)[~visual_mask]) + report("qwen_input_visual_rows", direct_embeds[visual_mask], expected_embeds.to(direct_embeds.device)[visual_mask]) + position_ids = mrope_position_ids(presentation.embeds_info, direct_embeds.shape[1], "cuda") + freqs = mrope_freqs_cis(position_ids) + hidden = (expected_embeds if args.oracle_qwen_input else direct_embeds).to(conditioner.encoder.dtype) + for index, layer in enumerate(conditioner.encoder.layers): + hidden = layer(hidden, freqs) + expected_layer = torch.load(qwen_capture / "qwen_layers" / f"{index:02d}.pt", map_location="cuda", weights_only=False) + report(f"qwen_layer_{index:02d}", hidden, expected_layer) + if index < len(direct_deepstack): + hidden[visual_mask] = hidden[visual_mask] + direct_deepstack[index].to(hidden) + expected_layer50 = torch.load(qwen_capture / "qwen_layer50.pt", map_location="cuda", weights_only=False) + report("qwen_layer50", presentation.text_states, expected_layer50) +vae_dtype = {"float16": torch.float16, "bfloat16": torch.bfloat16, "float32": torch.float32}[args.vae_dtype] +vae = MiniMaxH3VideoVAEEncoder.from_safetensors(args.vae, device="cuda", dtype=vae_dtype, tiling=not args.vae_no_tiling).eval() +cond_latents = [] +cond_images = [] +for keyframe in presentation.keyframes: + resized = resize_keyframe( + keyframe["image"], + width, + height, + crop="disabled" if keyframe["resolved_frame_index"] == 0 else "center", + ) + cond_images.append(resized) + pixels = resized.movedim(-1, 1).cuda().float().mul(2.0).sub(1.0) + cond_latents.append(vae.encode(pixels)) +if args.qwen_capture_dir: + for index, latent in enumerate(cond_latents): + captured_vae = torch.load(Path(args.qwen_capture_dir) / f"vae_keyframe_{index}.pt", map_location="cuda", weights_only=False) + print({"stage": f"vae_meta_{index}", **captured_vae.get("meta", {})}, flush=True) + report(f"vae_image_{index}", cond_images[index], captured_vae["image"]) + report(f"vae_keyframe_{index}", latent, captured_vae["latent"]) + +text = (early_refiner or H3TokenRefiner(checkpoint))(presentation.text_states) +packer = early_packer or H3PromptPacker(checkpoint) +sigma = beta_sigmas(12, device=video.device)[0] +hidden, times, segments, positions, _, _ = packer( + text, + video, + audio, + _model_sigma(sigma), + text_token_tags=presentation.text_token_tags, + cond_latents=cond_latents, + cond_frame_indices=[keyframe["resolved_frame_index"] for keyframe in presentation.keyframes], + frame_count=frame_count, + seed=seed, +) + +expected_hidden = capture["hidden"] +text_length = text.shape[1] +frame_rows = (video.shape[-2] // 2) * (video.shape[-1] // 2) +direct_first = hidden[text_length:text_length + frame_rows] +direct_last = hidden[text_length + frame_rows:text_length + 2 * frame_rows] +comfy_first = expected_hidden[text_length:text_length + frame_rows] +comfy_last = expected_hidden[text_length + frame_rows:text_length + 2 * frame_rows] + +print({"stage": "lengths", "text": text_length, "cond_each": frame_rows, "direct_total": hidden.shape[0], "comfy_total": expected_hidden.shape[0]}, flush=True) +report("text_rows", hidden[:text_length], expected_hidden[:text_length]) +report("cond_first_to_first", direct_first, comfy_first) +report("cond_first_to_last", direct_first, comfy_last) +report("cond_last_to_last", direct_last, comfy_last) +report("cond_last_to_first", direct_last, comfy_first) +report("timesteps", times, capture["timesteps"]) +report("positions", positions, capture["position_ids"]) +print({"stage": "segments", "direct": segments, "comfy": capture["segments"]}, flush=True) diff --git a/tools/compare_h3_keyframe_capture.py b/tools/compare_h3_keyframe_capture.py new file mode 100644 index 0000000..1249cbc --- /dev/null +++ b/tools/compare_h3_keyframe_capture.py @@ -0,0 +1,55 @@ +"""Replay a matched Comfy keyframe DiT capture through the direct H3 model.""" + +import argparse +from pathlib import Path + +import torch + +from h3_blackwell_runtime.checkpoint import H3Checkpoint +from h3_blackwell_runtime.denoiser import H3PackedDenoiser +from h3_blackwell_runtime.packing import unpatchify_video +from h3_blackwell_runtime.rope import h3_rope_rotation +from h3_blackwell_runtime.sampler import _unpack_audio + + +parser = argparse.ArgumentParser() +parser.add_argument("--capture", type=Path, required=True) +parser.add_argument("--model", default="/models/minimax_h3_fl2va_pruned_nvfp4.safetensors") +parser.add_argument("--attention", default="sage2") +args = parser.parse_args() + + +def report(name, actual, expected): + expected = expected.to(actual.device) + delta = (actual.float() - expected.float()).abs() + print({"stage": name, "shape": tuple(actual.shape), "mean_delta": float(delta.mean()), "max_delta": float(delta.max())}, flush=True) + + +captured_input = torch.load(args.capture / "input.pt", map_location="cuda", weights_only=False) +captured_output = torch.load(args.capture / "output.pt", map_location="cuda", weights_only=False) +checkpoint = H3Checkpoint(args.model) +model = H3PackedDenoiser.from_checkpoint(checkpoint, attention_backend=args.attention).eval() + +hidden = captured_input["hidden"].cuda() +timesteps = captured_input["timesteps"].cuda() +positions = captured_input["position_ids"].cuda() +segments = captured_input["segments"] +rotation = h3_rope_rotation(positions, model.backbone.inv_freq, hidden.dtype) +with torch.inference_mode(): + for index, (block, adaln) in enumerate(zip(model.backbone.blocks, model.backbone.adaln, strict=True)): + hidden = block(hidden, rotation, *adaln(timesteps), segments) + expected = torch.load(args.capture / "blocks" / f"{index:02d}.pt", map_location="cuda", weights_only=False) + report(f"block_{index:02d}", hidden, expected) + + video_rows, audio_rows = model.final_layer( + hidden, + timesteps, + tuple(captured_output["video_segment"]), + tuple(captured_output["audio_segment"]), + ) + expected_video = captured_output["video"].cuda() + expected_audio = captured_output["audio"].cuda() + video = unpatchify_video(video_rows, expected_video.shape[2], expected_video.shape[3], expected_video.shape[4]) + audio = _unpack_audio(audio_rows) + report("video_output", video, expected_video) + report("audio_output", audio, expected_audio) diff --git a/tools/compare_qwen_vision.py b/tools/compare_qwen_vision.py new file mode 100644 index 0000000..cb582a6 --- /dev/null +++ b/tools/compare_qwen_vision.py @@ -0,0 +1,212 @@ +"""Compare direct and Comfy Qwen3-VL vision outputs on one keyframe.""" + +import argparse +from pathlib import Path +import sys + +import numpy as np +from PIL import Image +from safetensors import safe_open +import torch +from torch.nn import functional as F + +from h3_blackwell_runtime.qwen3vl_vision import ( + Qwen3VL32BVision, + _apply_rope_vision, + process_image, + resize_keyframe, +) + + +parser = argparse.ArgumentParser() +parser.add_argument("--image", required=True) +parser.add_argument("--checkpoint", default="/text-encoders/qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors") +parser.add_argument("--width", type=int, default=384) +parser.add_argument("--height", type=int, default=384) +parser.add_argument("--comfy-path", default="/opt/ComfyUI") +parser.add_argument("--save-reference") +parser.add_argument("--dtype", choices=("float16", "bfloat16", "float32"), default="bfloat16") +parser.add_argument("--captured-reference") +args = parser.parse_args() + +sys.path.insert(0, args.comfy_path) +import comfy.ops # noqa: E402 +from comfy.ldm.modules.attention import optimized_attention_for_device # noqa: E402 +from comfy.text_encoders.qwen3vl import ( # noqa: E402 + QWEN3VL_VISION, + QWEN3VL_VISION_COMMON, + Qwen3VLVisionModel, +) +from comfy.text_encoders.qwen_vl import process_qwen2vl_images # noqa: E402 +from comfy.text_encoders.llama import apply_rope # noqa: E402 + + +def report(name, actual, expected): + actual = actual.detach() + expected = expected.detach().to(actual.device) + delta = (actual.float() - expected.float()).abs() + print({ + "stage": name, + "shape": tuple(actual.shape), + "actual_dtype": str(actual.dtype), + "expected_dtype": str(expected.dtype), + "mean_delta": float(delta.mean()), + "max_delta": float(delta.max()), + }, flush=True) + + +image = Image.open(args.image).convert("RGB") +pixels = torch.from_numpy(np.asarray(image).copy()).unsqueeze(0).cuda().float().div(255.0) +pixels = resize_keyframe(pixels, args.width, args.height) + +direct_flatten, direct_grid = process_image(pixels) +reference_flatten, reference_grid = process_qwen2vl_images( + pixels, + patch_size=16, + image_mean=[0.5, 0.5, 0.5], + image_std=[0.5, 0.5, 0.5], +) +report("flatten_patches", direct_flatten, reference_flatten) +print({"stage": "grid", "direct": direct_grid.tolist(), "reference": reference_grid.tolist()}, flush=True) + +dtype = {"float16": torch.float16, "bfloat16": torch.bfloat16, "float32": torch.float32}[args.dtype] +config = { + **QWEN3VL_VISION_COMMON, + **QWEN3VL_VISION["qwen3vl_32b"], + "out_hidden_size": 5120, +} +reference = Qwen3VLVisionModel( + config, + device="cuda", + dtype=dtype, + ops=comfy.ops.disable_weight_init, +).to("cuda").eval() +with safe_open(args.checkpoint, framework="pt", device="cuda") as checkpoint: + print({ + "stage": "checkpoint_dtypes", + "text_norm": str(checkpoint.get_tensor("model.layers.0.input_layernorm.weight").dtype), + "vision_norm": str(checkpoint.get_tensor("visual.blocks.0.norm1.weight").dtype), + "vision_patch": str(checkpoint.get_tensor("visual.patch_embed.proj.weight").dtype), + }, flush=True) + visual_state = { + name.removeprefix("visual."): checkpoint.get_tensor(name).to(dtype) + for name in checkpoint.keys() + if name.startswith("visual.") + } +reference.load_state_dict(visual_state, strict=True) +del visual_state + +direct = Qwen3VL32BVision(args.checkpoint, device="cuda", dtype=dtype).eval() +with torch.inference_mode(): + direct_x = direct.patch_embed(direct_flatten.cuda().to(dtype)) + direct_patch_embed = direct_x + reference_x = reference.patch_embed(reference_flatten.cuda().to(dtype)) + report("patch_embed", direct_x, reference_x) + direct_pos = direct.fast_pos_embed_interpolate(direct_grid).to(direct_x.device) + reference_pos = reference.fast_pos_embed_interpolate(reference_grid).to(reference_x.device) + report("position_embed", direct_pos, reference_pos) + direct_x = direct_x + direct_pos + direct_vision_input = direct_x + reference_x = reference_x + reference_pos + report("vision_input", direct_x, reference_x) + + direct_rotary = direct.rot_pos_emb(direct_grid.to(direct_x.device)).reshape(direct_x.shape[0], -1) + reference_rotary = reference.rot_pos_emb(reference_grid).to(reference_x.device).reshape(reference_x.shape[0], -1) + report("rotary", direct_rotary, reference_rotary) + + def position_tuple(rotary): + embedding = torch.cat((rotary, rotary), dim=-1) + cosine = embedding.cos().unsqueeze(-2) + sine = embedding.sin().unsqueeze(-2) + split = sine.shape[-1] // 2 + return cosine, sine[..., :split], -sine[..., split:] + + direct_position = position_tuple(direct_rotary) + reference_position = position_tuple(reference_rotary) + cu_seqlens = F.pad( + torch.repeat_interleave(direct_grid[:, 1] * direct_grid[:, 2], direct_grid[:, 0]).cumsum(0, dtype=torch.int32), + (1, 0), + value=0, + ) + optimized_attention = optimized_attention_for_device(reference_x.device, mask=False, small_input=True) + + direct_block0 = direct.blocks[0] + reference_block0 = reference.blocks[0] + direct_norm = F.layer_norm( + direct_x, + (direct_x.shape[-1],), + weight=direct_block0.norm1_weight, + bias=direct_block0.norm1_bias, + eps=1e-6, + ) + reference_norm = reference_block0.norm1(reference_x) + report("block0_norm1", direct_norm, reference_norm) + direct_qkv = F.linear(direct_norm, direct_block0.attn.qkv_weight, direct_block0.attn.qkv_bias) + reference_qkv = reference_block0.attn.qkv(reference_norm) + report("block0_qkv", direct_qkv, reference_qkv) + direct_q, direct_k, direct_v = direct_qkv.reshape(direct_x.shape[0], 3, 16, 72).permute(1, 0, 2, 3).unbind(0) + reference_q, reference_k, reference_v = reference_qkv.reshape(reference_x.shape[0], 3, 16, 72).permute(1, 0, 2, 3).unbind(0) + direct_q, direct_k = _apply_rope_vision(direct_q.float(), direct_k.float(), direct_position) + direct_q, direct_k = direct_q.to(dtype), direct_k.to(dtype) + reference_q, reference_k = apply_rope(reference_q, reference_k, reference_position) + report("block0_rope_q", direct_q, reference_q) + report("block0_rope_k", direct_k, reference_k) + direct_attention_heads = F.scaled_dot_product_attention( + direct_q.transpose(0, 1).unsqueeze(0), + direct_k.transpose(0, 1).unsqueeze(0), + direct_v.transpose(0, 1).unsqueeze(0), + ) + direct_attention = direct_attention_heads.transpose(1, 2).reshape(1, direct_x.shape[0], -1) + reference_attention = optimized_attention( + reference_q.transpose(0, 1).unsqueeze(0), + reference_k.transpose(0, 1).unsqueeze(0), + reference_v.transpose(0, 1).unsqueeze(0), + 16, + skip_reshape=True, + ) + report("block0_attention", direct_attention, reference_attention) + direct_projected = F.linear(direct_attention[0], direct_block0.attn.proj_weight, direct_block0.attn.proj_bias) + reference_projected = reference_block0.attn.proj(reference_attention)[0] + report("block0_projected", direct_projected, reference_projected) + + direct_deepstack = [] + direct_blocks = [] + reference_deepstack = [] + for index, (direct_block, reference_block) in enumerate(zip(direct.blocks, reference.blocks)): + direct_x = direct_block(direct_x, cu_seqlens, direct_position) + direct_blocks.append(direct_x) + reference_x = reference_block( + reference_x, + cu_seqlens, + reference_position, + optimized_attention=optimized_attention, + ) + report(f"block_{index:02d}", direct_x, reference_x) + if index in direct.deepstack_visual_indexes: + merger_index = direct.deepstack_visual_indexes.index(index) + direct_deepstack.append(direct.deepstack_merger_list[merger_index](direct_x)) + reference_deepstack.append(reference.deepstack_merger_list[merger_index](reference_x)) + direct_merged = direct.merger(direct_x) + reference_merged = reference.merger(reference_x) +report("merged", direct_merged, reference_merged) +for index, (actual, expected) in enumerate(zip(direct_deepstack, reference_deepstack)): + report(f"deepstack_{index}", actual, expected) +if args.captured_reference: + captured = torch.load(args.captured_reference, map_location="cuda", weights_only=False) + report("loaded_comfy_pixel_values", direct_flatten, captured["pixel_values"]) + print({"stage": "loaded_comfy_grid", "direct": direct_grid.tolist(), "expected": captured["grid"].tolist()}) + trace_path = Path(args.captured_reference).with_name(Path(args.captured_reference).name.replace("qwen_vision_", "qwen_vision_trace_")) + trace = torch.load(trace_path, map_location="cuda", weights_only=False) + report("loaded_comfy_patch_embed", direct_patch_embed, trace["patch_embed"]) + report("loaded_comfy_position_embed", direct_pos, trace["position_embed"]) + report("loaded_comfy_vision_input", direct_vision_input, trace["vision_input"]) + for index, block_output in enumerate(direct_blocks): + report(f"loaded_comfy_block_{index:02d}", block_output, trace[f"block_{index:02d}"]) + report("loaded_comfy_merged", direct_merged, captured["merged"]) + for index, (actual, expected) in enumerate(zip(direct_deepstack, captured["deepstack"])): + report(f"loaded_comfy_deepstack_{index}", actual, expected) +if args.save_reference: + torch.save({ + "merged": reference_merged.detach().cpu(), + "deepstack": [value.detach().cpu() for value in reference_deepstack], + }, args.save_reference) diff --git a/tools/compare_vae_decoder_clip.py b/tools/compare_vae_decoder_clip.py index 0e3fdae..71ef712 100644 --- a/tools/compare_vae_decoder_clip.py +++ b/tools/compare_vae_decoder_clip.py @@ -7,7 +7,7 @@ from pathlib import Path import torch from safetensors.torch import load_file -from h3_blackwell_runtime.vae_decoder import MiniMaxH3VideoVAE as DirectVAE +from h3_blackwell_runtime.vae_decoder import MiniMaxH3VideoVAE as DirectVAE, _conv3d parser = argparse.ArgumentParser() @@ -30,7 +30,7 @@ state = torch.load(args.latent, map_location="cuda", weights_only=False) latent = state["latent"].to("cuda") if isinstance(state, dict) else state.to("cuda") direct = DirectVAE.from_safetensors("/vae/minimax_h3_video_vae_fp16.safetensors", device="cuda", tiling=False).eval() -upstream = UpstreamVAE(tiling=False).to("cuda").eval() +upstream = UpstreamVAE(tiling=False).to("cuda", dtype=torch.float16).eval() upstream.load_state_dict(load_file("/vae/minimax_h3_video_vae_fp16.safetensors", device="cuda"), strict=True) with torch.inference_mode(): @@ -38,7 +38,7 @@ with torch.inference_mode(): z_u = z_d.clone().to(next(upstream.parameters()).dtype) z_d = z_d * direct.latents_std.view(1, -1, 1, 1, 1).to(z_d) + direct.latents_mean.view(1, -1, 1, 1, 1).to(z_d) z_u = z_u * upstream.latents_std.view(1, -1, 1, 1, 1).to(z_u) + upstream.latents_mean.view(1, -1, 1, 1, 1).to(z_u) - z_d = direct.post_quant_conv(z_d) + z_d = _conv3d(z_d, direct.post_quant_conv.weight, direct.post_quant_conv.bias) z_u = upstream.post_quant_conv(z_u) stats("post_quant_conv", z_d, z_u) diff --git a/tools/compare_vae_encoder.py b/tools/compare_vae_encoder.py new file mode 100644 index 0000000..61b1671 --- /dev/null +++ b/tools/compare_vae_encoder.py @@ -0,0 +1,110 @@ +"""Compare direct and upstream MiniMax H3 VAE encoding on one keyframe.""" + +import argparse +import os +import sys + +import numpy as np +from PIL import Image +import torch + +from h3_blackwell_runtime.vae_encoder import ( + MiniMaxH3VideoVAEEncoder as DirectVAEEncoder, + _causal_conv3d, + _downsample, + _group_norm_3d, + _resnet, +) + + +parser = argparse.ArgumentParser() +parser.add_argument("--image", required=True) +parser.add_argument("--vae", default="/vae/minimax_h3_video_vae_fp16.safetensors") +parser.add_argument("--comfy-path", default="/opt/ComfyUI") +parser.add_argument("--tiling", action="store_true") +args = parser.parse_args() + +sys.path.insert(0, args.comfy_path) +from h3_blackwell_runtime.upstream_vae import MiniMaxH3VideoVAE as UpstreamVAE # noqa: E402 + + +def load_checkpoint(path): + if os.getenv("H3_DISABLE_MMAP", "").lower() in {"1", "true", "yes", "on"}: + from safetensors.torch import load + + with open(path, "rb") as file: + return load(file.read()) + from safetensors.torch import load_file + + return load_file(path, device="cuda") + + +def report(name, actual, expected): + actual = actual.detach() + expected = expected.detach().to(actual.device) + delta = (actual.float() - expected.float()).abs() + print({ + "stage": name, + "shape": tuple(actual.shape), + "actual_min": float(actual.min()), + "actual_max": float(actual.max()), + "expected_min": float(expected.min()), + "expected_max": float(expected.max()), + "mean_delta": float(delta.mean()), + "max_delta": float(delta.max()), + }, flush=True) + + +image = Image.open(args.image).convert("RGB") +pixels = torch.from_numpy(np.asarray(image).copy()).permute(2, 0, 1).unsqueeze(0).cuda().float() +pixels = pixels.div(127.5).sub(1.0) + +direct = DirectVAEEncoder.from_safetensors(args.vae, device="cuda", tiling=args.tiling).eval() +upstream = UpstreamVAE(tiling=args.tiling).to("cuda").eval() +upstream = upstream.to(dtype=direct.compute_dtype) +state = load_checkpoint(args.vae) +upstream.load_state_dict(state, strict=True) +del state + +trace = {} +trace_names = {"encoder.conv_in", "encoder.norm_out", "encoder.conv_out", "quant_conv"} +for level in range(6): + trace_names.update(f"encoder.down.{level}.block.{block}" for block in range(2)) + if level < 4: + trace_names.add(f"encoder.down.{level}.downsample") +for name, module in upstream.named_modules(): + if name in trace_names: + module.register_forward_hook(lambda _module, _inputs, output, name=name: trace.__setitem__(name, output.detach().cpu())) + +report("quant_conv_weight", direct.quant_conv.weight, upstream.quant_conv.weight) +report("quant_conv_bias", direct.quant_conv.bias, upstream.quant_conv.bias) +with torch.inference_mode(): + direct_latent = direct.encode(pixels.clone()) + upstream_latent = upstream.encode(pixels.to(direct.compute_dtype)) +report("normalized_latent", direct_latent, upstream_latent) + +with torch.inference_mode(): + x = pixels.unsqueeze(2) + x = (x + 1.0) * 0.5 + x = (x - direct.pixel_mean.to(x)) / direct.pixel_std.to(x) + params = direct.W + x = _causal_conv3d(x, params["conv_in"][0], params["conv_in"][1], kernel_size=3, stride=(1, 1, 1), spatial_padding=1, temporal_causal=True) + report("encoder.conv_in", x, trace.pop("encoder.conv_in")) + for level_index, level in enumerate(params["down"]): + for block_index, block in enumerate(level["blocks"]): + x = _resnet(x, block) + name = f"encoder.down.{level_index}.block.{block_index}" + report(name, x, trace.pop(name)) + if level["down"] is not None: + x = _downsample(x, level["down"]) + name = f"encoder.down.{level_index}.downsample" + report(name, x, trace.pop(name)) + x = _group_norm_3d(x, params["norm_out_w"], params["norm_out_b"]) + report("encoder.norm_out", x, trace.pop("encoder.norm_out")) + x = torch.nn.functional.silu(x) + x = _causal_conv3d(x, params["conv_out"][0], params["conv_out"][1], kernel_size=3, stride=(1, 1, 1), spatial_padding=1, temporal_causal=True) + report("encoder.conv_out", x, trace.pop("encoder.conv_out")) + x = torch.nn.functional.conv3d(x, direct.quant_conv.weight, direct.quant_conv.bias) + report("quant_conv", x, trace.pop("quant_conv")) + report("latents_mean", direct.latents_mean, upstream.latents_mean) + report("latents_std", direct.latents_std, upstream.latents_std) diff --git a/tools/compare_vae_full_decode.py b/tools/compare_vae_full_decode.py index 936bc50..f228661 100644 --- a/tools/compare_vae_full_decode.py +++ b/tools/compare_vae_full_decode.py @@ -51,7 +51,7 @@ with torch.inference_mode(): torch.cuda.empty_cache() gc.collect() - upstream = UpstreamVAE().to("cuda").eval() + upstream = UpstreamVAE().to("cuda", dtype=torch.float16).eval() upstream.load_state_dict(load_file("/vae/minimax_h3_video_vae_fp16.safetensors", device="cuda"), strict=True) upstream_pixels = pixelize(upstream.decode(latent.to(next(upstream.parameters()).dtype)), frames) diff --git a/tools/direct_t2v_preview.py b/tools/direct_t2v_preview.py index 62793ac..c0b9992 100644 --- a/tools/direct_t2v_preview.py +++ b/tools/direct_t2v_preview.py @@ -118,7 +118,7 @@ if args.first_frame is not None or args.last_frame is not None: def load_image(path: Path) -> torch.Tensor: img = Image.open(path).convert("RGB") - return torch.from_numpy(np.array(img)).permute(2, 0, 1).unsqueeze(0).float() / 255.0 + return torch.from_numpy(np.array(img)).unsqueeze(0).float() / 255.0 first = load_image(args.first_frame) if args.first_frame is not None else None last = load_image(args.last_frame) if args.last_frame is not None else None @@ -127,7 +127,7 @@ if args.first_frame is not None or args.last_frame is not None: from h3_blackwell_runtime.qwen3vl_vision import Qwen3VL32BVision vision_tower = Qwen3VL32BVision( - "/text-encoders/qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors", device="cuda", dtype=torch.bfloat16 + "/text-encoders/qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors", device="cuda", dtype=torch.float32 ) report_memory("vision_tower_loaded") presentation = build_fl2va_presentation( @@ -145,7 +145,8 @@ if args.first_frame is not None or args.last_frame is not None: cond_latents = [] for kf in presentation.keyframes: resized = resize_keyframe(kf["image"].cuda(), args.width, args.height, crop="disabled" if kf["resolved_frame_index"] == 0 else "center") - cond_latents.append(vae_encoder.encode(resized.movedim(-1, 1).cuda().float())) + pixels = resized.movedim(-1, 1).cuda().float().mul(2.0).sub(1.0) + cond_latents.append(vae_encoder.encode(pixels)) report_memory("fl2va_conditioned") text = refiner(presentation.text_states) report_memory("text_conditioned") @@ -163,6 +164,7 @@ if args.first_frame is not None or args.last_frame is not None: seed=seed, text_token_tags=presentation.text_token_tags, cond_latents=cond_latents, + cond_frame_indices=[kf["resolved_frame_index"] for kf in presentation.keyframes], frame_count=frames, cache_mode=args.cache_mode, cache_threshold=args.cache_threshold, diff --git a/tools/patch_comfy_qwen_vision_capture.py b/tools/patch_comfy_qwen_vision_capture.py new file mode 100644 index 0000000..d249629 --- /dev/null +++ b/tools/patch_comfy_qwen_vision_capture.py @@ -0,0 +1,74 @@ +"""Capture actual loaded-Comfy Qwen merged and DeepStack vision tensors.""" + +from pathlib import Path + + +path = Path("/opt/ComfyUI/comfy/text_encoders/qwen3vl.py") +source = path.read_text(encoding="utf-8") +if "import os\n" not in source: + source = source.replace("import os\n", "import os\n", 1) if "import os\n" in source else "import os\n" + source +if "import traceback\n" not in source: + source = "import traceback\n" + source +old = ( + " merged, deepstack = self.visual(image.to(device, dtype=torch.float32), grid)\n" + " return merged, {\"grid\": grid, \"deepstack\": deepstack}\n" +) +new = ( + " merged, deepstack = self.visual(image.to(device, dtype=torch.float32), grid)\n" + " capture_dir = os.getenv(\"H3_CAPTURE_DIR\")\n" + " if capture_dir:\n" + " os.makedirs(capture_dir, exist_ok=True)\n" + " capture_index = getattr(self, \"_h3_vision_capture_index\", 0)\n" + " torch.save({\"merged\": merged.detach().cpu(), \"deepstack\": [value.detach().cpu() for value in deepstack], \"pixel_values\": image.detach().cpu(), \"grid\": grid.detach().cpu(), \"stack\": traceback.format_stack()}, os.path.join(capture_dir, f\"qwen_vision_{capture_index}.pt\"))\n" + " open(os.path.join(capture_dir, \"qwen_vision_stack.txt\"), \"w\", encoding=\"utf-8\").writelines(traceback.format_stack())\n" + " self._h3_vision_capture_index = capture_index + 1\n" + " return merged, {\"grid\": grid, \"deepstack\": deepstack}\n" +) +if source.count(old) != 1: + raise RuntimeError("Unable to locate Qwen3-VL vision preprocess block.") +path.write_text(source.replace(old, new), encoding="utf-8") + +vision_path = Path("/opt/ComfyUI/comfy/text_encoders/qwen35.py") +vision_source = vision_path.read_text(encoding="utf-8") +vision_source = vision_source.replace( + " x = self.patch_embed(x)\n pos_embeds = self.fast_pos_embed_interpolate(grid_thw).to(x.device)\n x = x + pos_embeds\n", + " x = self.patch_embed(x)\n capture_trace = {\"patch_embed\": x.detach().cpu()}\n pos_embeds = self.fast_pos_embed_interpolate(grid_thw).to(x.device)\n capture_trace[\"position_embed\"] = pos_embeds.detach().cpu()\n x = x + pos_embeds\n capture_trace[\"vision_input\"] = x.detach().cpu()\n", + 1, +) +vision_source = vision_source.replace( + " x = blk(x, cu_seqlens=cu_seqlens, position_embeddings=position_embeddings, optimized_attention=optimized_attention)\n", + " x = blk(x, cu_seqlens=cu_seqlens, position_embeddings=position_embeddings, optimized_attention=optimized_attention)\n capture_trace[f\"block_{layer_num:02d}\"] = x.detach().cpu()\n", + 1, +) +vision_source = vision_source.replace( + " merged = self.merger(x)\n if self.deepstack_merger_list is not None:\n", + " merged = self.merger(x)\n capture_dir = os.getenv(\"H3_CAPTURE_DIR\")\n if capture_dir:\n trace_index = getattr(self, \"_h3_vision_trace_index\", 0)\n torch.save(capture_trace, os.path.join(capture_dir, f\"qwen_vision_trace_{trace_index}.pt\"))\n self._h3_vision_trace_index = trace_index + 1\n if self.deepstack_merger_list is not None:\n", + 1, +) +vision_path.write_text(vision_source, encoding="utf-8") + +clip_path = Path("/opt/ComfyUI/comfy/sd1_clip.py") +clip_source = clip_path.read_text(encoding="utf-8") +if "import os\n" not in clip_source: + clip_source = "import os\n" + clip_source +clip_source = clip_source.replace( + " tokens_embed = torch.tensor([tokens_temp], device=device, dtype=torch.long)\n tokens_embed = self.transformer.get_input_embeddings()(tokens_embed, out_dtype=torch.float32)\n", + " tokens_embed = torch.tensor([tokens_temp], device=device, dtype=torch.long)\n capture_dir = os.getenv(\"H3_CAPTURE_DIR\")\n if capture_dir:\n torch.save(tokens_embed.detach().cpu(), os.path.join(capture_dir, \"qwen_compact_token_ids.pt\"))\n tokens_embed = self.transformer.get_input_embeddings()(tokens_embed, out_dtype=torch.float32)\n", + 1, +) +clip_path.write_text(clip_source, encoding="utf-8") + +nodes_path = Path("/opt/ComfyUI/comfy_extras/nodes_minimax_h3.py") +nodes_source = nodes_path.read_text(encoding="utf-8") +if "import os\n" not in nodes_source: + nodes_source = "import os\n" + nodes_source +vae_encode_line = "kf[\"latent\"] = vae.encode(kf.pop(\"image\"))" +if nodes_source.count(vae_encode_line) != 1: + raise RuntimeError("Unable to locate MiniMax H3 keyframe VAE encode call.") +nodes_source = nodes_source.replace( + vae_encode_line, + "keyframe_image = kf.pop(\"image\")\n kf[\"latent\"] = vae.encode(keyframe_image)\n capture_dir = os.getenv(\"H3_CAPTURE_DIR\")\n if capture_dir:\n capture_index = getattr(vae, \"_h3_vae_capture_index\", 0)\n vae_model = vae.first_stage_model\n vae_meta = {\"parameter_dtype\": str(next(vae_model.parameters()).dtype), \"tiling\": vae_model.tiling, \"tile_size\": vae_model.tile_size, \"tile_overlap_min\": vae_model.tile_overlap_min}\n torch.save({\"image\": keyframe_image.detach().cpu(), \"latent\": kf[\"latent\"].detach().cpu(), \"meta\": vae_meta}, os.path.join(capture_dir, f\"vae_keyframe_{capture_index}.pt\"))\n vae._h3_vae_capture_index = capture_index + 1", + 1, +) +nodes_path.write_text(nodes_source, encoding="utf-8") +print("Applied Qwen3-VL vision capture patch.") diff --git a/tools/submit_fl2va_keyframes.py b/tools/submit_fl2va_keyframes.py new file mode 100644 index 0000000..131ee20 --- /dev/null +++ b/tools/submit_fl2va_keyframes.py @@ -0,0 +1,46 @@ +"""Submit a matched Comfy FL2VA first/last-frame reference workflow.""" + +import argparse +import json +from urllib.request import Request, urlopen + + +parser = argparse.ArgumentParser() +parser.add_argument("--url", default="http://localhost:8188") +parser.add_argument("--first", default="fl2va_key_first.png") +parser.add_argument("--last", default="fl2va_key_last.png") +parser.add_argument("--prefix", default="fl2va-comfy-keyframes-seed440207") +args = parser.parse_args() + +prompt = { + "1": {"class_type": "UNETLoader", "inputs": {"unet_name": "minimax_h3_fl2va_pruned_nvfp4.safetensors", "weight_dtype": "default"}}, + "3": {"class_type": "CLIPLoader", "inputs": {"clip_name": "qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors", "type": "minimax"}}, + "4": {"class_type": "VAELoader", "inputs": {"vae_name": "minimax_h3_video_vae_fp16.safetensors"}}, + "5": {"class_type": "LoadImage", "inputs": {"image": args.first}}, + "6": {"class_type": "LoadImage", "inputs": {"image": args.last}}, + "8": {"class_type": "MiniMaxH3ImageToVideo", "inputs": { + "clip": ["3", 0], + "vae": ["4", 0], + "prompt": "A studio time-lapse of the same pink peony bud opening into the same fully bloomed pink peony, fixed camera, cream background.", + "width": 384, + "height": 384, + "length": 22, + "first_frame": ["5", 0], + "last_frame": ["6", 0], + }}, + "9": {"class_type": "BasicGuider", "inputs": {"model": ["1", 0], "conditioning": ["8", 0]}}, + "10": {"class_type": "RandomNoise", "inputs": {"noise_seed": 440207}}, + "11": {"class_type": "KSamplerSelect", "inputs": {"sampler_name": "res_multistep"}}, + "12": {"class_type": "BasicScheduler", "inputs": {"model": ["1", 0], "scheduler": "beta", "steps": 12, "denoise": 1.0}}, + "13": {"class_type": "SamplerCustomAdvanced", "inputs": {"noise": ["10", 0], "guider": ["9", 0], "sampler": ["11", 0], "sigmas": ["12", 0], "latent_image": ["8", 1]}}, + "14": {"class_type": "VAEDecode", "inputs": {"samples": ["13", 0], "vae": ["4", 0]}}, + "15": {"class_type": "SaveImage", "inputs": {"images": ["14", 0], "filename_prefix": args.prefix}}, +} + +request = Request( + args.url.rstrip("/") + "/prompt", + data=json.dumps({"prompt": prompt}).encode(), + headers={"Content-Type": "application/json"}, +) +with urlopen(request) as response: + print(response.read().decode())