[feat] GenRL: stabilize reward model compatibility (#1400)

Co-authored-by: Davids048 <jundasu@ucsd.edu>

Author: Abecid

GenRL

@@ -28,6 +28,129 @@
 
 # Global cache of HPSv3 inferencers keyed by device.
 _HPSV3_INFERENCERS: dict[str, Any] = {}
+_HPSV3_LOAD_PATCHED = False
+
+
+def _patch_transformers_video_input_alias() -> None:
+    """Keep HPSv3 compatible with newer transformers releases.
+
+    HPSv3 imports ``VideoInput`` from ``transformers.image_utils`` for type
+    annotations. Some transformers versions used by FastVideo no longer
+    export that alias, even though the runtime image utilities HPSv3 needs are
+    still present.
+    """
+    from transformers import image_utils
+
+    if not hasattr(image_utils, "VideoInput"):
+        image_utils.VideoInput = image_utils.ImageInput
+
+
+def _remap_hpsv3_state_dict(state_dict: dict[str, Any]) -> dict[str, Any]:
+    """Adapt HPSv3 checkpoints saved with older Qwen2-VL key names."""
+    remapped = {}
+    for key, value in state_dict.items():
+        if key.startswith("visual."):
+            key = f"model.{key}"
+        elif key.startswith("model.layers."):
+            key = f"model.language_model.{key[len('model.'):]}"
+        elif key.startswith("model.embed_tokens."):
+            key = f"model.language_model.{key[len('model.'):]}"
+        elif key.startswith("model.norm."):
+            key = f"model.language_model.{key[len('model.'):]}"
+
+        key = key.replace(
+            "base_model.model.visual.",
+            "base_model.model.model.visual.",
+            1,
+        )
+        key = key.replace(
+            "base_model.model.model.layers.",
+            "base_model.model.model.language_model.layers.",
+            1,
+        )
+        key = key.replace(
+            "base_model.model.model.embed_tokens.",
+            "base_model.model.model.language_model.embed_tokens.",
+            1,
+        )
+        key = key.replace(
+            "base_model.model.model.norm.",
+            "base_model.model.model.language_model.norm.",
+            1,
+        )
+        remapped[key] = value
+    return remapped
+
+
+def _walk_model_graph(model: Any):
+    """Yield common wrapper/base model objects without importing PEFT."""
+    stack = [model]
+    seen = set()
+    while stack:
+        current = stack.pop()
+        if current is None or id(current) in seen:
+            continue
+        seen.add(id(current))
+        yield current
+        for attr in ("base_model", "model"):
+            child = getattr(current, attr, None)
+            if child is not None:
+                stack.append(child)
+
+
+def _patch_load_state_dict(cls: Any) -> None:
+    """Patch a model class to accept old Qwen2-VL checkpoint keys."""
+    if getattr(cls, "_fastvideo_qwen2vl_key_remap", False):
+        return
+
+    original_load_state_dict = cls.load_state_dict
+
+    def load_state_dict_with_key_remap(
+        self,
+        state_dict,
+        strict=True,
+        assign=False,
+    ):
+        state_dict = _remap_hpsv3_state_dict(state_dict)
+        return original_load_state_dict(
+            self,
+            state_dict,
+            strict=strict,
+            assign=assign,
+        )
+
+    cls.load_state_dict = load_state_dict_with_key_remap
+    cls._fastvideo_qwen2vl_key_remap = True
+
+
+def _patch_hpsv3_state_dict_loader() -> None:
+    """Patch HPSv3 reward model loading for transformers key drift."""
+    global _HPSV3_LOAD_PATCHED
+    if _HPSV3_LOAD_PATCHED:
+        return
+
+    from hpsv3.model.qwen2vl_trainer import Qwen2VLRewardModelBT
+
+    _patch_load_state_dict(Qwen2VLRewardModelBT)
+    try:
+        from peft import PeftModel
+    except ImportError:
+        PeftModel = None
+    if PeftModel is not None:
+        _patch_load_state_dict(PeftModel)
+    _HPSV3_LOAD_PATCHED = True
+
+
+def _patch_hpsv3_runtime_model(model: Any) -> None:
+    """Add aliases expected by HPSv3's older Qwen2-VL forward."""
+    for candidate in _walk_model_graph(model):
+        language_model = getattr(candidate, "language_model", None)
+        if (
+            language_model is not None
+            and not hasattr(candidate, "embed_tokens")
+            and hasattr(language_model, "embed_tokens")
+        ):
+            candidate.__dict__["embed_tokens"] = language_model.embed_tokens
 
 
 def _normalize_device(device) -> str:
@@ -36,6 +159,19 @@ def _normalize_device(device) -> str:
     return str(torch.device(device))
 
 
+def _move_hpsv3_inferencer(inferencer: Any, device) -> None:
+    """Move an HPSv3 inferencer across devices.
+
+    HPSv3RewardInferencer does not expose ``.to()``, but it stores its torch
+    module on ``.model`` and reads ``.device`` when preparing batches.
+    """
+    device_str = _normalize_device(device)
+    model = getattr(inferencer, "model", None)
+    if model is not None and hasattr(model, "to"):
+        model.to(device)
+    inferencer.device = device_str
+
+
 def set_hpsv3_device(device) -> None:
     """Move cached HPSv3 inferencer to given device."""
     key = _normalize_device(device)
@@ -44,7 +180,7 @@ def set_hpsv3_device(device) -> None:
     # Move from any existing device.
     for old_key, inf in list(_HPSV3_INFERENCERS.items()):
         if old_key != key:
-            inf.to(device)
+            _move_hpsv3_inferencer(inf, device)
             _HPSV3_INFERENCERS[key] = inf
             del _HPSV3_INFERENCERS[old_key]
             return
@@ -55,15 +191,18 @@ def _get_hpsv3_inferencer(device):
     key = _normalize_device(device)
     if key not in _HPSV3_INFERENCERS:
         try:
+            _patch_transformers_video_input_alias()
             from hpsv3 import HPSv3RewardInferencer
+            _patch_hpsv3_state_dict_loader()
         except ImportError as exc:
             msg = (
-                "hpsv3 package not found. Ensure the "
-                "HPSv3 submodule is checked out under "
-                "fastvideo/train/methods/rl/reward/HPSv3"
+                "Failed to import HPSv3. Ensure the HPSv3 submodule is "
+                "checked out under fastvideo/train/methods/rl/reward/HPSv3 "
+                "and that its transformers dependencies are compatible."
             )
             raise ImportError(msg) from exc
         inf = HPSv3RewardInferencer(device=device)
+        _patch_hpsv3_runtime_model(inf.model)
         _HPSV3_INFERENCERS[key] = inf
     return _HPSV3_INFERENCERS[key]
 

fastvideo/train/methods/rl/reward/hpsv3.py

@@ -24,11 +24,20 @@ def prepare_images(
 
     if images.ndim == 4:
         # Image batch: (N, C, H, W) or (N, H, W, C)
-        if images.shape[1] in (1, 3):
+        if images.shape[-1] in (1, 3):
+            pass
+        elif images.shape[1] in (1, 3):
             images = images.transpose(0, 2, 3, 1)
     elif images.ndim == 5:
-        # Video batch: (N, F, C, H, W) or (N, C, F, H, W)
-        if images.shape[2] in (1, 3):
+        # Video batch: (N, F, H, W, C), (N, F, C, H, W),
+        # or (N, C, F, H, W). Check channel-last first because
+        # one-frame videos have shape[1] == 1.
+        if images.shape[-1] in (1, 3):
+            pass
+        elif images.shape[1] == 3 and images.shape[2] == 1:
+            # (N, C=3, F=1, H, W) -> (N, F, H, W, C)
+            images = images.transpose(0, 2, 3, 4, 1)
+        elif images.shape[2] in (1, 3):
             # (N, F, C, H, W) -> (N, F, H, W, C)
             images = images.transpose(0, 1, 3, 4, 2)
         elif images.shape[1] in (1, 3):