Update embed_lwm.py

embed_lwm.py

@@ -1,48 +1,114 @@
 import os
 import sys
-from typing import List, Tuple, Optional
+from typing import List, Optional, Tuple
 
 import torch
-from huggingface_hub import snapshot_download
 
+def _log(msg: str):
+    print(msg, flush=True)
 
-_LWM_MODEL = None
-_LWM_DIR = None
+def _candidate_repo_dirs():
+    return [
+        os.getenv("LWM_REPO_DIR", "").strip(),
+        "./LWM-v1.1",
+        "/home/user/app/LWM-v1.1",
+    ]
 
+def _ensure_repo_on_path() -> Optional[str]:
+    for d in _candidate_repo_dirs():
+        if d and os.path.isdir(d):
+            if d not in sys.path:
+                sys.path.insert(0, d)
+            return d
+    return None
+
+def _ensure_pretrained_model_shim(repo_dir: str) -> None:
+    """
+    Some LWM examples import: `from pretrained_model import lwm`
+    If the repo doesn't ship `pretrained_model.py`, but has `lwm_model.py` with class `LWM`,
+    we create a tiny shim so imports succeed.
+    """
+    shim_path = os.path.join(repo_dir, "pretrained_model.py")
+    lwm_path  = os.path.join(repo_dir, "lwm_model.py")
+    if os.path.isfile(shim_path):
+        return
+    if not os.path.isfile(lwm_path):
+        return  # nothing we can do
+
+    # Create a simple factory around LWM
+    shim_code = """# Auto-generated shim to satisfy `from pretrained_model import lwm`
+import torch
+try:
+    from lwm_model import LWM
+except Exception as e:
+    raise ImportError(f"Shim could not import LWM from lwm_model.py: {e}")
+
+def lwm():
+    # Build a default LWM encoder (adjust constructor args if your repo requires them)
+    return LWM()
+"""
+    try:
+        with open(shim_path, "w", encoding="utf-8") as f:
+            f.write(shim_code)
+        _log(f"[INFO] Created shim: {shim_path}")
+    except Exception as e:
+        _log(f"[WARN] Could not create pretrained_model shim: {e}")
+
+def _maybe_load_weights(model, repo_dir: str):
+    # Try common weight locations
+    candidates = [
+        os.path.join(repo_dir, "models", "model.pth"),
+        os.path.join(repo_dir, "model.pth"),
+    ]
+    for w in candidates:
+        if os.path.isfile(w):
+            try:
+                sd = torch.load(w, map_location="cpu")
+                # Sometimes saved as {'model': state_dict}
+                if isinstance(sd, dict) and "state_dict" in sd:
+                    sd = sd["state_dict"]
+                elif isinstance(sd, dict) and "model" in sd:
+                    sd = sd["model"]
+                model.load_state_dict(sd, strict=False)
+                _log(f"[INFO] Loaded LWM weights from {w}")
+                return
+            except Exception as e:
+                _log(f"[WARN] Failed to load weights from {w}: {e}")
+    _log("[WARN] No weights file found; using randomly-initialized LWM.")
 
 def get_lwm_encoder():
     """
-    Try to download & load wi-lab/lwm-v1.1 and create the encoder.
-    Returns a torch.nn.Module or None on failure.
+    Try to build an LWM encoder using the cloned repo.
+    Returns a torch.nn.Module or None.
     """
-    global _LWM_MODEL, _LWM_DIR
-    if _LWM_MODEL is not None:
-        return _LWM_MODEL
+    repo_dir = _ensure_repo_on_path()
+    if not repo_dir:
+        _log("[WARN] LWM repo not found; set LWM_REPO_DIR or clone to ./LWM-v1.1")
+        return None
+
+    # If the repo's modules expect `pretrained_model`, make sure it exists
+    _ensure_pretrained_model_shim(repo_dir)
+
+    # Try the most common entry point used in examples
     try:
-        _LWM_DIR = snapshot_download(repo_id="wi-lab/lwm-v1.1", local_dir="./LWM-v1.1", local_dir_use_symlinks=False)
-        if _LWM_DIR not in sys.path:
-            sys.path.append(_LWM_DIR)
-        from pretrained_model import lwm  # type: ignore
-        model = lwm()
-        # Try common ckpt filenames
-        cand = None
-        for fn in ["model_checkpoint.pth", "checkpoint.pth", "lwm_v1.1.pth"]:
-            p = os.path.join(_LWM_DIR, fn)
-            if os.path.exists(p):
-                cand = p
-                break
-        if cand:
-            state = torch.load(cand, map_location="cpu")
-            # handle optional "module." prefix
-            if any(k.startswith("module.") for k in state.keys()):
-                model.load_state_dict(state)
+        # Import order: prefer pretrained_model.lwm() if available
+        import pretrained_model  # type: ignore
+        if hasattr(pretrained_model, "lwm"):
+            model = pretrained_model.lwm()
+        else:
+            # Fallback: try lwm_model directly
+            import lwm_model  # type: ignore
+            if hasattr(lwm_model, "LWM"):
+                model = lwm_model.LWM()
+            elif hasattr(lwm_model, "build_model"):
+                model = lwm_model.build_model()
             else:
