From 838d34c42905f4a162d5ced16fdf867316e18b84 Mon Sep 17 00:00:00 2001 From: Bowen Xue Date: Tue, 6 Oct 2026 08:29:35 +0000 Subject: [PATCH 1/2] Encode prompts in BF16 without stalling on references, and cut references to the generated length (0.2.4) ComfyUI's text-encoder base ran the H3 Qwen3-VL encoder in FP32, so the NVFP4 weights were dequantized to FP32 matmuls. The official MiniMax pipeline runs it in BF16, and so does FreeVideo now. Encoding memory is planned from the sequence the language model really sees, checked against the dedicated memory actually free after loading (on Windows also the DXGI budget), and encoder weights move back to host memory until the forward fits. Inputs too long for one forward run the language model in blocks without a dense attention mask. DeepStack features are joined once, and each vision block returns its cache. On Windows a forward that starts to spill into shared GPU memory is stopped within seconds and retried with more room, and the room each kind of input needed is remembered per machine. References are cut to the generated length, at most 15 s, as the official pipeline and ComfyUI's native H3 node do, instead of refusing clips over 15 s and using shorter ones whole. The media card, the generation progress and the result say which clip was cut and how much was used. On an RTX 5060 Ti 16 GB (Windows, 832x480, 5 s), prompt encoding with one to three reference images takes 22-24 s instead of 35-41 s at a 12.3-12.6 GiB instead of a 15.0-15.4 GiB peak; a 14.5 s reference video with three images, which stalled for over ten minutes in shared memory, encodes in 30 s; two 15 s videos and three images at 1344x768, which failed out of memory, now complete. --- freevideo_engine/comfy_assets.py | 1 + freevideo_engine/comfy_bridge.py | 5 + freevideo_engine/comfy_media.py | 5 +- freevideo_engine/comfy_nodes.py | 2 +- freevideo_engine/encode_worker.py | 59 +++++- freevideo_engine/encoder_diagnostics.py | 2 +- freevideo_engine/encoder_lowmem.py | 180 ++++++++++++++++ freevideo_engine/encoder_memory.py | 33 ++- freevideo_engine/encoder_precision.py | 134 ++++++++++++ freevideo_engine/encoder_workspace.py | 271 ++++++++++++++++++++++++ freevideo_engine/generate.py | 2 + freevideo_engine/macos_encoder.py | 3 +- freevideo_engine/media_encoding.py | 81 +++++-- freevideo_engine/release_notes.json | 18 +- freevideo_engine/tuning.py | 1 + scripts/test_freevideo_entry.mjs | 9 +- web/freevideo.js | 22 +- web/generation_progress.css | 5 +- web/generation_progress.js | 31 ++- web/studio.js | 7 +- 20 files changed, 817 insertions(+), 54 deletions(-) create mode 100644 freevideo_engine/encoder_lowmem.py create mode 100644 freevideo_engine/encoder_precision.py create mode 100644 freevideo_engine/encoder_workspace.py diff --git a/freevideo_engine/comfy_assets.py b/freevideo_engine/comfy_assets.py index 0017d90..0badb58 100644 --- a/freevideo_engine/comfy_assets.py +++ b/freevideo_engine/comfy_assets.py @@ -152,6 +152,7 @@ def output_summary(report, relative_video): gpu_total_bytes=report.get('profile', {}).get('policy', {}).get('hardware', {}).get('vram_total'), geometry=report.get('geometry', {}), sampling_plan=report.get('sampling_plan'), + reference_trims=report.get('encoding', {}).get('reference_trims') or [], video=relative_video.as_posix(), report=relative_video.with_suffix( '.debug.json' if report.get('diagnostic_file') == relative_video.with_suffix('.debug.json').name else '.request.json').as_posix()) diff --git a/freevideo_engine/comfy_bridge.py b/freevideo_engine/comfy_bridge.py index 9c00f57..d75dcb3 100644 --- a/freevideo_engine/comfy_bridge.py +++ b/freevideo_engine/comfy_bridge.py @@ -355,6 +355,8 @@ def progress_message(event): 'encoder_compute': 'Encoding text and images', 'encoder_oom': 'Releasing encoder weights after insufficient GPU memory', 'encoder_retry': 'Retrying text encoding with more GPU workspace', + 'encoder_spill': 'Leaving more GPU memory for text encoding', + 'encoder_low_memory': 'Encoding long references in smaller blocks', 'encoder_conditioning_pack': 'Preparing prompt data', 'keyframe_vae': 'Encoding reference media', 'media_vae': 'Encoding reference media', @@ -414,6 +416,9 @@ def progress_message(event): elapsed_seconds=event.get('elapsed_seconds')) if name == 'media_encode_phase': return {'label': str(event.get('phase', 'Encoding input media')), 'timing_phase': 'encoding'} + if name == 'reference_trimmed': + return {'label': 'Preparing reference media', 'timing_phase': 'encoding', + 'reference_trimmed': {key: event.get(key) for key in ('kind', 'number', 'seconds', 'used_seconds')}} if name == 'prepared_blocks': return {'label': 'Loading cached video model · %s / 50 blocks' % event.get('blocks', '?'), 'timing_phase': 'load'} diff --git a/freevideo_engine/comfy_media.py b/freevideo_engine/comfy_media.py index 773f96d..f71300f 100644 --- a/freevideo_engine/comfy_media.py +++ b/freevideo_engine/comfy_media.py @@ -26,8 +26,6 @@ def _audio(value, path): or waveform.shape[1] not in (1, 2) or type(rate) is not int or not 8000 <= rate <= 192000 or not bool(torch.isfinite(waveform).all())): raise ValueError('Expected a finite mono/stereo Comfy AUDIO input') - if waveform.shape[-1] > rate * 15: - raise ValueError('Trim reference audio to 15 s or less before generation') samples = waveform[0].detach().float().cpu().contiguous().numpy() layout = 'mono' if len(samples) == 1 else 'stereo' with av.open(str(path), 'w', format='wav') as container: @@ -75,8 +73,7 @@ def export(run, canvas, *, first=None, last=None, references=None, loras=None, c path = path.with_suffix('.png') _image(value, path) elif kind == 'video': - if value.get_duration() > 15.05: - raise ValueError('Trim reference video to 15 s or less before generation') + # A clip longer than the generated video is shortened, and reported, by media encoding. path = path.with_suffix('.mp4') # Native save_to honors crops/trims and remuxes compatible file # inputs without expanding all frames into host float tensors. diff --git a/freevideo_engine/comfy_nodes.py b/freevideo_engine/comfy_nodes.py index e43c78f..5acc3ff 100644 --- a/freevideo_engine/comfy_nodes.py +++ b/freevideo_engine/comfy_nodes.py @@ -219,7 +219,7 @@ class FreeVideoReference(io.ComfyNode): def define_schema(cls): return io.Schema(node_id='FreeVideoReference', display_name='FreeVideo · Reference (experimental)', category='FreeVideo', is_experimental=True, - description='Append one image, video or audio reference in order. Clips must be trimmed to 15 s or less. ' + description='Append one image, video or audio reference in order. A clip is cut to the generated length, at most 15 s, and the cut is reported. ' 'Connect the stack to Generate. VDN reference quality is experimental.', inputs=[io.Image.Input('image', optional=True), io.Video.Input('video', optional=True), io.Audio.Input('audio', optional=True), References.Input('previous', optional=True)], diff --git a/freevideo_engine/encode_worker.py b/freevideo_engine/encode_worker.py index b10aa5c..40c64c5 100644 --- a/freevideo_engine/encode_worker.py +++ b/freevideo_engine/encode_worker.py @@ -211,6 +211,8 @@ def phase(name, **metrics): task = task_for(request['media']) if task != 't2va': normalized, media_kwargs = prepare(request['media'], request['geometry'], Path(request['output']).parent / 'media') + # Saved with the conditioning, so a reused input cache still reports them. + loading['reference_trims'] = [row['trimmed'] for row in normalized if row.get('trimmed')] if keyframes is not None: if not isinstance(keyframes, dict) or set(keyframes) != {'first', 'last'}: raise ValueError('FL2VA encoding requires keyframes.first and keyframes.last') @@ -281,6 +283,17 @@ def phase(name, **metrics): from .encoder_memory import token_counts loading['token_summary'] = token_counts(tokens) tokenize_seconds = time.perf_counter() - encode_started + from contextlib import contextmanager, nullcontext + from .encoder_workspace import (LOW_MEMORY_CHUNK, WHOLE_BLOCK_TOKENS, SharedMemorySpill, SpillGuard, Workspace, adapter_reader, + dedicated_room, input_size, make_room, plan) + size = loading['encoder_input'] = input_size(tokens) + workspace = Workspace.for_encoder(torch, path) + reader = adapter_reader() + guard = SpillGuard(torch, reader) + current = {} + # What this process may hold under FreeVideo's plan, beside the physical room. + budget_bytes = int(float(request['gpu_budget_gb']) * 1e9) if request.get('gpu_budget_gb') else None + if resident is not None and encoder_cache_hit: resident.encoder_room(clip, tokens, estimated) loading['resident_admission'] = list(resident.decisions) @@ -295,6 +308,10 @@ def timed_load(*values, **kwargs): phase('encoder_device_load', load_seconds=load_seconds) print(json.dumps({'event': 'encoder_device_reuse' if encoder_on_gpu(clip) else 'encoder_device_load_start'}), flush=True) tick = time.perf_counter() + import comfy.model_management as memory + if current.get('need') is not None: + # The native planner keeps 0.8 GiB beside the reserve; leave this attempt's room. + memory.EXTRA_RESERVED_VRAM = max(memory.EXTRA_RESERVED_VRAM, current['need'] - int(.8 * 2**30)) result = native_load(*values, **kwargs) torch.cuda.synchronize() transfer_seconds[0] += time.perf_counter() - tick @@ -328,19 +345,52 @@ def timed_load(*values, **kwargs): released = release_mapped_pages(checkpoint_paths) if released: print(json.dumps(dict(event='encoder_pages_released', **released)), flush=True) + if current.get('need') is not None: + # cudaMemGetInfo can promise room the WDDM budget does not have; check what is + # really free and move weights back to host memory until the forward fits. + room = make_room(clip.patcher, torch, current['need'], reader, budget_bytes) + loading.setdefault('encoder_workspace', []).append(dict(current['plan'], mode=current['mode'], **room)) + guard.arm() from .encoder_memory import snapshot - import comfy.model_management as memory phase('encoder_compute', load_seconds=load_seconds, device_load_seconds=transfer_seconds[0], gpu=snapshot(torch, clip, memory)) print(json.dumps({'event': 'encoder_compute_start'}), flush=True) return result + @contextmanager + def attempt(index, failure): + # Room is planned against what this device can give the encoder with none of its weights loaded. + capacity = dedicated_room(torch, reader, budget_bytes) + int(clip.patcher.loaded_size()) + mode, need, details = plan(workspace, size, capacity, index, current.get('mode') if failure else None) + current.update(mode=mode, need=need, plan=details) + from .encoder_lowmem import low_memory + if mode == 'blocks': + phase('encoder_low_memory', **details) + context = low_memory(clip.cond_stage_model, LOW_MEMORY_CHUNK) + elif size['sequence'] >= WHOLE_BLOCK_TOKENS: + # One block: the whole sequence at once, without the native dense T x T mask. + context = low_memory(clip.cond_stage_model, 1 << 30) + else: + context = nullcontext() + from .encoder_precision import bf16_language_model + try: + # BF16, as the official pipeline runs this encoder (ComfyUI's base runs it in FP32). + with context, bf16_language_model(clip.cond_stage_model): + yield + workspace.learn(mode, size, guard.used(), spilled=guard.spill is not None) + except SharedMemorySpill as error: + workspace.learn(mode, size, error.used_bytes, spilled=True) + raise + finally: + guard.disarm() + clip.load_model = timed_load try: from .encoder_memory import encode_with_recovery from .encoder_pinning import readonly_pinning import comfy.model_management as memory with readonly_pinning(torch, memory): - encoded, attempts = encode_with_recovery(clip, tokens, memory, torch, phase) + encoded, attempts = encode_with_recovery(clip, tokens, memory, torch, phase, max_attempts=4, + attempt=attempt) loading['encoder_attempts'] = attempts # Finish async forward work at the existing conditioning boundary, so # it is not attributed to packing/saving on the host. @@ -350,7 +400,10 @@ def timed_load(*values, **kwargs): phase('encoder_conditioning_pack', gpu=snapshot(torch, clip, memory)) finally: del clip.load_model - del timed_load, native_load + guard.disarm() + if reader is not None: + reader.close() + del timed_load, native_load, attempt del tokens task = 'fl2va' if images is not None else task value = to_cache(encoded, prompt, task=task) diff --git a/freevideo_engine/encoder_diagnostics.py b/freevideo_engine/encoder_diagnostics.py index b781411..3f48aea 100644 --- a/freevideo_engine/encoder_diagnostics.py +++ b/freevideo_engine/encoder_diagnostics.py @@ -17,7 +17,7 @@ 'encoder_options', 'encoder_cuda_setup', 'encoder_path_config', 'encoder_native_import', 'encoder_media_prepare', 'encoder_lookup', 'encoder_load', 'encoder_checkpoint_map', 'encoder_construct', 'encoder_tokenize', 'encoder_device_load', 'encoder_page_release', - 'encoder_compute', 'encoder_oom', 'encoder_retry', 'encoder_conditioning_pack', + 'encoder_compute', 'encoder_oom', 'encoder_spill', 'encoder_retry', 'encoder_low_memory', 'encoder_conditioning_pack', 'keyframe_vae', 'media_vae', 'encoder_save', 'encoder_idle_preload') diff --git a/freevideo_engine/encoder_lowmem.py b/freevideo_engine/encoder_lowmem.py new file mode 100644 index 0000000..5538859 --- /dev/null +++ b/freevideo_engine/encoder_lowmem.py @@ -0,0 +1,180 @@ +"""Run the H3 text encoder's language model in token blocks for very long inputs. + +The native forward keeps FP32 activations, builds one dense T x T causal mask +and runs every projection and MLP over the whole sequence at once. Measured on +the 32B encoder, a 15 s reference video (15.7k tokens) needed 9.5 GiB beyond +the weights and two videos plus three images (34k tokens) 22 GiB, more than a +16 GB card has. The same math fits in a fraction of that: + +- attention runs per block of queries, each against all earlier keys with the + causal limit as a lower-right bias, so no mask is ever materialized; the + memory-efficient kernel takes FP32 but not grouped K/V heads, so each K/V + head is broadcast (a view, no copy) to its group of query heads; +- the position-wise MLP runs per block of tokens; +- DeepStack features are dropped after the layers that use them. + +Only the language-model path the H3 encoder uses is replaced. Calls with KV +caches, attention masks or intermediate layer lists go to the native code. +""" +from contextlib import contextmanager +import types + +import torch +import torch.nn.functional as F + + +def grouped_causal_attention(q, k, v, start): + """Queries at positions start..start+len(q) over keys 0..start+len(q), causal, grouped K/V heads.""" + from torch.nn.attention.bias import causal_lower_right + batch, heads, rows, width = q.shape + groups = k.shape[1] + per_group = heads // groups + end = start + rows + bias = causal_lower_right(rows, end) + output = torch.empty((batch, heads, rows, width), device=q.device, dtype=q.dtype) + for group in range(groups): + first = group * per_group + keys = k[:, group:group + 1, :end].expand(batch, per_group, end, width) + values = v[:, group:group + 1, :end].expand(batch, per_group, end, width) + output[:, first:first + per_group] = F.scaled_dot_product_attention( + q[:, first:first + per_group], keys, values, attn_mask=bias) + return output + + +def rope(x, freqs_cis): + """The native apply_rope, for one tensor (queries or keys) at a time.""" + cos, sin, nsin = freqs_cis[0], freqs_cis[1], freqs_cis[2] + embedded = x * cos + half = embedded.shape[-1] // 2 + embedded[..., :half].addcmul_(x[..., half:], nsin) + embedded[..., half:].addcmul_(x[..., :half], sin) + return embedded.to(x.dtype) + + +def chunked_attention(attention, chunk): + original = attention.forward + + def forward(self, hidden_states, attention_mask=None, freqs_cis=None, optimized_attention=None, + past_key_value=None, sliding_window=None): + batch, length, _ = hidden_states.shape + if optimized_attention is not None: # a native caller: its mask and caches apply + return original(hidden_states, attention_mask=attention_mask, freqs_cis=freqs_cis, + optimized_attention=optimized_attention, past_key_value=past_key_value, + sliding_window=sliding_window) + if (attention_mask is not None or past_key_value is not None or sliding_window is not None + or freqs_cis is None or any(t.shape[-2] != length for t in freqs_cis[:3])): + raise ValueError('Block-wise encoder attention supports causal attention without caches only') + if self.merged_qkv: + fused = self.qkv_proj(hidden_states) + projected = dict(zip(('q', 'k', 'v'), fused.split((self.inner_size, self.kv_size, self.kv_size), dim=-1))) + else: + projected = {} + keys = (projected.get('k') if projected else self.k_proj(hidden_states)) + values = (projected.get('v') if projected else self.v_proj(hidden_states)) + keys = keys.view(batch, length, self.num_kv_heads, self.head_dim).transpose(1, 2) + values = values.view(batch, length, self.num_kv_heads, self.head_dim).transpose(1, 2) + if self.k_norm is not None: + keys = self.k_norm(keys) + keys = rope(keys, freqs_cis) + output = None + for start in range(0, length, chunk): + end = min(length, start + chunk) + queries = (projected['q'][:, start:end] if projected else self.q_proj(hidden_states[:, start:end])) + queries = queries.reshape(batch, end - start, self.num_heads, self.head_dim) + queries = queries.transpose(1, 2) + if self.q_norm is not None: + queries = self.q_norm(queries) + queries = rope(queries, tuple(t[..., start:end, :] for t in freqs_cis[:3])) + value = grouped_causal_attention(queries, keys, values, start) + del queries + value = self.o_proj(value.transpose(1, 2).reshape(batch, end - start, self.num_heads * self.head_dim)) + if output is None: + output = torch.empty((batch, length, value.shape[-1]), device=value.device, dtype=value.dtype) + output[:, start:end] = value + del value + return output, None + return types.MethodType(forward, attention) + + +def chunked_mlp(mlp, chunk): + original = mlp.forward + + def forward(self, x): + if x.shape[1] <= chunk: + return original(x) + output = torch.empty_like(x) + for start in range(0, x.shape[1], chunk): + output[:, start:start + chunk] = original(x[:, start:start + chunk]) + return output + return types.MethodType(forward, mlp) + + +def language_model(model): + """The Llama2_ module inside the native H3 text encoder.""" + for _, module in model.named_modules(): + if type(module).__name__ == 'Llama2_' and hasattr(module, 'layers'): + return module + raise ValueError('The native H3 text encoder has no Llama language model') + + +def chunked_forward(llama, original): + def forward(self, x, attention_mask=None, embeds=None, num_tokens=None, intermediate_output=None, + final_layer_norm_intermediate=True, dtype=None, position_ids=None, embeds_info=[], + past_key_values=None, input_ids=None, deepstack_embeds=None, visual_pos_masks=None): + if (attention_mask is not None or past_key_values is not None + or isinstance(intermediate_output, (list, str))): + return original(x, attention_mask=attention_mask, embeds=embeds, num_tokens=num_tokens, + intermediate_output=intermediate_output, + final_layer_norm_intermediate=final_layer_norm_intermediate, dtype=dtype, + position_ids=position_ids, embeds_info=embeds_info, past_key_values=past_key_values, + input_ids=input_ids, deepstack_embeds=deepstack_embeds, visual_pos_masks=visual_pos_masks) + x = embeds if embeds is not None else self.embed_tokens(x, out_dtype=dtype) + length = x.shape[1] + if position_ids is None: + position_ids = torch.arange(length, device=x.device).unsqueeze(0) + freqs_cis = self.compute_freqs_cis(position_ids, x.device) + if intermediate_output is not None and intermediate_output < 0: + intermediate_output = len(self.layers) + intermediate_output + deepstack = list(deepstack_embeds) if deepstack_embeds is not None else None + deepstack_embeds = None + intermediate = None + for index, layer in enumerate(self.layers): + # The block-wise attention installed on each layer needs no mask. + x, _ = layer(x=x, attention_mask=None, freqs_cis=freqs_cis, optimized_attention=None, + past_key_value=None) + # DeepStack: per-layer visual features at image positions (Qwen3-VL), as the native forward. + if deepstack is not None and index < len(deepstack): + x[visual_pos_masks] = x[visual_pos_masks] + deepstack[index].to(x) + deepstack[index] = None # used by this layer only + if index == intermediate_output: + intermediate = x.clone() + if self.norm is not None: + x = self.norm(x) + if intermediate is not None and final_layer_norm_intermediate and self.norm is not None: + intermediate = self.norm(intermediate) + return x, intermediate + return types.MethodType(forward, llama) + + +@contextmanager +def low_memory(model, chunk): + """Within the block, the encoder's language model runs attention and MLPs in `chunk`-token blocks.""" + llama = language_model(model) + patched = [] + + def patch(module, replacement): + patched.append((module, module.__dict__.get('forward'))) + module.forward = replacement + + try: + patch(llama, chunked_forward(llama, llama.forward)) + for layer in llama.layers: + patch(layer.self_attn, chunked_attention(layer.self_attn, chunk)) + patch(layer.mlp, chunked_mlp(layer.mlp, chunk)) + yield + finally: + for module, previous in reversed(patched): + if previous is None: + del module.forward # back to the class's own forward + else: + module.forward = previous diff --git a/freevideo_engine/encoder_memory.py b/freevideo_engine/encoder_memory.py index afcda3f..fc2a18a 100644 --- a/freevideo_engine/encoder_memory.py +++ b/freevideo_engine/encoder_memory.py @@ -148,14 +148,37 @@ def release_cast_buffers(manager, torch): released_reserved_bytes=max(0, reserved - torch.cuda.memory_reserved())) -def encode_with_recovery(clip, tokens, manager, torch, phase, *, max_attempts=3): +def encode_with_recovery(clip, tokens, manager, torch, phase, *, max_attempts=3, attempt=None): + """Encode, retrying out-of-memory and shared-memory spills with more room. + + `attempt(index, failure)` returns the context each forward runs in; it plans + that attempt's mode and room from the previous failure ('gpu_oom' or + 'shared_memory_spill', None at first). + """ + from contextlib import nullcontext from .adaptive import classify_failure + from .encoder_workspace import SharedMemorySpill original_reserve = manager.EXTRA_RESERVED_VRAM attempts = [] + previous = None try: for index in range(max_attempts): try: - return clip.encode_from_tokens_scheduled(tokens), attempts + with (attempt(index, previous) if attempt is not None else nullcontext()): + return clip.encode_from_tokens_scheduled(tokens), attempts + except SharedMemorySpill as error: + # Windows does not fail past the dedicated budget; the guard stopped + # this forward before it ran from shared system memory. + failure = dict(kind='shared_memory_spill', exception=[dict(type='SharedMemorySpill', message=str(error))]) + native_is_oom = None + failed = snapshot(torch, clip, manager, windows_memory=True) + failed.update(attempt=index + 1, kind=failure['kind'], used_bytes=error.used_bytes, + spilled_bytes=error.spilled_bytes) + attempts.append(failed) + phase('encoder_spill', gpu=failed, encoder_attempts=list(attempts)) + if index + 1 == max_attempts: + raise + retry_error = error except Exception as error: failure = classify_failure(error) native_is_oom = getattr(manager, 'is_oom', None) @@ -189,11 +212,12 @@ def encode_with_recovery(clip, tokens, manager, torch, phase, *, max_attempts=3) tb = tb.tb_next try: gc.collect() - native_oom = native_is_oom(retry_error) if callable(native_is_oom) else False + native_oom = (failure['kind'] == 'shared_memory_spill' + or (native_is_oom(retry_error) if callable(native_is_oom) else False)) except Exception as cleanup: retry_error.cleanup_errors = [repr(cleanup)] raise retry_error from cleanup - if failure['kind'] != 'gpu_oom' and not native_oom: + if failure['kind'] not in ('gpu_oom', 'shared_memory_spill') and not native_oom: raise retry_error try: manager.unload_all_models() @@ -202,6 +226,7 @@ def encode_with_recovery(clip, tokens, manager, torch, phase, *, max_attempts=3) retry_error.cleanup_errors = [repr(cleanup)] raise retry_error from cleanup del retry_error + previous = failure['kind'] manager.EXTRA_RESERVED_VRAM = original_reserve + (index + 1) * GiB phase('encoder_retry', gpu=snapshot(torch, clip, manager), encoder_attempts=list(attempts)) finally: diff --git a/freevideo_engine/encoder_precision.py b/freevideo_engine/encoder_precision.py new file mode 100644 index 0000000..95e1d71 --- /dev/null +++ b/freevideo_engine/encoder_precision.py @@ -0,0 +1,134 @@ +"""Run the H3 text encoder's language model in BF16 instead of FP32. + +ComfyUI's text-encoder base builds the input embeddings in FP32 and calls the +language model with dtype=float32, so every NVFP4 weight is dequantized to +FP32 and every matmul runs in FP32, without tensor cores. The official MiniMax +H3 pipeline runs this encoder in BF16, the model's own dtype, and conditions +the transformer on BF16 hidden states; FreeVideo saves them in BF16 too. BF16 +halves the working memory per token and uses tensor cores. +""" +from contextlib import contextmanager +import types + +import torch + +from .encoder_lowmem import language_model + + +def vision_model(model): + """The Qwen3-VL wrapper that prepares image positions and DeepStack features.""" + for _, module in model.named_modules(): + if callable(getattr(type(module), 'build_image_inputs', None)): + return module + return None + + +def single_concat_image_inputs(wrapper): + """build_image_inputs that joins each DeepStack layer once, in BF16. + + The native version grows every layer with one torch.cat per image or video + block, a full copy each time, and keeps every block's own FP32 features. For + two 15 s reference videos and three images (33 blocks) on Windows, where + PyTorch has no expandable segments, those growing copies fragmented the + allocator past 16 GiB before the language model started. The features are + cast to the hidden-state dtype where they are added, so BF16 here is exact. + """ + from comfy.text_encoders.qwen_vl import qwen2vl_mrope_position_ids + + def build_image_inputs(self, embeds, embeds_info): + images = sorted([e for e in embeds_info if e.get("type") == "image"], key=lambda e: e["index"]) + if not images: + return None, None, None + device, length = embeds.device, embeds.shape[1] + position_ids = qwen2vl_mrope_position_ids(embeds_info, length, device) + visual_pos_masks = torch.zeros((1, length), dtype=torch.bool, device=device) + for row in images: + visual_pos_masks[0, row["index"]:row["index"] + row["size"]] = True + for row in images: + row["extra"]["deepstack"] = list(row["extra"]["deepstack"]) + deepstack = [] + for layer in range(len(images[0]["extra"]["deepstack"])): + parts = [row["extra"]["deepstack"][layer] for row in images] + joined = torch.empty((sum(part.shape[0] for part in parts),) + tuple(parts[0].shape[1:]), + device=device, dtype=torch.bfloat16) + start = 0 + for index, part in enumerate(parts): + joined[start:start + part.shape[0]] = part + start += part.shape[0] + images[index]["extra"]["deepstack"][layer] = None # each block's own copy is no longer needed + deepstack.append(joined) + del parts + return position_ids, visual_pos_masks, deepstack + return types.MethodType(build_image_inputs, wrapper) + + +def compact_vision_outputs(wrapper): + """preprocess_embed that returns each block's cached memory before the next block. + + The vision tower runs once per image or video block. Without expandable + segments (Windows), PyTorch kept every block's freed activation memory + instead of reusing it: 33 blocks of two 15 s references grew reserved memory + from 12.7 to 26.5 GiB while live tensors grew by 2.7 GiB, and on a 16 GB + card the encode overflowed into shared memory before the language model + started. Returning the cache after each block kept the whole encode at 18.2 + GiB including 12.3 GiB of weights. DeepStack features are kept in BF16, the + dtype the language model casts them to anyway (exact). + """ + original = wrapper.preprocess_embed + + def preprocess_embed(self, embed, device): + merged, extra = original(embed, device) + if merged is None or not merged.is_cuda: + return merged, extra + if extra and extra.get('deepstack') is not None: + extra = dict(extra, deepstack=[value.to(torch.bfloat16) for value in extra['deepstack']]) + torch.cuda.empty_cache() + return merged, extra + return types.MethodType(preprocess_embed, wrapper) + + +@contextmanager +def bf16_language_model(model): + llama = language_model(model) + wrapper = vision_model(model) + wrapper_previous = wrapper.__dict__.get('build_image_inputs') if wrapper is not None else None + embed_previous = wrapper.__dict__.get('preprocess_embed') if wrapper is not None else None + if wrapper is not None: + wrapper.build_image_inputs = single_concat_image_inputs(wrapper) + wrapper.preprocess_embed = compact_vision_outputs(wrapper) + previous = llama.__dict__.get('forward') + current = llama.forward + + def forward(self, x, attention_mask=None, embeds=None, num_tokens=None, intermediate_output=None, + final_layer_norm_intermediate=True, dtype=None, position_ids=None, embeds_info=[], + past_key_values=None, input_ids=None, deepstack_embeds=None, visual_pos_masks=None): + if embeds is not None: + embeds = embeds.to(torch.bfloat16) + if deepstack_embeds is not None: + deepstack_embeds = [value.to(torch.bfloat16) for value in deepstack_embeds] + if attention_mask is not None and attention_mask.is_floating_point(): + attention_mask = attention_mask.to(torch.bfloat16) + result = current(x, attention_mask=attention_mask, embeds=embeds, num_tokens=num_tokens, + intermediate_output=intermediate_output, + final_layer_norm_intermediate=final_layer_norm_intermediate, dtype=torch.bfloat16, + position_ids=position_ids, embeds_info=embeds_info, past_key_values=past_key_values, + input_ids=input_ids, deepstack_embeds=deepstack_embeds, visual_pos_masks=visual_pos_masks) + return result + + llama.forward = types.MethodType(forward, llama) + try: + yield + finally: + if previous is None: + del llama.forward + else: + llama.forward = previous + if wrapper is not None: + if wrapper_previous is None: + del wrapper.build_image_inputs + else: + wrapper.build_image_inputs = wrapper_previous + if embed_previous is None: + del wrapper.preprocess_embed + else: + wrapper.preprocess_embed = embed_previous diff --git a/freevideo_engine/encoder_workspace.py b/freevideo_engine/encoder_workspace.py new file mode 100644 index 0000000..4589494 --- /dev/null +++ b/freevideo_engine/encoder_workspace.py @@ -0,0 +1,271 @@ +"""Dedicated VRAM for each text-encoder forward pass, sized from its real input. + +ComfyUI loads as many encoder weights as its reserve allows and streams the +rest. The forward then needs working memory that grows with the sequence the +language model actually sees: text tokens plus the expanded tokens of every +reference image and video block (the token list holds one entry per image). +When that room is short, Linux raises out-of-memory and the retry leaves more. +Windows does not: the driver places the overflow in shared system memory and +the encode runs many times slower. A reference request on a 16 GB RTX 4060 Ti +sat there for over ten minutes with zero dedicated VRAM free. + +Every encode therefore sizes its room from the input before loading, checks +the dedicated room really free after loading (moving weights back to host +memory until it fits), and switches long inputs to the block-wise language +model (encoder_lowmem) when the whole-sequence forward cannot fit at all. On +Windows a sampler stops a forward that still outgrows its room so recovery can +retry with more. What each forward used is kept per machine and input kind. +""" +import json +import math +from pathlib import Path +import sys +import threading +import time + +GiB, MiB = 2**30, 2**20 +LOW_MEMORY_CHUNK = 2048 +# From here the block-wise forward runs even when the whole sequence fits: the same +# math without the dense causal mask, 15.7k tokens in 7.2 s instead of 9.9 s in BF16. +WHOLE_BLOCK_TOKENS = 4096 +# Upper bounds of the measured working memory (allocator peak above the loaded +# weights) of the 32B H3 encoder running in BF16, partially loaded as on a 16 GB +# card, for text from 480 to 2880 tokens, one to eight reference images up to +# 1920x1080 and reference videos up to two 15 s clips (34k tokens). Per sequence +# token and per patch of the largest image or video block, which the vision +# tower runs alone. Measured peaks: whole 0.49-15.17 GiB, blocks 0.76-8.11 GiB. +FORMULA = {'whole': (1.5 * GiB, 0.45 * MiB, 0.30 * MiB), + 'blocks': (1.5 * GiB, 0.20 * MiB, 0.30 * MiB)} +HEADROOM, MARGIN = 1.1, 256 * MiB +SPILL_MARGIN = 256 * MiB +HISTORY_NAME = 'encoder-workspace.json' +KEEP = 32 +RECENT = 8 + + +class SharedMemorySpill(RuntimeError): + """The forward outgrew the dedicated VRAM that was free when it started.""" + + def __init__(self, used_bytes, spilled_bytes): + super().__init__('Text encoding spilled %.2f GiB into shared GPU memory' % (spilled_bytes / GiB)) + self.used_bytes, self.spilled_bytes = int(used_bytes), int(spilled_bytes) + + +def smart_grid(height, width, factor=32, min_pixels=3136, max_pixels=12845056, patch=16): + """Patch grid of one image or video block after the native Qwen3-VL resize.""" + h_bar, w_bar = round(height / factor) * factor, round(width / factor) * factor + if h_bar * w_bar > max_pixels: + beta = math.sqrt(height * width / max_pixels) + h_bar = max(factor, math.floor(height / beta / factor) * factor) + w_bar = max(factor, math.floor(width / beta / factor) * factor) + elif h_bar * w_bar < min_pixels: + beta = math.sqrt(min_pixels / (height * width)) + h_bar, w_bar = math.ceil(height * beta / factor) * factor, math.ceil(width * beta / factor) * factor + return (h_bar // patch) * (w_bar // patch) + + +def input_size(tokens): + """The language model's real sequence length and the largest vision-tower input.""" + entries = [] + if isinstance(tokens, dict): + batches = tokens.get('qwen3vl_32b') or [] + entries = [entry for batch in batches if isinstance(batch, (list, tuple)) for entry in batch] + text = vision = largest = blocks = 0 + for entry in entries: + value = entry[0] if isinstance(entry, (list, tuple)) and entry else entry + data = value.get('data') if isinstance(value, dict) else None + if data is not None and hasattr(data, 'shape') and len(data.shape) == 4: + grid = smart_grid(int(data.shape[1]), int(data.shape[2])) + vision += grid // 4 + largest = max(largest, grid) + blocks += 1 + else: + text += 1 + return dict(text_tokens=text, vision_tokens=vision, sequence=text + vision, largest_patches=largest, + vision_entries=blocks) + + +def formula(mode, size): + fixed, per_token, per_patch = FORMULA[mode] + return fixed + per_token * size['sequence'] + per_patch * size['largest_patches'] + + +class Workspace: + """Room estimates for this machine: the measured bound, raised by what forwards really used.""" + + def __init__(self, history=None, machine='default'): + self.history = Path(history) if history else None + self.machine = machine + self.data = {} + if self.history is not None: + try: + value = json.loads(self.history.read_text(encoding='utf-8')) + if value.get('schema') == 1 and isinstance(value.get('machines'), dict): + self.data = value + except (OSError, ValueError, AttributeError): + pass + self.data.setdefault('schema', 1) + self.data.setdefault('machines', {}) + + @classmethod + def for_encoder(cls, torch, checkpoint): + props = torch.cuda.get_device_properties(torch.cuda.current_device()) + path = Path(checkpoint) + machine = '|'.join((sys.platform, props.name, str(round(props.total_memory / GiB)), path.name, + str(path.stat().st_size))) + try: + from .paths import data_root + history = data_root() / HISTORY_NAME + except Exception: + history = None + return cls(history, machine) + + def scale(self, mode): + """How far this machine's recent forwards exceeded the measured bound's input-dependent part. + + Only the latest forwards count, so a driver or code change that fixed an + overflow stops reserving extra memory after a few encodes. + """ + rows = self.data['machines'].get(self.machine, {}).get(mode, [])[-RECENT:] + return max([1.] + [row[2] for row in rows if isinstance(row, list) and len(row) > 2 + and isinstance(row[2], (int, float)) and math.isfinite(row[2])]) + + def need(self, mode, size): + fixed = FORMULA[mode][0] + variable = (formula(mode, size) - fixed) * min(4., self.scale(mode)) + return int((fixed + variable) * HEADROOM + MARGIN) + + def learn(self, mode, size, used, spilled=False): + if not used or used <= 0: + return + fixed = FORMULA[mode][0] + variable = formula(mode, size) - fixed + ratio = (used - fixed) / variable if variable > 0 else 1. + rows = self.data['machines'].setdefault(self.machine, {}).setdefault(mode, []) + rows.append([int(size['sequence']), int(used), round(max(0., ratio), 4), bool(spilled), round(time.time())]) + del rows[:-KEEP] + if self.history is not None: + try: + from .monitoring import save + save(self.history, self.data) + except OSError: + pass # Learning only refines the bound; the encode itself does not depend on it. + + +def dedicated_room(torch, reader=None, budget=None): + """Bytes a new allocation can use without leaving dedicated VRAM or this request's budget. + + PyTorch's cached free blocks count. cudaMemGetInfo is not the WDDM budget: + after a load that left 1.9 GiB by its count, the 4060 Ti above reported no + free memory at all, so on Windows the DXGI budget bounds it as well. + `budget` is what this process may hold in total under FreeVideo's plan. + """ + free = torch.cuda.mem_get_info()[0] + if reader is not None: + try: + local = reader.sample()['local'] + free = min(free, max(0, local['budget_bytes'] - local['usage_bytes'])) + except (OSError, KeyError): + pass + room = free + torch.cuda.memory_reserved() - torch.cuda.memory_allocated() + if budget is not None: + room = min(room, budget - torch.cuda.memory_allocated()) + return int(max(0, room)) + + +def adapter_reader(): + if sys.platform != 'win32': + return None + try: + from .windows_gpu_memory import AdapterMemory + return AdapterMemory() + except Exception: + return None + + +def make_room(patcher, torch, need, reader=None, budget=None): + """Move loaded weights back to host memory until `need` bytes of dedicated VRAM are free.""" + before = dedicated_room(torch, reader, budget) + result = dict(need_bytes=int(need), room_before_bytes=before, loaded_before_bytes=int(patcher.loaded_size())) + if before < need: + unloaded = patcher.partially_unload(patcher.offload_device, need - before) + result.update(unloaded_bytes=int(unloaded or 0), room_after_bytes=dedicated_room(torch, reader, budget), + loaded_after_bytes=int(patcher.loaded_size())) + return result + + +def plan(workspace, size, capacity, attempt=0, failed=None): + """Mode and room for this attempt: the whole-sequence forward when it fits, else blocks. + + `capacity` is the dedicated room with no encoder weights on the device. After + a spill or an out-of-memory, whole-sequence forwards move to blocks and block + forwards ask for half again as much room. + """ + whole, blocks = workspace.need('whole', size), workspace.need('blocks', size) + mode = 'whole' if failed is None and whole <= capacity else 'blocks' + need = whole if mode == 'whole' else blocks + if failed == 'blocks': + need = int(need * 1.5 ** attempt) + return mode, min(need, max(capacity, 0)), dict(whole_need_bytes=whole, blocks_need_bytes=blocks, + capacity_bytes=int(capacity)) + + +class SpillGuard: + """Measure one forward's working memory; on Windows, stop it before it runs from shared memory. + + WDDM does not fail an allocation past the dedicated budget, so nothing + raises. A sampler compares the allocator's growth with the dedicated room + free at the start; once it is exceeded, the next module call raises + SharedMemorySpill and the recovery retries with more room. + """ + + def __init__(self, torch, reader=None, interval=.25): + self.torch, self.reader, self.interval = torch, reader, interval + self.hook = self.thread = None + self.spill = None + self.base_allocated = 0 + + def arm(self): + torch = self.torch + self.disarm() + self.spill = None + self.base_allocated = torch.cuda.memory_allocated() + self.base_reserved = torch.cuda.memory_reserved() + torch.cuda.reset_peak_memory_stats() + self.room = None + if self.reader is not None: + try: + local = self.reader.sample()['local'] + self.room = max(0, local['budget_bytes'] - local['usage_bytes']) + except (OSError, KeyError): + self.room = None + if self.room is None: + return + self.stop = threading.Event() + self.thread = threading.Thread(target=self._watch, name='freevideo-encoder-spill', daemon=True) + self.thread.start() + self.hook = torch.nn.modules.module.register_module_forward_pre_hook(self._check) + + def _watch(self): + while not self.stop.wait(self.interval): + grown = self.torch.cuda.memory_reserved() - self.base_reserved + if grown - self.room > SPILL_MARGIN: + self.spill = (self.used(), grown - self.room) + return + + def _check(self, module, arguments): + if self.spill is not None: + raise SharedMemorySpill(*self.spill) + + def used(self): + """Working memory beyond the loaded weights: cached blocks it reused plus new ones.""" + return max(0, int(self.torch.cuda.max_memory_reserved() - self.base_allocated)) + + def disarm(self): + if self.hook is not None: + self.hook.remove() + self.hook = None + if self.thread is not None: + self.stop.set() + self.thread.join() + self.thread = None diff --git a/freevideo_engine/generate.py b/freevideo_engine/generate.py index f373f2d..744b2c7 100644 --- a/freevideo_engine/generate.py +++ b/freevideo_engine/generate.py @@ -469,6 +469,8 @@ def interrupted(signum, frame): report['tuning']['conditioning_cache_hit'] = True report['input_cache']['conditioning_hit'] = True print(json.dumps({'event': 'conditioning_cache_hit'}), flush=True) + for trimmed in cached.get('reference_trims') or []: + print(json.dumps(dict(event='reference_trimmed', **trimmed)), flush=True) else: encoder_env = dict(env, PYTHONPATH=str(repo)) print(json.dumps({'event': 'encoding_start'}), flush=True) diff --git a/freevideo_engine/macos_encoder.py b/freevideo_engine/macos_encoder.py index ec2208b..9f195c3 100644 --- a/freevideo_engine/macos_encoder.py +++ b/freevideo_engine/macos_encoder.py @@ -221,7 +221,8 @@ def release_encoder(): info = encode_latents(normalized, value, base, canvas, device='mps') backend.synchronize() media_metrics.update(conditioning_info=info, media_vae_seconds=time.monotonic() - tick, - media_budget=media_policy) + media_budget=media_policy, + reference_trims=[row['trimmed'] for row in normalized if row.get('trimmed')]) output.parent.mkdir(parents=True, exist_ok=True) temporary = output.with_suffix('.partial') print(json.dumps(dict(event='encoder_phase', stage='encoder_save')), flush=True) diff --git a/freevideo_engine/media_encoding.py b/freevideo_engine/media_encoding.py index eaaa095..24cd264 100644 --- a/freevideo_engine/media_encoding.py +++ b/freevideo_engine/media_encoding.py @@ -2,7 +2,10 @@ Reference video pixels stay uint8 on disk; only the Qwen 2 fps sample and one VAE input are materialized. Neither ComfyUI nor a whole float RGB movie is kept -in the generation process. Reference clips must be explicitly trimmed to 15 s. +in the generation process. A reference clip contributes as much as the video +being generated, at most 15 s, as the official MiniMax H3 pipeline and ComfyUI's +native H3 node truncate references to the generated length. A shortened clip is +reported, never cut silently. """ import json import math @@ -10,6 +13,22 @@ from pathlib import Path +REFERENCE_FRAMES = 360 # 15 s at 24 fps, the longest H3 generation + + +def reference_frames(canvas): + """Frames of a reference clip the generation uses: its own length, at most 15 s.""" + frames = canvas.get('frames') + return min(REFERENCE_FRAMES, int(frames)) if frames else REFERENCE_FRAMES + + +def clip_seconds(container, stream): + """A clip's length from its metadata, or None when the file does not say.""" + if stream.duration and stream.time_base: + return float(stream.duration * stream.time_base) + return container.duration / 1e6 if container.duration else None + + def reference_size(width, height, canvas): scale = min(1., math.sqrt(canvas['width'] * canvas['height'] / (width * height))) return max(32, round(width * scale / 32) * 32), max(32, round(height * scale / 32) * 32) @@ -46,31 +65,39 @@ def encode_visual(vae, pixels, *, device='cuda'): return (latent - mean) / std -def _audio(path, destination): +def _audio(path, destination, seconds): + """The clip's first `seconds` at 32 kHz stereo, as the official pipeline truncates reference audio. + + Returns the saved path (None without an audio stream), whether the clip ran + more than 50 ms longer, and its length from the file's metadata. + """ import av import numpy as np chunks = [] + limit, count = int(seconds * 32000), 0 with av.open(str(path)) as container: if not container.streams.audio: - return None + return None, False, None + length = clip_seconds(container, container.streams.audio[0]) resampler = av.AudioResampler(format='fltp', layout='stereo', rate=32000) - count = 0 for frame in container.decode(audio=0): for item in resampler.resample(frame): data = item.to_ndarray() count += data.shape[1] - if count > 15 * 32000: - raise ValueError('Reference audio exceeds 15 s; trim it explicitly before generation') chunks.append(data) - for item in resampler.resample(None): - chunks.append(item.to_ndarray()) + if count > limit + 1600: + break + else: + for item in resampler.resample(None): + count += item.samples + chunks.append(item.to_ndarray()) if not chunks: raise ValueError('Reference audio is empty') - data = np.concatenate(chunks, axis=1) - if data.shape[1] > 15 * 32000 or not np.isfinite(data).all(): - raise ValueError('Invalid reference audio length or samples') + data = np.concatenate(chunks, axis=1)[:, :limit] + if not np.isfinite(data).all(): + raise ValueError('Invalid reference audio samples') np.save(destination, data, allow_pickle=False) - return str(destination) + return str(destination), count > limit + 1600, length def prepare(media, canvas, directory): @@ -82,6 +109,8 @@ def prepare(media, canvas, directory): directory = Path(directory) directory.mkdir(parents=True, exist_ok=True) result, key_images, ref_items = [], [], [] + limit = reference_frames(canvas) + numbers = {'image': 0, 'video': 0, 'audio': 0} for anchor in ('first', 'last'): if not media.get(anchor): continue @@ -95,6 +124,8 @@ def prepare(media, canvas, directory): for index, row in enumerate(media.get('references', [])): kind, path = row['kind'], row['path'] ref = {'kind': kind} + numbers[kind] += 1 + longer, length, used = False, None, limit / 24. target = directory / ('reference-%02d.npy' % index) if kind == 'image': with Image.open(path) as source: @@ -110,9 +141,8 @@ def prepare(media, canvas, directory): raise ValueError('Reference video has no video stream') stream = container.streams.video[0] width, height = reference_size(stream.width, stream.height, canvas) - if stream.duration and float(stream.duration * stream.time_base) > 15.05: - raise ValueError('Reference video exceeds 15 s; trim it explicitly') - pixels = np.lib.format.open_memmap(target, mode='w+', dtype=np.uint8, shape=(360, height, width, 3)) + length = clip_seconds(container, stream) + pixels = np.lib.format.open_memmap(target, mode='w+', dtype=np.uint8, shape=(limit, height, width, 3)) origin = None count = 0 previous = None @@ -122,11 +152,15 @@ def prepare(media, canvas, directory): if origin is None: origin = timestamp timestamp -= origin - if timestamp > 15.05: - raise ValueError('Reference video exceeds 15 s; trim it explicitly') + if count >= limit: + # Only the generated length is used; a frame 50 ms past it means the clip was longer. + if timestamp > limit / 24. + .05: + longer = True + break + continue image = frame.reformat(width=width, height=height, format='rgb24').to_ndarray() # Explicit 24 fps reference normalization, independent of output geometry. - while count / 24. <= timestamp + 1e-6 and count < 360: + while count / 24. <= timestamp + 1e-6 and count < limit: pixels[count] = image if previous is None or abs(count / 24. - timestamp) < .5 / fps else previous count += 1 previous = image @@ -136,7 +170,9 @@ def prepare(media, canvas, directory): raise ValueError('Reference video needs at least 22 normalized frames for H3/Qwen temporal encoding') pixels.flush() ref.update(pixels=str(target), frames=count, normalized_fps=24) - soundtrack = _audio(path, directory / ('reference-%02d-audio.npy' % index)) + used = count / 24. + soundtrack, longer_audio, _ = _audio(path, directory / ('reference-%02d-audio.npy' % index), limit / 24.) + longer = longer or longer_audio if soundtrack: ref['audio'] = soundtrack ref_items.append({'type': 'audio'}) @@ -145,10 +181,15 @@ def prepare(media, canvas, directory): ref_items.append({'type': 'video', 'data': sampled, 'timestamps': [i / 2. for i in range(len(selected))]}) del pixels else: - ref['audio'] = _audio(path, target) + ref['audio'], longer, length = _audio(path, target, limit / 24.) if ref['audio'] is None: raise ValueError('Reference audio file has no audio stream') ref_items.append({'type': 'audio'}) + if longer: + # Numbered per kind like the prompt's