Update embed_lwm.py

embed_lwm.py

@@ -1,69 +1,95 @@
 # embed_lwm.py
 import os
 import sys
-from typing import List, Optional, Tuple
+from typing import List, Tuple, Optional
 
 import torch
+from huggingface_hub import snapshot_download
 
-def _log(msg: str):
-    print(msg, flush=True)
+_LWM_MODEL = None
+_LWM_DIR = None
 
-def _maybe_add_lwm_repo_to_path():
+
+def _add_repo_to_path(path: str):
+    if path and os.path.isdir(path) and path not in sys.path:
+        sys.path.insert(0, path)
+
+
+def _load_state_dict_flex(model: torch.nn.Module, state):
     """
-    Ensure the HF-cloned LWM repo is importable.
-    You can override the location with env var LWM_REPO_DIR.
+    Load a variety of saved formats into `model`:
+      - plain state_dict
+      - {"model": state_dict}
+      - with or without "module." prefixes
     """
-    candidates = [
-        os.getenv("LWM_REPO_DIR", ""),      # user override
-        "./LWM-v1.1",                       # local default
-        "/home/user/app/LWM-v1.1",          # HF Space default path
-    ]
-    for c in candidates:
-        if c and os.path.isdir(c) and c not in sys.path:
-            sys.path.insert(0, c)
+    def _try(sd, strict=False):
+        try:
+            model.load_state_dict(sd, strict=strict)
+            return True
+        except Exception:
+            return False
+
+    # direct state dict?
+    if isinstance(state, dict) and all(isinstance(k, str) for k in state.keys()) and any(
+        torch.is_tensor(v) for v in state.values()
+    ):
+        sd = state
+    elif isinstance(state, dict) and "model" in state and isinstance(state["model"], dict):
+        sd = state["model"]
+    else:
+        raise ValueError("Unrecognized checkpoint format.")
+
+    # Try as-is
+    if _try(sd, strict=False):
+        return
+
+    # Try to add "module." prefix
+    if not any(k.startswith("module.") for k in sd.keys()):
+        sd_mod = {f"module.{k}": v for k, v in sd.items()}
+        if _try(sd_mod, strict=False):
+            return
+
+    # Try to strip "module." prefix
+    sd_strip = {k.replace("module.", "", 1): v for k, v in sd.items()}
+    if _try(sd_strip, strict=False):
+        return
+
+    # last resort strict=False on original again
+    model.load_state_dict(sd, strict=False)
+
 
 def get_lwm_encoder():
     """
-    Try to load the encoder from pretrained_model.py in the HF repo.
-    Returns a torch.nn.Module or None if it can’t be loaded.
+    Download & load wi-lab/lwm-v1.1 and create the encoder from lwm_model.py.
+    Returns a torch.nn.Module or None on failure.
     """
-    _maybe_add_lwm_repo_to_path()
-    try:
-        # HF repo exports a builder called `lwm`
-        from pretrained_model import lwm  # type: ignore
-    except Exception as e:
-        _log(f"[WARN] Failed to import pretrained_model.lwm: {e}")
-        return None
-
+    global _LWM_MODEL, _LWM_DIR
+    if _LWM_MODEL is not None:
+        return _LWM_MODEL
     try:
+        _LWM_DIR = snapshot_download(
+            repo_id="wi-lab/lwm-v1.1",
+            local_dir="./LWM-v1.1",
+            local_dir_use_symlinks=False,
+        )
+        _add_repo_to_path(_LWM_DIR)
+
+        # Import builder from the HF repo (it's named lwm_model.py)
+        from lwm_model import lwm  # type: ignore
         model = lwm()
-    except Exception as e:
-        _log(f"[WARN] pretrained_model.lwm() failed to build model: {e}")
-        return None
 
-    # Load weights if present
-    weights = None
-    for cand in ("models/model.pth", "./LWM-v1.1/models/model.pth"):
-        if os.path.isfile(cand):
-            weights = cand
-            break
+        # Load checkpoint from models/model.pth
+        ckpt_path = os.path.join(_LWM_DIR, "models", "model.pth")
+        if os.path.isfile(ckpt_path):
+            state = torch.load(ckpt_path, map_location="cpu")
+            _load_state_dict_flex(model, state)
 
-    if weights:
-        try:
-            sd = torch.load(weights, map_location="cpu")
-            try:
-                model.load_state_dict(sd)
-            except Exception:
-                # sometimes saved as {"model": state_dict}
-                if isinstance(sd, dict) and "model" in sd:
-                    model.load_state_dict(sd["model"])
-                else:
-                    raise
-        except Exception as e:
-            _log(f"[WARN] Could not load weights from {weights}: {e}")
-
-    model.eval()
-    return model
+        model.eval()
+        _LWM_MODEL = model
+        return _LWM_MODEL
+    except Exception as e:
+        print(f"[WARN] Failed to load LWM encoder: {e}", flush=True)
+        return None
 
 
 @torch.no_grad()
