cozyclay 2.1.0 → 3.0.0

package/tools/ardy/gvhmr_fastpath.py

@@ -1,168 +0,0 @@
-"""Quality-preserving preparation around the existing box-side GVHMR runner.
-
-No model, dtype, batch size, flip test, tracker, camera or motion math changes.
-Only two strictly checkpoint-loaded preprocessing models are cached on CPU.
-The upstream/checkpoint files remain untouched; all patches are scoped.
-"""
-import contextlib
-import gc
-import importlib
-import sys
-import time
-import weakref
-from pathlib import Path
-
-import torch
-
-
-@contextlib.contextmanager
-def checked_fast_load():
-    """Skip only random initialization that a complete checkpoint overwrites.
-
-    Audit live skipped tensors against tensors covered by load_state_dict.
-    Constants/zeros and non-persistent buffers are not skipped. An incomplete
-    checkpoint cannot leave uninitialized parameters silently in production.
-    """
-    saved = []
-    skipped = []
-    covered = []
-    loads = []
-    original_load = torch.load
-    original_state_load = torch.nn.Module.load_state_dict
-
-    def patch(owner, name, value):
-        saved.append((owner, name, getattr(owner, name)))
-        setattr(owner, name, value)
-
-    def skip(tensor, *args, **kwargs):
-        skipped.append(weakref.ref(tensor))
-        return tensor
-
-    def mmap_load(*args, **kwargs):
-        if "mmap" in kwargs:
-            return original_load(*args, **kwargs)
-        try:
-            return original_load(*args, **kwargs, mmap=True)
-        except (RuntimeError, ValueError) as exc:
-            # Old torch checkpoint serialization cannot be memory-mapped.
-            if "mmap" not in str(exc).lower():
-                raise
-            return original_load(*args, **kwargs)
-
-    def state_load(module, state, *args, **kwargs):
-        result = original_state_load(module, state, *args, **kwargs)
-        if result.missing_keys or result.unexpected_keys:
-            raise RuntimeError("gvhmr-fast-load-incomplete-checkpoint")
-        loads.append(module)
-        covered.extend(weakref.ref(v) for v in module.state_dict(keep_vars=True).values())
-        return result
-
-    try:
-        for name in ("uniform_", "normal_", "trunc_normal_", "kaiming_uniform_", "kaiming_normal_",
-                     "xavier_uniform_", "xavier_normal_", "orthogonal_", "sparse_"):
-            patch(torch.nn.init, name, skip)
-        # timm's imported trunc_normal_ delegates to this module-global helper.
-        for name in ("timm.layers.weight_init", "timm.models.layers.weight_init"):
-            module = sys.modules.get(name)
-            if module is not None and hasattr(module, "_trunc_normal_"):
-                patch(module, "_trunc_normal_", skip)
-        patch(torch, "load", mmap_load)
-        patch(torch.nn.Module, "load_state_dict", state_load)
-        yield
-        live = [ref() for ref in skipped if ref() is not None]
-        covered_storage = {ref().untyped_storage().data_ptr() for ref in covered if ref() is not None}
-        if not loads or any(t.untyped_storage().data_ptr() not in covered_storage for t in live):
-            raise RuntimeError("gvhmr-fast-load-uncovered-initialization")
-    finally:
-        for owner, name, value in reversed(saved):
-            setattr(owner, name, value)
-
-
-class FastRuntime:
-    def __init__(self, runner, enabled=True):
-        self.runner = runner
-        self.enabled = enabled
-        self.cache = {}
-        self.signature = None
-        self.stats = {}
-
-    def invalidate_changed_weights(self):
-        paths = [Path(self.runner.__file__), Path(".git/HEAD"),
-                 Path("inputs/checkpoints/vitpose/vitpose-h-multi-coco.pth"),
-                 Path("inputs/checkpoints/hmr2/epoch=10-step=25000.ckpt")]
-        signature = tuple((str(p.resolve()), p.stat().st_size, p.stat().st_mtime_ns) for p in paths if p.exists())
-        if self.signature is not None and signature != self.signature:
-            self.cache.clear()
-            gc.collect()
-        self.signature = signature
-
-    def release_gpu(self):
-        for item, attribute in self.cache.values():
-            getattr(item, attribute).cpu()
-        gc.collect()
-        torch.cuda.empty_cache()
-
-    def acquire(self, name, factory, attribute):
-        started = time.perf_counter()
-        if name in self.cache:
-            item = self.cache[name][0]
-            getattr(item, attribute).cuda().eval()
-            self.stats["modelCacheHits"].append(name)
-        else:
-            with checked_fast_load():
-                item = factory()
-            self.cache[name] = (item, attribute)
-        self.stats["modelLoadSeconds"][name] = time.perf_counter() - started
-        return item
-
-    @contextlib.contextmanager
-    def job(self, evidence=None):
-        self.invalidate_changed_weights()
-        self.stats = {"modelCacheHits": [], "modelLoadSeconds": {}, "decodeCropCalls": 0,
-                      "sharedCropHits": 0, "decodeCropSeconds": 0, "mode": "exact-fast" if self.enabled else "reference"}
-        runner = self.runner
-        vp = importlib.import_module("hmr4d.utils.preproc.vitpose")
-        vf = importlib.import_module("hmr4d.utils.preproc.vitfeat_extractor")
-        original_crop, original_vp_crop = vf.get_batch, vp.get_batch
-        original_preprocess, original_release = runner.preprocess, runner._release_gpu
-        original_vit, original_hmr = runner.VitPoseExtractor, runner.Extractor
-        crops = []
-
-        def shared_crop(path, boxes, *args, **kwargs):
-            for prior_path, prior_boxes, prior_args, prior_kwargs, output in crops:
-                if str(path) == prior_path and args == prior_args and kwargs == prior_kwargs and torch.equal(boxes, prior_boxes):
-                    self.stats["sharedCropHits"] += 1
-                    # Both consumers read their input. Cloning prevents a future
-                    # in-place consumer from contaminating the other model.
-                    return tuple(v.clone() for v in output)
-            started = time.perf_counter()
-            output = original_crop(path, boxes, *args, **kwargs)
-            self.stats["decodeCropCalls"] += 1
-            self.stats["decodeCropSeconds"] += time.perf_counter() - started
-            if self.enabled:
-                crops.append((str(path), boxes.clone(), args, dict(kwargs), tuple(v.clone() for v in output)))
-            return output
-
-        def preprocess(*args, **kwargs):
-            started = time.perf_counter()
-            data = original_preprocess(*args, **kwargs)
-            self.stats["preprocessSeconds"] = time.perf_counter() - started
-            if evidence:
-                torch.save(data, Path(evidence) / "preprocess.pt")
-            crops.clear()
-            return data
-
-        try:
-            vf.get_batch = vp.get_batch = shared_crop
-            runner.preprocess = preprocess
-            if self.enabled:
-                runner.VitPoseExtractor = lambda: self.acquire("vitpose", original_vit, "pose")
-                runner.Extractor = lambda: self.acquire("hmr2", original_hmr, "extractor")
-                runner._release_gpu = self.release_gpu
-            yield self.stats
-        finally:
-            runner.preprocess, runner._release_gpu = original_preprocess, original_release
-            runner.VitPoseExtractor, runner.Extractor = original_vit, original_hmr
-            vf.get_batch, vp.get_batch = original_crop, original_vp_crop
-            crops.clear()
-            self.release_gpu()

