diff --git a/__init__.py b/__init__.py index 8104589e..c3741389 100644 --- a/__init__.py +++ b/__init__.py @@ -20,7 +20,6 @@ if __package__: FunPackVideoStitch, ) from .model_management import FunPackApplyLoraWeights, FunPackLoraLoader - from .nvfp4_patch import FunPackNVFP4ModelLoader from .samplers import FunPackHybridEuler2SSampler, FunPackDistilledFlowSampler, FunPackLTXAVSceneChainSampler from .templates import FunPackRefinementKeyLoader try: @@ -63,10 +62,6 @@ else: except Exception: FunPackApplyLoraWeights = None FunPackLoraLoader = None - try: - from nvfp4_patch import FunPackNVFP4ModelLoader - except Exception: - FunPackNVFP4ModelLoader = None try: from samplers import FunPackHybridEuler2SSampler, FunPackDistilledFlowSampler, FunPackLTXAVSceneChainSampler except Exception: @@ -101,7 +96,6 @@ NODE_CLASS_MAPPINGS = { "FunPackLTXAVSceneChainSampler": FunPackLTXAVSceneChainSampler, "FunPackApplyLoraWeights": FunPackApplyLoraWeights, "FunPackLoraLoader": FunPackLoraLoader, - "FunPackNVFP4ModelLoader": FunPackNVFP4ModelLoader, "FunPackRefinementKeyLoader": FunPackRefinementKeyLoader, } NODE_CLASS_MAPPINGS = {name: cls for name, cls in NODE_CLASS_MAPPINGS.items() if cls is not None} @@ -127,7 +121,6 @@ NODE_DISPLAY_NAME_MAPPINGS = { "FunPackLTXAVSceneChainSampler": "FunPack LTXAV Scene Chain Sampler", "FunPackApplyLoraWeights": "FunPack Apply LoRA Weights", "FunPackLoraLoader": "FunPack LoRA Loader", - "FunPackNVFP4ModelLoader": "FunPack NVFP4 Model Loader (Experimental)", "FunPackRefinementKeyLoader": "FunPack Refinement Key Loader", } diff --git a/block_steer.py b/block_steer.py deleted file mode 100644 index ffa13a90..00000000 --- a/block_steer.py +++ /dev/null @@ -1,395 +0,0 @@ -"""Block Steering (EXPERIMENTAL): rating-learned per-block residual gains for LTX-AV. - -RLHF at the MODEL level instead of the conditioning level. The rating system already -attributes quality to sampling settings and conditioning; this adds the missing -coordinate system - which REGIONS OF THE MODEL were doing the work when a generation -came out good or bad - and then steers per-block application accordingly. - -Three stages, all riding existing FunPack infra: - -1. MEASURE - forward hooks on the transformer blocks record, per model call, each - block's relative residual delta on the VIDEO stream: ||vx_out - vx_in|| / ||vx_in||. - A block that barely moves the hidden state is coasting; one with a large relative - delta is doing the work for this prompt at this noise level. Hooks are tagged with - ltx_enhancements' _FUNPACK_HOOK_TAG so the existing leak-strip covers them. - -2. ATTRIBUTE - two signals, folded into a per-key `block_profile.pt` sidecar: - * run-level: the end-of-run mean activity fingerprint is staged (`block_snapshot.pt`, - same stage-then-pair convention as x0_snapshot) and paired with the user's rating - reward on the next Refiner cycle - liked / disliked / all-runs EMA profiles. - * step-level: when output_guidance's value function is trained, a model-function - wrapper scores each step's x0_hat (one tiny-MLP forward, no extra model pass) and - credits blocks by (value delta this step) x (their activity this step, centered). - Converges in a handful of runs - matters because keys are retrained per instance. - -3. STEER - per-block gains on the video residual only: vx_out' = vx_in + g_b*(vx_out - - vx_in), g_b = 1 + strength * score_b, scores zero-mean/unit-max so emphasis is - REDISTRIBUTED rather than globally inflated, clamped to +/-MAX_GAIN_DELTA. The audio - stream is never touched (residual-gain on the video half is the intervention family - that works on LTX-AV; hidden-state injection through joint attention is the one that - corrupts audio). Exact no-op until the profile has enough paired data. - -Everything is per-refinement-key and rebuilt from ratings in a few generations - a -mechanism upgrade, not a data store to migrate. -""" - -from __future__ import annotations - -import os - -import torch - -try: - from .ltx_enhancements import _FUNPACK_HOOK_TAG, _funpack_locate_blocks -except ImportError: - from ltx_enhancements import _FUNPACK_HOOK_TAG, _funpack_locate_blocks - -# Steering never moves a block's residual contribution by more than this, regardless of -# strength widget or how lopsided the learned scores are. -MAX_GAIN_DELTA = 0.10 -# Profile readiness: at least this many rating-paired runs (with 2, per-block correlation -# is degenerate — every block correlates ±1 trivially)… -MIN_RATED_RUNS = 3 -# …with at least this much reward spread. Attribution is the per-block CORRELATION between -# reward and activity across runs — if every run is rated the same there is nothing to -# attribute, no matter how many runs were paired (the v1 liked/disliked-EMA design failed -# exactly there: 16 same-pole ratings produced zero contrast and looked "not ready"). -# Correlation, not raw covariance: run-mean activity is nearly deterministic (averaged -# over ~1200 steps it varies in the 4th decimal run-to-run), so raw covariance is -# numerically dead ("flat activity across runs") — but a tiny variation that CONSISTENTLY -# tracks reward is exactly the signal we want, and correlation is scale-free. -MIN_REWARD_STD = 0.05 -# Confidence ramp: scores scale by min(1, (runs - 2) / 6) so a young profile steers -# gently (correlation over 3 runs is noisy) and reaches full strength around 8 runs. -CONFIDENCE_FULL_AT = 8 -# …or this many value-credit runs, before scores are considered meaningful. -MIN_CREDIT_RUNS = 2 -# Rolling history of (activity fingerprint, reward) pairs kept per key. -HISTORY_MAX = 32 -# Blend weights when both attribution signals are available. -RATING_WEIGHT = 0.6 -CREDIT_WEIGHT = 0.4 -_EPS = 1e-8 - - -def _video_stream(x): - """The video tensor of a block input/output: (vx, ax) tuple on LTX-AV, plain tensor - on video-only LTXV.""" - if isinstance(x, (tuple, list)) and len(x) >= 1: - return x[0] - return x if isinstance(x, torch.Tensor) else None - - -# --------------------------------------------------------------------------- -# Stage 1 - measurement -# --------------------------------------------------------------------------- - -class BlockActivityRecorder: - """Accumulates per-block relative residual deltas (video stream) across a run. - - `cur` always holds the most recent model call's per-block activity so the value- - credit wrapper (which runs after all blocks in the same call) can read it. - """ - - def __init__(self, n_blocks: int): - self.n_blocks = int(n_blocks) - self.sums = torch.zeros(self.n_blocks, dtype=torch.float32) - self.counts = torch.zeros(self.n_blocks, dtype=torch.float32) - self.cur = torch.zeros(self.n_blocks, dtype=torch.float32) - # value-credit accumulation (filled by the credit wrapper, if installed) - self.credit = torch.zeros(self.n_blocks, dtype=torch.float32) - self.credit_steps = 0 - self._prev_score = None - - def record(self, idx: int, rel_delta: float): - if 0 <= idx < self.n_blocks: - self.sums[idx] += rel_delta - self.counts[idx] += 1.0 - self.cur[idx] = rel_delta - - def fingerprint(self): - """Mean per-block activity over the run, or None if nothing was recorded.""" - if float(self.counts.max()) <= 0: - return None - return self.sums / self.counts.clamp(min=1.0) - - def credit_step(self, score: float): - """Fold one model call's value score into per-block credit: blocks that were - MORE active than average during a step where predicted quality rose (fell) get - positive (negative) credit.""" - if self._prev_score is not None: - delta = float(score) - self._prev_score - centered = self.cur - self.cur.mean() - self.credit += delta * centered - self.credit_steps += 1 - self._prev_score = float(score) - - def credit_vector(self): - if self.credit_steps <= 0: - return None - return self.credit / float(self.credit_steps) - - -def install_recorder(model, recorder: BlockActivityRecorder): - """Tagged forward hooks on every transformer block; returns handles (possibly []).""" - blocks = _funpack_locate_blocks(model) - if blocks is None: - return [] - handles = [] - for i, blk in enumerate(blocks): - def _hook(_module, args, output, _idx=i, _rec=recorder): - try: - v_in = _video_stream(args[0] if args else None) - v_out = _video_stream(output) - if v_in is None or v_out is None or v_in.shape != v_out.shape: - return output - with torch.no_grad(): - denom = float(v_in.norm()) + _EPS - _rec.record(_idx, float((v_out - v_in).norm()) / denom) - except Exception: - pass - return output - setattr(_hook, _FUNPACK_HOOK_TAG, True) - handles.append(blk.register_forward_hook(_hook)) - return handles - - -def make_credit_wrapper(old_wrapper, recorder: BlockActivityRecorder, value_fn): - """model_function_wrapper that scores each call's x0_hat with the (already trained) - output-space value function and folds the step-to-step delta into block credit. - Installed INNERMOST so it scores the base model's own prediction, before - embed_guidance / output_guidance corrections layer on top.""" - - def _call(apply_fn, a): - if old_wrapper is not None: - return old_wrapper(apply_fn, a) - return apply_fn(a["input"], a["timestep"], **a.get("c", {})) - - def _wrapper(apply_fn, args, _rec=recorder, _vf=value_fn): - denoised = _call(apply_fn, args) - try: - with torch.inference_mode(False), torch.no_grad(): - c = _vf.compress(denoised.detach().float()) - score = float(_vf(c).mean()) - _rec.credit_step(score) - except Exception: - pass - return denoised - - return _wrapper - - -# --------------------------------------------------------------------------- -# Stage 3 - steering -# --------------------------------------------------------------------------- - -def install_steer(model, gains): - """Tagged forward hooks scaling each block's VIDEO residual by its gain: - vx' = vx_in + g*(vx_out - vx_in). Audio passes through untouched. Returns handles; - blocks whose gain is ~1.0 get no hook at all (byte-identical native path).""" - blocks = _funpack_locate_blocks(model) - if blocks is None: - return [] - handles = [] - for i, blk in enumerate(blocks): - g = float(gains[i]) if i < len(gains) else 1.0 - if abs(g - 1.0) < 1e-4: - continue - def _hook(_module, args, output, _g=g): - try: - v_in = _video_stream(args[0] if args else None) - if isinstance(output, tuple) and len(output) == 2: - v_out, a_out = output - if v_in is None or v_in.shape != v_out.shape: - return output - return (v_in + _g * (v_out - v_in), a_out) - if isinstance(output, torch.Tensor) and v_in is not None and v_in.shape == output.shape: - return v_in + _g * (output - v_in) - except Exception: - pass - return output - setattr(_hook, _FUNPACK_HOOK_TAG, True) - handles.append(blk.register_forward_hook(_hook)) - return handles - - -def remove_handles(handles): - for h in handles or []: - try: - h.remove() - except Exception: - pass - - -# --------------------------------------------------------------------------- -# Stage 2 - persistence + attribution -# --------------------------------------------------------------------------- - -def _state_path(refinement_key, name): - try: - from .conditioning import refinement_state_path - except ImportError: - from conditioning import refinement_state_path - return refinement_state_path(refinement_key, name, prefix="refine_v2", extension="pt") - - -def save_run_snapshot(refinement_key, recorder: BlockActivityRecorder): - """End of run: stage this run's fingerprint (+ optional value credit) so the NEXT - rating cycle can pair it with a reward. Atomic write - a rating read never sees a - partial file. Same convention as the sampler's x0_snapshot.""" - if not refinement_key: - return False - fp = recorder.fingerprint() - if fp is None: - return False - # Store RELATIVE emphasis (mean-normalized): run-global scale differences (step - # counts, schedulers, scene counts) must not masquerade as per-block signal. - fp = fp / (fp.mean() + _EPS) - try: - path = _state_path(refinement_key, "block_snapshot") - os.makedirs(os.path.dirname(path), exist_ok=True) - tmp = path + ".tmp" - torch.save({ - "activity": fp.cpu(), - "credit": (recorder.credit_vector().cpu() if recorder.credit_vector() is not None else None), - "n_blocks": recorder.n_blocks, - }, tmp) - os.replace(tmp, path) - return True - except Exception as e: - print(f"[FunPackBlockSteer] snapshot save failed: {e}") - return False - - -def _ema(old, x, n): - a = max(0.2, 1.0 / (float(n) + 1.0)) - return x if old is None else old * (1.0 - a) + x * a - - -def update_profile_with_rating(refinement_key, reward): - """Pair the staged snapshot with a rating reward (called by the Refiner at the same - site that trains the output value function). Appends (activity fingerprint, reward) - to a rolling per-key history — EVERY admissible reward contributes with its actual - value (no liked/disliked pole thresholds, no dead band) — and folds the run's value - credit (rating-independent, self-supervised) into the credit EMA. Consumes the - snapshot so one run is never paired twice. Returns the number of rating-paired runs - in the history, or None when there was nothing to pair.""" - if not refinement_key: - return None - try: - snap_path = _state_path(refinement_key, "block_snapshot") - if not os.path.exists(snap_path): - return None - snap = torch.load(snap_path, map_location="cpu", weights_only=False) - activity = snap.get("activity") - if not isinstance(activity, torch.Tensor): - return None - prof_path = _state_path(refinement_key, "block_profile") - prof = {} - if os.path.exists(prof_path): - prof = torch.load(prof_path, map_location="cpu", weights_only=False) - # Start fresh on: different model depth, a v1 pole-EMA profile, or a v2 - # profile whose fingerprints predate mean-normalization (raw scales would - # poison the correlation). - if (int(prof.get("n_blocks", -1)) != activity.numel() - or ("history_act" not in prof and "liked_n" in prof) - or int(prof.get("fp_version", 1)) < 3): - prof = {} - prof["n_blocks"] = activity.numel() - prof["fp_version"] = 3 - - row = activity.float().reshape(1, -1) - rw = torch.tensor([float(reward)], dtype=torch.float32) - ha, hr = prof.get("history_act"), prof.get("history_rw") - prof["history_act"] = row if ha is None else torch.cat([ha, row], dim=0)[-HISTORY_MAX:] - prof["history_rw"] = rw if hr is None else torch.cat([hr, rw], dim=0)[-HISTORY_MAX:] - - credit = snap.get("credit") - if isinstance(credit, torch.Tensor) and credit.numel() == activity.numel(): - prof["credit"] = _ema(prof.get("credit"), credit, prof.get("credit_n", 0)) - prof["credit_n"] = int(prof.get("credit_n", 0)) + 1 - - tmp = prof_path + ".tmp" - torch.save(prof, tmp) - os.replace(tmp, prof_path) - os.remove(snap_path) # consumed - never pair one run with two ratings - return int(prof["history_rw"].numel()) - except Exception as e: - print(f"[FunPackBlockSteer] profile update failed: {e}") - return None - - -def _normalize(v): - v = v - v.mean() - m = float(v.abs().max()) - return v / m if m > _EPS else None - - -def block_scores_with_status(refinement_key): - """(scores, status): combined normalized per-block scores in [-1, 1] (zero-mean, - unit max-abs), or (None, why-not). Rating attribution is the per-block covariance - between reward and activity across the run history; value credit blends in when - trained. `status` is a user-facing sentence for the sampler log.""" - if not refinement_key: - return None, "no refinement key wired" - try: - prof_path = _state_path(refinement_key, "block_profile") - if not os.path.exists(prof_path): - return None, "no profile yet — rate a generation made with block_steer on" - prof = torch.load(prof_path, map_location="cpu", weights_only=False) - ha, hr = prof.get("history_act"), prof.get("history_rw") - credit, credit_n = prof.get("credit"), int(prof.get("credit_n", 0)) - - rating_part = None - peak_corr = 0.0 - n = int(hr.numel()) if isinstance(hr, torch.Tensor) else 0 - conf = min(1.0, max(0.0, (n - (MIN_RATED_RUNS - 1)) / float(CONFIDENCE_FULL_AT - (MIN_RATED_RUNS - 1)))) - if not isinstance(ha, torch.Tensor) or n < MIN_RATED_RUNS: - rating_status = f"needs {MIN_RATED_RUNS}+ rated runs for this key (have {n})" - elif float(hr.std()) < MIN_REWARD_STD: - rating_status = (f"{n} rated runs but all rated alike — rate both good AND bad " - "runs so blocks have something to contrast") - else: - # Per-block Pearson correlation between reward and activity. Scale-free: - # run-mean activity varies only in the 4th decimal run-to-run, so raw - # covariance is numerically dead — but consistency with reward isn't. - rc = hr - hr.mean() - ac = ha - ha.mean(dim=0) - corr = (rc.unsqueeze(1) * ac).sum(dim=0) / ( - (rc.square().sum().sqrt() * ac.square().sum(dim=0).sqrt()) + _EPS) - corr = corr.clamp(-1.0, 1.0) - peak_corr = float(corr.abs().max()) - rating_part = _normalize(corr) - rating_status = ("no block's activity varies at all across runs" - if rating_part is None else "ready") - - credit_part = _normalize(credit) if (credit is not None and credit_n >= MIN_CREDIT_RUNS) else None - - if rating_part is not None and credit_part is not None: - combined = RATING_WEIGHT * rating_part + CREDIT_WEIGHT * credit_part - else: - combined = rating_part if rating_part is not None else credit_part - if combined is None: - return None, rating_status - scores = _normalize(combined) - if scores is None: - return None, "flat combined scores" - # Confidence ramp: correlation over few runs is noisy — steer gently until the - # history has substance. Applied AFTER normalization so it survives into gains. - scores = scores * conf - return scores, (f"ready ({n} runs, peak |corr| {peak_corr:.2f}, " - f"confidence {conf:.2f})") - except Exception as e: - return None, f"score load failed: {e}" - - -def load_block_scores(refinement_key): - """Scores only (None while not ready) — see block_scores_with_status.""" - return block_scores_with_status(refinement_key)[0] - - -def gains_from_scores(scores, strength): - """Per-block gains 1 + strength*score, hard-clamped to 1 +/- MAX_GAIN_DELTA.""" - s = max(0.0, float(strength)) - g = 1.0 + (s * scores).clamp(-MAX_GAIN_DELTA, MAX_GAIN_DELTA) - return [float(x) for x in g] diff --git a/conditioning.py b/conditioning.py index 0e2d0cec..ba1cb550 100644 --- a/conditioning.py +++ b/conditioning.py @@ -11910,21 +11910,6 @@ class FunPackVideoRefinerV2(FunPackVideoRefiner): n_out = self._v2_train_output_value_function(refinement_key, float(learning_profile.get("reward", 0.0))) if n_out is not None: print(f"[FunPackRefiner] Output value function updated — {n_out} samples") - # Block Steering: the same reward pairs with the sampler's staged per-block - # activity fingerprint (refinement_key.block_snapshot.pt) — feeds the - # liked/disliked/credit block profile the Chain Sampler steers from. No-op - # when the sampler ran without block_steer (no snapshot staged). - try: - try: - from . import block_steer as _bsmod - except ImportError: - import block_steer as _bsmod - n_blocks = _bsmod.update_profile_with_rating( - refinement_key, float(learning_profile.get("reward", 0.0))) - if n_blocks is not None: - print(f"[FunPackRefiner] Block profile updated — {n_blocks} rated run(s)") - except Exception as _bse: - print(f"[FunPackRefiner] Block profile update failed: {_bse}") # Absolute store: the same rating also feeds the keyless, prompt-agnostic taste prior. # Runs even with no refinement_key (Absolute is global), so standalone runs still build it. # Skipped for Wrong-* repair ratings (skip_value_function): Absolute reads reward as pure diff --git a/movie_editor/frontend/engine_settings.js b/movie_editor/frontend/engine_settings.js index 3e1a706f..ddf26e97 100644 --- a/movie_editor/frontend/engine_settings.js +++ b/movie_editor/frontend/engine_settings.js @@ -254,8 +254,6 @@ { name: "score_slider_strength", label: "Slider strength (eta)", kind: "float", default: 1.0, min: 0, max: 3, step: 0.25, dependsOn: "score_slider" }, { name: "output_guidance", label: "Output guidance", kind: "bool", default: false }, { name: "output_guidance_strength", label: "Output guidance strength", kind: "float", default: 0.02, min: 0.005, max: 0.1, step: 0.005, dependsOn: "output_guidance" }, - { name: "block_steer", label: "Block steering", kind: "bool", default: false }, - { name: "block_steer_strength", label: "Block steer strength", kind: "float", default: 0.05, min: 0.01, max: 0.15, step: 0.005, dependsOn: "block_steer" }, { name: "decode_noise_scale", label: "Decode noise scale", kind: "float", default: 0.0, min: 0, max: 1, step: 0.01 }, { name: "decode_timestep", label: "Decode timestep", kind: "float", default: 0.05, min: 0, max: 1, step: 0.01 }, { name: "decode_tile_size", label: "Decode tile size", kind: "int", default: 0, min: 0, max: 4096, step: 64 }, @@ -408,7 +406,7 @@ { id: "chain_auto", title: "Auto continuity", defaultOpen: true, knobs: [] }, { id: "chain_timing", title: "Timing", defaultOpen: false, knobs: ["frame_overlap", "transition_duration", "use_same_seed"] }, { id: "chain_cont", title: "Manual continuity", defaultOpen: false, knobs: ["carry_i2v_guides"] }, - { id: "chain_guid", title: "Guidance", defaultOpen: false, knobs: ["cfg", "embed_guidance", "embed_guidance_source", "embed_guidance_strength", "score_slider", "score_slider_strength", "output_guidance", "output_guidance_strength", "block_steer", "block_steer_strength"] }, + { id: "chain_guid", title: "Guidance", defaultOpen: false, knobs: ["cfg", "embed_guidance", "embed_guidance_source", "embed_guidance_strength", "score_slider", "score_slider_strength", "output_guidance", "output_guidance_strength"] }, { id: "chain_dec", title: "Decode", defaultOpen: false, knobs: ["decode_noise_scale", "decode_timestep", "decode_tile_size"] }, { id: "chain_exp", title: "Experimental", defaultOpen: false, knobs: ["mid_scene_guide", "mid_scene_guide_strength", "joyai_memory", "joyai_memory_size", "joyai_fix_frames", "joyai_frame_select", "joyai_memory_strength", "joyai_audio_memory", "v2a_grad_scale", "alg_blur_guides", "bounded_attention_enabled"] }, ]; diff --git a/nvfp4_patch.py b/nvfp4_patch.py deleted file mode 100644 index 30bd349c..00000000 --- a/nvfp4_patch.py +++ /dev/null @@ -1,527 +0,0 @@ -"""FunPack NVFP4 on-the-fly model patching (EXPERIMENTAL). - -Blackwell's 5th-gen tensor cores run NVFP4 (e2m1 with per-16-element FP8-E4M3 block -scales + a per-tensor FP32 scale) at roughly 2x FP8 / 4x BF16 GEMM throughput. ComfyUI -core already ships the whole runtime for it (comfy.quant_ops layouts + MixedPrecisionOps -dispatch, dynamic activation quantization, LoRA fallback, dequant emulation on non-FP4 -GPUs) - but only for checkpoints that were quantized ahead of time. Stock loaders offer -fp8 at best, and on sm_120 the fp8 path is mostly storage, not tensor-core speedup. - -This module closes that gap: it quantizes an existing bf16/fp16 LTX-AV checkpoint to -NVFP4 *at load time* as a pure state-dict transform - selected 2D Linear weights are -re-serialized in core's exact quantized-checkpoint format (weight/weight_scale/ -weight_scale_2 + per-layer `comfy_quant` metadata), then handed to the normal -`load_diffusion_model_state_dict` path. Core's own detection routes the model through -MixedPrecisionOps; every layer we did NOT tag simply stays at compute dtype. No custom -forward, no monkeypatching, no new inference code paths to maintain. - -Zone-aware precision map (the part generic quantizers can't do): LTX-AV is a joint -audio-video transformer and its zones are not equally robust. By default we quantize -only the video-branch Linears (attn1/attn2/ff) of the middle blocks, keeping at full -precision: - - the audio branch (audio_attn1/audio_attn2/audio_ff) - joint attention is the part - that corrupts first (see the self-consistency post-block injection failure), - - the cross-modal bridges (audio_to_video_attn / video_to_audio_attn), - - the first/last N blocks (standard quantization practice; they carry the most - input/output-adjacent signal), - - everything that is not a plain 2D matmul (norms, ada modulation, embeddings). - -Scopes widen from there for experimentation. W4A4 without calibration WILL cost some -fidelity - this is an experimental speed/VRAM lever, off unless you wire the node in. - -Requires: ComfyUI with comfy.quant_ops (comfy_kitchen), torch cu13+ on Blackwell -(sm_120+) for real FP4 GEMMs. On anything older the model still loads and runs through -core's dequant emulation (no speedup, same quality) - useful for A/B sanity checks. -""" - -from __future__ import annotations - -import json -import logging -import os -import re - -import torch - -# --------------------------------------------------------------------------- -# Static input-scale calibration -# --------------------------------------------------------------------------- -# Core quantizes each layer's INPUT dynamically (amax reduction + quantize) on every -# call unless the layer has a static `input_scale` parameter. Live probe on sm_120: -# the dynamic path costs ~0.30 ms/layer-call and a provided scale cuts it to ~0.11 ms -# (amax dominates) - across ~800 quantized layers that is a top-line cost of every -# model call. Calibration is capture-then-freeze: a "calibrate" run records each -# quantized layer's input amax via forward-pre-hooks into a per-checkpoint sidecar; -# later loads inject `input_scale = amax * MARGIN / (448 * 6)` into the state dict, -# where core's loader picks it up natively. MARGIN leaves headroom because NVFP4 -# clips activations above the scale; per the disposable-keys philosophy the sidecar -# is per-machine and rebuilt in one run whenever absent. -CALIBRATION_MARGIN = 1.25 -_F8F4_MAX = 448.0 * 6.0 # F8_E4M3_MAX * F4_E2M1_MAX -INPUT_SCALE_MODES = ["dynamic", "auto-calibrate"] - - -def _scales_path(checkpoint_path): - base = os.path.join(os.path.dirname(os.path.abspath(__file__)), "refinements", "nvfp4_scales") - return os.path.join(base, os.path.basename(checkpoint_path) + ".json") - - -def load_calibration(checkpoint_path): - """{module_path: input_amax} recorded by a previous calibrate run, or None.""" - try: - p = _scales_path(checkpoint_path) - if not os.path.exists(p): - return None - with open(p, "r", encoding="utf-8") as f: - data = json.load(f) - amax = data.get("amax") or {} - return {str(k): float(v) for k, v in amax.items()} or None - except Exception as e: # noqa: BLE001 - logging.warning(f"[FunPack NVFP4] could not read calibration sidecar: {e}") - return None - - -def inject_input_scales(sd, target_weight_keys, calibration, margin=CALIBRATION_MARGIN): - """Add `.input_scale` next to each quantized weight whose module path has a - calibrated amax. Keys are matched by suffix so any checkpoint prefix - ('model.diffusion_model.' or none) works. Returns the number injected.""" - injected = 0 - for k in target_weight_keys: - prefix = k[: -len("weight")] # '...transformer_blocks.3.attn1.to_q.' - module_path = prefix.rstrip(".") - for cal_key, amax in calibration.items(): - if module_path == cal_key or module_path.endswith("." + cal_key): - sd[prefix + "input_scale"] = torch.tensor( - float(amax) * margin / _F8F4_MAX, dtype=torch.float32) - injected += 1 - break - return injected - - -def install_calibration_capture(model, checkpoint_path): - """Forward-pre-hooks on every quantized Linear recording the running input amax. - The sidecar is written on the first full sweep and refreshed periodically; hooks - live for the model instance's lifetime (a calibration run is expected to be - slightly SLOWER - it pays the amax it is trying to eliminate, once).""" - store = {"amax": {}, "writes": 0, "calls": 0} - path = _scales_path(checkpoint_path) - - def _flush(): - try: - os.makedirs(os.path.dirname(path), exist_ok=True) - tmp = path + ".tmp" - with open(tmp, "w", encoding="utf-8") as f: - json.dump({"amax": store["amax"], "margin": CALIBRATION_MARGIN}, f) - os.replace(tmp, path) - store["writes"] += 1 - except Exception as e: # noqa: BLE001 - logging.warning(f"[FunPack NVFP4] calibration flush failed: {e}") - - hooked = 0 - try: - for name, m in model.model.diffusion_model.named_modules(): - if getattr(m, "quant_format", None) != "nvfp4": - continue - - def _hook(_module, args, _name=name): - try: - x = args[0] if args else None - if isinstance(x, torch.Tensor): - a = float(x.detach().abs().amax()) - if a > store["amax"].get(_name, 0.0): - store["amax"][_name] = a - store["calls"] += 1 - if store["calls"] % 4096 == 0: - _flush() - except Exception: - pass - - m.register_forward_pre_hook(_hook) - hooked += 1 - except Exception as e: # noqa: BLE001 - logging.warning(f"[FunPack NVFP4] calibration capture install failed: {e}") - return 0 - - if hooked: - # First flush after one full model call so a single generation is enough. - def _first_flush_hook(_module, _args): - if store["writes"] == 0 and len(store["amax"]) >= hooked: - _flush() - try: - last = [m for _, m in model.model.diffusion_model.named_modules() - if getattr(m, "quant_format", None) == "nvfp4"][-1] - last.register_forward_pre_hook(_first_flush_hook) - except Exception: - pass - return hooked - -QUANTIZE_SCOPES = [ - "video blocks", # video-branch Linears in middle blocks (default, safest) - "video + cross-modal", # + audio_to_video_attn / video_to_audio_attn bridges - "all blocks", # every Linear inside transformer_blocks (incl. audio branch) - "all 2D layers", # + eligible Linears outside the blocks (most aggressive) -] - -# Video-branch submodules inside a BasicAVTransformerBlock (comfy/ldm/lightricks/av_model.py). -_VIDEO_MODULES = ("attn1.", "attn2.", "ff.") -_CROSS_MODULES = ("audio_to_video_attn.", "video_to_audio_attn.") - -# Matches "...transformer_blocks.." with any leading prefix -# ("model.diffusion_model." or none - load_diffusion_model_state_dict strips prefixes -# from ALL keys uniformly, so our metadata keys ride along). -_BLOCK_RE = re.compile(r"(?:^|\.)transformer_blocks\.(\d+)\.(.*)$") - -# Outside the blocks ("all 2D layers" scope) never touch input/output-adjacent or -# modulation weights even when they are 2D Linears. -_GENERIC_SKIP = ("proj_out", "adaln", "embed", "emb.", "norm", "scale_shift", "patchify") - -# Below this, quantization overhead outweighs any GEMM win and relative error is worst. -_MIN_DIM = 512 - - -def _block_ids(sd) -> list[int]: - ids = set() - for k in sd: - m = _BLOCK_RE.search(k) - if m: - ids.add(int(m.group(1))) - return sorted(ids) - - -def _selected(key_body: str, scope: str, keep: set[int]) -> bool: - """Decide whether the Linear at `key_body` (state-dict key minus '.weight') is - quantized under `scope`, given the high-precision `keep` block set.""" - m = _BLOCK_RE.search(key_body + ".") - if not m: - if scope != "all 2D layers": - return False - low = key_body.lower() - return not any(s in low for s in _GENERIC_SKIP) - idx, rest = int(m.group(1)), m.group(2) - if idx in keep: - return False - if scope in ("all blocks", "all 2D layers"): - return True - if any(rest.startswith(p) for p in _VIDEO_MODULES): - return True - if scope == "video + cross-modal" and any(rest.startswith(p) for p in _CROSS_MODULES): - return True - return False - - -def _eligible(key: str, w) -> bool: - return ( - key.endswith(".weight") - and isinstance(w, torch.Tensor) - and w.dim() == 2 - and w.shape[0] >= _MIN_DIM - and w.shape[1] >= _MIN_DIM - and w.shape[1] % 16 == 0 # NVFP4 blocks run along in_features - ) - - -def quantize_state_dict_nvfp4(sd, scope="video blocks", keep_first_blocks=2, - keep_last_blocks=2, device=None, progress_callback=None): - """Return (new_sd, quantized_layer_names, kept_block_ids). - - Pure transform: selected 2D weights are replaced by NVFP4 QuantizedTensor - serialization (weight uint8-packed, weight_scale fp8-e4m3 block scales, - weight_scale_2 fp32 tensor scale) plus a per-layer `comfy_quant` metadata key - - byte-for-byte the format `comfy.ops._load_quantized_module` consumes. Unselected - keys pass through untouched. - """ - import comfy.quant_ops as quant_ops - - if not quant_ops._CK_AVAILABLE: - raise RuntimeError( - "comfy_kitchen is unavailable in this ComfyUI - NVFP4 quantization needs it. " - "Update ComfyUI / reinstall requirements." - ) - if scope not in QUANTIZE_SCOPES: - raise ValueError(f"Unknown quantize scope: {scope}") - for k in sd: - if k.endswith(".comfy_quant") or k.endswith("scaled_fp8"): - raise ValueError( - "This checkpoint is already quantized (comfy_quant / scaled_fp8 metadata " - "found). NVFP4 patching needs the original bf16/fp16 weights." - ) - - ids = _block_ids(sd) - keep: set[int] = set() - if ids: - if keep_first_blocks > 0: - keep |= set(ids[:keep_first_blocks]) - if keep_last_blocks > 0: - keep |= set(ids[-keep_last_blocks:]) - - targets = [k for k, w in sd.items() if _eligible(k, w) and _selected(k[: -len(".weight")], scope, keep)] - - meta = torch.tensor(list(json.dumps({"format": "nvfp4"}).encode("utf-8")), dtype=torch.uint8) - out = {} - done = 0 - for k, w in sd.items(): - if k not in targets: - out[k] = w - continue - wq = w - if wq.dtype not in (torch.float32, torch.bfloat16): - # fp16 and fp8-stored checkpoints alike: quantize from bf16. - wq = wq.to(torch.bfloat16) - if device is not None: - wq = wq.to(device) - qt = quant_ops.QuantizedTensor.from_float(wq, "TensorCoreNVFP4Layout", scale="recalculate") - for kk, vv in qt.state_dict(k).items(): - out[kk] = vv.to("cpu") - out[k[: -len("weight")] + "comfy_quant"] = meta.clone() - del qt, wq - done += 1 - if progress_callback is not None: - progress_callback(done, len(targets)) - if done % 100 == 0: - logging.info(f"[FunPack NVFP4] quantized {done}/{len(targets)} layers…") - - return out, targets, keep - - -def log_nvfp4_diagnostics(model, device=None): - """Print everything needed to tell WHY NVFP4 is (or is not) fast on this machine. - - 'Loaded fine' and 'runs fast' are different claims. The three ways the fast path - silently degrades, in the order they bite: - 1. comfy_kitchen's CUDA backend is disabled when torch is built for CUDA < 13 - - quantized matmuls then fall through to the pure-PyTorch eager backend, which - is CORRECT but no faster than bf16 (the exact "speed didn't change" symptom). - 2. The GPU lacks FP4 tensor cores (pre-Blackwell) - core flags every nvfp4 layer - _full_precision_mm and dequantizes (visible in the layer census below). - 3. Everything is native but the quantized scope is too small a slice of the - total step time to move the needle. - A timed 4096x4096 GEMM probe settles it empirically: on a native FP4 path the - quantized matmul must beat bf16 clearly; parity or worse means emulation.""" - lines = [] - try: - import comfy_kitchen as ck - for name, info in ck.list_backends().items(): - if info.get("available"): - caps = info.get("capabilities") or [] - lines.append(f" ck backend '{name}': AVAILABLE, scaled_mm_nvfp4=" - f"{'yes' if 'scaled_mm_nvfp4' in caps else 'NO'}") - else: - lines.append(f" ck backend '{name}': unavailable ({info.get('unavailable_reason')})") - except Exception as e: # noqa: BLE001 - lines.append(f" comfy_kitchen introspection failed: {e}") - cap = torch.cuda.get_device_capability(device) if torch.cuda.is_available() else None - lines.append(f" torch {torch.__version__} | cuda build {torch.version.cuda} | device capability {cap}") - if torch.version.cuda is not None: - try: - if tuple(map(int, str(torch.version.cuda).split(".")))[0] < 13: - lines.append(" ⚠ torch is built for CUDA < 13 -> comfy_kitchen DISABLES its CUDA " - "backend -> nvfp4 matmuls run on the eager (pure PyTorch) backend: " - "correct but NO speedup. Install a cu13x torch build to unlock FP4 GEMMs.") - except Exception: # noqa: BLE001 - pass - - n_q = n_emul = 0 - try: - for m in model.model.diffusion_model.modules(): - if getattr(m, "quant_format", None) == "nvfp4": - n_q += 1 - if getattr(m, "_full_precision_mm", False): - n_emul += 1 - lines.append(f" nvfp4 layers in loaded model: {n_q}" - + (f" — {n_emul} flagged FULL-PRECISION EMULATION (no FP4 compute on this GPU)" - if n_emul else " — all on the quantized matmul path")) - except Exception as e: # noqa: BLE001 - lines.append(f" layer census failed: {e}") - - if torch.cuda.is_available(): - try: - import time - import comfy.quant_ops as quant_ops - d = torch.device(device) if device is not None else torch.device("cuda") - w = torch.randn(4096, 4096, dtype=torch.bfloat16, device=d) - qw = quant_ops.QuantizedTensor.from_float(w, "TensorCoreNVFP4Layout", scale="recalculate") - - def _timed(fn, n=30): - for _ in range(5): - fn() - torch.cuda.synchronize(d) - t0 = time.perf_counter() - for _ in range(n): - fn() - torch.cuda.synchronize(d) - return (time.perf_counter() - t0) / n - - # Decomposed probe at two token counts: M=1024 (small scene chunk) and - # M=16384 (realistic video token count). Times the input quantization and - # the FP4 GEMM SEPARATELY — "no speedup" has two very different fixes - # depending on which one is at fault: - # * GEMM alone beats bf16, quantization dominates -> per-call overhead - # problem (x1232 layers/step); needs static input scales or fusion. - # * GEMM alone loses to bf16 -> the FP4 kernel isn't engaging at all. - results = {} - for m in (1024, 16384): - x = torch.randn(m, 4096, dtype=torch.bfloat16, device=d) - t_bf16 = _timed(lambda: torch.nn.functional.linear(x, w)) - t_quant = _timed(lambda: quant_ops.QuantizedTensor.from_float(x, "TensorCoreNVFP4Layout")) - # Same quantization with a PREcomputed static scale: isolates the amax - # reduction (which per-layer input_scale calibration could remove) from - # tensor-construction/kernel-launch overhead (which it can't). - s = (x.abs().amax() / (448.0 * 6.0)).float() - t_quant_s = _timed(lambda: quant_ops.QuantizedTensor.from_float(x, "TensorCoreNVFP4Layout", scale=s)) - qx = quant_ops.QuantizedTensor.from_float(x, "TensorCoreNVFP4Layout") - t_gemm = _timed(lambda: torch.nn.functional.linear(qx, qw)) - total = t_quant + t_gemm - results[m] = (t_bf16, t_quant, t_quant_s, t_gemm) - lines.append( - f" probe 4096x4096 @ M={m}: bf16 {t_bf16 * 1e3:.3f} ms | " - f"input-quant {t_quant * 1e3:.3f} ms (static-scale {t_quant_s * 1e3:.3f} ms) " - f"+ fp4-gemm {t_gemm * 1e3:.3f} ms = {total * 1e3:.3f} ms " - f"({t_bf16 / max(total, 1e-9):.2f}x net, {t_bf16 / max(t_gemm, 1e-9):.2f}x gemm-only)") - del x, qx - - # Verdict from the LARGE-M run (realistic video token counts); small-M numbers - # are launch-overhead dominated by nature and must not drive the conclusion. - t_bf16, t_quant, t_quant_s, t_gemm = results[16384] - total = t_quant + t_gemm - if t_gemm < t_bf16 * 0.8 and total < t_bf16 * 0.8: - lines.append(" verdict: FAST PATH ACTIVE — FP4 GEMMs win at real token counts. " - "If end-to-end speed still looks flat, the input-quant overhead " - f"(~{t_quant * 1e3:.2f} ms/layer-call, largely size-independent) is " - "cancelling the win on SMALL layers: prefer scope 'video blocks' " - "(skips the audio branch, whose matmuls are too small to profit).") - if t_quant_s < t_quant * 0.6: - lines.append(" note: static input scales would reclaim most of the input-quant " - "cost (amax dominates) — calibration support is worth building.") - else: - lines.append(" note: static input scales would NOT help much " - "(overhead is construction/launch, not amax).") - elif t_gemm < t_bf16 * 0.8: - lines.append(" verdict: FP4 GEMM fast but INPUT QUANTIZATION dominates even at " - "large M — overhead-bound; scope down to the largest layers only.") - else: - lines.append(" verdict: FP4 GEMM not faster than bf16 even at M=16384 — the " - "kernel is not engaging; report these numbers upstream.") - except Exception as e: # noqa: BLE001 - lines.append(f" GEMM probe failed: {e}") - else: - lines.append(" GEMM probe skipped (no CUDA device)") - - print("[FunPack NVFP4] diagnostics:\n" + "\n".join(lines)) - return {"nvfp4_layers": n_q, "emulated_layers": n_emul} - - -class FunPackNVFP4ModelLoader: - """EXPERIMENTAL: load a diffusion model with on-the-fly NVFP4 quantization. - - Drop-in replacement for 'Load Diffusion Model' aimed at LTX-AV on Blackwell - (sm_120+): selected transformer Linears run real FP4 tensor-core GEMMs (~2x fp8 - throughput) and weight VRAM drops to ~1/4 of bf16. Audio branch, cross-modal - bridges, first/last blocks, and all non-matmul weights stay at full precision - by default. On GPUs without FP4 compute the model still works via emulation. - """ - - CATEGORY = "FunPack/Experimental" - RETURN_TYPES = ("MODEL",) - RETURN_NAMES = ("model",) - FUNCTION = "load" - DESCRIPTION = __doc__ - - @classmethod - def INPUT_TYPES(cls): - import folder_paths - return { - "required": { - "unet_name": (folder_paths.get_filename_list("diffusion_models"), - {"tooltip": "Original bf16/fp16 checkpoint - NOT a pre-quantized one."}), - "quantize_scope": (QUANTIZE_SCOPES, { - "default": "video blocks", - "tooltip": "Which Linears go NVFP4. 'video blocks' keeps the whole audio " - "path + cross-modal bridges at full precision (safest). Widen " - "only if quality holds.", - }), - "keep_first_blocks": ("INT", {"default": 2, "min": 0, "max": 16, - "tooltip": "First N transformer blocks stay full precision."}), - "keep_last_blocks": ("INT", {"default": 2, "min": 0, "max": 16, - "tooltip": "Last N transformer blocks stay full precision."}), - "input_scales": (INPUT_SCALE_MODES, { - "default": "dynamic", - "tooltip": "auto-calibrate: first run records each layer's input range " - "(slightly slower) into a per-checkpoint sidecar; every later " - "load injects static input scales, removing the per-call amax " - "cost (~0.2 ms x every quantized layer x every step). Reload " - "the model after the calibration run to activate.", - }), - } - } - - def load(self, unet_name, quantize_scope, keep_first_blocks, keep_last_blocks, - input_scales="dynamic"): - import comfy.model_management as mm - import comfy.sd - import comfy.utils - import folder_paths - - path = folder_paths.get_full_path_or_raise("diffusion_models", unet_name) - # return_metadata is NOT optional here: LTX-AV variants carry their architecture - # config in the safetensors metadata, and detection without it builds the WRONG - # model (6 vs 9 ada params, connector dims) -> size-mismatch wall on load. - sd, metadata = comfy.utils.load_torch_file(path, return_metadata=True) - - if mm.supports_nvfp4_compute(mm.get_torch_device()): - logging.info("[FunPack NVFP4] FP4 tensor-core compute available - native GEMM path.") - else: - logging.warning( - "[FunPack NVFP4] This GPU has no FP4 compute (needs Blackwell, sm_100+). " - "The model will load and run through dequant emulation: correct output, " - "no speedup. Use this only for A/B quality checks." - ) - - device = mm.get_torch_device() if torch.cuda.is_available() else None - pbar = comfy.utils.ProgressBar(1) - - def _tick(done, total): - if total: - pbar.total = total - pbar.update_absolute(done) - - sd, targets, keep = quantize_state_dict_nvfp4( - sd, scope=quantize_scope, - keep_first_blocks=keep_first_blocks, keep_last_blocks=keep_last_blocks, - device=device, progress_callback=_tick, - ) - if not targets: - raise RuntimeError( - "No layers matched the NVFP4 quantization scope - is this an LTX-AV " - "diffusion model checkpoint? (Expected transformer_blocks.* keys.)" - ) - logging.info( - f"[FunPack NVFP4] quantized {len(targets)} Linear layers " - f"(scope='{quantize_scope}', full-precision blocks: {sorted(keep) or 'none'})." - ) - - calibration = load_calibration(path) if input_scales == "auto-calibrate" else None - if calibration: - injected = inject_input_scales(sd, targets, calibration) - logging.info( - f"[FunPack NVFP4] static input scales ACTIVE on {injected}/{len(targets)} layers " - f"(calibrated {len(calibration)}, margin {CALIBRATION_MARGIN}x) — per-call input " - "amax eliminated. Delete the sidecar under refinements/nvfp4_scales/ to recalibrate." - ) - - model = comfy.sd.load_diffusion_model_state_dict(sd, metadata=metadata) - if model is None: - raise RuntimeError("ComfyUI could not detect the model type after NVFP4 patching.") - - if input_scales == "auto-calibrate" and not calibration: - hooked = install_calibration_capture(model, path) - logging.info( - f"[FunPack NVFP4] CALIBRATING input scales on {hooked} layers this session " - "(slightly slower). Run one full generation, then reload the model (change any " - "widget or restart) — static scales activate on the next load." - ) - - try: - log_nvfp4_diagnostics(model, device) - except Exception as e: # noqa: BLE001 - logging.warning(f"[FunPack NVFP4] diagnostics failed: {e}") - return (model,) diff --git a/samplers.py b/samplers.py index 718608bf..86734f63 100644 --- a/samplers.py +++ b/samplers.py @@ -2281,14 +2281,6 @@ class FunPackLTXAVSceneChainSampler: "default": 0.02, "min": 0.005, "max": 0.1, "step": 0.005, "tooltip": "Per-step correction strength applied to the model's predicted output. Same scale/units as embed_guidance_strength — start there and adjust.", }), - "block_steer": ("BOOLEAN", { - "default": False, - "tooltip": "EXPERIMENTAL: RLHF at the model level. Measures which transformer blocks do the work each run (per-block video-residual activity, ~free), pairs those fingerprints with your ratings (plus per-step value-function credit when output_guidance's VF is trained), and on later runs redistributes emphasis: blocks correlated with your Perfects get boosted, blocks correlated with Awfuls damped (video residual only — audio untouched, hard ±10% cap). Learns automatically while on; exact no-op until 2+ rated runs (or 2+ value-credit runs) for this key. Requires refinement_key_input.", - }), - "block_steer_strength": ("FLOAT", { - "default": 0.05, "min": 0.01, "max": 0.15, "step": 0.005, - "tooltip": "Max per-block residual gain deviation (scores are normalized to ±1, so 0.05 = up to ±5% per block; hard-capped at ±10% regardless). Start low — this reshapes the model's internal emphasis, not the conditioning.", - }), }, "hidden": { "unique_id": "UNIQUE_ID", @@ -3930,7 +3922,6 @@ class FunPackLTXAVSceneChainSampler: alg_blur_guides=False, bounded_attention_enabled=False, output_guidance=False, output_guidance_strength=0.02, - block_steer=False, block_steer_strength=0.05, unique_id=None, prompt=None): if not isinstance(positive, list) or not positive: raise ValueError("positive conditioning must contain at least one scene entry.") @@ -4036,65 +4027,6 @@ class FunPackLTXAVSceneChainSampler: print(f"[FunPackSceneChain] output_guidance: active ({_output_value_fn.n_trained} samples), " f"strength={output_guidance_strength}") - # Block Steering (see block_steer.py): run-level machinery, installed once for the - # whole chain. The recorder + steer hooks are tagged module hooks (the defensive - # strip above cleans any leak); the value-credit wrapper goes in INNERMOST so all - # per-scene wrappers (embed/score/output guidance, temporal) stack on top and see - # it as their old_wrapper — it scores the base model's own prediction. - _bs_recorder = None - _bs_handles = [] - _bs_wrapper_old = None - _bs_wrapper_installed = False - if block_steer: - try: - try: - from . import block_steer as _bsmod - except ImportError: - import block_steer as _bsmod - if not refinement_key_input: - print("[FunPackSceneChain] block_steer: no refinement_key_input wired — " - "nothing to learn into or steer from, skipping.") - else: - # Recover from a credit wrapper leaked by a crashed previous run. - _leaked = model.model_options.get("model_function_wrapper") - if getattr(_leaked, "_funpack_block_steer", False): - _old = getattr(_leaked, "_funpack_old", None) - if _old is not None: - model.model_options["model_function_wrapper"] = _old - else: - del model.model_options["model_function_wrapper"] - print("[FunPackSceneChain] block_steer: stripped a leaked credit wrapper from a previous run") - _bs_blocks = _bsmod._funpack_locate_blocks(model) - if _bs_blocks is None: - print("[FunPackSceneChain] block_steer: couldn't locate transformer blocks — skipping.") - else: - _bs_recorder = _bsmod.BlockActivityRecorder(len(_bs_blocks)) - _bs_handles = _bsmod.install_recorder(model, _bs_recorder) - _bs_scores, _bs_status = _bsmod.block_scores_with_status(refinement_key_input) - if _bs_scores is not None and _bs_scores.numel() != len(_bs_blocks): - _bs_scores, _bs_status = None, ( - f"profile was learned on a {_bs_scores.numel()}-block model, this one has " - f"{len(_bs_blocks)} — relearning") - if _bs_scores is not None: - _gains = _bsmod.gains_from_scores(_bs_scores, block_steer_strength) - _bs_handles += _bsmod.install_steer(model, _gains) - _dev = max(abs(g - 1.0) for g in _gains) if _gains else 0.0 - print(f"[FunPackSceneChain] block_steer: steering active " - f"({len(_bs_blocks)} blocks, max gain deviation ±{_dev:.3f}, {_bs_status})") - else: - print(f"[FunPackSceneChain] block_steer: learning only — {_bs_status}; " - "steering is a no-op this run.") - _bs_vf = _output_value_fn or self._load_output_value_function(refinement_key_input) - if _bs_vf is not None: - _bs_wrapper_old = model.model_options.get("model_function_wrapper") - _w = _bsmod.make_credit_wrapper(_bs_wrapper_old, _bs_recorder, _bs_vf) - _w._funpack_block_steer = True - _w._funpack_old = _bs_wrapper_old - model.model_options["model_function_wrapper"] = _w - _bs_wrapper_installed = True - except Exception as e: - print(f"[FunPackSceneChain] block_steer setup failed (continuing without): {e}") - # Phase timing + quantized-path census: "the matmuls got faster" and "the video # arrives sooner" are different claims — the report separates sampling, decode, # and everything else so an optimization's real ceiling is visible. The census @@ -4450,29 +4382,6 @@ class FunPackLTXAVSceneChainSampler: except Exception as e: print(f"[FunPackLTXAVSceneChainSampler] Failed to write sampler context: {e}") - # Block Steering teardown: drop hooks, restore the wrapper chain, and stage this - # run's activity fingerprint for the NEXT rating cycle (same stage-then-pair - # convention as the x0 snapshot below). - if _bs_recorder is not None: - try: - try: - from . import block_steer as _bsmod - except ImportError: - import block_steer as _bsmod - _bsmod.remove_handles(_bs_handles) - if _bs_wrapper_installed: - if _bs_wrapper_old is not None: - model.model_options["model_function_wrapper"] = _bs_wrapper_old - elif "model_function_wrapper" in model.model_options: - del model.model_options["model_function_wrapper"] - if _bsmod.save_run_snapshot(refinement_key_input, _bs_recorder): - _cs = _bs_recorder.credit_steps - print(f"[FunPackSceneChain] block_steer: staged activity fingerprint " - f"({int(_bs_recorder.counts.max())} calls" - + (f", value credit from {_cs} steps" if _cs else "") + ")") - except Exception as e: - print(f"[FunPackSceneChain] block_steer teardown failed: {e}") - # Snapshot the final output for output_guidance's NEXT training cycle (the rating that # scores THIS run pairs with this snapshot, not with denoised — see # _save_output_value_snapshot). Independent of the output_guidance toggle itself: a diff --git a/tests/test_block_steer.py b/tests/test_block_steer.py deleted file mode 100644 index a0aebff2..00000000 --- a/tests/test_block_steer.py +++ /dev/null @@ -1,239 +0,0 @@ -"""Block Steering: recorder/steer hooks on a toy AV model + profile/score/gain math. - -The toy model mirrors the LTX-AV calling convention the hooks rely on: blocks held in -`model.model.diffusion_model.transformer_blocks` (a ModuleList of >= 28 entries, which -is what _funpack_locate_blocks requires), each called as block((vx, ax), ...) and -returning (vx_out, ax_out). Persistence tests point the refinement store at a tmp dir. -""" -import sys -import types -from pathlib import Path - -import pytest -import torch - -sys.path.insert(0, str(Path(__file__).resolve().parents[1])) - -import block_steer as bs # noqa: E402 - -N_BLOCKS = 28 # _funpack_locate_blocks' minimum believable transformer depth -DIM = 8 - - -class _ToyAVBlock(torch.nn.Module): - def __init__(self, v_gain, a_gain=0.5): - super().__init__() - self.v_gain = v_gain - self.a_gain = a_gain - - def forward(self, x, **kwargs): - vx, ax = x - return vx + self.v_gain * torch.ones_like(vx), ax + self.a_gain * torch.ones_like(ax) - - -def _toy_model(v_gains): - diff = torch.nn.Module() - diff.transformer_blocks = torch.nn.ModuleList([_ToyAVBlock(g) for g in v_gains]) - inner = types.SimpleNamespace(diffusion_model=diff) - return types.SimpleNamespace(model=inner, model_options={}) - - -def _run(model, vx=None, ax=None): - vx = vx if vx is not None else torch.ones(1, 4, DIM) - ax = ax if ax is not None else torch.ones(1, 2, DIM) - for blk in model.model.diffusion_model.transformer_blocks: - vx, ax = blk((vx, ax)) - return vx, ax - - -@pytest.fixture -def store(tmp_path, monkeypatch): - """Redirect refinement sidecars into a tmp dir (no conditioning.py import needed).""" - def _path(key, name, prefix="refine_v2", extension="json"): - return str(tmp_path / f"{key}.{name}.{extension}") - monkeypatch.setattr(bs, "_state_path", lambda key, name: _path(key, name, extension="pt")) - return tmp_path - - -def test_recorder_ranks_blocks_by_video_activity(): - gains = [0.01] * N_BLOCKS - gains[5], gains[20] = 2.0, 1.0 # blocks 5 and 20 move the video stream hardest - model = _toy_model(gains) - rec = bs.BlockActivityRecorder(N_BLOCKS) - handles = bs.install_recorder(model, rec) - assert len(handles) == N_BLOCKS - _run(model) - _run(model) - bs.remove_handles(handles) - fp = rec.fingerprint() - assert fp is not None and int(rec.counts[0]) == 2 - assert int(fp.argmax()) == 5 - assert fp[20] > fp[3] - # hooks removed - a third run must not record - _run(model) - assert int(rec.counts[0]) == 2 - - -def test_hooks_carry_leak_strip_tag(): - """Recorder AND steer hooks must be tagged so the sampler's defensive - strip_funpack_block_hooks cleans them if a run crashes mid-sample.""" - model = _toy_model([1.0] * N_BLOCKS) - rec = bs.BlockActivityRecorder(N_BLOCKS) - h1 = bs.install_recorder(model, rec) - h2 = bs.install_steer(model, [1.05] * N_BLOCKS) - try: - for blk in model.model.diffusion_model.transformer_blocks: - fns = list(blk._forward_hooks.values()) - assert fns and all(getattr(f, bs._FUNPACK_HOOK_TAG, False) for f in fns) - finally: - bs.remove_handles(h1 + h2) - - -def test_steer_scales_video_residual_only(): - model = _toy_model([1.0] * N_BLOCKS) - v_ref, a_ref = _run(model) - gains = [1.0] * N_BLOCKS - gains[3] = 1.10 # boost block 3's video contribution by 10% - gains[7] = 0.90 # damp block 7's - handles = bs.install_steer(model, gains) - assert len(handles) == 2 # gain==1.0 blocks get NO hook (native path) - v_out, a_out = _run(model) - bs.remove_handles(handles) - # video: each block adds +1; two blocks changed by +/-0.1 -> net exactly equal - assert torch.allclose(v_out, v_ref, atol=1e-5) # +0.1 and -0.1 cancel - gains[7] = 1.0 - handles = bs.install_steer(model, gains) - v_out2, _ = _run(model) - bs.remove_handles(handles) - assert torch.allclose(v_out2 - v_ref, torch.full_like(v_out2, 0.10), atol=1e-5) - # audio stream byte-identical under steering - assert torch.equal(a_out, a_ref) - - -def test_credit_wrapper_scores_and_credits(): - model = _toy_model([1.0] * N_BLOCKS) - rec = bs.BlockActivityRecorder(N_BLOCKS) - handles = bs.install_recorder(model, rec) - - class _VF(torch.nn.Module): - def __init__(self): - super().__init__() - self.scores = iter([0.1, 0.5, 0.2]) - def compress(self, x): - return x.mean().unsqueeze(0) - def forward(self, c): - return torch.tensor([next(self.scores)]) - - def apply_fn(x, t, **c): - _run(model) # blocks fire -> recorder.cur populated - return x * 0.9 - - wrapper = bs.make_credit_wrapper(None, rec, _VF()) - args = {"input": torch.ones(1, 4, DIM), "timestep": torch.tensor([0.9]), "c": {}} - for _ in range(3): - wrapper(apply_fn, args) - bs.remove_handles(handles) - assert rec.credit_steps == 2 # first call only sets the baseline - cv = rec.credit_vector() - assert cv is not None and cv.shape == (N_BLOCKS,) - # In the toy model vx grows by +1 per block, so relative deltas DECREASE with depth: - # early blocks are the most active. Net score delta across the run is positive - # (0.1 -> 0.5 -> 0.2), so credit must follow the centered activity gradient - - # positive for the most active (early) blocks, negative for the least active (late). - assert float(cv[0]) > 0 > float(cv[-1]) - assert int(cv.argmax()) == 0 - - -def _rec_with(activity): - r = bs.BlockActivityRecorder(activity.numel()) - r.sums = activity.clone().float(); r.counts = torch.ones(activity.numel()) - return r - - -def test_snapshot_profile_scores_gains_roundtrip(store): - key = "toykey" - base = torch.full((N_BLOCKS,), 0.5) - hot = base.clone(); hot[10] = 1.5 # block 10 fires hard on the GOOD run - cold = base.clone(); cold[2] = 1.5 # block 2 fires hard on the BAD run - - assert bs.save_run_snapshot(key, _rec_with(hot)) - assert bs.update_profile_with_rating(key, 1.0) == 1 # Perfect - assert bs.update_profile_with_rating(key, 1.0) is None # snapshot consumed - no double pair - scores, status = bs.block_scores_with_status(key) - assert scores is None and "have 1" in status # 1 rated run < MIN_RATED_RUNS - - assert bs.save_run_snapshot(key, _rec_with(cold)) - assert bs.update_profile_with_rating(key, -0.9) == 2 # Awful - scores, status = bs.block_scores_with_status(key) - assert scores is None and "have 2" in status # correlation needs 3+ - - assert bs.save_run_snapshot(key, _rec_with(base)) - assert bs.update_profile_with_rating(key, 0.2) == 3 # neutral third run - scores, status = bs.block_scores_with_status(key) - assert status.startswith("ready") - assert "corr" in status and "confidence" in status - assert scores is not None and scores.shape == (N_BLOCKS,) - assert int(scores.argmax()) == 10 and int(scores.argmin()) == 2 - assert abs(float(scores.mean())) < 1e-5 # zero-mean: redistributes, no drift - # confidence ramp: 3 runs -> young profile steers gently, not at unit strength - conf = (3 - 2) / (bs.CONFIDENCE_FULL_AT - 2) - assert float(scores.abs().max()) == pytest.approx(conf, abs=1e-4) - - gains = bs.gains_from_scores(scores, 0.05) - # gains follow the (asymmetric, zero-mean) scores exactly; the largest deviation - # on either side is bounded by strength * confidence - assert max(gains) == pytest.approx(1.0 + 0.05 * float(scores.max()), abs=1e-4) - assert min(gains) == pytest.approx(1.0 + 0.05 * float(scores.min()), abs=1e-4) - assert max(abs(g - 1.0) for g in gains) == pytest.approx(0.05 * conf, abs=1e-4) - # hard cap regardless of strength widget - gains_hi = bs.gains_from_scores(scores, 50.0) - assert max(gains_hi) <= 1.0 + bs.MAX_GAIN_DELTA + 1e-6 - assert min(gains_hi) >= 1.0 - bs.MAX_GAIN_DELTA - 1e-6 - - -def test_same_pole_ratings_need_contrast_not_more_runs(store): - """The v1 regression: MANY rated runs, all rated alike -> the profile must say WHY - it can't steer (no reward contrast), not pretend it needs more runs.""" - key = "samepole" - for _ in range(16): - assert bs.save_run_snapshot(key, _rec_with(torch.rand(N_BLOCKS))) - assert bs.update_profile_with_rating(key, 1.0) is not None - scores, status = bs.block_scores_with_status(key) - assert scores is None - assert "rated alike" in status and "16" in status - # one contrasting rating unlocks attribution - assert bs.save_run_snapshot(key, _rec_with(torch.rand(N_BLOCKS))) - bs.update_profile_with_rating(key, -0.9) - scores, status = bs.block_scores_with_status(key) - assert scores is not None and status.startswith("ready") - - -def test_every_reward_value_contributes(store): - """No dead band: mild rewards like Missing action (0.05) / Wrong details (0.20) - count with their actual value instead of being dropped.""" - key = "mild" - base = torch.full((N_BLOCKS,), 0.5) - hi = base.clone(); hi[7] = 1.2 - lo = base.clone(); lo[21] = 1.2 - bs.save_run_snapshot(key, _rec_with(hi)); bs.update_profile_with_rating(key, 0.35) - bs.save_run_snapshot(key, _rec_with(lo)); bs.update_profile_with_rating(key, 0.05) - bs.save_run_snapshot(key, _rec_with(base)); bs.update_profile_with_rating(key, 0.20) - scores, status = bs.block_scores_with_status(key) - assert status.startswith("ready") - assert int(scores.argmax()) == 7 and int(scores.argmin()) == 21 - - -def test_depth_mismatch_and_v1_profile_reset(store): - key = "depth" - bs.save_run_snapshot(key, _rec_with(torch.rand(N_BLOCKS))) - bs.update_profile_with_rating(key, 1.0) - r2 = torch.rand(N_BLOCKS + 4) # different model depth - bs.save_run_snapshot(key, _rec_with(r2)) - assert bs.update_profile_with_rating(key, 1.0) == 1 # fresh profile, not a mixed one - prof = torch.load(bs._state_path(key, "block_profile"), weights_only=False) - assert prof["n_blocks"] == N_BLOCKS + 4 - # a v1 pole-EMA profile on disk (liked_n, no history) is discarded, not mixed in - torch.save({"n_blocks": N_BLOCKS, "liked": torch.rand(N_BLOCKS), "liked_n": 16}, - bs._state_path(key, "block_profile")) - bs.save_run_snapshot(key, _rec_with(torch.rand(N_BLOCKS))) - assert bs.update_profile_with_rating(key, 1.0) == 1 diff --git a/tests/test_nvfp4_patch.py b/tests/test_nvfp4_patch.py deleted file mode 100644 index c0aa0f69..00000000 --- a/tests/test_nvfp4_patch.py +++ /dev/null @@ -1,241 +0,0 @@ -"""NVFP4 on-the-fly patching: transform format + scope selection + core-load roundtrip. - -Uses the REAL comfy.quant_ops / comfy.ops from the local ComfyUI install (comfy_kitchen's -eager backend quantizes and matmuls NVFP4 on CPU), so these tests prove byte-for-byte -compatibility with the format `comfy.ops._load_quantized_module` consumes - not just our -own invariants. Skipped wholesale if ComfyUI or comfy_kitchen is unavailable. -""" -import os -import sys -from pathlib import Path - -import pytest -import torch - -sys.path.insert(0, str(Path(__file__).resolve().parents[1])) - -_COMFY_ROOT = os.environ.get("COMFYUI_ROOT", str(Path.home() / "Documents" / "ComfyUI")) -if _COMFY_ROOT not in sys.path and os.path.isdir(_COMFY_ROOT): - sys.path.insert(0, _COMFY_ROOT) - -quant_ops = pytest.importorskip("comfy.quant_ops") -if not quant_ops._CK_AVAILABLE: - pytest.skip("comfy_kitchen unavailable", allow_module_level=True) - -import comfy.ops # noqa: E402 - -from nvfp4_patch import ( # noqa: E402 - QUANTIZE_SCOPES, - _selected, - quantize_state_dict_nvfp4, -) - -D = 512 # >= _MIN_DIM and a multiple of 16 - - -def _w(out_f=D, in_f=D): - return torch.randn(out_f, in_f, dtype=torch.bfloat16) * 0.02 - - -def _fake_ltxav_sd(n_blocks=4, prefix=""): - """Minimal LTX-AV-shaped state dict: video attn/ff, audio branch, cross-modal - bridges, plus non-block layers that must never be touched by block scopes.""" - sd = {} - for i in range(n_blocks): - b = f"{prefix}transformer_blocks.{i}." - for mod in ("attn1", "attn2", "audio_attn1", "audio_to_video_attn", "video_to_audio_attn"): - for lin in ("to_q", "to_k", "to_v", "to_out.0"): - sd[f"{b}{mod}.{lin}.weight"] = _w() - sd[f"{b}ff.net.0.proj.weight"] = _w(2 * D, D) - sd[f"{b}ff.net.2.weight"] = _w(D, 2 * D) - sd[f"{b}audio_ff.net.0.proj.weight"] = _w(2 * D, D) - sd[f"{b}scale_shift_table"] = torch.randn(6, D, dtype=torch.bfloat16) # 2D but not .weight - sd[f"{b}attn1.norm_q.weight"] = torch.randn(D, dtype=torch.bfloat16) # 1D - sd[f"{prefix}patchify_proj.weight"] = _w() - sd[f"{prefix}proj_out.weight"] = _w() - sd[f"{prefix}adaln_single.linear.weight"] = _w(6 * D, D) - return sd - - -def _quant_layers(out_sd): - return {k[: -len(".comfy_quant")] for k in out_sd if k.endswith(".comfy_quant")} - - -def test_video_scope_protects_audio_cross_edges_and_nonblock(): - sd = _fake_ltxav_sd() - out, targets, keep = quantize_state_dict_nvfp4(sd, "video blocks", 1, 1) - q = _quant_layers(out) - assert keep == {0, 3} - # middle-block video linears quantized - assert "transformer_blocks.1.attn1.to_q" in q - assert "transformer_blocks.2.ff.net.2" in q - # first/last blocks untouched - assert not any(l.startswith("transformer_blocks.0.") or l.startswith("transformer_blocks.3.") for l in q) - # audio branch, cross-modal bridges, non-block layers untouched - assert not any("audio" in l for l in q) - assert not any(l.split(".")[-2:] == ["patchify_proj"] or "proj_out" in l or "adaln" in l for l in q) - # untouched keys pass through identically - assert torch.equal(out["transformer_blocks.0.attn1.to_q.weight"], - sd["transformer_blocks.0.attn1.to_q.weight"]) - assert len(targets) == len(q) - - -def test_scopes_widen_monotonically(): - sd = _fake_ltxav_sd() - counts = [] - for scope in QUANTIZE_SCOPES: - _, targets, _ = quantize_state_dict_nvfp4(sd, scope, 1, 1) - counts.append(len(targets)) - assert counts == sorted(counts) and counts[0] < counts[-1] - # cross-modal scope adds exactly the two bridges of the two middle blocks (4 linears each) - assert counts[1] - counts[0] == 2 * 2 * 4 - - -def test_all_2d_scope_still_skips_sensitive_generic_layers(): - sd = _fake_ltxav_sd() - out, _, _ = quantize_state_dict_nvfp4(sd, "all 2D layers", 0, 0) - q = _quant_layers(out) - for name in ("proj_out", "adaln_single.linear", "patchify_proj"): - assert not any(name in l for l in q) - assert any("audio_attn1" in l for l in q) # audio IS included in the widest scopes - - -def test_serialized_format_matches_core_loader_contract(): - sd = _fake_ltxav_sd(n_blocks=3) - out, _, _ = quantize_state_dict_nvfp4(sd, "video blocks", 1, 1) - layer = "transformer_blocks.1.attn1.to_q" - assert out[f"{layer}.weight"].dtype == torch.uint8 # packed fp4 pairs - assert out[f"{layer}.weight"].shape[1] == D // 2 - assert out[f"{layer}.weight_scale"].dtype == torch.float8_e4m3fn - assert out[f"{layer}.weight_scale_2"].dtype == torch.float32 - import json - conf = json.loads(bytes(out[f"{layer}.comfy_quant"].tolist())) - assert conf == {"format": "nvfp4"} - - -def test_prefix_and_refusals(): - sd = _fake_ltxav_sd(prefix="model.diffusion_model.") - out, targets, _ = quantize_state_dict_nvfp4(sd, "video blocks", 1, 1) - assert any(k.startswith("model.diffusion_model.") and k.endswith(".comfy_quant") for k in out) - assert targets - # already-quantized checkpoints are refused - with pytest.raises(ValueError): - quantize_state_dict_nvfp4(out, "video blocks", 1, 1) - with pytest.raises(ValueError): - quantize_state_dict_nvfp4({"scaled_fp8": torch.zeros(2)}, "video blocks", 1, 1) - - -def test_selector_never_matches_bias_or_small(): - sd = {"transformer_blocks.1.attn1.to_q.bias": torch.randn(D, dtype=torch.bfloat16), - "transformer_blocks.1.attn1.small.weight": _w(64, 64)} - _, targets, _ = quantize_state_dict_nvfp4(sd, "all blocks", 0, 0) - assert targets == [] - assert _selected("transformer_blocks.1.attn1.to_q", "video blocks", set()) - assert not _selected("transformer_blocks.1.audio_attn1.to_q", "video blocks", set()) - - -@pytest.mark.parametrize("disabled", [[], ["nvfp4"]]) -def test_core_mixed_precision_linear_loads_and_forwards(disabled): - """The emitted keys load through comfy.ops.mixed_precision_ops Linear and produce - output close to the bf16 reference - both on the eager NVFP4 matmul path and on - the dequant-emulation path used by GPUs without FP4 compute.""" - torch.manual_seed(7) - w = _w() - sd = {"transformer_blocks.1.attn1.to_q.weight": w, - "transformer_blocks.1.attn1.to_q.bias": torch.zeros(D, dtype=torch.bfloat16)} - out, targets, _ = quantize_state_dict_nvfp4(sd, "video blocks", 0, 0) - assert targets == ["transformer_blocks.1.attn1.to_q.weight"] - - ops = comfy.ops.mixed_precision_ops({}, torch.bfloat16, disabled=disabled) - lin = ops.Linear(D, D, bias=True, device="cpu", dtype=torch.bfloat16) - layer_sd = {k.replace("transformer_blocks.1.attn1.to_q.", ""): v for k, v in out.items()} - missing, unexpected = lin.load_state_dict(layer_sd, strict=False) - assert not missing and not unexpected - assert isinstance(lin.weight, quant_ops.QuantizedTensor) - assert lin.quant_format == "nvfp4" - - x = torch.randn(2, 8, D, dtype=torch.bfloat16) * 0.5 - with torch.no_grad(): - y = lin.forward(x) - y_ref = torch.nn.functional.linear(x, w) - assert y.shape == y_ref.shape - rel = (y.float() - y_ref.float()).abs().mean() / (y_ref.float().abs().mean() + 1e-8) - assert rel < 0.25, f"relative error too high: {rel:.3f} (disabled={disabled})" - - -def test_diagnostics_reports_layer_census(capsys): - """log_nvfp4_diagnostics counts nvfp4 layers + emulation flags and never raises - (GEMM probe skipped without CUDA).""" - import types - from nvfp4_patch import log_nvfp4_diagnostics - - class _Q: - quant_format = "nvfp4" - _full_precision_mm = False - - class _QEmul(_Q): - _full_precision_mm = True - - class _Diff(torch.nn.Module): - def modules(self): - return [self, _Q(), _Q(), _QEmul()] - - model = types.SimpleNamespace(model=types.SimpleNamespace(diffusion_model=_Diff())) - out = log_nvfp4_diagnostics(model, None) - assert out == {"nvfp4_layers": 3, "emulated_layers": 1} - text = capsys.readouterr().out - assert "nvfp4 layers in loaded model: 3" in text - assert "EMULATION" in text - assert "GEMM probe skipped" in text or "GEMM probe" in text - - -def test_input_scale_injection_and_core_load(tmp_path, monkeypatch): - """Calibration sidecar -> inject_input_scales matches by suffix across checkpoint - prefixes, and core's MixedPrecisionOps Linear registers the injected input_scale - and still forwards correctly.""" - from nvfp4_patch import inject_input_scales, _scales_path - import nvfp4_patch as np_mod - - ckpt = str(tmp_path / "model.safetensors") - sidecar = _scales_path(ckpt) - assert sidecar.endswith("model.safetensors.json") - - cal = {"transformer_blocks.1.attn1.to_q": 12.5} - w = _w() - sd = {"model.diffusion_model.transformer_blocks.1.attn1.to_q.weight": w, - "model.diffusion_model.transformer_blocks.1.attn1.to_q.bias": torch.zeros(D, dtype=torch.bfloat16)} - out, targets, _ = quantize_state_dict_nvfp4(sd, "video blocks", 0, 0) - n = inject_input_scales(out, targets, cal) - assert n == 1 - key = "model.diffusion_model.transformer_blocks.1.attn1.to_q.input_scale" - assert key in out and out[key].dtype == torch.float32 - assert float(out[key]) == pytest.approx(12.5 * np_mod.CALIBRATION_MARGIN / (448.0 * 6.0)) - - # core load consumes it: input_scale becomes a registered parameter - import comfy.ops - ops = comfy.ops.mixed_precision_ops({}, torch.bfloat16, disabled=[]) - lin = ops.Linear(D, D, bias=True, device="cpu", dtype=torch.bfloat16) - layer_sd = {k.replace("model.diffusion_model.transformer_blocks.1.attn1.to_q.", ""): v - for k, v in out.items()} - missing, unexpected = lin.load_state_dict(layer_sd, strict=False) - assert not missing and not unexpected - assert isinstance(getattr(lin, "input_scale", None), torch.nn.Parameter) - x = torch.randn(2, 4, D, dtype=torch.bfloat16) * 0.5 - with torch.no_grad(): - y = lin.forward(x) - y_ref = torch.nn.functional.linear(x, w) - rel = (y.float() - y_ref.float()).abs().mean() / (y_ref.float().abs().mean() + 1e-8) - assert y.shape == y_ref.shape and rel < 0.3 - - # no calibration entry for a layer -> nothing injected for it - sd2 = {"transformer_blocks.2.attn1.to_q.weight": _w()} - out2, targets2, _ = quantize_state_dict_nvfp4(sd2, "video blocks", 0, 0) - assert inject_input_scales(out2, targets2, cal) == 0 - - -def test_load_calibration_reads_sidecar(tmp_path, monkeypatch): - import nvfp4_patch as np_mod - monkeypatch.setattr(np_mod, "_scales_path", lambda p: str(tmp_path / "s.json")) - assert np_mod.load_calibration("whatever") is None - (tmp_path / "s.json").write_text('{"amax": {"a.b": 3.0}, "margin": 1.25}') - assert np_mod.load_calibration("whatever") == {"a.b": 3.0}