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