package/tools/ardy/gvhmr_trajectory.py

@@ -1,240 +0,0 @@
-"""Observation-anchored recovery of delayed vertical descents from GVHMR.
-
-This is NOT a ground-every-frame filter or a general dynamics solver. A
-stationary-camera take must contain a large, confident image-space descent
-and an observed landing plateau. A delayed world endpoint may be recovered
-from that observation when it never settles. Frames, rotations and XZ
-are never changed. Ambiguous events retain the original output.
-"""
-import contextlib
-import json
-import time
-from pathlib import Path
-
-import numpy as np
-
-
-def smooth(values, sigma):
-    radius = max(1, int(np.ceil(3 * sigma)))
-    x = np.arange(-radius, radius + 1)
-    weights = np.exp(-0.5 * (x / max(sigma, 0.25)) ** 2)
-    weights /= weights.sum()
-    return np.convolve(np.pad(values, (radius, radius), mode="edge"), weights, mode="valid")
-
-
-def center_of_mass(joints):
-    """Approximate segment-mass COM; normalized weights, not a pelvis proxy."""
-    segments = [(0, 12, .497), (12, 15, .081),
-                (1, 4, .100), (2, 5, .100), (4, 7, .0465), (5, 8, .0465),
-                (7, 10, .0145), (8, 11, .0145),
-                (16, 18, .028), (17, 19, .028), (18, 20, .016), (19, 21, .016),
-                (20, 20, .006), (21, 21, .006)]
-    return sum((joints[:, a] + joints[:, b]) * (weight * .5)
-               for a, b, weight in segments) / sum(s[2] for s in segments)
-
-
-def recover_vertical(root, joints, camera_joints, keypoints, intrinsics, camera_up, fps):
-    """Pure NumPy kernel. Returns Y corrections and inspectable diagnostics.
-
-    Units are meters and seconds. COCO hips 11/12 provide translation
-    observations; the ballistic proposal uses whole-body COM. FPS is always
-    the actual source FPS, never inferred from number of frames.
-    """
-    root = np.asarray(root, dtype=np.float64)
-    n = len(root)
-    delta = np.zeros(n, dtype=np.float64)
-    report = {"version": 2, "status": "unchanged", "events": [], "rejected": [],
-              "frames": n, "fps": float(fps), "changedFrames": 0}
-    arrays = [root, joints, camera_joints, keypoints, intrinsics, camera_up]
-    if not np.isfinite(fps) or not 5 <= fps <= 240 or n < 12 or any(not np.isfinite(a).all() for a in arrays):
-        report["reason"] = "invalid-input"
-        return delta, report
-    joints, camera_joints = np.asarray(joints), np.asarray(camera_joints)
-    keypoints, intrinsics = np.asarray(keypoints), np.asarray(intrinsics)
-    if root.shape != (n, 3) or joints.shape != (n, 22, 3) or camera_joints.shape != (n, 22, 3) or keypoints.shape != (n, 17, 3):
-        report["reason"] = "unsupported-shape"
-        return delta, report
-    k = intrinsics[0] if intrinsics.ndim == 3 else intrinsics
-    up = np.asarray(camera_up)
-    projected_up = np.array([k[0, 0] * up[0], k[1, 1] * up[1]])
-    norm = np.linalg.norm(projected_up)
-    depth = camera_joints[:, 0, 2]
-    if norm < .45 * min(k[0, 0], k[1, 1]) or np.median(depth) <= 0:
-        report["reason"] = "gravity-axis-not-observable"
-        return delta, report
-    if intrinsics.ndim == 3 and np.max(np.abs(intrinsics - k)) > .01 * np.max(k):
-        report["reason"] = "variable-intrinsics"
-        return delta, report
-    hip_conf = np.minimum(keypoints[:, 11, 2], keypoints[:, 12, 2])
-    hips = keypoints[:, [11, 12], :2].mean(axis=1)
-    # Fixed metric scale over the take avoids bbox/depth breathing as the
-    # person curls up. Only timing is recovered: endpoint world heights are
-    # retained, so this does not pretend to solve monocular absolute scale.
-    scale = np.median(depth) / norm
-    observed = smooth(hips @ (projected_up / norm) * scale, max(.6, fps * .04))
-    height = float(np.median(np.linalg.norm(joints[:, 12] - joints[:, 0], axis=1) +
-                             np.linalg.norm(joints[:, 0] - joints[:, 4], axis=1) +
-                             np.linalg.norm(joints[:, 4] - joints[:, 7], axis=1)) * 1.35)
-    height = max(.8, height)
-    velocity = np.gradient(observed) * fps
-    base = root[:, 1]
-    base_smooth = smooth(base, max(.6, fps * .04))
-    base_velocity = np.gradient(base_smooth) * fps
-    settle = max(4, round(.25 * fps))
-    com_offset = center_of_mass(joints)[:, 1] - base
-    i = 1
-    while i < n - settle:
-        if velocity[i] >= -.65 * height:
-            i += 1
-            continue
-        fast_start = i
-        while i < n - 1 and velocity[i] < -.3 * height:
-            i += 1
-        end_fast = i
-        start = fast_start
-        while start > 0 and velocity[start - 1] < -.1 * height and fast_start - start < .4 * fps:
-            start -= 1
-        # Landing evidence is a stable observed pelvis, not a local minimum
-        # of the already incorrect GVHMR root or a hard-coded pose category.
-        landing = None
-        for t in range(end_fast, min(n - settle, end_fast + round(.6 * fps)) + 1):
-            if np.max(observed[t:t + settle]) - np.min(observed[t:t + settle]) < .025 * height:
-                landing = t
-                break
-        if landing is None:
-            report["rejected"].append({"start": start, "reason": "no-observed-landing"})
-            continue
-        drop = observed[start] - np.median(observed[landing:landing + settle])
-        if drop < .35 * height or not .12 <= (landing - start) / fps <= 2.0:
-            continue
-        if np.quantile(hip_conf[start:landing + settle], .1) < .55:
-            report["rejected"].append({"start": start, "landing": landing, "reason": "uncertain-keypoints"})
-            continue
-        if np.ptp(depth[start:landing + settle]) > .35 * np.median(depth[start:landing + settle]):
-            report["rejected"].append({"start": start, "landing": landing, "reason": "uncertain-depth"})
-            continue
-        anchor = None
-        for t in range(landing, min(n - settle, landing + round(3 * fps)) + 1):
-            if np.max(np.abs(base_velocity[t:t + settle])) < .08 * height:
-                anchor = t + settle - 1
-                break
-        endpoint_source = "settled-world"
-        if anchor is None:
-            # A lagging prediction often keeps descending past the end of a
-            # short clip. Requiring that same prediction to settle prevented
-            # recovery exactly when it was needed. Use a persistent observed
-            # plateau and metric camera scale, NOT the erroneous last height.
-            anchor = n - 1
-            if anchor - landing < round(.4 * fps):
-                report["rejected"].append({"start": start, "landing": landing, "reason": "landing-observation-too-short"})
-                continue
-            endpoint_source = "observed-plateau"
-        # The image must remain on its plateau while the world catches up.
-        # This rejects another jump, stairs and camera motion during the tail.
-        if np.ptp(observed[landing:anchor + 1]) > .07 * height:
-            report["rejected"].append({"start": start, "landing": landing, "reason": "moving-landing"})
-            continue
-        tail_drift = base[landing] - base[anchor]
-        total_drop = base[start] - base[anchor]
-        if tail_drift < .12 * height or total_drop < .35 * height:
-            i = landing + settle
-            continue
-        stop_observed = float(np.median(observed[landing:anchor + 1]))
-        endpoint_y = base[anchor] if endpoint_source == "settled-world" else base[start] - (observed[start] - stop_observed)
-        progress = np.clip((observed[start] - observed[start:anchor + 1]) /
-                           max(observed[start] - stop_observed, 1e-8), 0, 1)
-        target = base[start] + progress * (endpoint_y - base[start])
-        # Keep the original endpoint exactly, with a short smooth residual
-        # bridge inside the observed plateau (no tail offset after the event).
-        end_blend = min(settle, len(target))
-        blend = np.linspace(0, 1, end_blend)
-        blend = blend * blend * (3 - 2 * blend)
-        target[-end_blend:] += (endpoint_y - target[-1]) * blend
-        airborne = landing - start + 1
-        observed_com = target[:airborne] + com_offset[start:landing + 1]
-        # The stable landing window includes a few post-impact frames.
-        # Do not force those frames to keep accelerating under gravity.
-        # Locate an impact endpoint near the plateau, after >=90% descent
-        # (the short image denoising kernel blurs the impact discontinuity).
-        ballistic_error, ballistic, gravity_end = float("inf"), None, None
-        for stop in range(max(4, airborne - round(.15 * fps)), airborne + 1):
-            if progress[stop - 1] < .9:
-                continue
-            times = np.arange(stop) / fps
-            duration = times[-1]
-            initial_velocity = (observed_com[stop - 1] - observed_com[0] + .5 * 9.81 * duration ** 2) / duration
-            proposal = observed_com[0] + initial_velocity * times - .5 * 9.81 * times ** 2
-            error = float(np.sqrt(np.mean((proposal - observed_com[:stop]) ** 2)))
-            if error < ballistic_error:
-                ballistic_error, ballistic, gravity_end = error, proposal, stop
-        gravity_applied = ballistic_error < .025 * height
-        if gravity_applied:
-            target[:gravity_end] = ballistic - com_offset[start:start + gravity_end]
-        change = target - base[start:anchor + 1]
-        if np.max(np.abs(change)) > 2.5 * height:
-            report["rejected"].append({"start": start, "reason": "excessive-correction"})
-            continue
-        delta[start:anchor + 1] = change
-        report["events"].append({"start": start, "landing": landing, "anchor": anchor,
-                                 "endpointSource": endpoint_source,
-                                 "sourceDropM": float(drop), "worldDropM": float(total_drop),
-                                 "tailDriftBeforeM": float(tail_drift),
-                                 "tailDriftAfterM": float(target[airborne - 1] - target[-1]),
-                                 "maxCorrectionM": float(np.max(np.abs(change))),
-                                 "gravityApplied": bool(gravity_applied),
-                                 "gravityFitRmseM": ballistic_error if np.isfinite(ballistic_error) else None,
-                                 "gravityIntervalEnd": start + gravity_end - 1 if gravity_applied else None,
-                                 "mode": "ballistic-com" if gravity_applied else "observed-timing",
-                                 "gravityReason": None if gravity_applied else "ballistic-curve-disagrees-with-video"})
-        i = anchor + 1
-    report["changedFrames"] = int(np.count_nonzero(np.abs(delta) > 1e-6))
-    report["status"] = "corrected" if report["events"] else "unchanged"
-    if not report["events"]:
-        report["reason"] = "no-confident-delayed-descent"
-    return delta, report
-
-
-@contextlib.contextmanager
-def trajectory_job(runner, video, output_root, *, enabled=True):
-    """Scoped post-inference hook; upstream files and original runner stay intact."""
-    stats = {"version": 1, "status": "disabled", "changedFrames": 0}
-    if not enabled:
-        yield stats
-        return
-    import torch
-    from pytorch3d.transforms import axis_angle_to_matrix
-    original = runner.DemoPL.predict
-    # The legacy NPZ exporter rounds FPS to an integer. Keep fractional
-    # source rates here so the physical proposal uses seconds, not that label.
-    fps = float(runner.iio.immeta(str(video), plugin="pyav").get("fps") or runner.video_fps(str(video)))
-
-    def predict(model, data, static_cam=False):
-        result = original(model, data, static_cam=static_cam)
-        began = time.perf_counter()
-        if not static_cam:
-            stats.update(status="unchanged", reason="moving-camera-not-supported")
-            return result
-        with torch.no_grad():
-            endecoder = model.pipeline.endecoder
-            gp, cp = result["smpl_params_global"], result["smpl_params_incam"]
-            w = endecoder.fk_v2(**{k: v[None] for k, v in gp.items()})[0]
-            c = endecoder.fk_v2(**{k: v[None] for k, v in cp.items()})[0]
-            r_c2w = axis_angle_to_matrix(gp["global_orient"][0]) @ axis_angle_to_matrix(cp["global_orient"][0]).T
-            cpu = lambda a: a.detach().cpu().numpy()
-            delta, report = recover_vertical(cpu(gp["transl"]), cpu(w), cpu(c), cpu(data["kp2d"]),
-                                             cpu(data["K_fullimg"]), cpu(r_c2w.T[:, 1]), fps)
-            if report["events"]:
-                corrected = gp["transl"].clone()
-                corrected[:, 1] += torch.as_tensor(delta, device=corrected.device, dtype=corrected.dtype)
-                gp["transl"] = corrected
-                result["net_outputs"]["pred_smpl_params_global"]["transl"] = corrected[None]
-            stats.update(report)
-            stats["seconds"] = time.perf_counter() - began
-            Path(output_root, "trajectory.json").write_text(json.dumps(stats, indent=2))
-        return result
-
-    runner.DemoPL.predict = predict
-    try:
-        yield stats
-    finally:
-        runner.DemoPL.predict = original

