diff --git a/GenRL b/GenRL deleted file mode 160000 index 82a867ed2..000000000 --- a/GenRL +++ /dev/null @@ -1 +0,0 @@ -Subproject commit 82a867ed24f6095a7d931ec76319aae20d9a01c4 diff --git a/fastvideo/train/methods/rl/reward/hpsv3.py b/fastvideo/train/methods/rl/reward/hpsv3.py index b01765335..7678bd073 100644 --- a/fastvideo/train/methods/rl/reward/hpsv3.py +++ b/fastvideo/train/methods/rl/reward/hpsv3.py @@ -28,6 +28,129 @@ if _HPSV3_ROOT.exists(): # 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] diff --git a/fastvideo/train/methods/rl/reward/utils.py b/fastvideo/train/methods/rl/reward/utils.py index 91f9cef02..151cc4cb9 100644 --- a/fastvideo/train/methods/rl/reward/utils.py +++ b/fastvideo/train/methods/rl/reward/utils.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): diff --git a/fastvideo/train/methods/rl/reward/videoalign.py b/fastvideo/train/methods/rl/reward/videoalign.py index abfc2d31c..d333a1ca8 100644 --- a/fastvideo/train/methods/rl/reward/videoalign.py +++ b/fastvideo/train/methods/rl/reward/videoalign.py @@ -7,6 +7,7 @@ from __future__ import annotations import os import sys import tempfile +from importlib import import_module, util from typing import Any import numpy as np @@ -29,6 +30,7 @@ if os.path.isdir(_VIDEOALIGN_ROOT): # Global cache of VideoAlign inferencers. _VIDEOALIGN_INFERENCERS: dict[str, Any] = {} +_VIDEOALIGN_PATCHED = False def _normalize_device_str(device) -> str: @@ -37,6 +39,283 @@ def _normalize_device_str(device) -> str: return str(torch.device(device)) +def _move_videoalign_inferencer(inferencer: Any, device) -> None: + """Move a VideoAlign inferencer across devices.""" + device_str = _normalize_device_str(device) + model = getattr(inferencer, "model", None) + if model is not None and hasattr(model, "to"): + model.to(device) + inferencer.device = device_str + + +def _remap_qwen2vl_state_dict_keys( + state_dict: dict[str, Any], +) -> dict[str, Any]: + """Adapt 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 VideoAlign 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_qwen2vl_state_dict_keys(state_dict) + if not assign: + try: + assign = any( + getattr(param, "is_meta", False) + for param in self.parameters() + ) + except Exception: + assign = False + 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 _select_videoalign_frame_indices( + vision_mod: Any, + ele: dict[str, Any], + total_frames: int, + video_fps: float, +) -> list[int]: + sample_type = ele.get("sample_type", "uniform") + if sample_type == "uniform": + nframes = vision_mod.smart_nframes( + ele, + total_frames=total_frames, + video_fps=video_fps, + ) + return torch.linspace( + 0, + total_frames - 1, + nframes, + ).round().long().tolist() + if sample_type == "multi_pts": + frames_each_pts = 6 + num_pts = 4 + fps = 8 + nframes = max( + frames_each_pts, + int(total_frames * fps // video_fps), + ) + frame_idx = torch.linspace( + 0, + total_frames - 1, + nframes, + ).round().long().tolist() + start_pt = int(frames_each_pts // 2) + end_pt = int(nframes - frames_each_pts // 2) + pts = torch.linspace( + start_pt, + end_pt, + num_pts, + ).round().long().tolist() + idx = [] + for pt in pts: + idx.extend( + frame_idx[ + pt - frames_each_pts // 2: + pt + frames_each_pts // 2 + ] + ) + return idx + raise ValueError(f"Unsupported VideoAlign sample_type: {sample_type}") + + +def _read_video_opencv( + vision_mod: Any, + ele: dict[str, Any], +) -> torch.Tensor: + """Read local MP4s without relying on torchvision.io.read_video.""" + import cv2 + + video_path = ele["video"] + if video_path.startswith("file://"): + video_path = video_path[7:] + cap = cv2.VideoCapture(video_path) + if not cap.isOpened(): + raise ValueError(f"Could not open video: {video_path}") + + video_fps = float(cap.get(cv2.CAP_PROP_FPS) or 30.0) + total_file_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT) or 0) + start_frame = max( + 0, + int(round(float(ele.get("video_start", 0.0) or 0.0) * video_fps)), + ) + end_sec = ele.get("video_end") + if end_sec is None: + end_frame = total_file_frames if total_file_frames > 0 else None + else: + end_frame = int(round(float(end_sec) * video_fps)) + if total_file_frames > 0: + end_frame = min(end_frame, total_file_frames) + + cap.set(cv2.CAP_PROP_POS_FRAMES, start_frame) + frames = [] + current_frame = start_frame + while end_frame is None or current_frame < end_frame: + ok, frame = cap.read() + if not ok: + break + frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)) + current_frame += 1 + cap.release() + + if not frames: + raise ValueError(f"No frames were read from video: {video_path}.") + + idx = _select_videoalign_frame_indices( + vision_mod, + ele, + total_frames=len(frames), + video_fps=video_fps, + ) + video = np.stack([frames[i] for i in idx], axis=0) + return torch.from_numpy(video).permute(0, 3, 1, 2) + + +def _torchvision_read_video_available() -> bool: + try: + torchvision_io = import_module("torchvision.io") + except Exception: + return False + return hasattr(torchvision_io, "read_video") + + +def _patch_videoalign_video_reader() -> None: + """Register an OpenCV reader for torchvision builds without read_video.""" + vision_mod = import_module("vision_process") + if "opencv" not in vision_mod.VIDEO_READER_BACKENDS: + + def read_video_opencv(ele): + return _read_video_opencv(vision_mod, ele) + + vision_mod.VIDEO_READER_BACKENDS["opencv"] = read_video_opencv + + if _torchvision_read_video_available(): + return + + vision_mod.FORCE_QWENVL_VIDEO_READER = "opencv" + if hasattr(vision_mod.get_video_reader_backend, "cache_clear"): + vision_mod.get_video_reader_backend.cache_clear() + + +def _patch_videoalign_modules() -> Any: + """Patch VideoAlign for the FastVideo dependency set.""" + global _VIDEOALIGN_PATCHED + + inference_mod = import_module("inference") + if _VIDEOALIGN_PATCHED: + return inference_mod + + train_reward_mod = import_module("train_reward") + trainer_mod = import_module("trainer") + _patch_videoalign_video_reader() + + if util.find_spec("flash_attn") is None: + for mod in (train_reward_mod, inference_mod): + original_create = mod.create_model_and_processor + + def create_model_and_processor_sdpa( + *args, + _original_create=original_create, + **kwargs, + ): + training_args = kwargs.get("training_args") + if training_args is not None: + training_args.disable_flash_attn2 = True + return _original_create(*args, **kwargs) + + mod.create_model_and_processor = create_model_and_processor_sdpa + + _patch_load_state_dict(trainer_mod.Qwen2VLRewardModelBT) + try: + peft_mod = import_module("peft") + except ImportError: + peft_mod = None + if peft_mod is not None: + _patch_load_state_dict(peft_mod.PeftModel) + + _VIDEOALIGN_PATCHED = True + return inference_mod + + +def _patch_videoalign_runtime_model(model: Any) -> None: + """Add aliases expected by VideoAlign'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 set_videoalign_device(device) -> None: """Move cached VideoAlign inferencers to device.""" key = _normalize_device_str(device) @@ -45,7 +324,7 @@ def set_videoalign_device(device) -> None: ): if old_key != key and old_key.split(":")[0] != key: new_key = inf._key_prefix + ":" + key - inf.to(device) + _move_videoalign_inferencer(inf, device) _VIDEOALIGN_INFERENCERS[new_key] = inf del _VIDEOALIGN_INFERENCERS[old_key] @@ -56,10 +335,13 @@ def _get_inferencer( ): """Get or create VideoAlign inferencer.""" if checkpoint_path is None: - checkpoint_path = os.path.join( - os.path.dirname(__file__), - "..", "..", "..", "..", "..", - "data", "VideoReward", + checkpoint_path = os.environ.get( + "VIDEOALIGN_CHECKPOINT_PATH", + os.path.join( + os.path.dirname(__file__), + "VideoAlign", + "checkpoints", + ), ) checkpoint_path = os.path.abspath(checkpoint_path) @@ -67,7 +349,8 @@ def _get_inferencer( cache_key = f"{checkpoint_path}:{key}" if cache_key not in _VIDEOALIGN_INFERENCERS: try: - from inference import VideoVLMRewardInference + inference_mod = _patch_videoalign_modules() + VideoVLMRewardInference = inference_mod.VideoVLMRewardInference except ImportError as exc: msg = ( "VideoAlign not found. Ensure the " @@ -81,6 +364,7 @@ def _get_inferencer( load_from_pretrained=checkpoint_path, device=device, ) + _patch_videoalign_runtime_model(inf.model) inf._key_prefix = checkpoint_path or "default" _VIDEOALIGN_INFERENCERS[cache_key] = inf return _VIDEOALIGN_INFERENCERS[cache_key] diff --git a/fastvideo/train/methods/rl/utils/rewards.py b/fastvideo/train/methods/rl/utils/rewards.py index a1e24b012..d7017c205 100644 --- a/fastvideo/train/methods/rl/utils/rewards.py +++ b/fastvideo/train/methods/rl/utils/rewards.py @@ -3,6 +3,7 @@ from __future__ import annotations +import gc import importlib import inspect import time @@ -20,9 +21,11 @@ from fastvideo.train.methods.rl.reward import ( videoalign_ta_score, ) from fastvideo.train.methods.rl.reward.hpsv3 import ( + _HPSV3_INFERENCERS, set_hpsv3_device, ) from fastvideo.train.methods.rl.reward.videoalign import ( + _VIDEOALIGN_INFERENCERS, set_videoalign_device, ) @@ -154,6 +157,28 @@ def move_reward_models(reward_cfg, device) -> None: set_videoalign_device(device) +def clear_reward_models(reward_cfg) -> None: + """Drop cached GPU-backed reward models before PPO training.""" + cleared = False + if _has_reward( + reward_cfg, + {"hpsv3_general", "hpsv3_percentile"}, + ): + _HPSV3_INFERENCERS.clear() + cleared = True + if _has_reward( + reward_cfg, + {"videoalign_mq", "videoalign_ta"}, + ): + _VIDEOALIGN_INFERENCERS.clear() + cleared = True + if cleared: + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + torch.cuda.ipc_collect() + + @contextmanager def reward_models_on_device(reward_cfg, device): """Temporarily move reward models to device.""" @@ -173,10 +198,9 @@ def reward_models_on_device(reward_cfg, device): _t2 = time.perf_counter() move_reward_models(reward_cfg, "cpu") if use_cuda: - import gc - gc.collect() torch.cuda.empty_cache() + torch.cuda.ipc_collect() _t3 = time.perf_counter() logger.info( "[rewards] move_to_cpu+gc=%.1fs",