@@ -72,101 +98,91 @@ 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]]]:
+):
     """
-    Build per-dataset embeddings using the LWM encoder.
-
+    Build embeddings with the LWM encoder.
     Strategy:
-    1) If `utils.tokenizer` exists in the repo, try tokenizing each channel sample
-       and pass the tokenized tensor to the model.
-    2) If that fails, try feeding a flattened real-valued vector to the model.
-    3) If the forward still fails, fall back to using the flattened vector as the “embedding”.
+      1) Try repo's tokenizer if available (utils.tokenizer), feed to model.
+      2) Else try feeding flattened real vectors to the model.
+      3) If forward fails, fall back to using flattened vectors as embeddings.
 
     Returns:
       embs: [D, n, d]
-      labels_per_ds (optional)
+      labels_per_ds: Optional[List[Tensor]]
     """
-    _maybe_add_lwm_repo_to_path()
-
-    # Try to import tokenizer if present; fall back to identity
-    def _identity(x): return x
+    # Try optional tokenizer
     try:
         from utils import tokenizer as lwm_tokenizer  # type: ignore
     except Exception:
-        lwm_tokenizer = _identity  # type: ignore
+        lwm_tokenizer = None
 
-    all_feats = []
-    labels_per_ds = [] if label_aware else None
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model = model.to(device).eval()
 
-    try:
-        params = list(model.parameters())
-        device = next(p.device for p in params) if params else torch.device("cpu")
-    except Exception:
-        device = torch.device("cpu")
+    all_embs = []
+    labels_per_ds = [] if label_aware else None
 
-    model = model.to(device)
-    model.eval()
+    for ch, y, _name in datasets:
+        N = int(ch.shape[0])
+        n = min(int(n_per_dataset), N)
+        idx = torch.randperm(N)[:n]
+        Xi = ch[idx]
 
-    for chs, labels, _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 = []
+        feats = []
+        for x in Xi:
+            x2 = x
+            if x2.ndim > 2:
+                x2 = x2.squeeze(0)
 
-        for x in sub:
-            # Ensure 2D (e.g., [N_ant, SC]) if possible
-            x_proc = x
-            if x_proc.ndim > 2:
-                x_proc = x_proc.squeeze(0)
+            # 1) tokenizer path
+            if lwm_tokenizer is not None:
+                try:
+                    tok = lwm_tokenizer(x2)
+                    tok = tok.to(device)
+                    out = model(tok)
+                    out = torch.as_tensor(out).reshape(1, -1).detach().cpu()
+                    feats.append(out)
+                    continue
+                except Exception:
+                    pass
 
-            # First, try tokenizer-based forward
-            did_forward = False
+            # 2) flattened forward path
             try:
-                tok = lwm_tokenizer(x_proc)  # repo-specific; often returns a tensor
-                tok = tok.to(device)
-                y = model(tok)
-                y = torch.as_tensor(y).reshape(1, -1).detach().cpu()
-                feats_this.append(y)
-                did_forward = True
+                vec = x2.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)
+                out = model(vec)
+                out = torch.as_tensor(out).reshape(1, -1).detach().cpu()
+                feats.append(out)
+                continue
             except Exception:
-                # If tokenizer-based call fails, try flat-vector forward
                 pass
 
-            if not did_forward:
-                try:
-                    # Flatten to real vector
-                    vec = x_proc.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]
-                    y2 = model(vec)
-                    y2 = torch.as_tensor(y2).reshape(1, -1).detach().cpu()
-                    feats_this.append(y2)
-                    did_forward = True
-                except Exception:
-                    # Last resort: use the flattened vector as the embedding
-                    vec = x_proc.reshape(-1)
-                    if torch.is_complex(vec):
-                        vec = torch.cat([vec.real, vec.imag], dim=0)
-                    vec = vec.to(torch.float32).unsqueeze(0).cpu()
-                    feats_this.append(vec)
+            # 3) fallback: use flattened vector directly
+            vec = x2.reshape(-1)
+            if torch.is_complex(vec):
+                vec = torch.cat([vec.real, vec.imag], dim=0)
+            vec = vec.to(torch.float32).unsqueeze(0).cpu()
+            feats.append(vec)
 
-        embs_this = torch.cat(feats_this, dim=0)  # [n, d’]
-        all_feats.append(embs_this)
+        Zi = torch.cat(feats, dim=0)  # [n, d]
+        all_embs.append(Zi)
 
-        if label_aware and labels is not None and labels.numel() > 0:
-            labels_per_ds.append(labels[idx].clone())
+        if label_aware:
+            if y is not None and len(y) >= n:
+                labels_per_ds.append(y[idx].clone())
+            else:
+                labels_per_ds.append(torch.empty((0,), dtype=torch.long))
 
-    # Pad to common dimension
-    max_d = max(t.shape[1] for t in all_feats)
+    # Pad to common dim
+    max_d = max(t.shape[1] for t in all_embs)
     padded = []
-    for t in all_feats:
+    for t in all_embs:
         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
+    embs = torch.stack(padded, dim=0)  # [D, n, d]
+    return embs, labels_per_ds