package/tools/ardy/merge-text-encoder.py

@@ -1,160 +0,0 @@
-#!/usr/bin/env python3
-"""merge-text-encoder.py - bake the MNTP LoRA into the base weights.
-
-Run under the ARDY text-encoder venv (needs torch, peft, transformers, and
-the `ardy` package importable — i.e. from the ARDY checkout):
-
-  cd ~/ardy && .venv/bin/python merge-text-encoder.py --dest ~/cclay-text-encoders
-
-Why this exists: transformers v5 removed the automatic PEFT-adapter
-resolution that `from_pretrained` did in v4 (`find_adapter_config_file` is
-gone from modeling_utils). ARDY's vendored LLM2Vec loader relies on that
-behavior — it points `from_pretrained` at the MNTP *adapter* repo and lets
-transformers chase `adapter_config.json` to the (gated) base. Under v5 that
-path simply fails with "no model.safetensors found".
-
-So instead of an adapter directory, the local MNTP path becomes a full
-model: this script loads the ungated base (NousResearch mirror, already
-hash-verified by setup-text-encoder.py), applies the MNTP LoRA with peft,
-`merge_and_unload()`s it — the exact math upstream LLM2Vec performs at load
-time — and saves the merged weights to the path ARDY's TEXT_ENCODERS_DIR
-override resolves. Plain full-model loading works identically under
-transformers v4 and v5; the supervised LoRA keeps loading through peft
-directly, which both versions support.
-
-Layout before:  <dest>/McGill-NLP/...-mntp        (downloaded adapter files)
-Layout after:   <dest>/McGill-NLP/...-mntp-src    (the downloads, kept for
-                                                   provenance + re-merges)
-                <dest>/McGill-NLP/...-mntp        (merged full model +
-                                                   tokenizer + MERGE_MANIFEST)
-
-The merged config.json keeps `_name_or_path` =
-"meta-llama/Meta-Llama-3-8B-Instruct": it is a metadata string (no download
-or gate involved) that LLM2Vec's `prepare_for_tokenization` matches to pick
-the Llama-3 chat template — losing it would silently change every prompt
-embedding.
-"""
-
-import argparse
-import hashlib
-import json
-import os
-import shutil
-import sys
-
-BASE = "NousResearch/Meta-Llama-3-8B-Instruct"
-MNTP = "McGill-NLP/LLM2Vec-Meta-Llama-3-8B-Instruct-mntp"
-MNTP_SRC = MNTP + "-src"
-NAME_OR_PATH = "meta-llama/Meta-Llama-3-8B-Instruct"
-TOKENIZER_FILES = ["tokenizer.json", "tokenizer_config.json", "special_tokens_map.json"]
-
-
-def log(msg):
-    print("merge-text-encoder: %s" % msg, flush=True)
-
-
-def fail(msg):
-    print("merge-text-encoder: ERROR: %s" % msg, file=sys.stderr, flush=True)
-    sys.exit(1)
-
-
-def sha256_file(path):
-    h = hashlib.sha256()
-    with open(path, "rb") as f:
-        for block in iter(lambda: f.read(1024 * 1024), b""):
-            h.update(block)
-    return h.hexdigest()
-
-
-def main():
-    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
-    ap.add_argument("--dest", default=os.environ.get(
-        "TEXT_ENCODERS_DIR", os.path.expanduser("~/cclay-text-encoders")))
-    args = ap.parse_args()
-    dest = os.path.abspath(os.path.expanduser(args.dest))
-    base_dir = os.path.join(dest, BASE)
-    mntp_dir = os.path.join(dest, MNTP)
-    src_dir = os.path.join(dest, MNTP_SRC)
-
-    if not os.path.isfile(os.path.join(base_dir, "model.safetensors.index.json")):
-        fail("base model missing under %s — run setup-text-encoder.py first" % base_dir)
-
-    # Migrate a pre-merge tree: the download step used to land the adapter
-    # files directly at the mntp path.
-    if os.path.isdir(mntp_dir) and not os.path.isdir(src_dir) \
-            and os.path.isfile(os.path.join(mntp_dir, "adapter_model.safetensors")) \
-            and not os.path.isfile(os.path.join(mntp_dir, "model.safetensors.index.json")):
-        os.rename(mntp_dir, src_dir)
-        log("moved downloaded adapter layout to %s" % MNTP_SRC)
-
-    if not os.path.isfile(os.path.join(src_dir, "adapter_model.safetensors")):
-        fail("MNTP adapter missing under %s — run setup-text-encoder.py first" % src_dir)
-
-    adapter_sha = sha256_file(os.path.join(src_dir, "adapter_model.safetensors"))
-    manifest_path = os.path.join(mntp_dir, "MERGE_MANIFEST.json")
-    if os.path.isfile(manifest_path):
-        with open(manifest_path, "r", encoding="utf-8") as f:
-            manifest = json.load(f)
-        if manifest.get("mntp_adapter_sha256") == adapter_sha \
-                and os.path.isfile(os.path.join(mntp_dir, "model.safetensors.index.json")):
-            log("merged model already present and matches the adapter hash — nothing to do")
-            return
-
-    import torch  # deferred: keep --help fast and errors readable
-    import peft
-    import transformers
-    from peft import PeftModel
-    from ardy.model.llm2vec.llm2vec import LLM2Vec
-
-    log("loading base (bfloat16, cpu) from %s" % base_dir)
-    model_class = LLM2Vec._get_model_class("LlamaConfig", enable_bidirectional=True)
-    model = model_class.from_pretrained(base_dir, torch_dtype=torch.bfloat16)
-
-    log("applying + merging the MNTP LoRA from %s" % src_dir)
-    model = PeftModel.from_pretrained(model, src_dir)
-    model = model.merge_and_unload()
-
-    tmp_out = mntp_dir + ".merging"
-    if os.path.isdir(tmp_out):
-        shutil.rmtree(tmp_out)
-    log("saving merged model to %s" % mntp_dir)
-    model.save_pretrained(tmp_out, safe_serialization=True)
-    del model
-
-    for name in TOKENIZER_FILES:
-        src = os.path.join(src_dir, name)
-        if os.path.isfile(src):
-            shutil.copy2(src, os.path.join(tmp_out, name))
-
-    # Preserve the chat-template key and the source architecture label.
-    cfg_path = os.path.join(tmp_out, "config.json")
-    with open(cfg_path, "r", encoding="utf-8") as f:
-        cfg = json.load(f)
-    with open(os.path.join(src_dir, "config.json"), "r", encoding="utf-8") as f:
-        src_cfg = json.load(f)
-    cfg["_name_or_path"] = NAME_OR_PATH
-    if "architectures" in src_cfg:
-        cfg["architectures"] = src_cfg["architectures"]
-    with open(cfg_path, "w", encoding="utf-8") as f:
-        json.dump(cfg, f, indent=2)
-        f.write("\n")
-
-    with open(os.path.join(tmp_out, "MERGE_MANIFEST.json"), "w", encoding="utf-8") as f:
-        json.dump({
-            "base": BASE,
-            "mntp_adapter_sha256": adapter_sha,
-            "torch": torch.__version__,
-            "transformers": transformers.__version__,
-            "peft": peft.__version__,
-            "name_or_path": NAME_OR_PATH,
-        }, f, indent=2)
-        f.write("\n")
-
-    if os.path.isdir(mntp_dir):
-        shutil.rmtree(mntp_dir)
-    os.rename(tmp_out, mntp_dir)
-    log("done — %s is a full merged model; the supervised LoRA stays a peft adapter" % MNTP)
-
-
-if __name__ == "__main__":
-    main()