-                model.load_state_dict({f"module.{k}": v for k, v in state.items()}, strict=False)
+                raise ImportError("No LWM builder found in lwm_model or pretrained_model")
+        _maybe_load_weights(model, repo_dir)
         model.eval()
-        _LWM_MODEL = model
-        return _LWM_MODEL
+        return model
     except Exception as e:
-        print(f"[WARN] Failed to load LWM encoder: {e}", flush=True)
+        _log(f"[WARN] Failed to load LWM encoder: {e}")
         return None
 
 
@@ -52,36 +118,62 @@ def build_lwm_embeddings(
     datasets: List[Tuple[torch.Tensor, Optional[torch.Tensor], str]],
     n_per_dataset: int,
     label_aware: bool
-):
+) -> Tuple[torch.Tensor, Optional[List[torch.Tensor]]]:
     """
-    Minimal: flatten inputs and pass through model if it accepts tensors directly.
-    If the repo expects a tokenizer or different forward, adapt here.
+    Generic embedding builder:
+    - Flattens each complex channel (concat real/imag),
+    - Forwards through the model if it accepts a flat vector,
+    - Pads to a common embedding dim.
+    If forward fails, falls back to the raw flattened vector.
     """
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    model = model.to(device)
-
-    embs = []
+    all_feats = []
     labels_per_ds = [] if label_aware else None
 
-    for ch, y, _ in datasets:
-        N = ch.shape[0]
-        n = min(int(n_per_dataset), int(N))
-        idx = torch.randperm(N)[:n]
-        Xi = ch[idx]
-        Xi = Xi.reshape(n, -1)
-        if torch.is_complex(Xi):
-            Xi = torch.cat([Xi.real, Xi.imag], dim=-1)
-        Xi = Xi.to(torch.float32).to(device)
-        # naive forward: assume model(X)->[n, d]
-        try:
-            Zi = model(Xi)
-        except Exception:
-            # fallback: identity if forward signature differs
-            Zi = Xi
-        Zi = Zi.detach().to("cpu")
-        embs.append(Zi)
-        if label_aware:
-            labels_per_ds.append(y[idx].to(torch.long) if (y is not None and len(y) >= n) else torch.empty((0,), dtype=torch.long))
-
-    embs = torch.stack(embs, dim=0)  # [D, n, d]
-    return embs, labels_per_ds
+    try:
+        device = next(model.parameters()).device
+    except StopIteration:
+        device = torch.device("cpu")
+    model = model.to(device).eval()
+
+    for chs, y, _name in datasets:
+        n = min(int(n_per_dataset), int(chs.shape[0]))
+        idx = torch.randperm(chs.shape[0])[:n]
+        sub = chs[idx]
+        feats_this = []
+
+        for x in sub:
+            if x.ndim > 2:
+                x = x.squeeze(0)
+            vec = x.reshape(-1)
+            if torch.is_complex(vec):
+                vec = torch.cat([vec.real, vec.imag], dim=0)
+            vec = vec.to(torch.float32).unsqueeze(0).to(device)  # [1, d]
+
+            try:
+                out = model(vec)  # adapt here if your model expects another shape
+                out = out.reshape(1, -1).detach().cpu()
+            except Exception:
+                # If the model forward signature mismatches, use the raw vector
+                out = vec.detach().cpu()
+
+            feats_this.append(out)
+
+        embs_this = torch.cat(feats_this, dim=0)  # [n, d’]
+        all_feats.append(embs_this)
+
+        if label_aware and y is not None and y.numel() > 0:
+            labels_per_ds.append(y[idx].clone())
+
+    # Pad to common dim
+    max_d = max(t.shape[1] for t in all_feats)
+    padded = []
+    for t in all_feats:
+        if t.shape[1] < max_d:
+            pad = torch.zeros((t.shape[0], max_d - t.shape[1]), dtype=t.dtype)
+            t = torch.cat([t, pad], dim=1)
+        padded.append(t)
+
+    embs = torch.stack(padded, dim=0)  # [D, n, d]
+    if label_aware:
+        return embs, labels_per_ds if labels_per_ds is not None else []
+    return embs, None