38 lines
933 B
Python
38 lines
933 B
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Reusable reward models for training methods."""
|
|
|
|
from fastvideo.train.methods.rl.rewards.frame_rewards import (
|
|
ClipScoreScorer,
|
|
PickScoreScorer,
|
|
)
|
|
from fastvideo.train.methods.rl.rewards.media import (
|
|
MultiRewardScorer,
|
|
RewardScorer,
|
|
select_first_frame,
|
|
)
|
|
|
|
|
|
def build_multi_reward_scorer(
|
|
reward_weights,
|
|
*,
|
|
device="cuda",
|
|
scorers: dict[str, RewardScorer] | None = None,
|
|
) -> MultiRewardScorer:
|
|
available: dict[str, RewardScorer] = dict(scorers or {})
|
|
if not available:
|
|
available = {
|
|
"pickscore": PickScoreScorer(device=device),
|
|
"clipscore": ClipScoreScorer(device=device),
|
|
}
|
|
return MultiRewardScorer(reward_weights, scorers=available)
|
|
|
|
|
|
__all__ = [
|
|
"ClipScoreScorer",
|
|
"MultiRewardScorer",
|
|
"PickScoreScorer",
|
|
"RewardScorer",
|
|
"build_multi_reward_scorer",
|
|
"select_first_frame",
|
|
]
|