Add gearcut_infer.py

gearcut_infer.py

@@ -0,0 +1,82 @@
+"""
+GearCut — inference bridge.
+Natural-language instruction (+ current timeline) -> editing operations (JSON).
+
+Loads the trained gc_editor checkpoint and the GearCut tokenizer. If the files
+aren't found locally, they are downloaded from the AMFORGE HuggingFace repos
+(needs HF_TOKEN in the environment for private repos).
+"""
+import os, json
+import torch
+import sentencepiece as spm
+from gearcut_model import SparseMind, Config
+from gearcut_ground import ground_operations
+
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+class GearCutModel:
+    def __init__(self,
+                 checkpoint="gc_editor.pt",
+                 tokenizer="gearcut_tok.model",
+                 hf_repo_model="AMFORGE/gc_editor",
+                 hf_repo_tok="AMFORGE/gearcut_tok",
+                 hf_token=None):
+        tok_path, ckpt_path = tokenizer, checkpoint
+        if not os.path.exists(tok_path) or not os.path.exists(ckpt_path):
+            from huggingface_hub import hf_hub_download
+            token = hf_token or os.environ.get("HF_TOKEN")
+            if not os.path.exists(tok_path):
+                tok_path = hf_hub_download(hf_repo_tok, "gearcut_tok.model", token=token)
+            if not os.path.exists(ckpt_path):
+                ckpt_path = hf_hub_download(hf_repo_model, "gc_editor.pt", token=token)
+
+        self.sp = spm.SentencePieceProcessor(); self.sp.Load(tok_path)
+
+        ckpt = torch.load(ckpt_path, map_location=DEVICE, weights_only=False)
+        fields = set(Config.__dataclass_fields__)
+        self.cfg = Config(**{k: v for k, v in ckpt["config"].items() if k in fields})
+        self.model = SparseMind(self.cfg).to(DEVICE)
+        self.model.load_state_dict(ckpt["model"])
+        self.model.eval()
+        print(f"gc_editor loaded ({self.model.n_params:,} params, vocab={self.cfg.vocab_size})")
+
+    @staticmethod
+    def context(clips):
+        # clips: list of (id, name, in, out) -> "clips: c1=name(0.0-8.0) c2=..."
+        return "clips: " + " ".join(
+            f"{cid}={name}({i:.1f}-{o:.1f})" for cid, name, i, o in clips)
+
+    @torch.no_grad()
+    def operations(self, instruction, clips, max_new=160):
+        """Return (ops, repairs, raw). ops is grounded against the instruction;
+        repairs lists what the safety layer changed/rejected; ops is None if JSON failed."""
+        prompt = f"{self.context(clips)} | {instruction} =>"
+        ids = torch.tensor([self.sp.EncodeAsIds(prompt)], device=DEVICE)
+        out = self.model.generate(ids, max_new=max_new, temp=1.0, top_k=1)
+        gen_ids = out[0, ids.shape[1]:].tolist()
+        eos = self.cfg.eos_id
+        if eos in gen_ids:
+            gen_ids = gen_ids[:gen_ids.index(eos)]
+        text = self.sp.DecodeIds(gen_ids).strip()
+        if "[" in text and "]" in text:
+            text = text[text.index("["): text.rindex("]") + 1]
+        try:
+            raw_ops = json.loads(text)
+        except Exception:
+            return None, [], text
+        clip_ids = [c[0] for c in clips]
+        grounded, repairs = ground_operations(raw_ops, instruction, clip_ids)
+        return grounded, repairs, text
+
+
+if __name__ == "__main__":
+    m = GearCutModel()
+    demo_clips = [("c1", "intro.mp4", 0.0, 8.0)]
+    for instr in ["remove the first 3 seconds of the clip",
+                  "keep only 2 to 6 of clip 1",
+                  "export as final.mp4 in tiktok_vertical"]:
+        ops, repairs, raw = m.operations(instr, demo_clips)
+        print(f"\n> {instr}\n  {ops if ops is not None else '[parse failed] ' + raw}")
+        for r in repairs:
+            print("   ~", r)