diff --git a/experimental/__init__.py b/experimental/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/experimental/refl/README.md b/experimental/refl/README.md new file mode 100644 index 000000000..9c38aaeb9 --- /dev/null +++ b/experimental/refl/README.md @@ -0,0 +1,56 @@ +# experimental/refl — WAN ReFL/BPTT (differentiable reward backprop) + +The first self-contained UniRL recipe: direct reward backprop (ReFL / DRaFT) +for WAN video models. Two colocated roles — a `ReflActorRole` (FSDP WAN + +grad BPTT sampling + optimizer) and a frozen differentiable reward +(`RewardService`) — run a 3-RPC step under the distributed `enable_grad()` +context: generate → score → backward, then optimizer step. No advantages, +no replay, no rollout engine, no weight sync. + +## Launch + +One command per config; a Ray cluster must be up (`ray start --head`). + +```bash +# WAN 2.1 T2V + VideoAlign (Qwen2-VL VQ/MQ/TA) reward +export PRETRAINED_MODEL=/path/to/Wan2.1-T2V-1.3B-Diffusers \ + VIDEOALIGN_MODEL_PATH=/path/to/VideoReward \ + DATA_PATH=/path/to/prompts.txt +RAY_ADDRESS=auto python -m experimental.refl.run --config-name=wan21_t2v_videoalign_refl num_devices=8 + +# WAN 2.2 I2V + Face-identity reward (first frame via (image, condition) +# MediaRef; face reference via per-sample metadata ref_video_path) +pip install -r experimental/refl/reward/face/requirements.txt +export PRETRAINED_MODEL=/path/to/Wan2.2-I2V-A14B-Diffusers \ + FACE_MODEL_PATH=/path/to/antelodev2 \ + DATA_PATH=/path/to/i2v_prompts.jsonl +RAY_ADDRESS=auto python -m experimental.refl.run --config-name=wan22_i2v_face_refl num_devices=8 +``` + +## Layout + +| Path | What | +|---|---| +| `trainer.py` | `REFLTrainer(BaseTrainer)` — driver: wiring + the 3-RPC train step | +| `roles.py` | `ReflActorRole(Remote)` — family-agnostic actor (`pipeline_target` + `model_config`) | +| `models/` | Per-model BPTT adaptations subclassing the core pipelines (`types.py` defines the contract): `wan21.py`, `wan22.py` — mirrors `unirl/models/` (graduates into the matching model packages) | +| `reward/` | Package-local differentiable rewards (VideoAlign, Face), each with an additive-only `requirements.txt` — mirrors `unirl/reward/` (graduates into it) | +| `examples/` | Flat Hydra configs (repo-wide schema) — mirrors the top-level `examples/` (graduates into it) | + +## Environment + +Targets the **locked core stack only** (`transformers>=5.6,<5.7`, +`peft>=0.20` — see `pyproject.toml`). There are deliberately no +version-compat branches: a wrong environment fails loudly; align the +environment, not the code. Reward and actor share one Python process, so +recipe `requirements.txt` files may only ADD packages, never re-pin the +core stack. + +## Verification + +| Config | Hardware | Head | Status | +| --- | --- | --- | --- | +| `wan21_t2v_videoalign_refl` (835 rollouts) | 8xH20 | pre-adjustment (`e3c6b940` lineage) | contributor long run — reward curve in PR #210 | +| `wan21_t2v_videoalign_refl` (2-rollout smoke, full 81f/480x832 geometry) | 8xH20, fleet image | `40b3f4c9` | PASS — grads flow reward → VAE → DiT LoRA | +| VideoAlign load + differentiable fwd/bwd on transformers 5.6.2 + peft 0.20 | 8xH20 (isolated venv) | `9087a671` | PASS — `grad_abs_mean=3.5e-3` | +| `wan22_i2v_face_refl` | 8xH20 | current head | pending (needs face assets + I2V dataset) | diff --git a/experimental/refl/__init__.py b/experimental/refl/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/experimental/refl/examples/wan21_t2v_videoalign_refl.yaml b/experimental/refl/examples/wan21_t2v_videoalign_refl.yaml new file mode 100644 index 000000000..e229d3435 --- /dev/null +++ b/experimental/refl/examples/wan21_t2v_videoalign_refl.yaml @@ -0,0 +1,149 @@ +# @package _global_ +# REFL WAN 2.1 T2V — VideoAlign (Qwen2-VL VQ/MQ/TA) reward. +# +# Two roles, always (same shape as examples/diffusion/refl_sd3.yaml): a +# ReflActorRole (FSDP WAN 2.1 + grad BPTT sampling + optimizer) and the frozen +# differentiable VideoAlign reward, colocated on the actor's worker slab so +# decoded video never leaves the GPU. GradContext backprops reward → VAE → +# DiT LoRA across the sibling RPC boundary. + +num_devices: 8 +batch_size: 8 +num_rollouts: 1000 +save_interval: 100 +save_dir: ${oc.env:OUTPUT_DIR,outputs/wan21_t2v_videoalign_refl} +save_mode: adapter +max_grad_norm: 1.0 + +actor: + _target_: experimental.refl.roles.ReflActorRole + # Family selector — swap pipeline_target + model_config for another family, + # no code changes (mirrors ReFLPolicy's pipeline_target contract). + pipeline_target: experimental.refl.models.wan21.Wan21ReflPipeline + block_class_names: ["WanTransformerBlock"] + # REFL loss: -(reward - baseline) / scale * weight + kl_weight * KL. + reward_weight: 0.25 + reward_baseline: 0.0 + reward_scale: 1.0 + kl_weight: 0.0 + strategy: + # sampling.eta=0.0 reduces FlowSDE to the deterministic ODE — REFL wants a + # deterministic transition on the differentiable path. + _target_: unirl.sde.kernels.FlowSDEStrategy + model_config: + _target_: unirl.models.wan21.config.WAN21PipelineConfig + pretrained_model_ckpt_path: ${oc.env:PRETRAINED_MODEL} + model_precision: bf16 + autocast_precision: bf16 + trajectory_precision: bf16 + logprob_precision: fp32 + shift: 5.0 + max_sequence_length: 512 + fsdp_cfg: + _target_: unirl.train.configs.FSDPConfig + param_dtype: bf16 + # fp32 LoRA master over the bf16 base: AdamW steps (~lr=5e-6) are below + # bf16 ULP at lora_A's magnitude, so a bf16 master silently freezes A. + # Needs the pinned torch (>=2.11): older FSDP2 asserts uniform dtype over + # ALL params in a group; the pinned family checks trainables only. + master_dtype: fp32 + cpu_offload: false + mixed_precision: true + fsdp_mode: full + reshard_after_forward: true + # BPTT keeps the full mid→final grad window alive; activation + # checkpointing is mandatory for any non-trivial + # num_frames × num_inference_steps product. + activation_checkpointing: true + use_torch_compile: false + optimizer_cfg: + _target_: unirl.train.backend.base.OptimizerConfig + learning_rate: 5.0e-6 + adam_beta1: 0.9 + adam_beta2: 0.999 + adam_epsilon: 1.0e-8 + weight_decay: 0.0 + scheduler_cfg: + _target_: unirl.train.backend.base.LrSchedulerConfig + type: linear_warmup + warmup_steps: 10 + total_steps: ${num_rollouts} + lora_cfg: + _target_: unirl.train.configs.LoraConfig + rank: 64 + alpha: 128 + dropout: 0.0 + bias: none + task_type: FEATURE_EXTRACTION + # Peft path takes suffix matches → these six suffixes catch both + # attn1 (self-attention) and attn2 (cross-attention) LoRA targets + # inside every WanTransformerBlock. + target_modules: + - to_q + - to_k + - to_v + - to_out.0 + - ffn.net.0.proj + - ffn.net.2 + +reward: + _target_: unirl.reward.service.RewardService + backend: + _target_: experimental.refl.reward.videoalign.VideoAlignRewardScorer + base_device: cuda + config: + _target_: experimental.refl.reward.videoalign.VideoAlignSpec + reward_model_path: ${oc.env:VIDEOALIGN_MODEL_PATH} + device: cuda + batch_size: 1 + resize_height: 336 + resize_width: 588 + micro_batch_size: 1 + reward_num_frames: 36 # 81 → 36 uniform frames before scoring + use_norm: true + w_vq: 1.0 + w_mq: 1.0 + w_ta: 1.0 + differentiable: true # REFL requires grad through the reward forward + +data_source: + _target_: unirl.data.data_source.MultimodalRLDataSource + args: + run: + data_path: ${oc.env:DATA_PATH} + eval_data_path: ${oc.env:EVAL_DATA_PATH,${oc.env:DATA_PATH}} + seed: 42 + shuffle: false + algorithm: + prompts_per_rollout: ${batch_size} + +sampling: + _target_: unirl.types.sampling.DiffusionSamplingParams + # 25 inference steps with mid_timestep = final_timestep = 24 + # → DRaFT-1 (only the final step is differentiable). + num_inference_steps: 25 + guidance_scale: 5.0 + height: 480 + width: 832 + # (num_frames - 1) % 4 == 0 is the WAN 2.1 T2V hard constraint. + num_frames: 81 + eta: 0.0 + samples_per_prompt: 1 + seed: 42 + init_same_noise: false + sampler_kwargs: + negative_prompt: "blurry, low quality, distorted, oversaturated" + # Truncated BPTT window: no_grad until ``mid_timestep``, grad on from + # ``mid_timestep`` .. ``final_timestep`` (inclusive). 24/24/25 → only + # the final step traverses the reward backward pass (DRaFT-1 / ReFL). + mid_timestep: 24 + final_timestep: 24 + # KL on/off + weight live on actor.kl_weight (single knob; 0 here → + # single-pass forward per step, no LoRA-disabled reference recompute). + +logging: + report_to_wandb: ${oc.decode:${oc.env:REPORT_TO_WANDB,true}} + project_name: ${oc.env:WANDB_PROJECT,unirl-refl} + run_name: ${oc.env:WANDB_RUN_NAME,wan21_t2v_videoalign_refl} + tags: ["wan21", "t2v", "refl", "bptt", "videoalign"] + log_media: false diff --git a/experimental/refl/examples/wan22_i2v_face_refl.yaml b/experimental/refl/examples/wan22_i2v_face_refl.yaml new file mode 100644 index 000000000..2693ae5ae --- /dev/null +++ b/experimental/refl/examples/wan22_i2v_face_refl.yaml @@ -0,0 +1,138 @@ +# @package _global_ +# REFL WAN 2.2 I2V — Face-identity reward. +# +# Two roles, always (same shape as examples/diffusion/refl_sd3.yaml): a +# ReflActorRole (FSDP WAN 2.2 dual-DiT + grad BPTT sampling + optimizer, +# LoRA restricted to the low_noise DiT) and the frozen differentiable Face +# reward, colocated on the actor's worker slab. The I2V first frame rides the +# data source's (image, condition) MediaRef; the face reference video path +# rides per-sample metadata (``ref_video_path``) into the reward's records. + +num_devices: 8 +batch_size: 8 +num_rollouts: 1000 +save_interval: 100 +save_dir: ${oc.env:OUTPUT_DIR,outputs/wan22_i2v_face_refl} +save_mode: adapter +max_grad_norm: 1.0 + +actor: + _target_: experimental.refl.roles.ReflActorRole + pipeline_target: experimental.refl.models.wan22.Wan22ReflPipeline + block_class_names: ["WanTransformerBlock"] + # REFL loss: -(reward - baseline) / scale * weight + kl_weight * KL. + reward_weight: 0.1 + reward_baseline: 0.54 + reward_scale: 0.16 + kl_weight: 1.0 + strategy: + # sampling.eta=0.0 reduces FlowSDE to the deterministic ODE — REFL wants a + # deterministic transition on the differentiable path. + _target_: unirl.sde.kernels.FlowSDEStrategy + model_config: + _target_: unirl.models.wan22.config.WAN22PipelineConfig + pretrained_model_ckpt_path: ${oc.env:PRETRAINED_MODEL} + model_precision: bf16 + autocast_precision: bf16 + trajectory_precision: bf16 + logprob_precision: fp32 + shift: 5.0 + max_sequence_length: 512 + boundary_ratio: 0.9 + num_train_timesteps: 1000 + fsdp_cfg: + _target_: unirl.train.configs.FSDPConfig + param_dtype: bf16 + # fp32 LoRA master over the bf16 base: AdamW steps (~lr=5e-6) are below + # bf16 ULP at lora_A's magnitude, so a bf16 master silently freezes A. + # Needs the pinned torch (>=2.11): older FSDP2 asserts uniform dtype over + # ALL params in a group; the pinned family checks trainables only. + master_dtype: fp32 + cpu_offload: false + mixed_precision: true + fsdp_mode: full + reshard_after_forward: true + activation_checkpointing: true + use_torch_compile: false + optimizer_cfg: + _target_: unirl.train.backend.base.OptimizerConfig + learning_rate: 2.5e-6 + adam_beta1: 0.9 + adam_beta2: 0.999 + adam_epsilon: 1.0e-8 + weight_decay: 0.0 + scheduler_cfg: + _target_: unirl.train.backend.base.LrSchedulerConfig + type: linear_warmup + warmup_steps: 10 + total_steps: ${num_rollouts} + lora_cfg: + _target_: unirl.train.configs.LoraConfig + rank: 32 + alpha: 32 + dropout: 0.0 + bias: none + task_type: FEATURE_EXTRACTION + # Inject LoRA only into the low-noise DiT (the final-timestep expert the + # BPTT window trains); the high-noise DiT stays frozen. + module_prefix: low_noise + target_modules: + - to_q + - to_k + - to_v + - to_out.0 + - ffn.net.0.proj + - ffn.net.2 + +reward: + _target_: unirl.reward.service.RewardService + backend: + _target_: experimental.refl.reward.face.FaceRewardScorer + base_device: cuda + config: + _target_: experimental.refl.reward.face.FaceRewardSpec + model_path: ${oc.env:FACE_MODEL_PATH,/path/to/antelodev2_face_ckpt} + device: cuda + batch_size: 1 + image_size: 112 + ref_max_frames: 81 + ref_max_pixels: 230400 + differentiable: true + +data_source: + _target_: unirl.data.data_source.MultimodalRLDataSource + args: + run: + data_path: ${oc.env:DATA_PATH,/path/to/wan22_i2v_face_refl_prompts.jsonl} + eval_data_path: ${oc.env:EVAL_DATA_PATH,${oc.env:DATA_PATH,/path/to/wan22_i2v_face_refl_prompts.jsonl}} + seed: 42 + shuffle: false + algorithm: + prompts_per_rollout: ${batch_size} + +sampling: + _target_: unirl.types.sampling.DiffusionSamplingParams + num_inference_steps: 8 + guidance_scale: 1.0 + guidance_scale_2: 1.0 + height: 352 + width: 640 + num_frames: 81 + eta: 0.0 + samples_per_prompt: 1 + seed: 42 + init_same_noise: false + sampler_kwargs: + # Truncated BPTT window: grad on from mid_timestep..final_timestep — the + # last 4 of 8 steps, i.e. the low-noise DiT's sigma range. The per-step + # KL against the LoRA-disabled reference is switched/weighted by the + # single actor.kl_weight knob. + mid_timestep: 4 + final_timestep: 7 + +logging: + report_to_wandb: ${oc.decode:${oc.env:REPORT_TO_WANDB,true}} + project_name: ${oc.env:WANDB_PROJECT,unirl-refl} + run_name: ${oc.env:WANDB_RUN_NAME,wan22_i2v_face_refl} + tags: ["wan22", "i2v", "refl", "bptt", "face"] + log_media: false diff --git a/experimental/refl/models/__init__.py b/experimental/refl/models/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/experimental/refl/models/types.py b/experimental/refl/models/types.py new file mode 100644 index 000000000..096c8f2cf --- /dev/null +++ b/experimental/refl/models/types.py @@ -0,0 +1,40 @@ +"""Recipe-local BPTT stage contract for the refl recipe. + +``diffuse_with_grad`` is deliberately NOT part of the core +:class:`~unirl.models.types.diffusion.DiffusionStage` protocol: concrete +stages inherit that Protocol *explicitly*, so a protocol-level stub would +become a real ``None``-returning method on every diffusion stage in the +repo and make ``hasattr``-based capability checks meaningless. While REFL +is the only BPTT consumer, the contract lives here; if a second consumer +appears outside ``experimental/refl``, promote it to core as a separate opt-in +``@runtime_checkable`` protocol (the ``DifferentiableReward`` / +``LatentShapeProvider`` idiom), not as a method on ``DiffusionStage``. + +Implementors (``Wan21ReflDiffusionStage`` / ``Wan22ReflDiffusionStage``) +provide:: + + diffuse_with_grad(conditions, *, schedule, params, initial_latents=None) + -> DiffuseWithGradResult +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch + + +@dataclass +class DiffuseWithGradResult: + """Output of a recipe stage's ``diffuse_with_grad``. + + ``kl_loss`` is per-sample ``[B]`` (zeros when the KL branch is off) so + DP-scattered consumers round-trip each shard's own KL, never a + cross-shard aggregate. + """ + + z_final: torch.Tensor + kl_loss: torch.Tensor + + +__all__ = ["DiffuseWithGradResult"] diff --git a/experimental/refl/models/wan21.py b/experimental/refl/models/wan21.py new file mode 100644 index 000000000..508289c67 --- /dev/null +++ b/experimental/refl/models/wan21.py @@ -0,0 +1,412 @@ +"""Recipe-local WAN 2.1 T2V step + stage + pipeline for REFL BPTT. + +Mirrors ``experimental.refl.models.wan22`` but targets the WAN 2.1 T2V +single-DiT stack. The REFL-specific pieces live here: + +- :class:`Wan21ReflDiffusionStep` — strict recipe-local single-branch + CFG predictor used only by REFL BPTT. +- :class:`Wan21ReflDiffusionStage` — subclass of the mainline + :class:`WAN21DiffusionStage` with an extra :meth:`diffuse_with_grad` + method implementing single-branch CFG + BPTT windowing (mid/final + timestep) + optional per-step KL against the LoRA-disabled reference. +- :class:`Wan21ReflPipeline` — subclass of the mainline + :class:`WAN21Pipeline` that post-swaps ``self.diffusion`` for the REFL + variant above, keeping every other stage (text/vae/scheduler) intact. + +The mainline WAN 2.1 code path is unchanged. REFL-specific single-branch +CFG prediction lives in recipe-local :class:`Wan21ReflDiffusionStep`. +""" + +from __future__ import annotations + +import dataclasses +from contextlib import nullcontext +from typing import Any, Dict, Optional, Tuple + +import torch + +from experimental.refl.models.types import DiffuseWithGradResult +from unirl.models.wan21.bundle import WAN21Bundle +from unirl.models.wan21.clip_vision_encode import WAN21CLIPVisionEncodeStage +from unirl.models.wan21.conditions import WAN21Conditions +from unirl.models.wan21.diffusion import WAN21DiffusionStage, WAN21DiffusionStep +from unirl.models.wan21.image_encode import WAN21ImageLatentEncodeStage +from unirl.models.wan21.pipeline import WAN21Pipeline +from unirl.train.lora import adapters_disabled +from unirl.types.primitives import Images, Texts +from unirl.types.sampling import DiffusionSamplingParams + +# Matches the mainline module-level constant in unirl/models/wan21/diffusion.py. +# Not imported directly to keep the recipe decoupled from the mainline's +# private surface. +_WAN_TIMESTEP_SCALE: float = 1000.0 + +# Inclusive max for torch.Generator.manual_seed and torch initial_seed conventions. +MAX_TORCH_SEED = (1 << 63) - 1 + + +class Wan21ReflDiffusionStep(WAN21DiffusionStep): + """Recipe-local WAN 2.1 single-branch denoising for REFL BPTT.""" + + def predict_noise( # pyright: ignore[reportIncompatibleMethodOverride] + self, + model: WAN21Bundle, + sample: torch.Tensor, + sigma: torch.Tensor, + conditions: WAN21Conditions, + *, + branch: str, + ) -> torch.Tensor: + """Run exactly one CFG branch through the WAN 2.1 transformer.""" + if branch not in {"cond", "uncond"}: + raise ValueError(f"Wan21ReflDiffusionStep.predict_noise: unknown branch={branch!r}") + if conditions.text is None or conditions.text.embeds is None: + raise ValueError("Wan21ReflDiffusionStep.predict_noise: conditions.text.embeds is None") + + prompt_embeds = conditions.text.embeds + if branch == "cond": + branch_embeds = prompt_embeds + else: + neg = conditions.negative_text + branch_embeds = ( + neg.embeds if neg is not None and neg.embeds is not None else torch.zeros_like(prompt_embeds) + ) + + batch_size = int(sample.shape[0]) + timestep = sigma * _WAN_TIMESTEP_SCALE + if timestep.dim() == 0: + timestep = timestep.expand(batch_size) + elif int(timestep.shape[0]) != batch_size: + timestep = timestep.expand(batch_size) + + embeds_dtype = branch_embeds.dtype + sample_cast = sample.to(dtype=embeds_dtype) + + image_latent = conditions.image_latent + if image_latent is not None and image_latent.latents is not None: + sample_cat = torch.cat( + [sample_cast, image_latent.latents.to(device=sample_cast.device, dtype=embeds_dtype)], + dim=1, + ) + else: + sample_cat = sample_cast + + extra: Dict[str, Any] = {} + image_embed = conditions.image_embed + image_embeds = image_embed.embeds if image_embed is not None and image_embed.embeds is not None else None + if image_embeds is not None: + extra["encoder_hidden_states_image"] = image_embeds.to(device=sample_cast.device, dtype=embeds_dtype) + + return model.transformer( + hidden_states=sample_cat, + encoder_hidden_states=branch_embeds, + timestep=timestep, + return_dict=False, + **extra, + )[0] + + +class Wan21ReflDiffusionStage(WAN21DiffusionStage): + """WAN 2.1 T2V diffusion stage + REFL BPTT sampling override. + + Adds ``diffuse_with_grad`` that keeps the autograd graph alive on the + returned final latent, honouring the REFL-specific BPTT knobs + (``mid_timestep`` / ``final_timestep`` / ``kl_weight``) via + ``params.sampler_kwargs``. Everything else is inherited from the + mainline stage — ``diffuse`` / ``replay`` / ``predict_noise`` / + ``trainable_module`` are untouched. + """ + + def generate_latents( + self, + batch_size: int, + latent_shape: Tuple[int, ...], + device: torch.device, + dtype: torch.dtype = torch.float32, + base_seed: Optional[int] = None, + ) -> torch.Tensor: + """ + High-level function for generating initial latents. + + Args: + batch_size: Total number of samples + latent_shape: Shape of a single latent (C, H, W) or (C, T, H, W) + device: Device for the tensor + dtype: Data type for the tensor + + Returns: + Latent tensor [batch_size, *latent_shape] + """ + if base_seed is not None: + generator = torch.Generator(device=device) + generator.manual_seed(int(base_seed) % (MAX_TORCH_SEED + 1)) + return torch.randn( + batch_size, + *latent_shape, + device=device, + dtype=dtype, + generator=generator, + ) + return torch.randn( + batch_size, + *latent_shape, + device=device, + dtype=dtype, + ) + + def diffuse_with_grad( + self, + conditions: WAN21Conditions, + *, + schedule: torch.Tensor, + params: DiffusionSamplingParams, + initial_latents: Optional[torch.Tensor] = None, + ) -> DiffuseWithGradResult: + """Differentiable WAN 2.1 T2V sampling for REFL-style BPTT training. + + Returns :class:`DiffuseWithGradResult` with the live-grad + ``z_final`` + per-sample ``kl_loss`` ``[B]``. + + BPTT knobs (read from ``params.sampler_kwargs``): + + - ``mid_timestep`` (int, default 0): step index at which the forward + switches from ``torch.no_grad`` to grad-enabled — implements + truncated BPTT (only the last ``T - mid_timestep`` steps + participate in backward). DRaFT-1 sets ``mid_timestep = T - 1``. + - ``final_timestep`` (int, default ``num_inference_steps - 1``): + early stop. The loop breaks once ``i >= final_timestep``. + - ``kl_weight`` (float, default 0.0): when non-zero, per-step KL + ``(pred - ref_pred)**2 / (2 * sigma**2)`` is reduced per sample + and accumulated into the ``[B]`` ``kl_loss``. Per-sample (not a + scalar) so DP_SCATTER merge/re-shard round-trips each shard's own + KL. The actor multiplies it by its ``kl_weight`` at the loss site. + + ``initial_latents`` follows the same contract as :meth:`diffuse`: + when provided, used verbatim and the internal RNG path is + bypassed. + """ + if conditions.text is None or conditions.text.embeds is None: + raise ValueError("Wan21ReflDiffusionStage.diffuse_with_grad: conditions.text.embeds is None") + prompt_embeds = conditions.text.embeds + device = prompt_embeds.device + batch_size = int(prompt_embeds.shape[0]) + T = int(params.num_inference_steps) + schedule = schedule.to(device) + if int(schedule.shape[0]) != T + 1: + raise ValueError( + f"Wan21ReflDiffusionStage.diffuse_with_grad: schedule length {schedule.shape[0]} != T+1={T + 1}" + ) + self.strategy.init_schedule(schedule) + + latent_shape = self._latent_shape( + num_frames=int(params.num_frames), + height=int(params.height), + width=int(params.width), + ) + if initial_latents is not None: + if int(initial_latents.shape[0]) != batch_size: + raise ValueError( + f"Wan21ReflDiffusionStage.diffuse_with_grad: initial_latents.shape[0]=" + f"{int(initial_latents.shape[0])} != batch_size={batch_size}." + ) + if tuple(initial_latents.shape[1:]) != tuple(latent_shape): + raise ValueError( + f"Wan21ReflDiffusionStage.diffuse_with_grad: initial_latents.shape[1:]=" + f"{tuple(initial_latents.shape[1:])} != expected {tuple(latent_shape)} " + f"for num_frames={int(params.num_frames)}, " + f"height={int(params.height)}, width={int(params.width)}." + ) + latents = initial_latents.to(device=device, dtype=self.trajectory_dtype) + else: + latents = self.generate_latents( + batch_size=batch_size, + latent_shape=latent_shape, + device=device, + dtype=self.trajectory_dtype, + base_seed=int(params.seed), + ) + + # BPTT knobs. + sk: Dict[str, Any] = dict(getattr(params, "sampler_kwargs", {}) or {}) + mid_timestep = int(sk.get("mid_timestep", 0)) + final_timestep = int(sk.get("final_timestep", T - 1)) + kl_weight = float(sk.get("kl_weight", 0.0)) + if not (0 <= mid_timestep <= final_timestep < T): + raise ValueError( + f"Wan21ReflDiffusionStage.diffuse_with_grad: require 0 <= mid_timestep <= " + f"final_timestep < num_inference_steps, got mid={mid_timestep} " + f"final={final_timestep} T={T}." + ) + + autocast_ctx = ( + torch.autocast("cuda", self.autocast_dtype) + if device.type == "cuda" and self.autocast_dtype in (torch.float16, torch.bfloat16) + else nullcontext() + ) + sigma_max = float(schedule[1].item()) if int(schedule.shape[0]) > 1 else 0.99 + + transformer = self.model.transformer + kl_total = torch.zeros(batch_size, device=device, dtype=torch.float32) + kl_steps = 0 + + guidance_scale = float(params.guidance_scale) + use_cfg = guidance_scale > 1.0 + step = self.step + if not isinstance(step, Wan21ReflDiffusionStep): + raise TypeError( + f"Wan21ReflDiffusionStage.diffuse_with_grad requires Wan21ReflDiffusionStep, got {type(step).__name__}." + ) + + for i in range(T): + sigma = schedule[i].to(device) + sigma_next = schedule[i + 1].to(device) + grad_enabled = i >= mid_timestep + + if use_cfg: + # REFL parity: the conditional branch follows the BPTT grad + # window, while the uncond/negative branch is always + # stop-grad. Batched CFG would incorrectly add a + # ``(1 - guidance_scale) * d(uncond)/dθ`` term. + cond_ctx = nullcontext() if grad_enabled else torch.no_grad() + with cond_ctx, autocast_ctx: + noise_pred_cond = step.predict_noise( + self.model, + latents, + sigma, + conditions, + branch="cond", + ) + with torch.no_grad(), autocast_ctx: + noise_pred_uncond = step.predict_noise( + self.model, + latents, + sigma, + conditions, + branch="uncond", + ) + noise_pred_cond = noise_pred_cond.float() + noise_pred_uncond = noise_pred_uncond.float() + noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond) + kl_pred = noise_pred_cond + else: + pred_ctx = nullcontext() if grad_enabled else torch.no_grad() + with pred_ctx, autocast_ctx: + noise_pred = step.predict_noise( + self.model, + latents, + sigma, + conditions, + branch="cond", + ) + noise_pred = noise_pred.float() + kl_pred = noise_pred + + # Per-step KL against the LoRA-disabled reference (REFL-style). + # When CFG is enabled, this intentionally computes KL on the + # conditional prediction before the stop-grad uncond CFG mix. + if kl_weight != 0.0 and grad_enabled: + with torch.no_grad(), autocast_ctx, adapters_disabled(transformer): + ref_pred = step.predict_noise( + self.model, + latents, + sigma, + conditions, + branch="cond", + ) + sigma_f32 = sigma.to(dtype=torch.float32) + kl_step = ((kl_pred.float() - ref_pred.float()) ** 2 / (2.0 * sigma_f32**2)).flatten(1).mean(dim=1) + kl_total = kl_total + kl_step + kl_steps += 1 + + transition_ctx = nullcontext() if grad_enabled else torch.no_grad() + with transition_ctx: + new_latents, _, _ = step.forward( + strategy=self.strategy, + noise_pred=noise_pred, + sample=latents, + sigma=sigma, + sigma_next=sigma_next, + eta=float(params.eta), + sigma_max=sigma_max, + step_index=i, + ) + latents = new_latents.to(dtype=self.trajectory_dtype) + + if i >= final_timestep: + break + + kl_loss = kl_total / max(kl_steps, 1) if kl_steps > 0 else kl_total + return DiffuseWithGradResult(z_final=latents, kl_loss=kl_loss) + + +class Wan21ReflPipeline(WAN21Pipeline): + """WAN 2.1 T2V pipeline for the REFL recipe. + + Reuses the mainline :class:`WAN21Pipeline` construction (text encode / + condition build / VAE decode / schedule) and post-swaps + ``self.diffusion`` for the REFL-flavoured stage. The parent already + validates every constructor argument and wires the strategy / step / + precision policy through — we only need to rebuild the diffusion + stage with the same underlying components. + """ + + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + old = self.diffusion + assert isinstance(old, WAN21DiffusionStage), ( + f"Wan21ReflPipeline expects parent to build WAN21DiffusionStage, got {type(old).__name__}" + ) + self.diffusion = Wan21ReflDiffusionStage( + model=old.model, + step=Wan21ReflDiffusionStep(), + strategy=old.strategy, + autocast_precision=old.autocast_dtype, + trajectory_precision=old.trajectory_dtype, + logprob_precision=old.logprob_dtype, + ) + + def build_refl_conditions( + self, + texts: Texts, + *, + images: Optional[Images] = None, + params: DiffusionSamplingParams, + ) -> WAN21Conditions: + """Full REFL conditioning: text + CFG negative + optional I2V image. + + Mirrors the condition assembly of :meth:`WAN21Pipeline.generate` + (text via the public ``build_conditions``, image slots attached with + the diffusion geometry from ``params``), plus the REFL convention + that an explicit negative prompt rides + ``params.sampler_kwargs['negative_prompt']``. + """ + guidance = float(params.guidance_scale) + negative_prompt = (params.sampler_kwargs or {}).get("negative_prompt") + negatives = ( + Texts(texts=[str(negative_prompt)] * len(texts.texts)) if negative_prompt and guidance > 1.0 else None + ) + conds = self.build_conditions(texts, negatives=negatives, guidance_scale=guidance) + + if images is not None: + if images.pixels is None or int(images.pixels.shape[0]) != len(texts.texts): + raise ValueError( + f"Wan21ReflPipeline.build_refl_conditions: image count " + f"{None if images.pixels is None else int(images.pixels.shape[0])} " + f"!= text count {len(texts.texts)}" + ) + image_latent = WAN21ImageLatentEncodeStage( + self.bundle, + num_frames=int(params.num_frames), + height=int(params.height), + width=int(params.width), + ).encode(images) + image_embed = ( + WAN21CLIPVisionEncodeStage(self.bundle).encode(images) + if getattr(self.bundle, "uses_clip_vision", False) + else None + ) + conds = dataclasses.replace(conds, image_latent=image_latent, image_embed=image_embed) + return conds + + +__all__ = ["Wan21ReflDiffusionStep", "Wan21ReflDiffusionStage", "Wan21ReflPipeline"] diff --git a/experimental/refl/models/wan22.py b/experimental/refl/models/wan22.py new file mode 100644 index 000000000..74d865d1a --- /dev/null +++ b/experimental/refl/models/wan22.py @@ -0,0 +1,391 @@ +"""Recipe-local WAN 2.2 step + stage + pipeline for REFL BPTT.""" + +from __future__ import annotations + +import dataclasses +from contextlib import nullcontext +from typing import Any, Dict, Optional, Tuple + +import torch + +from experimental.refl.models.types import DiffuseWithGradResult +from unirl.models.wan21.clip_vision_encode import WAN21CLIPVisionEncodeStage +from unirl.models.wan21.conditions import WAN21Conditions +from unirl.models.wan21.image_encode import WAN21ImageLatentEncodeStage +from unirl.models.wan22.bundle import WAN22Bundle +from unirl.models.wan22.diffusion import WAN22DiffusionStage, WAN22DiffusionStep +from unirl.models.wan22.pipeline import WAN22Pipeline +from unirl.train.lora import adapters_disabled +from unirl.types.primitives import Images, Texts +from unirl.types.sampling import DiffusionSamplingParams + +# Matches the mainline module-level constant in unirl/models/wan22/diffusion.py. +# We do not import the private symbol so the recipe stays decoupled from the +# mainline's private surface. +_WAN_TIMESTEP_SCALE: float = 1000.0 + +# Inclusive max for torch.Generator.manual_seed and torch initial_seed conventions. +MAX_TORCH_SEED = (1 << 63) - 1 + + +class Wan22ReflDiffusionStep(WAN22DiffusionStep): + """Recipe-local WAN 2.2 single-branch denoising for REFL BPTT.""" + + def predict_noise( # pyright: ignore[reportIncompatibleMethodOverride] + self, + model: WAN22Bundle, + sample: torch.Tensor, + sigma: torch.Tensor, + conditions: WAN21Conditions, + *, + branch: str, + use_high_noise: Optional[bool] = None, + ) -> torch.Tensor: + """Run exactly one CFG branch through the currently active WAN 2.2 DiT.""" + if branch not in {"cond", "uncond"}: + raise ValueError(f"Wan22ReflDiffusionStep.predict_noise: unknown branch={branch!r}") + if conditions.text is None or conditions.text.embeds is None: + raise ValueError("Wan22ReflDiffusionStep.predict_noise: conditions.text.embeds is None") + + prompt_embeds = conditions.text.embeds + if branch == "cond": + branch_embeds = prompt_embeds + else: + neg = conditions.negative_text + branch_embeds = ( + neg.embeds if neg is not None and neg.embeds is not None else torch.zeros_like(prompt_embeds) + ) + + if use_high_noise is None: + use_high_noise, _ = self._select_for_sigma( + sigma, + 1.0, + None, + boundary_ratio=model.boundary_ratio, + ) + + batch_size = int(sample.shape[0]) + timestep = sigma * _WAN_TIMESTEP_SCALE + if timestep.dim() == 0: + timestep = timestep.expand(batch_size) + elif int(timestep.shape[0]) != batch_size: + timestep = timestep.expand(batch_size) + + embeds_dtype = branch_embeds.dtype + sample_cast = sample.to(dtype=embeds_dtype) + + image_latent = conditions.image_latent + if image_latent is not None and image_latent.latents is not None: + sample_cat = torch.cat( + [sample_cast, image_latent.latents.to(device=sample_cast.device, dtype=embeds_dtype)], + dim=1, + ) + else: + sample_cat = sample_cast + + extra: Dict[str, Any] = {} + image_embed = conditions.image_embed + image_embeds = image_embed.embeds if image_embed is not None and image_embed.embeds is not None else None + if image_embeds is not None: + extra["encoder_hidden_states_image"] = image_embeds.to(device=sample_cast.device, dtype=embeds_dtype) + + return model.transformer( + use_high_noise=bool(use_high_noise), + hidden_states=sample_cat, + encoder_hidden_states=branch_embeds, + timestep=timestep, + return_dict=False, + **extra, + )[0] + + +class Wan22ReflDiffusionStage(WAN22DiffusionStage): + """WAN 2.2 diffusion stage + REFL BPTT sampling override.""" + + def generate_latents( + self, + batch_size: int, + latent_shape: Tuple[int, ...], + device: torch.device, + dtype: torch.dtype = torch.float32, + base_seed: Optional[int] = None, + ) -> torch.Tensor: + """ + High-level function for generating initial latents. + + Args: + batch_size: Total number of samples + latent_shape: Shape of a single latent (C, H, W) or (C, T, H, W) + device: Device for the tensor + dtype: Data type for the tensor + + Returns: + Latent tensor [batch_size, *latent_shape] + """ + if base_seed is not None: + generator = torch.Generator(device=device) + generator.manual_seed(int(base_seed) % (MAX_TORCH_SEED + 1)) + return torch.randn( + batch_size, + *latent_shape, + device=device, + dtype=dtype, + generator=generator, + ) + return torch.randn( + batch_size, + *latent_shape, + device=device, + dtype=dtype, + ) + + def diffuse_with_grad( + self, + conditions: WAN21Conditions, + *, + schedule: torch.Tensor, + params: DiffusionSamplingParams, + initial_latents: Optional[torch.Tensor] = None, + ) -> DiffuseWithGradResult: + """Differentiable WAN 2.2 sampling for REFL-style BPTT training. + + Returns :class:`DiffuseWithGradResult` with the live-grad + ``z_final`` + per-sample ``kl_loss`` ``[B]``. + """ + + if conditions.text is None or conditions.text.embeds is None: + raise ValueError("Wan22ReflDiffusionStage.diffuse_with_grad: conditions.text.embeds is None") + prompt_embeds = conditions.text.embeds + device = prompt_embeds.device + batch_size = int(prompt_embeds.shape[0]) + T = int(params.num_inference_steps) + schedule = schedule.to(device) + if int(schedule.shape[0]) != T + 1: + raise ValueError( + f"Wan22ReflDiffusionStage.diffuse_with_grad: schedule length {schedule.shape[0]} != T+1={T + 1}" + ) + self.strategy.init_schedule(schedule) + + latent_shape = self._latent_shape( + num_frames=int(params.num_frames), + height=int(params.height), + width=int(params.width), + ) + if initial_latents is not None: + if int(initial_latents.shape[0]) != batch_size: + raise ValueError( + f"Wan22ReflDiffusionStage.diffuse_with_grad: initial_latents.shape[0]=" + f"{int(initial_latents.shape[0])} != batch_size={batch_size}." + ) + if tuple(initial_latents.shape[1:]) != tuple(latent_shape): + raise ValueError( + f"Wan22ReflDiffusionStage.diffuse_with_grad: initial_latents.shape[1:]=" + f"{tuple(initial_latents.shape[1:])} != expected {tuple(latent_shape)} " + f"for num_frames={int(params.num_frames)}, " + f"height={int(params.height)}, width={int(params.width)}." + ) + latents = initial_latents.to(device=device, dtype=self.trajectory_dtype) + else: + latents = self.generate_latents( + batch_size=batch_size, + latent_shape=latent_shape, + device=device, + dtype=self.trajectory_dtype, + base_seed=int(params.seed), + ) + + # BPTT knobs. + sk: Dict[str, Any] = dict(getattr(params, "sampler_kwargs", {}) or {}) + mid_timestep = int(sk.get("mid_timestep", 0)) + final_timestep = int(sk.get("final_timestep", T - 1)) + kl_weight = float(sk.get("kl_weight", 0.0)) + if not (0 <= mid_timestep <= final_timestep < T): + raise ValueError( + f"Wan22ReflDiffusionStage.diffuse_with_grad: require 0 <= mid_timestep <= " + f"final_timestep < num_inference_steps, got mid={mid_timestep} " + f"final={final_timestep} T={T}." + ) + + autocast_ctx = ( + torch.autocast("cuda", self.autocast_dtype) + if device.type == "cuda" and self.autocast_dtype in (torch.float16, torch.bfloat16) + else nullcontext() + ) + sigma_max = float(schedule[1].item()) if int(schedule.shape[0]) > 1 else 0.99 + + guidance_scale = float(params.guidance_scale) + guidance_scale_2 = ( + params.guidance_scale_2 if params.guidance_scale_2 is not None else self.model.guidance_scale_2 + ) + boundary_ratio = float(self.model.boundary_ratio) + step = self.step + if not isinstance(step, Wan22ReflDiffusionStep): + raise TypeError( + f"Wan22ReflDiffusionStage.diffuse_with_grad requires Wan22ReflDiffusionStep, got {type(step).__name__}." + ) + + dual = self.model.transformer + kl_total = torch.zeros(batch_size, device=device, dtype=torch.float32) + kl_steps = 0 + + for i in range(T): + sigma = schedule[i].to(device) + sigma_next = schedule[i + 1].to(device) + + use_high_noise, active_guidance = step._select_for_sigma( + sigma, + guidance_scale, + guidance_scale_2, + boundary_ratio=boundary_ratio, + ) + use_cfg = active_guidance > 1.0 + grad_enabled = i >= mid_timestep + + if use_cfg: + cond_ctx = nullcontext() if grad_enabled else torch.no_grad() + with cond_ctx, autocast_ctx: + noise_pred_cond = step.predict_noise( + self.model, + latents, + sigma, + conditions, + branch="cond", + use_high_noise=use_high_noise, + ) + with torch.no_grad(), autocast_ctx: + noise_pred_uncond = step.predict_noise( + self.model, + latents, + sigma, + conditions, + branch="uncond", + use_high_noise=use_high_noise, + ) + noise_pred_cond = noise_pred_cond.float() + noise_pred_uncond = noise_pred_uncond.float() + noise_pred = noise_pred_uncond + active_guidance * (noise_pred_cond - noise_pred_uncond) + kl_pred = noise_pred_cond + else: + pred_ctx = nullcontext() if grad_enabled else torch.no_grad() + with pred_ctx, autocast_ctx: + noise_pred = step.predict_noise( + self.model, + latents, + sigma, + conditions, + branch="cond", + use_high_noise=use_high_noise, + ) + noise_pred = noise_pred.float() + kl_pred = noise_pred + + # Per-step KL against the LoRA-disabled reference, on the + # conditional prediction of the currently active branch. When CFG + # is enabled, this intentionally excludes the stop-grad uncond mix. + if kl_weight != 0.0 and grad_enabled: + with torch.no_grad(), autocast_ctx, adapters_disabled(dual): + ref_pred = step.predict_noise( + self.model, + latents, + sigma, + conditions, + branch="cond", + use_high_noise=use_high_noise, + ) + sigma_f32 = sigma.to(dtype=torch.float32) + kl_step = ((kl_pred.float() - ref_pred.float()) ** 2 / (2.0 * sigma_f32**2)).flatten(1).mean(dim=1) + kl_total = kl_total + kl_step + kl_steps += 1 + + # Keep the transition outside autocast and under the same BPTT + # window used by WAN21: grad-enabled for train steps, no-grad + # before ``mid_timestep``. + transition_ctx = nullcontext() if grad_enabled else torch.no_grad() + with transition_ctx: + new_latents, _, _ = step.forward( + strategy=self.strategy, + noise_pred=noise_pred, + sample=latents, + sigma=sigma, + sigma_next=sigma_next, + eta=float(params.eta), + sigma_max=sigma_max, + step_index=i, + ) + latents = new_latents.to(dtype=self.trajectory_dtype) + + if i >= final_timestep: + break + + kl_loss = kl_total / max(kl_steps, 1) if kl_steps > 0 else kl_total + return DiffuseWithGradResult(z_final=latents, kl_loss=kl_loss) + + +class Wan22ReflPipeline(WAN22Pipeline): + """WAN 2.2 pipeline for the REFL recipe (Scheme A: post-swap ``self.diffusion``).""" + + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + old = self.diffusion + assert isinstance(old, WAN22DiffusionStage), ( + f"Wan22ReflPipeline expects parent to build WAN22DiffusionStage, got {type(old).__name__}" + ) + self.diffusion = Wan22ReflDiffusionStage( + model=old.model, + step=Wan22ReflDiffusionStep(), + strategy=old.strategy, + autocast_precision=old.autocast_dtype, + trajectory_precision=old.trajectory_dtype, + logprob_precision=old.logprob_dtype, + ) + + def build_refl_conditions( + self, + texts: Texts, + *, + images: Optional[Images] = None, + params: DiffusionSamplingParams, + ) -> WAN21Conditions: + """Full REFL conditioning: text + CFG negative + optional I2V image. + + Mirrors the condition assembly of :meth:`WAN22Pipeline.generate` — + the negative branch is encoded against the **effective** guidance + ``max(guidance_scale, guidance_scale_2)`` (WAN22 routes CFG by sigma + across the low/high-noise DiTs), and image slots are attached with + the diffusion geometry from ``params``. An explicit negative prompt + rides ``params.sampler_kwargs['negative_prompt']``. + """ + primary_g = float(params.guidance_scale) + low_g = float(params.guidance_scale_2) if params.guidance_scale_2 is not None else primary_g + effective_guidance = max(primary_g, low_g) + negative_prompt = (params.sampler_kwargs or {}).get("negative_prompt") + negatives = ( + Texts(texts=[str(negative_prompt)] * len(texts.texts)) + if negative_prompt and effective_guidance > 1.0 + else None + ) + conds = self.build_conditions(texts, negatives=negatives, guidance_scale=effective_guidance) + + if images is not None: + if images.pixels is None or int(images.pixels.shape[0]) != len(texts.texts): + raise ValueError( + f"Wan22ReflPipeline.build_refl_conditions: image count " + f"{None if images.pixels is None else int(images.pixels.shape[0])} " + f"!= text count {len(texts.texts)}" + ) + image_latent = WAN21ImageLatentEncodeStage( + self.bundle, + num_frames=int(params.num_frames), + height=int(params.height), + width=int(params.width), + ).encode(images) + image_embed = ( + WAN21CLIPVisionEncodeStage(self.bundle).encode(images) + if getattr(self.bundle, "uses_clip_vision", False) + else None + ) + conds = dataclasses.replace(conds, image_latent=image_latent, image_embed=image_embed) + return conds + + +__all__ = ["Wan22ReflDiffusionStep", "Wan22ReflDiffusionStage", "Wan22ReflPipeline"] diff --git a/experimental/refl/reward/__init__.py b/experimental/refl/reward/__init__.py new file mode 100644 index 000000000..8eedd85f6 --- /dev/null +++ b/experimental/refl/reward/__init__.py @@ -0,0 +1 @@ +"""Recipe-local reward implementations.""" diff --git a/experimental/refl/reward/face/__init__.py b/experimental/refl/reward/face/__init__.py new file mode 100644 index 000000000..d4b9c0826 --- /dev/null +++ b/experimental/refl/reward/face/__init__.py @@ -0,0 +1,5 @@ +"""Recipe-local face reward for REFL.""" + +from .scorer import FaceRewardScorer, FaceRewardSpec + +__all__ = ["FaceRewardScorer", "FaceRewardSpec"] diff --git a/experimental/refl/reward/face/face_tools.py b/experimental/refl/reward/face/face_tools.py new file mode 100644 index 000000000..eb52d52d0 --- /dev/null +++ b/experimental/refl/reward/face/face_tools.py @@ -0,0 +1,421 @@ +"""A simple, flexible implementation of a face analysis tool.""" + +import math +import os + +import onnx +import torch +import torch.nn.functional as F +import torchvision.ops as ops +from onnx2torch import convert +from skimage import transform as trans +from torchvision.transforms.functional import resize + +arcface_dst = torch.tensor( + [[38.2946, 51.6963], [73.5318, 51.5014], [56.0252, 71.7366], [41.5493, 92.3655], [70.7299, 92.2041]] +).float() + + +def distance2bbox(points, distance, max_shape=None): + """Decode distance prediction to bounding box.""" + x1 = points[:, 0] - distance[:, 0] + y1 = points[:, 1] - distance[:, 1] + x2 = points[:, 0] + distance[:, 2] + y2 = points[:, 1] + distance[:, 3] + if max_shape is not None: + x1 = x1.clamp(min=0, max=max_shape[1]) + y1 = y1.clamp(min=0, max=max_shape[0]) + x2 = x2.clamp(min=0, max=max_shape[1]) + y2 = y2.clamp(min=0, max=max_shape[0]) + return torch.stack([x1, y1, x2, y2], axis=-1) + + +def distance2kps(points, distance, max_shape=None): + """Decode distance prediction to keypoints.""" + preds = [] + for i in range(0, distance.shape[1], 2): + px = points[:, i % 2] + distance[:, i] + py = points[:, i % 2 + 1] + distance[:, i + 1] + if max_shape is not None: + px = px.clamp(min=0, max=max_shape[1]) + py = py.clamp(min=0, max=max_shape[0]) + preds.append(px) + preds.append(py) + return torch.stack(preds, axis=-1) + + +def face_transform(data, center, output_size, scale, rotation, device): + def to_homogeneous(mat): + return torch.vstack([mat, torch.tensor([0.0, 0.0, 1.0])]) + + scale_ratio = scale + rot = float(rotation) * math.pi / 180.0 + cx = center[0] * scale_ratio + cy = center[1] * scale_ratio + + C, H, W = data.shape + + t1 = to_homogeneous(torch.tensor([[scale_ratio, 0, 0], [0, scale_ratio, 0]])).float() + t2 = to_homogeneous(torch.tensor([[1, 0, -cx], [0, 1, -cy]])).float() + cos_theta = math.cos(rot) + sin_theta = math.sin(rot) + t3 = to_homogeneous(torch.tensor([[cos_theta, -sin_theta, 0], [sin_theta, cos_theta, 0]])).float() + t4 = to_homogeneous(torch.tensor([[1, 0, output_size / 2], [0, 1, output_size / 2]])).float() + M_homogeneous = t4 @ t3 @ t2 @ t1 + M = M_homogeneous[:2, :] + T = torch.tensor([[2 / W, 0, -1], [0, 2 / H, -1], [0, 0, 1]]) + theta = torch.inverse(T @ M_homogeneous @ torch.inverse(T)) + theta = theta[:2, :].unsqueeze(0).to(device) + grid = F.affine_grid(theta, data.unsqueeze(0).size(), align_corners=True) + transformed = F.grid_sample(data.unsqueeze(0), grid, align_corners=True) + cropped = transformed[0] + cropped = cropped[:, :output_size, :output_size] + return cropped.unsqueeze(0), M + + +def trans_points2d(pts, M): + ones = torch.ones((pts.shape[0], 1), dtype=pts.dtype, device=pts.device) + points_hom = torch.cat([pts, ones], dim=1) + points_hom = points_hom.unsqueeze(-1) + transformed_hom = torch.matmul(M, points_hom) + transformed = transformed_hom[:, :2, :].squeeze(-1) + return transformed + + +def estimate_norm(lmk, image_size=112, mode="arcface"): + assert lmk.shape == (5, 2) + assert image_size % 112 == 0 or image_size % 128 == 0 + if image_size % 112 == 0: + ratio = float(image_size) / 112.0 + diff_x = 0 + else: + ratio = float(image_size) / 128.0 + diff_x = 8.0 * ratio + dst = arcface_dst * ratio + dst[:, 0] += diff_x + tform = trans.SimilarityTransform() + tform.estimate(lmk, dst) + M = torch.from_numpy(tform.params).float() + return M + + +def norm_crop(img, landmark, image_size=112, mode="arcface"): + """Align an image into ArcFace canonical 112x112 using a similarity transform. + + Differentiable in ``img`` — the only non-grad path is ``landmark``, which + comes from the no-grad SCRFD detector. This is what carries the REFL + gradient from cosine reward all the way back to the generated pixels. + """ + M_homogeneous = estimate_norm(landmark, image_size, mode) + C, H, W = img.shape + img = img.unsqueeze(0) + T = torch.tensor([[2 / W, 0, -1], [0, 2 / H, -1], [0, 0, 1]]) + T_inv = torch.inverse(T) + theta = torch.inverse(T @ M_homogeneous @ T_inv) + theta = theta[:2, :].unsqueeze(0).to(img.device) + grid = F.affine_grid(theta, img.size(), align_corners=True) + transformed = F.grid_sample(img, grid, align_corners=True) + cropped = transformed[0] + warped = cropped[:, :image_size, :image_size] + return warped + + +def invert_affine_transform(matrix): + L = matrix[..., :2] + T = matrix[..., 2:] + a, b = L[..., 0, 0], L[..., 0, 1] + c, d = L[..., 1, 0], L[..., 1, 1] + det = a * d - b * c + inv_det = 1.0 / det + inv_L = torch.stack( + [torch.stack([d * inv_det, -b * inv_det], dim=-1), torch.stack([-c * inv_det, a * inv_det], dim=-1)], dim=-2 + ) + inv_T = -torch.matmul(inv_L, T) + inv_matrix = torch.cat([inv_L, inv_T], dim=-1) + return inv_matrix + + +class Face(dict): + def __init__(self, d=None, **kwargs): + if d is None: + d = {} + if kwargs: + d.update(**kwargs) + for k, v in d.items(): + setattr(self, k, v) + + def __setattr__(self, name, value): + if isinstance(value, (list, tuple)): + value = [self.__class__(x) if isinstance(x, dict) else x for x in value] + elif isinstance(value, dict) and not isinstance(value, self.__class__): + value = self.__class__(value) + super(Face, self).__setattr__(name, value) + super(Face, self).__setitem__(name, value) + + __setitem__ = __setattr__ + + def __getattr__(self, name): + return None + + @property + def embedding_norm(self): + if self.embedding is None: + return None + return torch.norm(self.embedding) + + @property + def normed_embedding(self): + if self.embedding is None: + return None + return self.embedding / self.embedding_norm + + +class SCRFD: + def __init__(self, model_file=None, device="cuda"): + self.model_file = model_file + self.device = device + self.center_cache = {} + model = onnx.load(self.model_file) + self.torch_model = convert(model) + self.torch_model.eval() + self.torch_model.requires_grad_(False) + self.torch_model.to(self.device) + self.use_kps = True + self.fmc = 3 + self._num_anchors = 2 + self._feat_stride_fpn = [8, 16, 32] + self.input_size = (640, 640) + + def forward(self, det_img, threshold=0.5): + input_height = det_img.shape[2] + input_width = det_img.shape[3] + scores_list = [] + bboxes_list = [] + kpss_list = [] + net_outs = self.torch_model(det_img.float()) + + for idx, stride in enumerate(self._feat_stride_fpn): + scores = net_outs[idx].cpu() + bbox_preds = net_outs[idx + self.fmc].cpu() + bbox_preds = bbox_preds * stride + if self.use_kps: + kps_preds = net_outs[idx + self.fmc * 2].cpu() * stride + + height = input_height // stride + width = input_width // stride + key = (height, width, stride) + if key in self.center_cache: + anchor_centers = self.center_cache[key] + else: + rows = torch.arange(height) + cols = torch.arange(width) + grid_y, grid_x = torch.meshgrid(rows, cols, indexing="ij") + anchor_centers = torch.stack([grid_x, grid_y], dim=-1).float() + anchor_centers = (anchor_centers * stride).reshape((-1, 2)) + if self._num_anchors > 1: + anchor_centers = torch.stack([anchor_centers] * self._num_anchors, axis=1).reshape((-1, 2)) + if len(self.center_cache) < 100: + self.center_cache[key] = anchor_centers + + scores = scores.reshape(-1, 1) + bbox_preds = bbox_preds.reshape(-1, 4) + if self.use_kps: + kps_preds = kps_preds.reshape(-1, kps_preds.shape[-1]) + pos_mask = scores[:, 0] >= threshold + bboxes = distance2bbox(anchor_centers, bbox_preds) + pos_scores = scores[pos_mask] + pos_bboxes = bboxes[pos_mask] + scores_list.append(pos_scores) + bboxes_list.append(pos_bboxes) + if self.use_kps: + kpss = distance2kps(anchor_centers, kps_preds) + kpss = kpss.reshape((kpss.shape[0], -1, 2)) + pos_kpss = kpss[pos_mask] + kpss_list.append(pos_kpss) + + return scores_list, bboxes_list, kpss_list + + @torch.no_grad() + def detect(self, image, input_size=None, max_num=0, metric="default", nms_thresh=0.4, det_thresh=0.5): + assert input_size is not None or self.input_size is not None + input_size = self.input_size if input_size is None else input_size + + im_ratio = float(image.shape[1]) / image.shape[2] + model_ratio = float(input_size[1]) / input_size[0] + if im_ratio > model_ratio: + new_height = input_size[1] + new_width = int(new_height / im_ratio) + else: + new_width = input_size[0] + new_height = int(new_width * im_ratio) + det_scale = float(new_height) / image.shape[1] + resized_img = resize(image, (new_height, new_width), antialias=False) + det_img = torch.zeros((3, input_size[1], input_size[0]), device=self.device) + det_img[:, :new_height, :new_width] = resized_img + det_img = det_img.unsqueeze(0) + scores_list, bboxes_list, kpss_list = self.forward(det_img, det_thresh) + + scores = torch.vstack(scores_list) + scores_ravel = scores.flatten() + order = torch.argsort(scores_ravel, descending=True) + bboxes = torch.vstack(bboxes_list) / det_scale + if self.use_kps: + kpss = torch.vstack(kpss_list) / det_scale + + pre_det = torch.cat((bboxes, scores), dim=1).float() + pre_det = pre_det[order] + keep = self.nms(pre_det, nms_thresh) + det = pre_det[keep, :] + + if self.use_kps: + kpss = kpss[order, :, :] + kpss = kpss[keep, :, :] + else: + kpss = None + return det, kpss + + def nms(self, dets, nms_thresh): + boxes = dets[:, :4] + scores = dets[:, 4] + keep = ops.nms(boxes, scores, iou_threshold=nms_thresh) + return keep.tolist() + + +class ArcFace: + def __init__(self, model_file=None, device="cuda"): + self.model_file = model_file + self.device = device + model = onnx.load(self.model_file) + self.torch_model = convert(model) + self.torch_model.eval() + self.torch_model.to(self.device) + # Frozen weights — gradient flows through inputs only. + self.torch_model.requires_grad_(False) + self.taskname = "recognition" + self.input_size = (112, 112) + + def get(self, img, face, input_size=(112, 112)): + aimg = norm_crop(img, landmark=face.kps, image_size=self.input_size[0]) + im_ratio = float(aimg.shape[1]) / aimg.shape[2] + model_ratio = float(input_size[1]) / input_size[0] + if im_ratio > model_ratio: + new_height = input_size[1] + new_width = int(new_height / im_ratio) + else: + new_width = input_size[0] + new_height = int(new_width * im_ratio) + resized_img = resize(aimg, (new_height, new_width), antialias=False) + face.embedding = self.get_feat(resized_img.unsqueeze(0)).flatten() + return face.embedding + + def compute_sim(self, feat1, feat2): + feat1 = feat1.ravel() + feat2 = feat2.ravel() + sim = torch.dot(feat1, feat2) / (torch.norm(feat1) * torch.norm(feat2)) + return sim + + def get_feat(self, imgs): + imgs = imgs[:, [2, 1, 0], :, :] + net_out = self.torch_model(imgs) + return net_out + + +class Landmark: + def __init__(self, model_file=None, device="cuda"): + self.model_file = model_file + self.device = device + model = onnx.load(self.model_file) + self.torch_model = convert(model) + self.torch_model.eval() + self.torch_model.to(device) + self.torch_model.requires_grad_(False) + self.lmk_dim = 2 + self.lmk_num = 106 + self.taskname = "landmark_%dd_%d" % (self.lmk_dim, self.lmk_num) + self.input_size = (192, 192) + + def get(self, img, face, input_size=(192, 192)): + bbox = face.bbox + w, h = (bbox[2] - bbox[0]), (bbox[3] - bbox[1]) + center = (bbox[2] + bbox[0]) / 2, (bbox[3] + bbox[1]) / 2 + rotate = 0 + _scale = self.input_size[0] / (max(w, h) * 1.5) + aimg, M = face_transform(img, center, self.input_size[0], _scale, rotate, img.device) + aimg = (aimg + 1) / 2 * 255.0 + aimg = aimg[:, [2, 1, 0], :, :] + + input_size = self.input_size if input_size is None else input_size + im_ratio = float(aimg.shape[2]) / aimg.shape[3] + model_ratio = float(input_size[1]) / input_size[0] + if im_ratio > model_ratio: + new_height = input_size[1] + new_width = int(new_height / im_ratio) + else: + new_width = input_size[0] + new_height = int(new_width * im_ratio) + resized_img = resize(aimg, (new_height, new_width), antialias=False) + det_img = torch.zeros((aimg.shape[0], 3, input_size[1], input_size[0]), device=self.device) + det_img[:, :, :new_height, :new_width] = resized_img + + pred = self.torch_model(det_img)[0] + pred = pred.reshape((-1, 2)) + if self.lmk_num < pred.shape[0]: + pred = pred[self.lmk_num * -1 :, :] + pred[:, 0:2] += 1 + pred[:, 0:2] *= self.input_size[0] // 2 + + IM = invert_affine_transform(M).to(img.device) + pred = trans_points2d(pred, IM) + face[self.taskname] = pred + return pred + + +class FaceAnalysis: + def __init__(self, root="~/.insightface", device="cuda"): + self.root = root + self.device = device + self.detection_root = os.path.join(root, "scrfd_10g_bnkps.onnx") + self.landmark_root = os.path.join(root, "2d106det.onnx") + self.arcface_root = os.path.join(root, "glintr100.onnx") + self.detection_model = SCRFD(self.detection_root, self.device) + self.landmark_model = Landmark(self.landmark_root, self.device) + self.arcface_model = ArcFace(self.arcface_root, self.device) + + def landmark_loss(self, id_landmark=None, gt_landmark=None, mask=None): + mask = mask.unsqueeze(-1).unsqueeze(-1) + error = torch.abs(id_landmark - gt_landmark) * mask + valid_frame_count = mask.sum() + 1e-8 + loss = error.sum() / valid_frame_count / id_landmark.shape[-2] + return loss + + def embedding_loss(self, id_embedding=None, gt_embedding=None, mask=None): + cos_sim = F.cosine_similarity(id_embedding, gt_embedding, dim=2) + cos_loss = (1 - cos_sim) * mask + valid_frame_count = mask.sum() + 1e-8 + loss = cos_loss.sum() / valid_frame_count + return loss + + def pool_embedding_loss(self, id_embedding=None, gt_embedding=None, id_mask=None): + """Pool-style cosine similarity between gen frames and any ref frame. + + Returns a scalar reward (per-call). The scorer calls this once per + sample and stacks the result into a per-sample reward tensor. + """ + id_emb_expanded = id_embedding.unsqueeze(2) + gt_emb_expanded = gt_embedding.unsqueeze(1) + gt_mask = torch.ones(gt_embedding.shape[0], gt_embedding.shape[1]).to(id_mask.device) + if gt_mask.shape[1] > 1: + gt_mask[:, 0] = 0 + is_all_zero = (gt_embedding == 0).all(dim=-1) + gt_mask[is_all_zero] = 0 + + cos_sim_all = F.cosine_similarity(id_emb_expanded, gt_emb_expanded, dim=3) + valid_mask = id_mask.unsqueeze(2) * gt_mask.unsqueeze(1) + + gt_valid_count = gt_mask.sum(dim=1) + 1e-8 + weight_matrix = valid_mask / (gt_valid_count.unsqueeze(1).unsqueeze(2) + 1e-8) + mean_similarities = (cos_sim_all * weight_matrix).sum(dim=2) + cos_loss = mean_similarities * id_mask + valid_frame_count = id_mask.sum() + 1e-8 + loss = cos_loss.sum() / valid_frame_count + return loss diff --git a/experimental/refl/reward/face/requirements.txt b/experimental/refl/reward/face/requirements.txt new file mode 100644 index 000000000..329432f90 --- /dev/null +++ b/experimental/refl/reward/face/requirements.txt @@ -0,0 +1,6 @@ +imageio>=2.31 +imageio-ffmpeg>=0.4 +onnx>=1.14 +onnx2torch>=1.5 +scikit-image>=0.21 +scipy>=1.11 diff --git a/experimental/refl/reward/face/scorer.py b/experimental/refl/reward/face/scorer.py new file mode 100644 index 000000000..c1e28a285 --- /dev/null +++ b/experimental/refl/reward/face/scorer.py @@ -0,0 +1,315 @@ +"""Face similarity reward scorer (REFL-compatible / BPTT-differentiable).""" + +from __future__ import annotations + +import logging +from collections import OrderedDict +from dataclasses import dataclass +from typing import List, Optional + +import torch + +from unirl.reward.base import BaseRewardComponentSpec +from unirl.reward.local.base import LocalRewardBackend +from unirl.reward.local.device import resolve_device +from unirl.types.reward import RewardRequest + +from .face_tools import Face, FaceAnalysis + +logger = logging.getLogger(__name__) + +# Reference-embedding LRU bound (see FaceRewardScorer._ref_cache). +_REF_CACHE_MAX = 64 + + +# --------------------------------------------------------------------------- +# Reference-video loader (imageio + resize-to-cover + center_crop) +# --------------------------------------------------------------------------- + + +def _load_ref_video_frames( + video_path: str, + *, + max_frames: int = 81, + max_pixels: int = 480 * 480, + height_div: int = 16, + width_div: int = 16, +) -> torch.Tensor: + """Load a reference video as ``(C, T, H, W)`` float32 in ``[-1, 1]``. + + Frame-for-frame preprocessing so the reference-side stays aligned with the + reference-loader used to train the REFL data pipeline. + """ + import imageio + import PIL.Image + import torchvision.transforms.functional as TF + + reader = imageio.get_reader(video_path) + total = int(reader.count_frames()) # type: ignore[attr-defined] + + # time_division_factor=4, remainder=1 + nf = min(max_frames, total) + while nf > 1 and nf % 4 != 1: + nf -= 1 + + first = PIL.Image.fromarray(reader.get_data(0)) + fw, fh = first.size + if fw * fh > max_pixels: + scale = (fw * fh / max_pixels) ** 0.5 + fh, fw = int(fh / scale), int(fw / scale) + target_h = fh // height_div * height_div + target_w = fw // width_div * width_div + + frames = [] + for i in range(nf): + frame = PIL.Image.fromarray(reader.get_data(i)) + iw, ih = frame.size + scale = max(target_w / iw, target_h / ih) + frame = TF.resize( + frame, + (round(ih * scale), round(iw * scale)), + interpolation=TF.InterpolationMode.BILINEAR, + ) + frame = TF.center_crop(frame, (target_h, target_w)) + t = TF.to_tensor(frame) # (C, H, W) [0, 1] + frame = TF.to_pil_image(t) # round-trip PIL to match training-time preprocessing + frames.append(TF.to_tensor(frame) * 2.0 - 1.0) + reader.close() + + # (T, C, H, W) -> (C, T, H, W) + return torch.stack(frames).permute(1, 0, 2, 3).contiguous() + + +# --------------------------------------------------------------------------- +# Reward scorer +# --------------------------------------------------------------------------- + + +class FaceRewardScorer(LocalRewardBackend): + """SCRFD detection + ArcFace embedding + pool cosine similarity for REFL.""" + + canonical_model_name = "face" + input_kind = "video" + + def __init__(self, *, config: "FaceRewardSpec", base_device: str) -> None: + super().__init__( + model_name=self.canonical_model_name, + device=resolve_device(config.device, base_device), + batch_size=int(config.batch_size), + model_path=config.model_path, + image_size=config.image_size, + ref_max_frames=config.ref_max_frames, + ref_max_pixels=config.ref_max_pixels, + differentiable=config.differentiable, + ) + + def _load_model(self) -> None: + model_path = self.model_kwargs["model_path"] + self.model = FaceAnalysis(root=model_path, device=self.device) + self._image_size = int(self.model_kwargs.get("image_size", 112)) + self._ref_max_frames = int(self.model_kwargs.get("ref_max_frames", 81)) + self._ref_max_pixels = int(self.model_kwargs.get("ref_max_pixels", 480 * 480)) + self._differentiable = bool(self.model_kwargs.get("differentiable", True)) + # Bounded LRU: reference embeddings are small, but a many-identity + # dataset must not grow GPU residency without bound. + self._ref_cache: OrderedDict[str, tuple[torch.Tensor, torch.Tensor]] = OrderedDict() + + def _compute_model_rewards(self, request: RewardRequest) -> List[float]: + raise NotImplementedError("FaceRewardScorer is REFL-only; use compute_rewards_differentiable().") + + # ------------------------------------------------------------------ + # Per-video face embedding (recipe-local REFL implementation) + # ------------------------------------------------------------------ + + def _extract_face_embeddings( + self, + video: torch.Tensor, + *, + with_grad: bool, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Extract per-frame face embeddings. + + Args: + video: ``(C, T, H, W)`` float in ``[-1, 1]``. + with_grad: When True, the ArcFace forward is kept inside autograd + (used for the GENERATED video on the BPTT path); when False, + the whole computation runs under ``torch.no_grad`` (used for + the REFERENCE video, which never receives gradient). + + Returns: + embeddings: ``(1, T, 512)`` on ``self.device``. Differentiable in + ``video`` when ``with_grad=True``. + mask: ``(1, T)`` on ``self.device`` (1 = face found). + """ + fa = self.model + fa.detection_model.torch_model.to(self.device) + fa.arcface_model.torch_model.to(self.device) + + T = int(video.shape[1]) + embeddings: List[torch.Tensor] = [] + mask: List[int] = [] + zero_emb = None + + autograd_ctx = torch.enable_grad if with_grad else torch.no_grad + + for t in range(T): + frame = video[:, t].float().to(self.device) + bboxes, kpss = fa.detection_model.detect(frame) + if bboxes.shape[0] > 0: + indexed = [(i, x) for i, x in enumerate(bboxes)] + sorted_bboxes = sorted( + indexed, + key=lambda item: (item[1][2] - item[1][0]) * (item[1][3] - item[1][1]), + ) + max_index, _ = sorted_bboxes[-1] + face = Face( + bbox=bboxes[max_index][0:4], + kps=kpss[max_index], + det_score=bboxes[max_index][4], + ) + with autograd_ctx(): + emb = fa.arcface_model.get(frame, face) + embeddings.append(emb) + mask.append(1) + else: + if zero_emb is None: + zero_emb = torch.zeros(512, device=self.device) + embeddings.append(zero_emb) + mask.append(0) + + emb_stack = torch.stack(embeddings).unsqueeze(0) # (1, T, 512) + mask_tensor = torch.tensor(mask, device=self.device).unsqueeze(0).float() # (1, T) + return emb_stack, mask_tensor + + # ------------------------------------------------------------------ + # Reference-video caching + # ------------------------------------------------------------------ + + def _get_ref_embeddings(self, ref_video_path: str) -> tuple[torch.Tensor, torch.Tensor]: + cached = self._ref_cache.get(ref_video_path) + if cached is not None: + self._ref_cache.move_to_end(ref_video_path) + return cached + with torch.no_grad(): + video = _load_ref_video_frames( + ref_video_path, + max_frames=self._ref_max_frames, + max_pixels=self._ref_max_pixels, + ).to(self.device) + ref_emb, ref_mask = self._extract_face_embeddings(video, with_grad=False) + # Detach (defence-in-depth) before caching. + cached = (ref_emb.detach(), ref_mask.detach()) + self._ref_cache[ref_video_path] = cached + while len(self._ref_cache) > _REF_CACHE_MAX: + self._ref_cache.popitem(last=False) + return cached + + # ------------------------------------------------------------------ + # Differentiable REFL reward entry point + # ------------------------------------------------------------------ + + def compute_rewards_differentiable( + self, + media_tensor: torch.Tensor, + prompts: List[str], + records: Optional[List[dict]] = None, + ) -> torch.Tensor: + """Score generated videos ``[B,C,T,H,W]`` against reference-video metadata.""" + if media_tensor.ndim != 5: + raise ValueError(f"FaceRewardScorer expects [B,C,T,H,W], got {tuple(media_tensor.shape)}") + B = int(media_tensor.shape[0]) + if len(prompts) != B: + raise ValueError(f"FaceRewardScorer: prompts length {len(prompts)} != batch size {B}.") + metadata = records or [None] * B + if len(metadata) != B: + raise ValueError(f"FaceRewardScorer: records length {len(metadata)} != batch size {B}.") + + rewards: List[torch.Tensor] = [] + for i in range(B): + md = metadata[i] or {} + ref_path: Optional[str] = md.get("ref_video_path") or md.get("ref_video") + if not ref_path: + raise ValueError( + f"FaceRewardScorer: sample {i} is missing metadata['ref_video_path']. " + "Make sure the JSONL row was written with absolute ref video paths." + ) + + ref_emb, ref_mask = self._get_ref_embeddings(str(ref_path)) + + gen_video = torch.clamp(media_tensor[i], -1, 1).to(self.device) + gen_emb, gen_mask = self._extract_face_embeddings(gen_video, with_grad=self._differentiable) + + if int(gen_mask.sum().item()) == 0: + rewards.append(gen_video.sum() * 0.0) + continue + + score = self.model.pool_embedding_loss( + id_embedding=gen_emb, + gt_embedding=ref_emb, + id_mask=gen_mask, + ) + rewards.append(score) + + return torch.stack(rewards).float() + + # ------------------------------------------------------------------ + # Lifecycle hooks + # ------------------------------------------------------------------ + + def offload(self) -> None: + self._ref_cache.clear() + fa = self.model + if fa is not None: + for sub in (fa.detection_model, fa.landmark_model, fa.arcface_model): + if sub is not None and hasattr(sub.torch_model, "cpu"): + sub.torch_model.cpu() + torch.cuda.empty_cache() + + def onload(self) -> None: + fa = self.model + if fa is not None: + for sub in (fa.detection_model, fa.landmark_model, fa.arcface_model): + if sub is not None and hasattr(sub.torch_model, "to"): + sub.torch_model.to(self.device) + + def is_available(self) -> bool: + return bool(self._is_loaded) + + def dispose(self) -> None: + self.offload() + + +@dataclass +class FaceRewardSpec(BaseRewardComponentSpec): + """Typed config for :class:`FaceRewardScorer`. + + Args: + model_path: Directory containing ONNX model files + (scrfd_10g_bnkps.onnx, 2d106det.onnx, glintr100.onnx). Passed + straight to :class:`FaceAnalysis` as its model root. + device: "auto" / "cuda" / "cuda:N" — resolved against ``base_device``. + batch_size: kept for parity with sibling specs; the scorer iterates + per-sample internally because each sample needs its own ref video. + image_size: ArcFace alignment size, must be a multiple of 112 or 128. + ref_max_frames: Frame cap for ref-video decoding (default 81, matches + the REFL training-time cap). + ref_max_pixels: Pixel cap for ref-video decoding (default 480*480 — + keeps SCRFD's short-side input near its trained resolution). + differentiable: Keep autograd on the generated ArcFace forward (default + True — required for REFL). Set False to score under no_grad for + the historical GRPO/replay path. + """ + + model_path: str = "" + device: str = "auto" + batch_size: int = 1 + image_size: int = 112 + ref_max_frames: int = 81 + ref_max_pixels: int = 480 * 480 + differentiable: bool = True + + +__all__ = [ + "FaceRewardScorer", + "FaceRewardSpec", +] diff --git a/experimental/refl/reward/videoalign/__init__.py b/experimental/refl/reward/videoalign/__init__.py new file mode 100644 index 000000000..8b1b998fc --- /dev/null +++ b/experimental/refl/reward/videoalign/__init__.py @@ -0,0 +1,17 @@ +"""Recipe-local VideoAlign reward (Qwen2-VL VQ/MQ/TA scorer). + +Self-contained port of the VideoAlign reward family for the REFL WAN recipe. +The reward model code (``Qwen2VLRewardModelBT``), prompt templates, config +dataclasses and checkpoint-loading helpers all live under +:mod:`experimental.refl.reward.videoalign.model`, matching the recipe-local +layout used by the WAN22 face reward. + +Public API +---------- +- :class:`VideoAlignRewardScorer` — REFL-compatible reward backend. +- :class:`VideoAlignSpec` — typed config dataclass. +""" + +from .scorer import VideoAlignRewardScorer, VideoAlignSpec + +__all__ = ["VideoAlignRewardScorer", "VideoAlignSpec"] diff --git a/experimental/refl/reward/videoalign/model/__init__.py b/experimental/refl/reward/videoalign/model/__init__.py new file mode 100644 index 000000000..ddd5a7684 --- /dev/null +++ b/experimental/refl/reward/videoalign/model/__init__.py @@ -0,0 +1,42 @@ +"""Inference-only port of the VideoAlign reward model. + +This subpackage contains everything needed to *load and forward* the +VideoAlign Qwen2-VL reward checkpoint, but **none** of the training-time +code (no ``VideoVLMRewardTrainer``, no GSB CSV data loader, no Bradley-Terry +loss). The original training pipeline lives at the upstream +``TIGER-AI-Lab/VideoScore`` / ``Tencent-Hunyuan/VideoAlign`` repos; mmrl +ships a vendored snapshot for its own RL training. UniRL does not need any +of that — at REFL rollout time we only call ``forward`` on the reward model +to read out three scalars per (video, prompt) pair. + +Files +----- +- :mod:`prompt_template` — ``build_prompt`` + the four prompt variants used + by the published checkpoints. +- :mod:`configs` — ``ModelConfig`` / ``PEFTLoraConfig`` / + ``TrainingConfig`` (the *inference-relevant* subset of fields; loaded from + the checkpoint's ``model_config.json``). +- :mod:`reward_model` — ``Qwen2VLRewardModelBT`` — Qwen2-VL with an + ``rm_head`` linear projection to (VQ, MQ, TA) scalars. +- :mod:`checkpoint` — ``load_model_from_checkpoint`` (full / LoRA + branches, plus the transformers>=5 key-remap). +- :mod:`factory` — ``create_model_and_processor`` — builds the + model, processor and (optionally) wraps with PEFT LoRA. +""" + +from .checkpoint import load_model_from_checkpoint +from .configs import ModelConfig, PEFTLoraConfig, TrainingConfig +from .factory import create_model_and_processor +from .prompt_template import DIMENSION_DESCRIPTIONS, build_prompt +from .reward_model import Qwen2VLRewardModelBT + +__all__ = [ + "DIMENSION_DESCRIPTIONS", + "ModelConfig", + "PEFTLoraConfig", + "Qwen2VLRewardModelBT", + "TrainingConfig", + "build_prompt", + "create_model_and_processor", + "load_model_from_checkpoint", +] diff --git a/experimental/refl/reward/videoalign/model/checkpoint.py b/experimental/refl/reward/videoalign/model/checkpoint.py new file mode 100644 index 000000000..b592177ad --- /dev/null +++ b/experimental/refl/reward/videoalign/model/checkpoint.py @@ -0,0 +1,162 @@ +"""Checkpoint loader for VideoAlign reward checkpoints. + +Supports both layouts produced by the upstream trainer: + +1. **Full state dict** (``save_full_model=True``) — ``checkpoint-K/model.pth``. + Loaded directly into the freshly-constructed + :class:`Qwen2VLRewardModelBT`. + +2. **LoRA split** (the default for the published VideoAlign release) — + ``checkpoint-K/adapter_model.safetensors`` (LoRA weights) plus + ``checkpoint-K/non_lora_state_dict.pth`` (rm_head + special-token + embeddings + any other ``requires_grad`` non-LoRA tensor). Both pieces + are merged into the PEFT-wrapped model. + +Includes the transformers>=5 compatibility shim +``base_model.model.model.*`` → ``base_model.model.model.language_model.*`` +when the target model uses the new Qwen2-VL submodule layout but the +checkpoint was saved against the old one. The remap is a no-op when both +sides agree. +""" + +from __future__ import annotations + +import glob +import logging +import os +from typing import Dict, Optional, Tuple + +import safetensors.torch +import torch + +logger = logging.getLogger(__name__) + + +def _insert_adapter_name_into_state_dict( + state_dict: Dict[str, torch.Tensor], + adapter_name: str, + parameter_prefix: str, +) -> Dict[str, torch.Tensor]: + """Rewrite raw LoRA keys to match peft's ``{module}.lora_{A,B}.{adapter}.*`` layout.""" + peft_model_state_dict: Dict[str, torch.Tensor] = {} + for key, val in state_dict.items(): + if parameter_prefix in key: + suffix = key.split(parameter_prefix)[1] + if "." in suffix: + suffix_to_replace = ".".join(suffix.split(".")[1:]) + key = key.replace(suffix_to_replace, f"{adapter_name}.{suffix_to_replace}") + else: + key = f"{key}.{adapter_name}" + peft_model_state_dict[key] = val + else: + peft_model_state_dict[key] = val + return peft_model_state_dict + + +def _pick_checkpoint_dir(checkpoint_dir: str, checkpoint_step: Optional[int]) -> str: + """Return the absolute path of the ``checkpoint-`` subdir to load. + + ``checkpoint_step is None`` or ``-1`` selects the largest step found. + A specific step falls back to the latest with a warning if missing. + """ + candidates = glob.glob(os.path.join(checkpoint_dir, "checkpoint-*")) + if not candidates: + raise FileNotFoundError( + f"No 'checkpoint-*' subdirectory under {checkpoint_dir!r}. " + "Make sure you point ``reward_model_path`` at the directory that " + "contains ``model_config.json`` alongside one or more ``checkpoint-K`` dirs." + ) + candidates.sort(key=lambda x: int(x.split("-")[-1]), reverse=True) + + if checkpoint_step is None or checkpoint_step == -1: + chosen = candidates[0] + logger.info("Using latest VideoAlign checkpoint: %s", chosen) + return chosen + + explicit = os.path.join(checkpoint_dir, f"checkpoint-{checkpoint_step}") + if explicit in candidates: + logger.info("Using requested VideoAlign checkpoint: %s", explicit) + return explicit + + chosen = candidates[0] + logger.warning( + "Requested VideoAlign checkpoint-%s not found; falling back to latest %s", + checkpoint_step, + chosen, + ) + return chosen + + +def load_model_from_checkpoint( + model: torch.nn.Module, + checkpoint_dir: str, + checkpoint_step: Optional[int], +) -> Tuple[torch.nn.Module, str]: + """Load VideoAlign weights into ``model`` from a ``checkpoint-K`` subdir. + + Args: + model: an already-constructed reward model (PEFT-wrapped or not). + checkpoint_dir: parent dir containing ``checkpoint-*`` subdirs. + checkpoint_step: which step to load; ``None`` / ``-1`` → latest. + + Returns: + ``(model, step_str)`` — the same model object (mutated in place by + ``load_state_dict``) and the step number of the loaded checkpoint + as a string (e.g. ``"3000"``), for downstream logging. + """ + checkpoint_path = _pick_checkpoint_dir(checkpoint_dir, checkpoint_step) + loaded_step = checkpoint_path.split("checkpoint-")[-1].split("/")[0] + + full_ckpt = os.path.join(checkpoint_path, "model.pth") + lora_ckpt = os.path.join(checkpoint_path, "adapter_model.safetensors") + non_lora_ckpt = os.path.join(checkpoint_path, "non_lora_state_dict.pth") + + # Upstream checkpoints use the old Qwen2VL layout (LM at + # base_model.model.model.*, vision at base_model.model.visual.*); + # transformers 5.6 nests them under model.language_model / model.visual. + # Applied to every loaded dict — LoRA keys embed module paths too — and + # idempotent for already-new-layout keys. + def _remap_qwen_layout(state_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: + new_state_dict: Dict[str, torch.Tensor] = {} + for key, value in state_dict.items(): + if key.startswith("base_model.model.model.language_model.") or key.startswith( + "base_model.model.model.visual." + ): + new_state_dict[key] = value # already the new layout + elif key.startswith("base_model.model.model"): + new_state_dict["base_model.model.model.language_model" + key[len("base_model.model.model") :]] = value + elif key.startswith("base_model.model.visual"): + new_state_dict["base_model.model.model.visual" + key[len("base_model.model.visual") :]] = value + else: + new_state_dict[key] = value + return new_state_dict + + if os.path.exists(full_ckpt): + model_state_dict = torch.load(full_ckpt, map_location="cpu", weights_only=True) + model.load_state_dict(_remap_qwen_layout(model_state_dict)) + else: + # LoRA branch — merge LoRA adapter + non-LoRA tensors. + if not os.path.exists(lora_ckpt) or not os.path.exists(non_lora_ckpt): + raise FileNotFoundError( + f"Neither {full_ckpt!r} nor the LoRA pair " + f"({lora_ckpt!r} + {non_lora_ckpt!r}) was found under {checkpoint_path!r}." + ) + + lora_state_dict = _remap_qwen_layout(safetensors.torch.load_file(lora_ckpt)) + non_lora_state_dict = _remap_qwen_layout(torch.load(non_lora_ckpt, map_location="cpu")) + + lora_state_dict = _insert_adapter_name_into_state_dict( + lora_state_dict, + adapter_name="default", + parameter_prefix="lora_", + ) + + model_state_dict = model.state_dict() + model_state_dict.update(non_lora_state_dict) + model_state_dict.update(lora_state_dict) + model.load_state_dict(model_state_dict) + + return model, loaded_step + + +__all__ = ["load_model_from_checkpoint"] diff --git a/experimental/refl/reward/videoalign/model/configs.py b/experimental/refl/reward/videoalign/model/configs.py new file mode 100644 index 000000000..f2b32003b --- /dev/null +++ b/experimental/refl/reward/videoalign/model/configs.py @@ -0,0 +1,151 @@ +"""Dataclass configs for the VideoAlign reward model. + +Inference-only subset of the originals at +``mmrl/recipes/rewards/videoalign/vendor/videoalign/utils.py``. We keep the +exact field names + default values that appear in checkpoints' +``model_config.json``, so ``ModelConfig(**dict_from_json)`` / +``PEFTLoraConfig(**dict_from_json)`` continue to round-trip every public +VideoAlign release. Training-only fields (``vision_lr``, ``merger_lr``, +``conduct_eval``, ``logging_epochs`` …) live on :class:`TrainingConfig` and +are accepted-but-ignored by the inference path — they only end up here +because ``model_config.json`` was written by the trainer. + +Why we DON'T inherit from ``transformers.TrainingArguments`` +----------------------------------------------------------- +The mmrl vendor's ``TrainingConfig`` extends ``TrainingArguments`` which +runs an ``__post_init__`` validating distributed launch fields +(``local_rank``, ``deepspeed`` …). At reward *inference* time the call site +is ``TrainingConfig(load_from_pretrained=..., bf16=..., output_dir="")`` and +the only fields we actually read downstream are ``bf16`` / ``fp16`` / +``gradient_checkpointing`` / ``disable_flash_attn2``. So we re-declare a +plain ``@dataclass`` with just those fields plus a generous ``**kwargs`` +catch-all (``__init__`` ignores unknown keys via ``__init_subclass__`` — +no, simpler: we filter in the loader). That keeps load-time light and +removes the transformers private-symbol coupling entirely. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import List, Literal, Optional + + +@dataclass +class TrainingConfig: + """Inference-relevant slice of the trainer's :class:`TrainingArguments`. + + Only ``bf16`` / ``fp16`` / ``gradient_checkpointing`` / + ``disable_flash_attn2`` are actually consumed by the inference path + (see :func:`experimental.refl.reward.videoalign.model.factory.create_model_and_processor`). + The other fields are kept so ``TrainingConfig(**model_config_json["training_args"])`` + still works when someone wants to introspect the saved config. + """ + + # Inference-time knobs (read by ``create_model_and_processor``) + output_dir: str = "" + bf16: bool = False + fp16: bool = False + gradient_checkpointing: bool = False + disable_flash_attn2: bool = False + + # Train-only — kept as no-op defaults so __init__ accepts them. + max_length: Optional[int] = None + dataset_num_proc: Optional[int] = None + center_rewards_coefficient: Optional[float] = None + vision_lr: Optional[float] = None + merger_lr: Optional[float] = None + special_token_lr: Optional[float] = None + conduct_eval: Optional[bool] = True + load_from_pretrained: Optional[str] = None + load_from_pretrained_step: Optional[int] = None + logging_epochs: Optional[float] = None + eval_epochs: Optional[float] = None + save_epochs: Optional[float] = None + remove_unused_columns: Optional[bool] = False + save_full_model: Optional[bool] = False + + +@dataclass +class PEFTLoraConfig: + """LoRA wiring for the reward model. + + For inference we typically want ``lora_enable=True`` only if the + checkpoint was saved as a LoRA split (``adapter_model.safetensors`` + + ``non_lora_state_dict.pth``) — :func:`load_model_from_checkpoint` + auto-detects which branch to take. The other fields drive the LoRA + target-module discovery inside :func:`create_model_and_processor`. + """ + + lora_enable: bool = False + vision_lora: bool = False + lora_r: int = 16 + lora_alpha: int = 32 + lora_dropout: float = 0.05 + lora_target_modules: Optional[List[str]] = None + lora_namespan_exclude: Optional[List[str]] = None + lora_modules_to_save: Optional[List[str]] = None + lora_task_type: str = "CAUSAL_LM" + use_rslora: bool = False + num_lora_modules: int = -1 + + def __post_init__(self) -> None: + # Mirror the upstream normalisation: a single-element list is + # flattened to a scalar so peft's ``target_modules`` accepts it. + if isinstance(self.lora_target_modules, list) and len(self.lora_target_modules) == 1: + self.lora_target_modules = self.lora_target_modules[0] + if isinstance(self.lora_namespan_exclude, list) and len(self.lora_namespan_exclude) == 1: + self.lora_namespan_exclude = self.lora_namespan_exclude[0] + + +@dataclass +class ModelConfig: + """Backbone + reward-head configuration. + + Read from ``model_config.json::model_config``. Most fields just pass + through to :meth:`Qwen2VLRewardModelBT.from_pretrained`. + """ + + model_name_or_path: Optional[str] = None + model_revision: str = "main" + + # Output dimensionality of the reward head: 3 for joint (VQ, MQ, TA) + # heads, 1 for single-attribute checkpoints. + output_dim: int = 1 + + # Whether the checkpoint introduces ``<|VQ_reward|>`` / ``<|MQ_reward|>`` + # / ``<|TA_reward|>`` special tokens (matches the + # ``detailed_special`` prompt template). When True, the reward is read + # from those token positions instead of the last token. + use_special_tokens: bool = False + + freeze_vision_tower: bool = field(default=False) + freeze_llm: bool = field(default=False) + tune_merger: bool = field(default=False) + + torch_dtype: Optional[Literal["auto", "bfloat16", "float16", "float32"]] = None + trust_remote_code: bool = False + attn_implementation: Optional[str] = None + load_in_8bit: bool = False + load_in_4bit: bool = False + bnb_4bit_quant_type: Literal["fp4", "nf4"] = "nf4" + use_bnb_nested_quant: bool = False + + # Where the reward is read from the hidden-state sequence: + # ``"last"`` — last non-pad token (standard). + # ``"mean"`` — masked mean over the prompt. + # ``"special"`` — special-token positions (set automatically when + # ``use_special_tokens=True`` and the checkpoint + # declares ``additional_special_tokens``). + reward_token: Literal["last", "mean", "special"] = "last" + + # Bradley-Terry / regression flavour — train-only. Inference ignores + # this entirely; kept for round-trip compatibility with + # ``model_config.json``. + loss_type: Literal["bt", "reg", "btt", "margin", "constant_margin", "scaled", "regular"] = "regular" + + def __post_init__(self) -> None: + if self.load_in_8bit and self.load_in_4bit: + raise ValueError("You can't use 8 bit and 4 bit precision at the same time") + + +__all__ = ["ModelConfig", "PEFTLoraConfig", "TrainingConfig"] diff --git a/experimental/refl/reward/videoalign/model/factory.py b/experimental/refl/reward/videoalign/model/factory.py new file mode 100644 index 000000000..1afce4df0 --- /dev/null +++ b/experimental/refl/reward/videoalign/model/factory.py @@ -0,0 +1,162 @@ +"""Build the VideoAlign reward model + Qwen2-VL processor. + +Inference-only counterpart to ``create_model_and_processor`` in +``mmrl/recipes/rewards/videoalign/vendor/videoalign/train_reward.py``. +Specifically: + +- Drops the ``trl.get_kbit_device_map`` / ``trl.get_quantization_config`` + dependency — we never 4-/8-bit quantise a reward model at REFL time. +- Drops the optimiser / loss / dataset side entirely. +- Keeps the optional LoRA wrapping and the optional ``<|VQ_reward|>`` / + ``<|MQ_reward|>`` / ``<|TA_reward|>`` special-token registration, since + both are needed to construct a *load-able* parameter graph for the + published checkpoints. +""" + +from __future__ import annotations + +import logging +from typing import Optional, Tuple + +import torch +from transformers import AutoProcessor + +from .configs import ModelConfig, PEFTLoraConfig, TrainingConfig +from .reward_model import Qwen2VLRewardModelBT + +logger = logging.getLogger(__name__) + + +def _find_target_linear_names( + model: torch.nn.Module, + num_lora_modules: int = -1, + lora_namespan_exclude: Optional[list] = None, +) -> list: + """Discover LoRA target modules by introspecting the model graph. + + Mirrors the upstream selection rule: every ``nn.Linear`` / ``nn.Embedding`` + whose qualified name doesn't contain any excluded namespace keyword. + ``num_lora_modules > 0`` truncates to the last N modules (the upstream + knob for cheaply LoRA-tuning only the top of the network). + """ + lora_namespan_exclude = lora_namespan_exclude or [] + if isinstance(lora_namespan_exclude, str): + lora_namespan_exclude = [lora_namespan_exclude] + + target_classes = (torch.nn.Linear, torch.nn.Embedding) + out = [] + for name, module in model.named_modules(): + if any(keyword in name for keyword in lora_namespan_exclude): + continue + if isinstance(module, target_classes): + out.append(name) + + if num_lora_modules > 0: + out = out[-num_lora_modules:] + return out + + +def create_model_and_processor( + model_config: ModelConfig, + peft_lora_config: PEFTLoraConfig, + training_args: TrainingConfig, + cache_dir: Optional[str] = None, +) -> Tuple[torch.nn.Module, AutoProcessor, object]: + """Build the Qwen2-VL reward model + matching processor. + + Args: + model_config: backbone configuration (matches the saved + ``model_config`` block in ``model_config.json``). + peft_lora_config: LoRA wiring (matches ``peft_lora_config``). + training_args: dtype / flash-attn knobs only; the dataset/optimiser + fields of :class:`TrainingConfig` are not used here. + cache_dir: optional HF cache directory. + + Returns: + ``(model, processor, peft_config)`` — same tuple shape as the + upstream helper. ``peft_config`` is ``None`` when LoRA is disabled. + """ + # Resolve the torch dtype string to a real dtype. + torch_dtype = ( + model_config.torch_dtype + if model_config.torch_dtype in ["auto", None] + else getattr(torch, model_config.torch_dtype) + ) + + # Build processor + (optional) special tokens. + processor = AutoProcessor.from_pretrained( + model_config.model_name_or_path, + padding_side="right", + cache_dir=cache_dir, + ) + + special_token_ids = None + if model_config.use_special_tokens: + special_tokens = ["<|VQ_reward|>", "<|MQ_reward|>", "<|TA_reward|>"] + processor.tokenizer.add_special_tokens({"additional_special_tokens": special_tokens}) + special_token_ids = processor.tokenizer.convert_tokens_to_ids(special_tokens) + + # Build the reward model. Quantisation is intentionally not supported + # here — the reward path expects full-precision (or bf16/fp16) weights. + model = Qwen2VLRewardModelBT.from_pretrained( + model_config.model_name_or_path, + output_dim=model_config.output_dim, + reward_token=model_config.reward_token, + special_token_ids=special_token_ids, + torch_dtype=torch_dtype, + attn_implementation=("flash_attention_2" if not training_args.disable_flash_attn2 else "sdpa"), + cache_dir=cache_dir, + revision=getattr(model_config, "model_revision", "main"), + ) + # transformers 5.x forwards unknown from_pretrained kwargs to the model + # ctor (4.x absorbed config fields like use_cache) — set it on the config. + model.config.use_cache = bool(training_args.gradient_checkpointing) + + if model_config.use_special_tokens: + model.resize_token_embeddings(len(processor.tokenizer)) + + if training_args.bf16: + model.to(torch.bfloat16) + if training_args.fp16: + model.to(torch.float16) + + # Optional LoRA wrapping — required to *load* a LoRA-split checkpoint. + if peft_lora_config.lora_enable: + # peft is an optional dep for the non-LoRA inference path; import + # lazily so users who only ever load full-state-dict ckpts don't + # need it installed. + from peft import LoraConfig, get_peft_model + + namespan_exclude = list(peft_lora_config.lora_namespan_exclude or []) + if isinstance(peft_lora_config.lora_namespan_exclude, str): + namespan_exclude = [peft_lora_config.lora_namespan_exclude] + # Mirror upstream: when vision_lora is off, exclude the visual tower. + if not peft_lora_config.vision_lora and "visual" not in namespan_exclude: + namespan_exclude.append("visual") + + target_modules = _find_target_linear_names( + model, + num_lora_modules=peft_lora_config.num_lora_modules, + lora_namespan_exclude=namespan_exclude, + ) + peft_config = LoraConfig( + target_modules=target_modules, + r=peft_lora_config.lora_r, + lora_alpha=peft_lora_config.lora_alpha, + lora_dropout=peft_lora_config.lora_dropout, + task_type=peft_lora_config.lora_task_type, + use_rslora=peft_lora_config.use_rslora, + bias="none", + modules_to_save=peft_lora_config.lora_modules_to_save, + ) + model = get_peft_model(model, peft_config) + else: + peft_config = None + + model.config.tokenizer_padding_side = processor.tokenizer.padding_side + model.config.pad_token_id = processor.tokenizer.pad_token_id + + return model, processor, peft_config + + +__all__ = ["create_model_and_processor"] diff --git a/experimental/refl/reward/videoalign/model/prompt_template.py b/experimental/refl/reward/videoalign/model/prompt_template.py new file mode 100644 index 000000000..7687af0da --- /dev/null +++ b/experimental/refl/reward/videoalign/model/prompt_template.py @@ -0,0 +1,187 @@ +"""Prompt templates for the VideoAlign reward model. + +Verbatim port of ``mmrl/recipes/rewards/videoalign/vendor/videoalign/prompt_template.py``. +The string constants are part of the *training contract* of the published +checkpoints — changing them would silently shift the reward distribution +(the Qwen2-VL chat template was fed exactly these strings during reward-model +training). Do not edit unless you are retraining from scratch. +""" + +from __future__ import annotations + +from typing import List, Union + +VIDEOSCORE_QUERY_PROMPT = """ +Suppose you are an expert in judging and evaluating the quality of AI-generated videos, +please watch the frames of a given video and see the text prompt for generating the video, +then give scores based on its {dimension_name}, i.e., {dimension_description}. +Output a float number from 1.0 to 5.0 for this dimension, +the higher the number is, the better the video performs in that sub-score, +the lowest 1.0 means Bad, the highest 5.0 means Perfect/Real (the video is like a real video). +The text prompt used for generation is "{text_prompt}". +""" + +DIMENSION_DESCRIPTIONS = { + "VQ": [ + "visual quality", + "the quality of the video in terms of clearness, resolution, brightness, and color", + ], + "TA": [ + "text-to-video alignment", + "the alignment between the text prompt and the video content and motion", + ], + "MQ": [ + "motion quality", + "the quality of the motion in terms of consistency, smoothness, and completeness", + ], + "Overall": [ + "Overall Performance", + "the overall performance of the video in terms of visual quality, text-to-video alignment, and motion quality", + ], +} + +SIMPLE_PROMPT = """ +Please evaluate the {dimension_name} of a generated video. Consider {dimension_description}. +The text prompt used for generation is "{text_prompt}". +""" + +DETAILED_PROMPT_WITH_SPECIAL_TOKEN = """ +You are tasked with evaluating a generated video based on three distinct criteria: Visual Quality, Motion Quality, and Text Alignment. Please provide a rating from 0 to 10 for each of the three categories, with 0 being the worst and 10 being the best. Each evaluation should be independent of the others. + +**Visual Quality:** +Evaluate the overall visual quality of the video, with a focus on static factors. The following sub-dimensions should be considered: +- **Reasonableness:** The video should not contain any significant biological or logical errors, such as abnormal body structures or nonsensical environmental setups. +- **Clarity:** Evaluate the sharpness and visibility of the video. The image should be clear and easy to interpret, with no blurring or indistinct areas. +- **Detail Richness:** Consider the level of detail in textures, materials, lighting, and other visual elements (e.g., hair, clothing, shadows). +- **Aesthetic and Creativity:** Assess the artistic aspects of the video, including the color scheme, composition, atmosphere, depth of field, and the overall creative appeal. The scene should convey a sense of harmony and balance. +- **Safety:** The video should not contain harmful or inappropriate content, such as political, violent, or adult material. If such content is present, the image quality and satisfaction score should be the lowest possible. + +Please provide the ratings of Visual Quality: <|VQ_reward|> +END + +**Motion Quality:** +Assess the dynamic aspects of the video, with a focus on dynamic factors. Consider the following sub-dimensions: +- **Stability:** Evaluate the continuity and stability between frames. There should be no sudden, unnatural jumps, and the video should maintain stable attributes (e.g., no fluctuating colors, textures, or missing body parts). +- **Naturalness:** The movement should align with physical laws and be realistic. For example, clothing should flow naturally with motion, and facial expressions should change appropriately (e.g., blinking, mouth movements). +- **Aesthetic Quality:** The movement should be smooth and fluid. The transitions between different motions or camera angles should be seamless, and the overall dynamic feel should be visually pleasing. +- **Fusion:** Ensure that elements in motion (e.g., edges of the subject, hair, clothing) blend naturally with the background, without obvious artifacts or the feeling of cut-and-paste effects. +- **Clarity of Motion:** The video should be clear and smooth in motion. Pay attention to any areas where the video might have blurry or unsteady sections that hinder visual continuity. +- **Amplitude:** If the video is largely static or has little movement, assign a low score for motion quality. + +Please provide the ratings of Motion Quality: <|MQ_reward|> +END + +**Text Alignment:** +Assess how well the video matches the textual prompt across the following sub-dimensions: +- **Subject Relevance** Evaluate how accurately the subject(s) in the video (e.g., person, animal, object) align with the textual description. The subject should match the description in terms of number, appearance, and behavior. +- **Motion Relevance:** Evaluate if the dynamic actions (e.g., gestures, posture, facial expressions like talking or blinking) align with the described prompt. The motion should match the prompt in terms of type, scale, and direction. +- **Environment Relevance:** Assess whether the background and scene fit the prompt. This includes checking if real-world locations or scenes are accurately represented, though some stylistic adaptation is acceptable. +- **Style Relevance:** If the prompt specifies a particular artistic or stylistic style, evaluate how well the video adheres to this style. +- **Camera Movement Relevance:** Check if the camera movements (e.g., following the subject, focus shifts) are consistent with the expected behavior from the prompt. + +Textual prompt - {text_prompt} +Please provide the ratings of Text Alignment: <|TA_reward|> +END +""" + +DETAILED_PROMPT = """ +You are tasked with evaluating a generated video based on three distinct criteria: Visual Quality, Motion Quality, and Text Alignment. Please provide a rating from 0 to 10 for each of the three categories, with 0 being the worst and 10 being the best. Each evaluation should be independent of the others. + +**Visual Quality:** +Evaluate the overall visual quality of the video, with a focus on static factors. The following sub-dimensions should be considered: +- **Reasonableness:** The video should not contain any significant biological or logical errors, such as abnormal body structures or nonsensical environmental setups. +- **Clarity:** Evaluate the sharpness and visibility of the video. The image should be clear and easy to interpret, with no blurring or indistinct areas. +- **Detail Richness:** Consider the level of detail in textures, materials, lighting, and other visual elements (e.g., hair, clothing, shadows). +- **Aesthetic and Creativity:** Assess the artistic aspects of the video, including the color scheme, composition, atmosphere, depth of field, and the overall creative appeal. The scene should convey a sense of harmony and balance. +- **Safety:** The video should not contain harmful or inappropriate content, such as political, violent, or adult material. If such content is present, the image quality and satisfaction score should be the lowest possible. + +**Motion Quality:** +Assess the dynamic aspects of the video, with a focus on dynamic factors. Consider the following sub-dimensions: +- **Stability:** Evaluate the continuity and stability between frames. There should be no sudden, unnatural jumps, and the video should maintain stable attributes (e.g., no fluctuating colors, textures, or missing body parts). +- **Naturalness:** The movement should align with physical laws and be realistic. For example, clothing should flow naturally with motion, and facial expressions should change appropriately (e.g., blinking, mouth movements). +- **Aesthetic Quality:** The movement should be smooth and fluid. The transitions between different motions or camera angles should be seamless, and the overall dynamic feel should be visually pleasing. +- **Fusion:** Ensure that elements in motion (e.g., edges of the subject, hair, clothing) blend naturally with the background, without obvious artifacts or the feeling of cut-and-paste effects. +- **Clarity of Motion:** The video should be clear and smooth in motion. Pay attention to any areas where the video might have blurry or unsteady sections that hinder visual continuity. +- **Amplitude:** If the video is largely static or has little movement, assign a low score for motion quality. + + +**Text Alignment:** +Assess how well the video matches the textual prompt across the following sub-dimensions: +- **Subject Relevance** Evaluate how accurately the subject(s) in the video (e.g., person, animal, object) align with the textual description. The subject should match the description in terms of number, appearance, and behavior. +- **Motion Relevance:** Evaluate if the dynamic actions (e.g., gestures, posture, facial expressions like talking or blinking) align with the described prompt. The motion should match the prompt in terms of type, scale, and direction. +- **Environment Relevance:** Assess whether the background and scene fit the prompt. This includes checking if real-world locations or scenes are accurately represented, though some stylistic adaptation is acceptable. +- **Style Relevance:** If the prompt specifies a particular artistic or stylistic style, evaluate how well the video adheres to this style. +- **Camera Movement Relevance:** Check if the camera movements (e.g., following the subject, focus shifts) are consistent with the expected behavior from the prompt. + +Textual prompt - {text_prompt} +Please provide the ratings of Visual Quality, Motion Quality, and Text Alignment. +""" + +SIMPLE_PROMPT_NO_PROMPT = """ +Please evaluate the {dimension_name} of a generated video. Consider {dimension_description}. +""" + + +def build_prompt( + prompt: str, + dimension: Union[str, List[str]], + template_type: str, +) -> str: + """Render the user-message text for the Qwen2-VL chat template. + + Args: + prompt: the text-to-video user prompt that produced the candidate + generation. May be ``""`` for the ``"none"`` template (then the + user message is empty — the model still reads the video frames). + dimension: ``"VQ"`` / ``"MQ"`` / ``"TA"`` / ``"Overall"`` (or a list + for multi-dim prompts). Drives the ``{dimension_name}`` / + ``{dimension_description}`` placeholders. + template_type: must match what the checkpoint was trained with. Read + this from ``model_config.json::data_config.prompt_template_type``. + One of ``"none"`` / ``"simple"`` / ``"video_score"`` / + ``"detailed_special"`` / ``"detailed"``. + + Returns: + The fully rendered user message string. + """ + if isinstance(dimension, list) and len(dimension) > 1: + dimension_name = ", ".join(DIMENSION_DESCRIPTIONS[d][0] for d in dimension) + dimension_name = f"overall performance({dimension_name})" + dimension_description = "the overall performance of the video" + else: + if isinstance(dimension, list): + dimension = dimension[0] + dimension_name = DIMENSION_DESCRIPTIONS[dimension][0] + dimension_description = DIMENSION_DESCRIPTIONS[dimension][1] + + if template_type == "none": + return prompt + elif template_type == "simple": + return SIMPLE_PROMPT.format( + dimension_name=dimension_name, + dimension_description=dimension_description, + text_prompt=prompt, + ) + elif template_type == "video_score": + return VIDEOSCORE_QUERY_PROMPT.format( + dimension_name=dimension_name, + dimension_description=dimension_description, + text_prompt=prompt, + ) + elif template_type == "detailed_special": + return DETAILED_PROMPT_WITH_SPECIAL_TOKEN.format(text_prompt=prompt) + elif template_type == "detailed": + return DETAILED_PROMPT.format(text_prompt=prompt) + else: + raise ValueError(f"Invalid template type: {template_type!r}") + + +__all__ = [ + "DIMENSION_DESCRIPTIONS", + "VIDEOSCORE_QUERY_PROMPT", + "SIMPLE_PROMPT", + "SIMPLE_PROMPT_NO_PROMPT", + "DETAILED_PROMPT", + "DETAILED_PROMPT_WITH_SPECIAL_TOKEN", + "build_prompt", +] diff --git a/experimental/refl/reward/videoalign/model/reward_model.py b/experimental/refl/reward/videoalign/model/reward_model.py new file mode 100644 index 000000000..b5179a097 --- /dev/null +++ b/experimental/refl/reward/videoalign/model/reward_model.py @@ -0,0 +1,233 @@ +"""Qwen2-VL reward model with a 3-dim (VQ, MQ, TA) regression head. + +Verbatim port of ``Qwen2VLRewardModelBT`` from +``mmrl/recipes/rewards/videoalign/vendor/videoalign/trainer.py`` — the +*only* class from that file we actually need at inference time. Training- +specific machinery (``VideoVLMRewardTrainer``, ``PartialEmbeddingUpdateCallback``, +``compute_multi_attr_accuracy``) lived in the same file but are not part of +the inference contract and have been intentionally dropped to remove the +``transformers.trainer`` private-symbol coupling that breaks on +transformers>=5. + +The forward signature is preserved so the same checkpoint state_dict loads +cleanly (no key remapping needed beyond the standard transformers>=5 +``base_model.model.model.language_model.*`` shim handled by +:func:`experimental.refl.reward.videoalign.model.checkpoint.load_model_from_checkpoint`). +""" + +from __future__ import annotations + +from typing import Any, List, Optional + +import torch +import torch.nn as nn +from transformers import Qwen2VLForConditionalGeneration + + +def _cfg_get(config: Any, name: str) -> Any: + """Read a ``Qwen2VLConfig`` field from wherever 5.6 keeps it: media + token ids live on the top-level config, LM fields like ``hidden_size`` + under the nested ``text_config``.""" + if hasattr(config, name): + val = getattr(config, name) + if val is not None: + return val + text_cfg = getattr(config, "text_config", None) + if text_cfg is not None and hasattr(text_cfg, name): + return getattr(text_cfg, name) + # Fall back to whatever the top-level had (possibly None) so callers + # that tolerate ``None`` (e.g. ``pad_token_id``) keep working. + if hasattr(config, name): + return getattr(config, name) + raise AttributeError(f"{type(config).__name__} has no attribute {name!r} (checked top-level and .text_config).") + + +class Qwen2VLRewardModelBT(Qwen2VLForConditionalGeneration): + """Qwen2-VL backbone + ``nn.Linear`` reward head (Bradley-Terry / regression). + + Differences vs ``Qwen2VLForConditionalGeneration``: + + - Replaces the LM-head pathway with an ``rm_head: Linear(hidden, output_dim)`` + whose output is pooled per-sample (last / mean / special-token). + - ``forward`` returns ``{"logits": Tensor[B, output_dim]}`` rather than a + ``CausalLMOutputWithPast``. + + Output dim conventions used by VideoAlign checkpoints: + + - ``output_dim=3`` + ``reward_token="last"`` → joint head; final-token + logits are the (VQ, MQ, TA) scalars directly. + - ``output_dim=3`` + ``reward_token="special"`` + 3 special-token IDs → + each special token contributes its own row of the 3×3 head; we + diagonal-extract to (VQ from <|VQ_reward|>, MQ from <|MQ_reward|>, …). + - ``output_dim=1`` → single-attribute heads + (one ckpt per dimension, rarely used by the public release). + """ + + def __init__( + self, + config, + output_dim: int = 4, + reward_token: str = "last", + special_token_ids: Optional[List[int]] = None, + ) -> None: + super().__init__(config) + self.output_dim = output_dim + hidden_size = _cfg_get(config, "hidden_size") + self.rm_head = nn.Linear(hidden_size, output_dim, bias=False) + self.reward_token = reward_token + + self.special_token_ids = special_token_ids + # When special tokens are configured, the pooling mode is forced — + # otherwise the trainer-set ``reward_token`` setting on disk wins. + if self.special_token_ids is not None: + self.reward_token = "special" + + # The forward code below addresses ``self.visual``; transformers 5.6 + # owns the vision tower at ``self.model.visual``. Read via ``_modules`` + # to avoid recursing into ``nn.Module.__getattr__``, which would + # re-trigger this property with a misleading "no attribute" error. + @property + def visual(self): # type: ignore[override] + inner = self._modules.get("model", None) + visual = getattr(inner, "visual", None) if inner is not None else None + if visual is None: + raise AttributeError( + "Qwen2VLRewardModelBT: vision tower not found at " + "self.model.visual — this code targets the locked " + "transformers 5.6 stack; align the environment." + ) + return visual + + # ------------------------------------------------------------------ + # Forward + # ------------------------------------------------------------------ + + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, # noqa: ARG002 — kept for API parity + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + pixel_values: Optional[torch.Tensor] = None, + pixel_values_videos: Optional[torch.FloatTensor] = None, + image_grid_thw: Optional[torch.LongTensor] = None, + video_grid_thw: Optional[torch.LongTensor] = None, + rope_deltas: Optional[torch.LongTensor] = None, # noqa: ARG002 — kept for API parity + ): + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # Visual + text token fusion. Identical to the upstream + # Qwen2VLForConditionalGeneration.forward except we don't compute + # an LM head — only the regression head over the final hidden + # states. + # + # transformers 5.6: the vision tower returns BaseModelOutputWithPooling + # whose pooler_output holds the merged features that fill the media + # placeholders (what get_video_features reads). last_hidden_state is + # the PRE-merger states (4x tokens, vision dim) — never usable here. + def _as_tensor(visual_out): + t = getattr(visual_out, "pooler_output", None) + if t is None: + raise TypeError( + "Qwen2-VL vision tower returned " + f"{type(visual_out).__name__} with no pooler_output. This " + "code targets the locked transformers 5.6 stack — align " + "the environment instead of widening this path." + ) + return t + + if inputs_embeds is None: + inputs_embeds = self.get_input_embeddings()(input_ids) + if pixel_values is not None: + pixel_values = pixel_values.type(self.visual.get_dtype()) + image_embeds = _as_tensor(self.visual(pixel_values, grid_thw=image_grid_thw)) + image_token_id = _cfg_get(self.config, "image_token_id") + image_mask = (input_ids == image_token_id).unsqueeze(-1).expand_as(inputs_embeds) + image_embeds = image_embeds.to(inputs_embeds.device, inputs_embeds.dtype) + inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) + + if pixel_values_videos is not None: + pixel_values_videos = pixel_values_videos.type(self.visual.get_dtype()) + video_embeds = _as_tensor(self.visual(pixel_values_videos, grid_thw=video_grid_thw)) + video_token_id = _cfg_get(self.config, "video_token_id") + video_mask = (input_ids == video_token_id).unsqueeze(-1).expand_as(inputs_embeds) + video_embeds = video_embeds.to(inputs_embeds.device, inputs_embeds.dtype) + inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds) + + if attention_mask is not None: + attention_mask = attention_mask.to(inputs_embeds.device) + + outputs = self.model( + input_ids=None, + position_ids=position_ids, + attention_mask=attention_mask, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + hidden_states = outputs[0] # (B, L, D) + logits = self.rm_head(hidden_states) # (B, L, output_dim) + + if input_ids is not None: + batch_size = input_ids.shape[0] + else: + batch_size = inputs_embeds.shape[0] + + # Locate per-sample sequence length so we can pool the right token. + try: + pad_token_id = _cfg_get(self.config, "pad_token_id") + except AttributeError: + pad_token_id = None + if pad_token_id is None and batch_size != 1: + raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.") + if pad_token_id is None: + sequence_lengths = -1 + else: + if input_ids is not None: + # Find the first pad token; previous token is the last real + # token. ``%`` keeps the index in range for ONNX export. + sequence_lengths = torch.eq(input_ids, pad_token_id).int().argmax(-1) - 1 + sequence_lengths = sequence_lengths % input_ids.shape[-1] + sequence_lengths = sequence_lengths.to(logits.device) + else: + sequence_lengths = -1 + + if self.reward_token == "last": + pooled_logits = logits[torch.arange(batch_size, device=logits.device), sequence_lengths] + elif self.reward_token == "mean": + valid_lengths = torch.clamp(sequence_lengths, min=0, max=logits.size(1) - 1) + pooled_logits = torch.stack([logits[i, : valid_lengths[i]].mean(dim=0) for i in range(batch_size)]) + elif self.reward_token == "special": + special_token_mask = torch.zeros_like(input_ids, dtype=torch.bool) + for special_token_id in self.special_token_ids: + special_token_mask = special_token_mask | (input_ids == special_token_id) + pooled_logits = logits[special_token_mask, ...] + # Each sample has exactly 3 special tokens (VQ/MQ/TA), each + # producing an ``output_dim``-row of logits. Reshape → + # (B, 3, output_dim); when output_dim=3 we keep the diagonal + # (matching the special-token training contract). + pooled_logits = pooled_logits.view(batch_size, 3, -1) + if self.output_dim == 3: + pooled_logits = pooled_logits.diagonal(dim1=1, dim2=2) + pooled_logits = pooled_logits.view(batch_size, -1) + else: + raise ValueError(f"Invalid reward_token: {self.reward_token!r}") + + return {"logits": pooled_logits} + + +__all__ = ["Qwen2VLRewardModelBT"] diff --git a/experimental/refl/reward/videoalign/requirements.txt b/experimental/refl/reward/videoalign/requirements.txt new file mode 100644 index 000000000..bcd1c63c0 --- /dev/null +++ b/experimental/refl/reward/videoalign/requirements.txt @@ -0,0 +1,9 @@ +# VideoAlign reward extras — ADDITIVE ONLY on top of the UniRL core stack. +# +# The reward shares one Python process with the actor (colocated Remote +# siblings), so pins here cannot "isolate" versions — an overlapping pin +# would downgrade the core environment itself. Core already provides +# transformers / peft / safetensors / huggingface-hub / einops (see +# pyproject.toml); never re-pin those here. Attention runs SDPA — flash-attn +# is not part of the locked stack. +torchvision # frame preprocessing transforms; matches the engine extra's torch build diff --git a/experimental/refl/reward/videoalign/scorer.py b/experimental/refl/reward/videoalign/scorer.py new file mode 100644 index 000000000..f19fcc94e --- /dev/null +++ b/experimental/refl/reward/videoalign/scorer.py @@ -0,0 +1,232 @@ +"""VideoAlign reward scorer (REFL-compatible / BPTT-differentiable). + +Wraps :class:`VideoRewardWrapper` (Qwen2-VL-based reward model producing +three scalar scores per (video, prompt) pair: VQ / MQ / TA) and exposes a +recipe-local differentiable REFL entry point. + +Reward = ``w_vq * VQ + w_mq * MQ + w_ta * TA`` (defaults to 1 / 1 / 1). + +Gradient flow +------------- +The Qwen2-VL vision encoder is differentiable w.r.t. the input pixels when +the *fast* image processor is used (the wrapper force-installs +``Qwen2VLImageProcessorFast`` on construction). The generated video arrives +via ``compute_rewards_differentiable`` as ``[B, C, T, H, W]`` float in +``[-1, 1]`` with a live ``grad_fn`` (BPTT path); we forward into the wrapper under +``torch.enable_grad`` so the linear combination of VQ/MQ/TA traces back +through the vision tower into the diffusion graph. + +Self-containment +---------------- +This scorer no longer requires the sibling ``mmrl`` repo on disk. The +Qwen2-VL reward backbone, prompt template, checkpoint loader and +inference wrapper all live under +:mod:`experimental.refl.reward.videoalign.model` / :mod:`...wrapper`. The +``mmrl_repo_root`` Spec field has been removed; ``MMRL_REPO_ROOT`` env +var is now irrelevant. +""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass +from typing import List, Optional + +import torch + +from unirl.reward.base import BaseRewardComponentSpec +from unirl.reward.local.base import LocalRewardBackend +from unirl.reward.local.device import resolve_device +from unirl.types.reward import RewardRequest + +from .wrapper import VideoRewardWrapper + +logger = logging.getLogger(__name__) + + +class VideoAlignRewardScorer(LocalRewardBackend): + """Qwen2-VL VideoAlign reward — VQ + MQ + TA → scalar per sample.""" + + canonical_model_name = "videoalign" + input_kind = "video" + + def __init__(self, *, config: "VideoAlignSpec", base_device: str) -> None: + super().__init__( + model_name=self.canonical_model_name, + device=resolve_device(config.device, base_device), + batch_size=int(config.batch_size), + reward_model_path=config.reward_model_path, + resize_height=config.resize_height, + resize_width=config.resize_width, + micro_batch_size=config.micro_batch_size, + reward_num_frames=config.reward_num_frames, + use_norm=config.use_norm, + w_vq=config.w_vq, + w_mq=config.w_mq, + w_ta=config.w_ta, + differentiable=config.differentiable, + ) + + # ------------------------------------------------------------------ + # Model loading + # ------------------------------------------------------------------ + + def _load_model(self) -> None: + reward_model_path = str(self.model_kwargs["reward_model_path"]) + if not reward_model_path: + raise ValueError( + "VideoAlignRewardScorer: ``reward_model_path`` must be set " + "(the directory containing ``model_config.json`` + the " + "``checkpoint-*`` subdir)." + ) + + self._w_vq = float(self.model_kwargs["w_vq"]) + self._w_mq = float(self.model_kwargs["w_mq"]) + self._w_ta = float(self.model_kwargs["w_ta"]) + self._use_norm = bool(self.model_kwargs["use_norm"]) + self._reward_num_frames = int(self.model_kwargs["reward_num_frames"]) + self._differentiable = bool(self.model_kwargs["differentiable"]) + + logger.info( + "VideoAlignRewardScorer: loading VideoRewardWrapper from %s", + reward_model_path, + ) + self.model = VideoRewardWrapper( + checkpoint_dir=reward_model_path, + device=self.device, + dtype=torch.bfloat16, + use_norm=self._use_norm, + resize_height=int(self.model_kwargs["resize_height"]), + resize_width=int(self.model_kwargs["resize_width"]), + micro_batch_size=int(self.model_kwargs["micro_batch_size"]), + ) + + def _compute_model_rewards(self, request: RewardRequest) -> List[float]: + raise NotImplementedError("VideoAlignRewardScorer is REFL-only; use compute_rewards_differentiable().") + + # ------------------------------------------------------------------ + # Differentiable REFL reward entry point + # ------------------------------------------------------------------ + + def compute_rewards_differentiable( + self, + media_tensor: torch.Tensor, + prompts: List[str], + records: Optional[List[dict]] = None, + ) -> torch.Tensor: + del records + if media_tensor.ndim != 5: + raise ValueError(f"VideoAlignRewardScorer expects [B,C,T,H,W], got {tuple(media_tensor.shape)}") + if len(prompts) != int(media_tensor.shape[0]): + raise ValueError( + f"VideoAlignRewardScorer: prompts length {len(prompts)} != batch size {int(media_tensor.shape[0])}." + ) + + # The wrapper expects per-sample [T, C, H, W] in [-1, 1]. + # NOTE: ``.clamp(-1.0, 1.0)`` mirrors mmrl's ``role.score`` — VAE + # decode can produce slightly out-of-range pixels (e.g. -1.02 / + # 1.03), and ``_pixels_neg1_to_255`` only clamps the [0, 1] + # midpoint afterwards, so values can still spill above 255 or + # below 0 without this safeguard. Required for numeric parity + # with the mmrl baseline. + per_sample_videos: List[torch.Tensor] = [] + for v in media_tensor: + v = v.to(self.device).permute(1, 0, 2, 3).clamp(-1.0, 1.0).contiguous() # → (T, C, H, W) + per_sample_videos.append(v) + + if self._reward_num_frames > 0: + ds: List[torch.Tensor] = [] + for v in per_sample_videos: + if v.shape[0] > self._reward_num_frames: + idx = torch.linspace( + 0, + v.shape[0] - 1, + self._reward_num_frames, + device=v.device, + ).long() + v = v[idx] + ds.append(v) + per_sample_videos = ds + + autograd_ctx = torch.enable_grad if self._differentiable else torch.no_grad + with autograd_ctx(): + scores = self.model.forward_scores( + per_sample_videos, + prompts, + use_norm=self._use_norm, + ) + + reward = self._w_vq * scores["VQ"] + self._w_mq * scores["MQ"] + self._w_ta * scores["TA"] + return reward.float() + + # ------------------------------------------------------------------ + # Lifecycle hooks (CPU offload between rollouts to free VRAM) + # ------------------------------------------------------------------ + + def offload(self) -> None: + if self.model is not None and getattr(self.model, "model", None) is not None: + self.model.model.cpu() + torch.cuda.empty_cache() + + def onload(self) -> None: + if self.model is not None and getattr(self.model, "model", None) is not None: + self.model.model.to(self.device) + + def is_available(self) -> bool: + return bool(self._is_loaded) + + def dispose(self) -> None: + self.offload() + + +# --------------------------------------------------------------------------- +# Spec +# --------------------------------------------------------------------------- + + +@dataclass +class VideoAlignSpec(BaseRewardComponentSpec): + """Typed config for :class:`VideoAlignRewardScorer`. + + Args: + reward_model_path: Directory containing ``model_config.json`` and + the ``checkpoint-*`` subdir (with ``model.pth`` or LoRA split). + Required. + device: ``"auto"`` / ``"cuda"`` / ``"cuda:N"`` — resolved against + ``base_device``. + batch_size: kept for parity with sibling specs. + resize_height / resize_width: bicubic target before the Qwen2-VL + vision encoder. Defaults (336 × 588) match the published + checkpoints. + micro_batch_size: max samples per reward forward (peak-VRAM knob). + reward_num_frames: temporal downsample to this many uniformly + spaced frames before scoring; ``<= 0`` disables. + use_norm: z-score normalise each dimension using the means / stds + stored under ``inference_config`` in ``model_config.json``. + w_vq, w_mq, w_ta: linear combination weights into the final scalar. + differentiable: keep autograd on the reward forward (default True — + required for REFL). Set False for historical GRPO / replay. + """ + + reward_model_path: str = "" + + device: str = "auto" + batch_size: int = 1 + + resize_height: int = 336 + resize_width: int = 588 + micro_batch_size: int = 1 + reward_num_frames: int = 36 + + use_norm: bool = True + w_vq: float = 1.0 + w_mq: float = 1.0 + w_ta: float = 1.0 + + differentiable: bool = True + + +__all__ = [ + "VideoAlignRewardScorer", + "VideoAlignSpec", +] diff --git a/experimental/refl/reward/videoalign/wrapper.py b/experimental/refl/reward/videoalign/wrapper.py new file mode 100644 index 000000000..e5ebc2397 --- /dev/null +++ b/experimental/refl/reward/videoalign/wrapper.py @@ -0,0 +1,321 @@ +"""Differentiable forward over a loaded VideoAlign reward model. + +Self-contained re-implementation — does NOT touch ``sys.path`` . The Qwen2-VL reward +backbone, its prompt template and the checkpoint loader all live under +:mod:`experimental.refl.reward.videoalign.model`. + +Public API +---------- +``VideoRewardWrapper(checkpoint_dir, device, ...) -> .forward_scores(...)`` + +The ``forward_scores`` signature is preserved verbatim from the mmrl +wrapper so the same callsite — ``self.model.forward_scores(per_sample_videos, +prompts, use_norm=...)`` — keeps returning ``{"VQ": Tensor[B], "MQ": +Tensor[B], "TA": Tensor[B], "Overall": Tensor[B]}`` with autograd-live +graphs when the input videos have ``grad_fn`` set. + +Gradient flow +------------- +The Qwen2-VL vision encoder is fully differentiable w.r.t. its pixel input +through the *fast* (tensor-native) image processor, which transformers 5.6 +loads by default — the slow PIL-based variant routes through ``numpy`` and +silently cuts the graph. +""" + +from __future__ import annotations + +# Runs on the shared core stack — reward and actor share one process, so +# there is no separate env to pin. The 5.6 processor injects +# ``mm_token_type_ids``, which the reward backbone does not accept; +# ``compute_scores`` pops it explicitly. Never filter by forward-signature: +# under PEFT that resolves to ``LoraModel.forward(*args, **kwargs)`` and +# silently drops the video tensors. +import json +import logging +import os +from collections.abc import Mapping +from typing import Dict, List, Optional + +import torch +import torchvision.transforms.functional as TF +from torchvision.transforms import InterpolationMode + +from .model import ( + ModelConfig, + PEFTLoraConfig, + TrainingConfig, + build_prompt, + create_model_and_processor, + load_model_from_checkpoint, +) + +logger = logging.getLogger(__name__) + + +def _load_configs_from_json(config_path: str): + """Parse ``model_config.json`` into the four typed configs. + + Drops trainer-only fields that may carry absolute filesystem paths + (they would break on a different machine and are never read at + inference). + """ + with open(config_path, "r", encoding="utf-8") as f: + config_dict = json.load(f) + + data_config = dict(config_dict["data_config"]) + data_config.pop("meta_data", None) + data_config.pop("data_dir", None) + + return ( + data_config, + config_dict["model_config"], + config_dict["peft_lora_config"], + config_dict.get("inference_config", None), + ) + + +class VideoRewardWrapper: + """Frozen VideoAlign reward model with a differentiable forward. + + Constructor responsibilities: + + 1. Parse ``model_config.json`` into typed configs. + 2. Build the Qwen2-VL reward model + processor via + :func:`create_model_and_processor`. + 3. Load weights from ``checkpoint-K`` via + :func:`load_model_from_checkpoint`. + 4. Move to the requested device + dtype, ``eval()`` + freeze + parameters (RL gradients flow *through* the activations, not into + the reward weights). + """ + + def __init__( + self, + checkpoint_dir: str, + device: str = "cuda", + dtype: torch.dtype = torch.bfloat16, + use_norm: bool = True, + resize_height: int = 336, + resize_width: int = 588, + micro_batch_size: int = 1, + ) -> None: + self.device = device + self.dtype = dtype + self.use_norm = use_norm + self.resize_height = resize_height + self.resize_width = resize_width + self.micro_batch_size = max(1, int(micro_batch_size)) + + config_path = os.path.join(checkpoint_dir, "model_config.json") + if not os.path.exists(config_path): + raise FileNotFoundError( + f"VideoRewardWrapper: expected ``model_config.json`` at " + f"{config_path!r}. Make sure ``reward_model_path`` points at " + "the directory containing both ``model_config.json`` and " + "one or more ``checkpoint-K`` subdirs (the layout produced " + "by the upstream VideoAlign trainer)." + ) + + data_config_dict, model_config_dict, peft_lora_config_dict, inference_config = _load_configs_from_json( + config_path + ) + + # We only need two fields out of the data_config block — the + # template type and the eval-dim list. Stash them directly without + # constructing the full dataclass. + self.prompt_template_type = data_config_dict.get("prompt_template_type", "none") + self.eval_dim = data_config_dict.get("eval_dim", "VQ") + self.inference_config = inference_config + self.build_prompt = build_prompt + + model_config = ModelConfig(**model_config_dict) + peft_lora_config = PEFTLoraConfig(**peft_lora_config_dict) + training_args = TrainingConfig( + load_from_pretrained=checkpoint_dir, + load_from_pretrained_step=-1, + gradient_checkpointing=False, + # sdpa — deterministic on the locked stack; flash-attn 2 is not + # part of it. + disable_flash_attn2=True, + bf16=(dtype == torch.bfloat16), + fp16=(dtype == torch.float16), + output_dir="", + ) + + # transformers 5.6 loads the fast (tensor-native) image processor by + # default; gradient flow requires it — the slow variant round-trips + # through PIL and severs autograd. + model, processor, _ = create_model_and_processor( + model_config=model_config, + peft_lora_config=peft_lora_config, + training_args=training_args, + ) + model, _ = load_model_from_checkpoint(model, checkpoint_dir, -1) + + model.to(self.device) + model.eval() + model.requires_grad_(False) + + self.model = model + self.processor = processor + self.data_config = data_config_dict # raw dict, kept for debugging + + logger.info( + "VideoRewardWrapper loaded: ckpt=%s device=%s dtype=%s resize=%dx%d micro_bs=%d use_norm=%s template=%s", + checkpoint_dir, + self.device, + self.dtype, + self.resize_height, + self.resize_width, + self.micro_batch_size, + self.use_norm, + self.prompt_template_type, + ) + + # ------------------------------------------------------------------ + # Internals + # ------------------------------------------------------------------ + + def _prepare_input(self, data): + if isinstance(data, Mapping): + return type(data)({k: self._prepare_input(v) for k, v in data.items()}) + if isinstance(data, (tuple, list)): + return type(data)(self._prepare_input(v) for v in data) + if isinstance(data, torch.Tensor): + return data.to(device=self.device) + return data + + def _prepare_inputs(self, inputs): + inputs = self._prepare_input(inputs) + if len(inputs) == 0: + raise ValueError("VideoReward inputs must not be empty.") + return inputs + + @staticmethod + def _pixels_neg1_to_255(video: torch.Tensor) -> torch.Tensor: + """[-1, 1] float → [0, 255] float — stays differentiable.""" + return ((video * 0.5 + 0.5).clamp(0, 1) * 255.0).float() + + def _resize_for_reward(self, video_tchw: torch.Tensor) -> torch.Tensor: + return TF.resize( # type: ignore[no-any-return] + video_tchw, + [self.resize_height, self.resize_width], + interpolation=InterpolationMode.BICUBIC, + antialias=True, + ).float() + + # ------------------------------------------------------------------ + # Public interface + # ------------------------------------------------------------------ + + def prepare_batch_from_frames( + self, + video_tensors: List[torch.Tensor], + prompts: List[str], + ): + """Build a Qwen2-VL processor batch from ``[T,C,H,W]`` videos + prompts. + + Auto-transposes ``[C,T,H,W]`` callers when ``C==3``. + """ + chat_data = [ + [ + { + "role": "user", + "content": [ + {"type": "video", "video": "file://dummy_path"}, + { + "type": "text", + "text": self.build_prompt( + prompt, + self.eval_dim, + self.prompt_template_type, + ), + }, + ], + } + ] + for prompt in prompts + ] + + processed: List[torch.Tensor] = [] + for video in video_tensors: + if video.dim() != 4: + raise ValueError(f"Expected video tensor shape [T,C,H,W], got {tuple(video.shape)}") + # Auto-transpose if caller passed [C,T,H,W] with T > 3. + if video.shape[0] == 3 and video.shape[1] > 3: + video = video.permute(1, 0, 2, 3) + video = self._pixels_neg1_to_255(video) + video = self._resize_for_reward(video) + processed.append(video) + + batch = self.processor( + text=self.processor.apply_chat_template( + chat_data, + tokenize=False, + add_generation_prompt=True, + ), + images=None, + videos=processed, + padding=True, + return_tensors="pt", + videos_kwargs={"do_rescale": True}, + ) + return self._prepare_inputs(batch) + + def _norm( + self, + vq: torch.Tensor, + mq: torch.Tensor, + ta: torch.Tensor, + ): + if self.inference_config is None: + return vq, mq, ta + vq = (vq - self.inference_config["VQ_mean"]) / self.inference_config["VQ_std"] + mq = (mq - self.inference_config["MQ_mean"]) / self.inference_config["MQ_std"] + ta = (ta - self.inference_config["TA_mean"]) / self.inference_config["TA_std"] + return vq, mq, ta + + def forward_scores( + self, + video_tensors: List[torch.Tensor], + prompts: List[str], + use_norm: Optional[bool] = None, + ) -> Dict[str, torch.Tensor]: + """Differentiable scoring → per-sample (VQ, MQ, TA, Overall) scalars. + + Args: + video_tensors: list of ``(T, C, H, W)`` tensors, pixels in + ``[-1, 1]``. May carry ``grad_fn`` (REFL path). + prompts: per-sample user prompt text. + use_norm: override ``self.use_norm`` for this call. + + Returns: + ``{"VQ": Tensor[B], "MQ": Tensor[B], "TA": Tensor[B], "Overall": Tensor[B]}``. + Gradients flow back into ``video_tensors`` when they require grad. + """ + if len(video_tensors) != len(prompts): + raise ValueError("video_tensors and prompts must have the same batch size.") + use_norm = self.use_norm if use_norm is None else use_norm + + all_vq, all_mq, all_ta = [], [], [] + for start in range(0, len(video_tensors), self.micro_batch_size): + end = start + self.micro_batch_size + batch = self.prepare_batch_from_frames(video_tensors[start:end], prompts[start:end]) + # 5.x processors inject mm_token_type_ids; the backbone doesn't + # accept it (see module NOTE — pop, never signature-filter). + batch.pop("mm_token_type_ids", None) + logits = self.model(**batch, return_dict=True)["logits"] # (B, 3) + vq, mq, ta = logits[:, 0], logits[:, 1], logits[:, 2] + if use_norm: + vq, mq, ta = self._norm(vq, mq, ta) + all_vq.append(vq) + all_mq.append(mq) + all_ta.append(ta) + + vq = torch.cat(all_vq, dim=0) + mq = torch.cat(all_mq, dim=0) + ta = torch.cat(all_ta, dim=0) + return {"VQ": vq, "MQ": mq, "TA": ta, "Overall": vq + mq + ta} + + +__all__ = ["VideoRewardWrapper"] diff --git a/experimental/refl/roles.py b/experimental/refl/roles.py new file mode 100644 index 000000000..2265695d6 --- /dev/null +++ b/experimental/refl/roles.py @@ -0,0 +1,254 @@ +"""ReflActorRole — family-agnostic REFL/BPTT actor Remote for the refl recipe. + +Mirrors :class:`unirl.train.refl.policy.ReFLPolicy` (the SD3 image-ReFL actor): +a **config-chosen** ``Pipeline`` (``pipeline_target`` + ``model_config``, no +per-family imports), FSDP-wrapped in place via ``FSDPBackend``, driven by three +driver RPCs per step under the distributed ``enable_grad()`` context:: + + gen = actor.generate_samples(texts=…, images=…, params=…) # grad BPTT sample + VAE decode + rew = reward.score_differentiable(gen.decoded, prompts, records) + actor.forward_backward_loss(rewards=rew, kl_loss=gen.kl_loss) # loss seed → local backward + actor.step(max_grad_norm=…) # ctx exit routed grads → optimizer + +The pipeline named by ``pipeline_target`` must expose the recipe contract +(see ``experimental.refl.models``): ``build_refl_conditions(texts, images=…, +params=…)`` plus a ``diffusion`` stage with ``diffuse_with_grad`` and a +``vae_decode`` stage with ``decode_with_grad``. +""" + +from __future__ import annotations + +import dataclasses +import logging +from dataclasses import dataclass +from typing import Any, Dict, List, Optional, Tuple + +import torch +import torch.distributed as dist +from hydra.utils import get_class + +from unirl.distributed.group.dispatch import Dispatch, Execute, distributed +from unirl.distributed.group.remote import Remote +from unirl.distributed.tensor.batch import Batch, concat_field +from unirl.sde.runtime import get_sigma_schedule +from unirl.train.backend.base import LrSchedulerConfig, OptimizerConfig +from unirl.train.backend.fsdp import FSDPBackend +from unirl.train.configs import FSDPConfig, LoraConfig +from unirl.types.primitives import Images, Texts +from unirl.types.sampling import DiffusionSamplingParams + +logger = logging.getLogger(__name__) + + +@dataclass +class REFLGenerated(Batch): + """Generated BPTT payload: live-grad decoded pixels + per-sample KL. + + Both fields are batch-aligned concat fields (``decoded`` ``[B, …]``, + ``kl_loss`` ``[B]``) so DP_SCATTER merge/re-shard keeps each DP shard's + own values — the KL fed back into ``forward_backward_loss`` is the KL + that shard actually accumulated. + """ + + decoded: torch.Tensor = concat_field(default_factory=lambda: torch.empty(0)) + kl_loss: torch.Tensor = concat_field(default_factory=lambda: torch.empty(0)) + + +@dataclass +class REFLLossMetrics(Batch): + """Per-DP-shard REFL scalar metrics.""" + + loss: List[float] = concat_field(default_factory=list) + reward_loss: List[float] = concat_field(default_factory=list) + kl_loss: List[float] = concat_field(default_factory=list) + reward_mean: List[float] = concat_field(default_factory=list) + + +class ReflActorRole(Remote): + """Family-agnostic REFL actor: config-chosen pipeline + FSDP + BPTT loss.""" + + def __init__( + self, + *, + pipeline_target: str, + model_config: Any, + fsdp_cfg: FSDPConfig, + optimizer_cfg: OptimizerConfig, + scheduler_cfg: LrSchedulerConfig, + lora_cfg: Optional[LoraConfig] = None, + strategy: Optional[Any] = None, + block_class_names: Tuple[str, ...] = ("WanTransformerBlock",), + reward_weight: float = 1.0, + reward_baseline: float = 0.0, + reward_scale: float = 1.0, + kl_weight: float = 0.0, + ) -> None: + super().__init__() + self._pipeline_target = str(pipeline_target) + self._model_config = model_config + self._fsdp_cfg = fsdp_cfg + self._optimizer_cfg = optimizer_cfg + self._scheduler_cfg = scheduler_cfg + self._lora_cfg = lora_cfg + self._strategy = strategy + self._block_class_names = tuple(block_class_names) + self.reward_weight = float(reward_weight) + self.reward_baseline = float(reward_baseline) + self.reward_scale = float(reward_scale) + self.kl_weight = float(kl_weight) + + def initialize(self) -> None: + torch.cuda.set_device(self.device) + # Default PG over the actor role's workers (env:// from Remote.setup's + # dist_env); FSDP2 fully_shard wraps over it. Same order as ReFLPolicy. + if self.rank_info is not None and int(self.rank_info.world_size) > 1 and not dist.is_initialized(): + dist.init_process_group(backend="nccl") + + try: + self._model_config.device = self.device # runtime device injection + except Exception: + pass + + pipeline_cls = get_class(self._pipeline_target) + self.pipeline = pipeline_cls.from_config(self._model_config, strategy=self._strategy) + + for stage_attr, method in (("diffusion", "diffuse_with_grad"), ("vae_decode", "decode_with_grad")): + stage = getattr(self.pipeline, stage_attr, None) + if not hasattr(stage, method): + raise TypeError( + f"ReflActorRole: pipeline {self._pipeline_target} .{stage_attr} lacks {method}(...); " + f"use a experimental.refl.models pipeline (or implement the REFL contract)." + ) + if not hasattr(self.pipeline, "build_refl_conditions"): + raise TypeError( + f"ReflActorRole: pipeline {self._pipeline_target} lacks build_refl_conditions(...); " + f"use a experimental.refl.models pipeline." + ) + + # FSDP-wrap pipeline.bundle.transformer in place + LoRA + optimizer. + # The pipeline's stages reference the same bundle, so sampling runs + # through the wrapped trainable transformer. + self.backend = FSDPBackend( + bundle=self.pipeline.bundle, + block_class_names=self._block_class_names, + trainable_attr="transformer", + fsdp_cfg=self._fsdp_cfg, + optimizer_cfg=self._optimizer_cfg, + scheduler_cfg=self._scheduler_cfg, + device=self.device, + rank=int(self.rank_info.rank) if self.rank_info is not None else 0, + lora_cfg=self._lora_cfg, + ) + logger.info( + "ReflActorRole initialized: pipeline=%s reward_weight=%.3f kl_weight=%.3f", + self._pipeline_target, + self.reward_weight, + self.kl_weight, + ) + + # ------------------------------------------------------------------ + # Grad chain (run under the driver's enable_grad() context) + # ------------------------------------------------------------------ + + @distributed(dispatch_mode=Dispatch.DP_SCATTER) + def generate_samples( + self, + *, + texts: Texts, + images: Optional[Images] = None, + params: DiffusionSamplingParams, + ) -> REFLGenerated: + """Grad-enabled BPTT sampling + in-graph VAE decode. + + ``texts`` / ``images`` are the data-source conditioning primitives + (``images`` is the I2V first frame; ``None`` for T2V). Negative + prompts ride ``params.sampler_kwargs['negative_prompt']`` and are + expanded by the pipeline's ``build_refl_conditions``. + """ + self.backend.model.train() + self.backend.zero_grad() + + # Single KL knob: the actor owns kl_weight (loss-side), the sampling + # config owns only sampling-shape knobs (mid/final window). A stale + # sampler_kwargs.kl_weight is an error, not a silent override. + sampler_kwargs = dict(params.sampler_kwargs or {}) + if "kl_weight" in sampler_kwargs: + raise ValueError( + "ReflActorRole: sampling.sampler_kwargs.kl_weight is no longer read; " + "configure actor.kl_weight instead (single source of truth)." + ) + sampler_kwargs["kl_weight"] = self.kl_weight + + # Fixed init noise across rollouts and ranks — the contributor's + # verified DRaFT regime (params.seed used verbatim by the stage). + # Varying noise per rollout is a training-semantics change; do not + # introduce it here without its own evidence. + params = dataclasses.replace(params, sampler_kwargs=sampler_kwargs) + + conditions = self.pipeline.build_refl_conditions(texts, images=images, params=params) + schedule = get_sigma_schedule( + int(params.num_inference_steps), + shift=float(getattr(self.pipeline, "shift", 5.0)), + device=self.pipeline.bundle.device, + ) + result = self.pipeline.diffusion.diffuse_with_grad(conditions, schedule=schedule, params=params) + pixels = self.pipeline.vae_decode.decode_with_grad(result.z_final) + return REFLGenerated(decoded=pixels, kl_loss=result.kl_loss) + + @distributed(dispatch_mode=Dispatch.DP_SCATTER) + def forward_backward_loss( + self, + *, + rewards: torch.Tensor, + kl_loss: Optional[torch.Tensor] = None, + ) -> REFLLossMetrics: + """Assemble the REFL loss and backward on this shard's actor graph. + + ``rewards`` and ``kl_loss`` are RPC inputs marked as grad leaves by the + framework; on ``enable_grad()`` exit their grads route back through the + reward / generate calls into the sampling graph in ONE backward pass. + """ + reward = rewards.to(dtype=torch.bfloat16) + reward_loss = (-(reward - self.reward_baseline) / self.reward_scale * self.reward_weight).mean() + if kl_loss is not None and self.kl_weight != 0.0: + kl_term = self.kl_weight * kl_loss.float().mean() + else: + kl_term = torch.zeros((), device=reward.device, dtype=reward_loss.dtype) + loss = reward_loss + kl_term + loss.backward() + return REFLLossMetrics( + loss=[float(loss.detach().item())], + reward_loss=[float(reward_loss.detach().item())], + kl_loss=[float(kl_term.detach().item())], + reward_mean=[float(reward.detach().float().mean().item())], + ) + + # ------------------------------------------------------------------ + # Optimizer / checkpoint (delegate to the composed FSDPBackend) + # ------------------------------------------------------------------ + + @distributed(dispatch_mode=Dispatch.BROADCAST, execute_mode=Execute.ALL) + def step(self, *, max_grad_norm: float) -> Dict[str, float]: + """Clip + one optimizer step; returns grad_norm and current lr.""" + grad_norm = float(self.backend.optimizer_step(max_grad_norm=float(max_grad_norm))) + lr = 0.0 + sched = getattr(self.backend, "scheduler", None) + if sched is not None: + last = sched.get_last_lr() + lr = float(last[0]) if last else 0.0 + return {"grad_norm": grad_norm, "lr": lr} + + @distributed(dispatch_mode=Dispatch.BROADCAST, execute_mode=Execute.ALL) + def save(self, path: str, step: Optional[int] = None, mode: str = "adapter") -> None: + self.backend.save(path, step=step, mode=mode) + + @distributed(dispatch_mode=Dispatch.BROADCAST, execute_mode=Execute.ALL) + def load(self, path: str) -> int: + return self.backend.load(path) + + @distributed(dispatch_mode=Dispatch.BROADCAST, execute_mode=Execute.ALL) + def wait_for_checkpoint(self) -> None: + self.backend.wait_for_checkpoint() + + +__all__ = ["ReflActorRole", "REFLGenerated", "REFLLossMetrics"] diff --git a/experimental/refl/run.py b/experimental/refl/run.py new file mode 100755 index 000000000..2e28a1328 --- /dev/null +++ b/experimental/refl/run.py @@ -0,0 +1,19 @@ +#!/usr/bin/env python +"""Hydra entry point for the experimental.refl REFL recipe.""" + +from __future__ import annotations + +import hydra +from omegaconf import DictConfig + +from experimental.refl.trainer import REFLTrainer + + +@hydra.main(version_base=None, config_path="examples", config_name="wan22_i2v_face_refl") +def main(cfg: DictConfig) -> None: + trainer = REFLTrainer(cfg=cfg) + trainer.train() + + +if __name__ == "__main__": + main() diff --git a/experimental/refl/trainer.py b/experimental/refl/trainer.py new file mode 100644 index 000000000..3217fb3d8 --- /dev/null +++ b/experimental/refl/trainer.py @@ -0,0 +1,192 @@ +"""REFLTrainer — recipe driver for WAN REFL/BPTT (video reward backprop). + +The video sibling of :class:`unirl.trainer.refl.RewardBackpropTrainer`: two +roles, always — a :class:`experimental.refl.roles.ReflActorRole` (FSDP WAN + +grad BPTT sampling + optimizer) and a frozen differentiable video reward +(:class:`unirl.reward.service.RewardService`), colocated on the same worker +slab so decoded video never leaves the GPU. Each step runs, under the +distributed ``enable_grad()`` context:: + + gen = actor.generate_samples(texts, images, params) # grad through BPTT window + VAE + rew = reward.score_differentiable(gen.decoded, prompts, records) + actor.forward_backward_loss(rewards=rew, kl_loss=gen.kl_loss) + # ctx exit → grads route reward → decode → DiT LoRA params + actor.step(max_grad_norm) + +No advantages / replay / ratio / rollout-engine / weight-sync — REFL is not a +PG-RL loop. Success signal: the reward curve rises. +""" + +from __future__ import annotations + +import logging +import time +from typing import Any, Dict, List, Optional + +import numpy as np +from hydra.utils import instantiate +from omegaconf import DictConfig + +from unirl.distributed.group.placement import placement +from unirl.distributed.tensor.grad_context import enable_grad +from unirl.trainer.base import BaseTrainer, build_sampling_dict +from unirl.types.primitives import Images, Texts +from unirl.types.sample import Sample +from unirl.types.sampling import DiffusionSamplingParams, total_samples_per_prompt +from unirl.utils.hydra import remote_hydra + +logger = logging.getLogger(__name__) + + +def _text_inputs(inputs: Sample) -> Texts: + """Read the text root from a data-source input ``Sample``.""" + prompts = inputs.parts[0].primitives.get("text") if inputs.parts else None + if not isinstance(prompts, Texts): + raise TypeError(f"REFL data-source root requires Texts, got {type(prompts).__name__}.") + return prompts + + +def _image_inputs(inputs: Sample) -> Optional[Images]: + """Read the optional I2V first-frame conditioning (a chained input Part).""" + for part in inputs.parts[1:]: + image = part.primitives.get("image") + if isinstance(image, Images): + return image + return None + + +def _records(inputs: Sample) -> Optional[List[Optional[Dict[str, Any]]]]: + """Per-sample data-source metadata (e.g. ``ref_video_path`` for Face reward).""" + metadata = inputs.parts[0].metadata if inputs.parts else None + return list(metadata) if metadata else None + + +class REFLTrainer(BaseTrainer): + """REFL/BPTT recipe driver: actor + differentiable reward, 3-RPC step.""" + + def __init__(self, *, cfg: DictConfig) -> None: + super().__init__(cfg=cfg, logging_cfg=cfg.get("logging")) + self.cfg = cfg # train() reads run-length/checkpoint defaults from it + self.batch_size = int(cfg.batch_size) + self.max_grad_norm = float(cfg.get("max_grad_norm", 1.0)) + self.data_source = instantiate(cfg.data_source) + + sampling = build_sampling_dict(cfg.sampling) + params = sampling.get("diffusion") + if not isinstance(params, DiffusionSamplingParams): + raise TypeError("REFLTrainer requires cfg.sampling to build DiffusionSamplingParams.") + if total_samples_per_prompt(sampling) != 1: + raise ValueError( + "REFLTrainer trains one sample per prompt (no advantage groups); set sampling.samples_per_prompt=1." + ) + self.sampling_params = params + + # Reward shares the actor's workers: decoded video stays on-GPU for + # scoring. (Cross-slab reward placement — DiffusionTrainer's + # reward_fraction — is deliberately not supported here; add it only + # when a recipe actually cannot colocate.) + with placement(self.pool, fraction=1.0, shared_workers=True): + self.actor = remote_hydra(cfg.actor) + self.reward = remote_hydra(cfg.reward) + self.actor.initialize() + # BaseTrainer.maybe_save/load_checkpoint operate on ``self.backend``. + self.backend = self.actor + + adp, rdp = self.actor.dp_size, self.reward.dp_size + if self.batch_size % adp or self.batch_size % rdp: + raise ValueError(f"batch_size={self.batch_size} must be divisible by actor dp={adp} and reward dp={rdp}") + logger.info( + "REFLTrainer ready: actor dp=%d reward dp=%d batch=%d max_grad_norm=%.2f", + adp, + rdp, + self.batch_size, + self.max_grad_norm, + ) + + def train_step(self, inputs: Sample) -> Dict[str, float]: + """One enable_grad() generate → score → backward, then optimizer step.""" + t0 = time.perf_counter() + texts = _text_inputs(inputs) + images = _image_inputs(inputs) + records = _records(inputs) + with enable_grad(): + gen = self.actor.generate_samples( + texts=texts, + images=images, + params=self.sampling_params, + ) + rewards = self.reward.score_differentiable(gen.decoded, list(texts.texts), records) + loss_metrics = self.actor.forward_backward_loss(rewards=rewards, kl_loss=gen.kl_loss) + step_result = self.actor.step(max_grad_norm=self.max_grad_norm) + if isinstance(step_result, list): # BROADCAST → one result per worker + step_result = step_result[0] + + return { + "loss": float(np.mean(loss_metrics.loss)), + "reward_loss": float(np.mean(loss_metrics.reward_loss)), + "kl_loss": float(np.mean(loss_metrics.kl_loss)), + "reward_mean": float(np.mean(loss_metrics.reward_mean)), + "grad_norm": float(step_result.get("grad_norm", 0.0)), + "lr": float(step_result.get("lr", 0.0)), + "step_time_s": time.perf_counter() - t0, + } + + def train( + self, + *, + num_rollouts: Optional[int] = None, + save_interval: Optional[int] = None, + save_dir: Optional[str] = None, + load_dir: Optional[str] = None, + save_mode: Optional[str] = None, + ) -> None: + cfg = self.cfg + num_rollouts = int(num_rollouts if num_rollouts is not None else cfg.get("num_rollouts", 100)) + save_interval = int(save_interval if save_interval is not None else cfg.get("save_interval", 0)) + save_dir = save_dir if save_dir is not None else cfg.get("save_dir") + load_dir = load_dir if load_dir is not None else cfg.get("load_dir") + save_mode = str(save_mode if save_mode is not None else cfg.get("save_mode", "adapter")) + + start = self.maybe_load_checkpoint(load_dir, num_rollouts=num_rollouts) + for _ in range(start): # fast-forward the data stream on resume + self.data_source.get_samples(self.batch_size) + self._init_wandb(num_rollouts=num_rollouts) + try: + for rollout_id in range(start, num_rollouts): + inputs = self.data_source.get_samples(self.batch_size) + metrics = self.train_step(inputs) + logger.info( + "rollout %d/%d reward=%.4f loss=%.4f kl=%.4f grad_norm=%.4f %.1fs", + rollout_id + 1, + num_rollouts, + metrics["reward_mean"], + metrics["loss"], + metrics["kl_loss"], + metrics["grad_norm"], + metrics["step_time_s"], + ) + self.wandb_logger.log_step( + rollout_id + 1, + { + "rollout/mean_reward": metrics["reward_mean"], + "train/loss": metrics["loss"], + "train/reward_loss": metrics["reward_loss"], + "train/kl_loss": metrics["kl_loss"], + "train/grad_norm": metrics["grad_norm"], + "train/lr": metrics["lr"], + "perf/step_time_s": metrics["step_time_s"], + }, + prefix="", + ) + self.maybe_save_checkpoint( + rollout_id, + num_rollouts, + save_interval=save_interval, + save_dir=save_dir, + save_mode=save_mode, + ) + finally: + self._finish_wandb() + + +__all__ = ["REFLTrainer"] diff --git a/pyproject.toml b/pyproject.toml index 9eded6e3c..6428210fe 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -29,7 +29,10 @@ dependencies = [ # in transformer_boogu.py / attention_processor.py / rope.py. Pure-python, # engine-agnostic. "einops>=0.7", - "peft>=0.14.0", # exclude_modules support used by LoRA/NFT adapter injection + # 0.20 floor: older peft imports transformers cache symbols (HybridCache) + # removed in 5.x and fails at import against the transformers pin above. + # exclude_modules support (>=0.14) used by LoRA/NFT adapter injection. + "peft>=0.20", "safetensors>=0.4", "Pillow>=10", "requests>=2.31", diff --git a/scripts/check_recipe_targets.py b/scripts/check_recipe_targets.py index ceb14ba3f..eeb1a7d69 100755 --- a/scripts/check_recipe_targets.py +++ b/scripts/check_recipe_targets.py @@ -23,11 +23,11 @@ ROOT = Path(__file__).resolve().parents[1] # YAML trees that hold recipes / stage configs with ``_target_`` entries. -SCAN_DIRS = ["examples", "CPPO", "DRPO", "FlowDPPO", "unirl"] +SCAN_DIRS = ["examples", "experimental", "CPPO", "DRPO", "FlowDPPO", "unirl"] # Vendored / sub-project trees kept byte-pristine (mirror .pre-commit-config exclude). SKIP_PARTS = {".git", "vendor"} -_TARGET_RE = re.compile(r"""^\s*_target_:\s*['"]?(unirl\.[A-Za-z0-9_.]+)['"]?\s*$""") +_TARGET_RE = re.compile(r"""^\s*_target_:\s*['"]?((?:unirl|experimental)\.[A-Za-z0-9_.]+)['"]?\s*$""") @lru_cache(maxsize=None) diff --git a/scripts/verify_refl_kl_batching.py b/scripts/verify_refl_kl_batching.py new file mode 100755 index 000000000..cd76e979b --- /dev/null +++ b/scripts/verify_refl_kl_batching.py @@ -0,0 +1,166 @@ +#!/usr/bin/env python3 +"""Verify KL DP-batching semantics for the refl recipe (CPU, no Ray, no GPU). + +Regression verification for the P1 review finding on PR #210: the original +``REFLGenerated.kl_loss`` was a per-shard *scalar* ``shared_field`` — DP +collect kept only rank 0's KL and re-broadcast it to every actor rank, which +happened to "work" only on the verified ``batch_size == actor_dp == 8`` +topology (logs duplicated rank 0; other B/dp splits risked a hard shape +mismatch between the routed KL grad and each rank's saved scalar). + +The fix makes ``kl_loss`` a batch-aligned per-sample ``[B]`` concat field. +This script pins the invariants at the exact wire layer DP dispatch uses +(``pytree_chunk`` / ``pytree_cat`` / ``infer_batch_size``) across the +topologies called out in review: B == dp, B > dp, non-power-of-two, dp == 1, +and unequal actor/reward dp. + +Standalone by design (repo policy after #99/#267 is no unenforced test tree): +run it directly whenever the refl recipe's KL/reward wire types change:: + + python scripts/verify_refl_kl_batching.py +""" + +from __future__ import annotations + +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +import torch # noqa: E402 + +from experimental.refl.roles import REFLGenerated, REFLLossMetrics # noqa: E402 +from unirl.distributed.tensor.pytree import infer_batch_size, pytree_cat, pytree_chunk # noqa: E402 + +# (batch_size, dp) — B == dp (the only previously-verified shape), B > dp, +# non-power-of-two, and the degenerate dp == 1. +TOPOLOGIES = [(8, 8), (8, 2), (6, 3), (12, 4), (4, 1)] + + +def _generated(batch: int) -> REFLGenerated: + decoded = torch.arange(batch * 6, dtype=torch.float32).reshape(batch, 2, 3) + kl = torch.arange(batch, dtype=torch.float32) + 1.0 # distinct per sample + return REFLGenerated(decoded=decoded, kl_loss=kl) + + +def check_refl_generated_chunk_cat_roundtrip(batch: int, dp: int) -> None: + """Each DP shard sees exactly its own KL rows; merge restores the batch.""" + gen = _generated(batch) + assert infer_batch_size((gen,), {}) == batch + + shards = pytree_chunk(gen, dp, batch) + assert len(shards) == dp + per = batch // dp + for rank, shard in enumerate(shards): + expect = gen.kl_loss[rank * per : (rank + 1) * per] + assert torch.equal(shard.kl_loss, expect), f"rank {rank} got foreign KL rows" + assert torch.equal(shard.decoded, gen.decoded[rank * per : (rank + 1) * per]) + + merged = pytree_cat(shards) + assert torch.equal(merged.kl_loss, gen.kl_loss) + assert torch.equal(merged.decoded, gen.decoded) + + +def check_forward_backward_payload_alignment(batch: int, dp: int) -> None: + """rewards and kl_loss chunk in lockstep — the forward_backward_loss wire. + + The old scalar-shared KL made this payload rewards=[B] + kl=[1]; chunking + by the inferred batch could not keep the two aligned off the B == dp + topology. Per-sample KL makes both first-class batch columns. + """ + kwargs = { + "rewards": torch.randn(batch), + "kl_loss": torch.arange(batch, dtype=torch.float32), + } + assert infer_batch_size((), kwargs) == batch + shards = pytree_chunk(kwargs, dp, batch) + per = batch // dp + for rank, shard in enumerate(shards): + assert shard["rewards"].shape == (per,) + assert shard["kl_loss"].shape == (per,) + assert torch.equal(shard["kl_loss"], kwargs["kl_loss"][rank * per : (rank + 1) * per]) + + +def check_per_shard_backward_grad_shape(batch: int, dp: int) -> None: + """The KL grad each rank produces matches its saved generate output rows. + + Mirrors ReflActorRole.forward_backward_loss on one shard: the KL input is + a grad leaf of shape [B/dp]; after backward its .grad must be the same + shape, because GradContext routes it as out_grads onto the SAME rank's + saved kl tensor from generate_samples. With the old scalar KL this pairing + was [broadcast scalar] vs [rank-local scalar] and only lined up by luck. + """ + per = batch // dp + for _rank in range(dp): + rewards = torch.randn(per, requires_grad=True) + kl = torch.rand(per, requires_grad=True) + reward_loss = (-(rewards.to(torch.bfloat16) - 0.5) / 0.25 * 1.0).mean() + loss = reward_loss + 1.0 * kl.float().mean() + loss.backward() + assert kl.grad is not None and kl.grad.shape == kl.shape + assert rewards.grad is not None and rewards.grad.shape == rewards.shape + + +def check_unequal_actor_reward_dp_stays_aligned() -> None: + """B-length columns survive reward-dp merge → actor-dp re-chunk. + + The recipe colocates actor and reward (equal dp by construction), but the + wire contract must not depend on that: scoring merged at reward dp=2 and + re-scattered at actor dp=4 must hand every actor rank the reward/KL rows + of its own samples. + """ + batch, rdp, adp = 8, 2, 4 + rewards = torch.arange(batch, dtype=torch.float32) + reward_shards = pytree_chunk({"r": rewards}, rdp, batch) + merged = pytree_cat(reward_shards)["r"] + assert torch.equal(merged, rewards) + + kwargs = {"rewards": merged, "kl_loss": torch.arange(batch, dtype=torch.float32) * 10.0} + actor_shards = pytree_chunk(kwargs, adp, batch) + per = batch // adp + for rank, shard in enumerate(actor_shards): + assert torch.equal(shard["rewards"], rewards[rank * per : (rank + 1) * per]) + assert torch.equal(shard["kl_loss"], kwargs["kl_loss"][rank * per : (rank + 1) * per]) + + +def check_per_sample_kl_equals_legacy_scalar_mean() -> None: + """Per-sample reduction then batch-mean == the legacy global scalar mean.""" + torch.manual_seed(0) + kl_pred = torch.randn(3, 4, 2, 5, 5) + ref = torch.randn(3, 4, 2, 5, 5) + sigma = torch.tensor(0.7) + per_sample = ((kl_pred - ref) ** 2 / (2.0 * sigma**2)).flatten(1).mean(dim=1) + legacy = ((kl_pred - ref) ** 2 / (2.0 * sigma**2)).mean() + assert per_sample.shape == (3,) + assert torch.allclose(per_sample.mean(), legacy, atol=1e-6) + + +def check_loss_metrics_concat_keeps_every_shard() -> None: + """REFLLossMetrics concat lists must surface every shard's scalars.""" + shards = [ + REFLLossMetrics(loss=[0.1], reward_loss=[0.2], kl_loss=[0.3], reward_mean=[0.4]), + REFLLossMetrics(loss=[1.1], reward_loss=[1.2], kl_loss=[1.3], reward_mean=[1.4]), + ] + merged = pytree_cat(shards) + assert merged.loss == [0.1, 1.1] + assert merged.kl_loss == [0.3, 1.3] + + +def main() -> int: + for batch, dp in TOPOLOGIES: + check_refl_generated_chunk_cat_roundtrip(batch, dp) + check_forward_backward_payload_alignment(batch, dp) + check_per_shard_backward_grad_shape(batch, dp) + print(f"[ok] topology B={batch} dp={dp}: chunk/cat, payload lockstep, backward shapes") + check_unequal_actor_reward_dp_stays_aligned() + print("[ok] unequal actor/reward dp (rdp=2 → adp=4) stays row-aligned") + check_per_sample_kl_equals_legacy_scalar_mean() + print("[ok] per-sample KL batch-mean equals legacy scalar mean") + check_loss_metrics_concat_keeps_every_shard() + print("[ok] REFLLossMetrics keeps every shard's scalars") + print("verify-refl-kl-batching: ALL CHECKS PASSED") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/unirl/data/data_source.py b/unirl/data/data_source.py index 20ac93080..411becbc7 100644 --- a/unirl/data/data_source.py +++ b/unirl/data/data_source.py @@ -310,12 +310,14 @@ def __init__(self, args): args: Hydra ``cfg`` (DictConfig) with: - run.data_path: Path to data file (JSON, JSONL, or TXT) - run.seed: Random seed + - run.shuffle: Whether to shuffle the training prompts (default: True) - algorithm.prompts_per_rollout: Batch size """ self.args = args self.data_path = args.run.data_path self.eval_data_path = args.run.eval_data_path self.seed = args.run.seed + self.shuffle = bool(getattr(args.run, "shuffle", True)) self.prompts_per_rollout = int(args.algorithm.prompts_per_rollout) self.drop_last = True @@ -401,12 +403,13 @@ def _create_dataloader(self) -> None: f"prompts_per_rollout={self.prompts_per_rollout})." ) + should_shuffle = sampler is None and self.shuffle self._dataloader = DataLoader( self.train_dataset, batch_size=self.prompts_per_rollout, sampler=sampler, - shuffle=(sampler is None), # Only shuffle if not using custom sampler - generator=self._shuffle_generator if sampler is None else None, + shuffle=should_shuffle, # Only shuffle if not using custom sampler + generator=self._shuffle_generator if should_shuffle else None, num_workers=0, # Keep simple for Ray collate_fn=self._collate_text, drop_last=True, diff --git a/unirl/distributed/group/worker.py b/unirl/distributed/group/worker.py index 5f16d68e4..1949d8aa6 100644 --- a/unirl/distributed/group/worker.py +++ b/unirl/distributed/group/worker.py @@ -309,9 +309,14 @@ def resolve(o): resolved_kwargs = map_tree(kwargs, resolve) if grad_mode: - # get_batch returns detached views/copies, so resolved tensors are - # fresh objects that don't alias store contents — mark them directly. - tensors = collect_leaves(resolved_args, Tensor) + collect_leaves(tuple(resolved_kwargs.values()), Tensor) + # Cross-RPC autograd can only propagate gradients back to controller-side + # TensorRef inputs recorded by Handle as input_metas. + tensors = [fetched[str(i)] for i in range(len(in_metas))] + if len(tensors) != len(in_metas): + raise RuntimeError( + f"Worker.call grad input mismatch for {method_name} call_id={call_id}: " + f"saved {len(tensors)} tensors for {len(in_metas)} TensorRef inputs" + ) for t in tensors: t.requires_grad_(True) t.retain_grad() diff --git a/unirl/models/wan21/bundle.py b/unirl/models/wan21/bundle.py index 2a5c29ccf..a5c157228 100644 --- a/unirl/models/wan21/bundle.py +++ b/unirl/models/wan21/bundle.py @@ -81,14 +81,13 @@ def uses_clip_vision(self) -> bool: def from_config(cls, config: WAN21PipelineConfig) -> "WAN21Bundle": """Load all WAN 2.1 components from a HuggingFace checkpoint.""" try: - from diffusers import AutoencoderKLWan, WanTransformer3DModel + from diffusers import WanTransformer3DModel except ImportError: # Fallback for older diffusers: ``AutoModel`` does dynamic # dispatch on the checkpoint config. Matches the fallback in # legacy ``models/wan21.py``. from diffusers import AutoModel - AutoencoderKLWan = AutoModel WanTransformer3DModel = AutoModel try: from transformers import AutoTokenizer, UMT5EncoderModel @@ -136,7 +135,28 @@ def from_config(cls, config: WAN21PipelineConfig) -> "WAN21Bundle": vae: Optional[nn.Module] = None if config.load_vae: - vae = AutoencoderKLWan.from_pretrained(vae_path, subfolder="vae", torch_dtype=vae_dtype).to(device).eval() + from .wan_video_vae import WanVideoVAE + + # ``load_from_diffusers`` reads local files only. Hub repo ids + # (e.g. ``Wan-AI/Wan2.1-T2V-1.3B-Diffusers``, the mainline WAN + # config default) resolve through the HF cache first, preserving + # the loading semantics of the previous + # ``AutoencoderKLWan.from_pretrained`` path. + vae_src = vae_path + if not os.path.isdir(vae_src): + from huggingface_hub import snapshot_download + + vae_src = snapshot_download(repo_id=vae_src, allow_patterns=["vae/*"]) + + vae = ( + WanVideoVAE.load_from_diffusers( + vae_src, + use_nested_grad_checkpoint=True, + use_act_grad_only_conv=True, + ) + .to(device=device, dtype=vae_dtype) + .eval() + ) vae.requires_grad_(False) text_encoder = ( diff --git a/unirl/models/wan21/image_encode.py b/unirl/models/wan21/image_encode.py index 9a591e660..1767aebde 100644 --- a/unirl/models/wan21/image_encode.py +++ b/unirl/models/wan21/image_encode.py @@ -109,20 +109,6 @@ def encode(self, p: Images) -> ImageLatentCondition: latent_condition = latent_condition.to(device=device, dtype=dtype) - # Per-channel normalization — strict inverse of ``wan21/vae.py:78-87`` - # decode: ``latent_decoded = stored * std + mean``. - vae_config = vae.config - latents_mean = getattr(vae_config, "latents_mean", None) - latents_std = getattr(vae_config, "latents_std", None) - if latents_mean is not None and latents_std is not None: - z_dim = int(getattr(vae_config, "z_dim", latent_condition.shape[1])) - mean = torch.tensor(latents_mean, device=device, dtype=dtype).view(1, z_dim, 1, 1, 1) - std = torch.tensor(latents_std, device=device, dtype=dtype).view(1, z_dim, 1, 1, 1) - latent_condition = (latent_condition - mean) / std - else: - scaling_factor = float(getattr(vae_config, "scaling_factor", 1.0)) - latent_condition = latent_condition * scaling_factor - # 4-channel first-frame mask (mirrors diffusers # ``pipeline_wan_i2v.py:468-479``): 1.0 at first pixel-time slot, # 0.0 elsewhere; then view-reshape + transpose so the temporal diff --git a/unirl/models/wan21/vae.py b/unirl/models/wan21/vae.py index 49e65238e..a8b8b7f45 100644 --- a/unirl/models/wan21/vae.py +++ b/unirl/models/wan21/vae.py @@ -1,31 +1,24 @@ """WAN21VAEDecodeStage — LatentSegment → Videos via 3D VAE decode. +Framework-level ``DecodeStage`` implementation (no REFL-only knobs). Both +GRPO and REFL recipes route their final ``LatentSegment`` through this stage +— the memory / autograd optimizations in ``WanVideoVAE`` (below) are what +make BPTT decode feasible, so pushing this into a recipe would fork the +decode contract across recipes. Kept in the framework. + Implements ``DecodeStage[LatentSegment, Videos]``. Reads the final stored position from ``LatentSegment.latents[:, -1]`` (the clean latent at -``T``, which ``WAN21DiffusionStage`` always stores), denormalizes using -either per-channel ``latents_mean`` / ``latents_std`` (when the VAE -config carries them, as recent diffusers ``AutoencoderKLWan`` does) or -the scalar ``scaling_factor`` fallback, runs VAE decode in fp32 (bf16 is -unsupported by most VAE implementations), and packs the 5D -``[B, C, T, H, W]`` output into a varlen-batched ``Videos`` primitive. - -**Why per-channel mean/std support (vs the legacy scalar-only path):** -diffusers' canonical Wan VAE ships with ``latents_mean`` / -``latents_std`` arrays in its config; using only ``scaling_factor`` -yields off-distribution decodes. Legacy ``models/wan21.py::decode_latents`` -uses the scalar fallback alone — a known latent-norm bug. The new path -follows the diffusers spec. +``T``, which ``WAN21DiffusionStage`` always stores), then dispatches to +the bundle's ``WanVideoVAE.decode`` . All un-normalization +(per-channel ``mean`` / ``std``), spatial tiling, nested gradient +checkpointing inside the decoder, and ``Conv3dActGradOnly`` are owned +by ``WanVideoVAE`` itself; this stage is just the +``LatentSegment → Videos`` adapter. VAE encode for I2V's image-condition latent lives in -``WAN21ImageLatentEncodeStage`` (``image_encode.py``) — it shares the -same per-channel norm helper path but as a sibling stage so the decode -side stays focused on its single job. There is intentionally no -generic ``WAN21VAEEncodeStage`` here; the I2V encode is the only -encode path we need. - -Decode math derived from diffusers' ``AutoencoderKLWan`` reference -and ``unirl/models/wan21.py:325-342, 450-462`` (do NOT import -legacy code). +``WAN21ImageLatentEncodeStage`` (``image_encode.py``) — it goes through +the same ``WanVideoVAE.encode(x).latent_dist.mode()`` diffusers-style +contract; there is intentionally no generic ``WAN21VAEEncodeStage``. """ from __future__ import annotations @@ -83,44 +76,42 @@ def decode(self, s: LatentSegment, *, grad: bool = False, activation_checkpoint: f"[B, C, T_lat, H_lat, W_lat], got {tuple(clean.shape)}" ) - vae = self.bundle.vae - vae_config = vae.config - latents_mean = getattr(vae_config, "latents_mean", None) - latents_std = getattr(vae_config, "latents_std", None) - scaling_factor = getattr(vae_config, "scaling_factor", 1.0) - - device = clean.device - - def _decode(lat: torch.Tensor) -> torch.Tensor: - latents_f32 = lat.to(dtype=torch.float32) - if latents_mean is not None and latents_std is not None: - # diffusers Wan VAE spec: latent ↦ latent * std + mean - # (un-normalize). Reshape to [1, C, 1, 1, 1] to broadcast - # over the (B, T, H, W) axes. - z_dim = int(getattr(vae_config, "z_dim", lat.shape[1])) - mean = torch.tensor(latents_mean, device=device, dtype=torch.float32).view(1, z_dim, 1, 1, 1) - std = torch.tensor(latents_std, device=device, dtype=torch.float32).view(1, z_dim, 1, 1, 1) - latents_f32 = latents_f32 * std + mean - else: - latents_f32 = latents_f32 / float(scaling_factor) - - return vae.to(torch.float32).decode(latents_f32).sample - with nullcontext() if grad else torch.no_grad(): - if grad and activation_checkpoint and clean.requires_grad: - from torch.utils.checkpoint import checkpoint - - decoded = checkpoint(_decode, clean, use_reentrant=False) - else: - decoded = _decode(clean) + decoded = self._vae_decode(clean) # Decoded layout is [B, C, T_dec, H_dec, W_dec] in [-1, 1]. # Normalize to [0, 1] and clamp before packing. decoded = ((decoded + 1.0) / 2.0).clamp(0.0, 1.0) - # Pack into the varlen ``Videos`` primitive: ``Video.frames`` is - # ``[T, C, H, W]`` so we permute each sample from (C, T, H, W) → - # (T, C, H, W) and let ``Videos.from_list`` concat along T. + return self._pack_videos(decoded) + + # ------------------------------------------------------------------ + # BPTT path (REFL): differentiable decode of a live grad latent. + # ------------------------------------------------------------------ + + def decode_with_grad(self, z_final: torch.Tensor) -> torch.Tensor: + """Differentiable VAE decode: ``z_final → pixels`` with autograd alive.""" + if z_final.ndim != 5: + raise ValueError( + f"WAN21VAEDecodeStage.decode_with_grad: expected 5D z_final " + f"[B, C, T_lat, H_lat, W_lat], got {tuple(z_final.shape)}" + ) + return self._vae_decode(z_final) + + # ------------------------------------------------------------------ + # Shared decode kernel (used by both `decode` and `decode_with_grad`). + # ------------------------------------------------------------------ + + def _vae_decode(self, clean: torch.Tensor) -> torch.Tensor: + """Run ``WanVideoVAE.decode`` and return pixels in the VAE's native ``[-1, 1]``.""" + vae = self.bundle.vae + vae_dtype = vae.dtype + latents = clean.to(dtype=vae_dtype) + decoded = vae.decode(latents, device=latents.device, tiled=True) + return decoded + + def _pack_videos(self, decoded: torch.Tensor) -> Videos: + """Pack a dense ``[B, C, T_dec, H_dec, W_dec]`` pixel tensor into ``Videos``.""" videos = [Video(frames=decoded[i].permute(1, 0, 2, 3).contiguous()) for i in range(int(decoded.shape[0]))] return Videos.from_list(videos) diff --git a/unirl/models/wan21/wan_video_vae.py b/unirl/models/wan21/wan_video_vae.py new file mode 100644 index 000000000..22b246266 --- /dev/null +++ b/unirl/models/wan21/wan_video_vae.py @@ -0,0 +1,1288 @@ +"""WanVideoVAE — Wan 2.x Video VAE with training optimizations.""" + +from __future__ import annotations + +import os +from collections import OrderedDict + +import torch +import torch.distributed as dist +import torch.nn as nn +import torch.nn.functional as F +import torch.utils.checkpoint +from einops import rearrange, repeat +from torch.nn.modules.utils import _triple + +CACHE_T = 2 + + +# ── Memory-efficient Conv layers ── + + +class Conv3dActGradOnlyFunction(torch.autograd.Function): + """Conv3d that only computes input gradient, not weight gradient. + + For frozen VAE: weights never update, so no need to save activations + for weight grad computation. Saves significant GPU memory. + """ + + @staticmethod + def forward(ctx, input, weight, bias, stride, padding, dilation, groups): + ctx.save_for_backward(weight) + ctx.input_info = (input.size(), input.stride(), input.dtype) + ctx.stride = stride + ctx.padding = padding + ctx.dilation = dilation + ctx.groups = groups + return F.conv3d(input, weight, bias, stride, padding, dilation, groups) + + @staticmethod + def backward(ctx, grad_output): + (weight,) = ctx.saved_tensors + grad_input = None + if ctx.needs_input_grad[0]: + input = torch.empty_strided( + ctx.input_info[0], ctx.input_info[1], dtype=ctx.input_info[2], device=grad_output.device + ) + grad_input = torch.ops.aten.convolution_backward( + grad_output, + input, + weight, + None, + _triple(ctx.stride), + _triple(ctx.padding), + _triple(ctx.dilation), + False, + [0, 0, 0], + ctx.groups, + (True, False, False), + )[0] + return grad_input, None, None, None, None, None, None + + +class Conv2dActGradOnlyFunction(torch.autograd.Function): + """Conv2d that only computes input gradient.""" + + @staticmethod + def forward(ctx, input, weight, bias, stride, padding, dilation, groups): + ctx.save_for_backward(weight) + ctx.input_shape = input.shape + ctx.stride = stride + ctx.padding = padding + ctx.dilation = dilation + ctx.groups = groups + return F.conv2d(input, weight, bias, stride, padding, dilation, groups) + + @staticmethod + def backward(ctx, grad_output): + (weight,) = ctx.saved_tensors + grad_input = None + if ctx.needs_input_grad[0]: + from torch.nn.grad import conv2d_input + + grad_input = conv2d_input( + ctx.input_shape, weight, grad_output, ctx.stride, ctx.padding, ctx.dilation, ctx.groups + ) + return grad_input, None, None, None, None, None, None + + +class Conv2dActGradOnly(nn.Conv2d): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.weight.requires_grad = False + if self.bias is not None: + self.bias.requires_grad = False + + def forward(self, x): + return Conv2dActGradOnlyFunction.apply( + x, self.weight, self.bias, self.stride, self.padding, self.dilation, self.groups + ) + + +class CausalConv3d(nn.Conv3d): + """Causal 3D convolution with temporal padding.""" + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._padding = (self.padding[2], self.padding[2], self.padding[1], self.padding[1], 2 * self.padding[0], 0) + self.padding = (0, 0, 0) + self.time_kernel_size = self.kernel_size[0] + + def forward_pad(self, x, cache_x=None): + padding = list(self._padding) + if cache_x is not None and self._padding[4] > 0: + cache_x = cache_x.to(x.device) + x = torch.cat([cache_x, x], dim=2) + padding[4] -= cache_x.shape[2] + x = F.pad(x, padding) + return x + + def forward(self, x, cache_x=None): + x = self.forward_pad(x, cache_x) + return super().forward(x) + + +class CausalConv3dActGradOnly(CausalConv3d): + """CausalConv3d with act-grad-only optimization.""" + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.weight.requires_grad = False + if self.bias is not None: + self.bias.requires_grad = False + + def forward(self, x, cache_x=None): + x = self.forward_pad(x, cache_x) + return Conv3dActGradOnlyFunction.apply( + x, self.weight, self.bias, self.stride, self.padding, self.dilation, self.groups + ) + + +# ── Building blocks ── + + +def check_is_instance(model, module_class): + if isinstance(model, module_class): + return True + if hasattr(model, "module") and isinstance(model.module, module_class): + return True + return False + + +class RMS_norm(nn.Module): + def __init__(self, dim, channel_first=True, images=True, bias=False): + super().__init__() + broadcastable_dims = (1, 1, 1) if not images else (1, 1) + shape = (dim, *broadcastable_dims) if channel_first else (dim,) + self.channel_first = channel_first + self.scale = dim**0.5 + self.gamma = nn.Parameter(torch.ones(shape)) + self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0.0 + + def forward(self, x): + return F.normalize(x, dim=(1 if self.channel_first else -1)) * self.scale * self.gamma + self.bias + + +class Upsample(nn.Upsample): + def forward(self, x): + return super().forward(x.float()).type_as(x) + + +class Resample(nn.Module): + def __init__(self, dim, mode): + assert mode in ("none", "upsample2d", "upsample3d", "downsample2d", "downsample3d") + super().__init__() + self.dim = dim + self.mode = mode + + if mode == "upsample2d": + self.resample = nn.Sequential( + Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"), nn.Conv2d(dim, dim // 2, 3, padding=1) + ) + elif mode == "upsample3d": + self.resample = nn.Sequential( + Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"), nn.Conv2d(dim, dim // 2, 3, padding=1) + ) + causal_conv3d = CausalConv3d(dim, dim * 2, (3, 1, 1), padding=(1, 0, 0)) + if hasattr(self, "decoder") and getattr(self, "decoder"): + setattr(causal_conv3d, "decoder", True) + self.time_conv = causal_conv3d + elif mode == "downsample2d": + self.resample = nn.Sequential(nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2))) + elif mode == "downsample3d": + self.resample = nn.Sequential(nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2))) + causal_conv3d = CausalConv3d(dim, dim, (3, 1, 1), stride=(2, 1, 1), padding=(0, 0, 0)) + if hasattr(self, "decoder") and getattr(self, "decoder"): + setattr(causal_conv3d, "decoder", True) + self.time_conv = causal_conv3d + else: + self.resample = nn.Identity() + + def forward(self, x, feat_cache=None, feat_idx=None): + b, c, t, h, w = x.size() + if self.mode == "upsample3d": + if hasattr(self, "decoder") and getattr(self, "decoder"): + x1 = x[:, :, 1:] + x1 = self.time_conv(x1) + x1 = x1.reshape(b, 2, c, t - 1, h, w) + x1 = torch.stack((x1[:, 0], x1[:, 1]), 3) + x1 = x1.reshape(b, c, (t - 1) * 2, h, w) + x = torch.cat([x[:, :, 0:1], x1], dim=2) + elif feat_cache is not None: + idx = feat_idx[0] + if feat_cache[idx] is None: + feat_cache[idx] = "Rep" + feat_idx[0] += 1 + else: + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx] != "Rep": + cache_x = torch.cat( + [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2 + ) + if cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx] == "Rep": + cache_x = torch.cat([torch.zeros_like(cache_x).to(cache_x.device), cache_x], dim=2) + if feat_cache[idx] == "Rep": + x = self.time_conv(x) + else: + x = self.time_conv(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + x = x.reshape(b, 2, c, t, h, w) + x = torch.stack((x[:, 0], x[:, 1]), 3) + x = x.reshape(b, c, t * 2, h, w) + + t = x.shape[2] + x = rearrange(x, "b c t h w -> (b t) c h w") + x = self.resample(x) + x = rearrange(x, "(b t) c h w -> b c t h w", t=t) + + if self.mode == "downsample3d": + if hasattr(self, "decoder") and getattr(self, "decoder"): + x = self.time_conv(x) + elif feat_cache is not None: + idx = feat_idx[0] + if feat_cache[idx] is None: + feat_cache[idx] = x.clone() + feat_idx[0] += 1 + else: + cache_x = x[:, :, -1:, :, :].clone() + x = self.time_conv(torch.cat([feat_cache[idx][:, :, -1:, :, :], x], 2)) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + + if feat_cache is None or feat_idx is None: + return x + return x, feat_cache, feat_idx + + +def patchify(x, patch_size): + if patch_size == 1: + return x + if x.dim() == 5: + return rearrange(x, "b c f (h q) (w r) -> b (c r q) f h w", q=patch_size, r=patch_size) + if x.dim() == 4: + return rearrange(x, "b c (h q) (w r) -> b (c r q) h w", q=patch_size, r=patch_size) + raise ValueError(f"Invalid input shape: {x.shape}") + + +def unpatchify(x, patch_size): + if patch_size == 1: + return x + if x.dim() == 5: + return rearrange(x, "b (c r q) f h w -> b c f (h q) (w r)", q=patch_size, r=patch_size) + if x.dim() == 4: + return rearrange(x, "b (c r q) h w -> b c (h q) (w r)", q=patch_size, r=patch_size) + return x + + +class ResidualBlock(nn.Module): + def __init__(self, in_dim, out_dim, dropout=0.0): + super().__init__() + self.in_dim = in_dim + self.out_dim = out_dim + + causal_conv3d_block1 = CausalConv3d(in_dim, out_dim, 3, padding=1) + causal_conv3d_block2 = CausalConv3d(out_dim, out_dim, 3, padding=1) + causal_conv3d_shortcut = CausalConv3d(in_dim, out_dim, 1) + self.residual = nn.Sequential( + RMS_norm(in_dim, images=False), + nn.SiLU(), + causal_conv3d_block1, + RMS_norm(out_dim, images=False), + nn.SiLU(), + nn.Dropout(dropout), + causal_conv3d_block2, + ) + self.shortcut = causal_conv3d_shortcut if in_dim != out_dim else nn.Identity() + + def forward(self, x, feat_cache=None, feat_idx=None): + h = self.shortcut(x) + for layer in self.residual: + if hasattr(self, "decoder") and getattr(self, "decoder"): + x = layer(x) + elif check_is_instance(layer, CausalConv3d) and feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat( + [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2 + ) + x = layer(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = layer(x) + + if feat_cache is None or feat_idx is None: + return x + h + return x + h, feat_cache, feat_idx + + +class AttentionBlock(nn.Module): + """Causal self-attention with a single head.""" + + def __init__(self, dim): + super().__init__() + self.dim = dim + self.norm = RMS_norm(dim) + self.to_qkv = nn.Conv2d(dim, dim * 3, 1) + self.proj = nn.Conv2d(dim, dim, 1) + nn.init.zeros_(self.proj.weight) + + def forward(self, x): + identity = x + b, c, t, h, w = x.size() + x = rearrange(x, "b c t h w -> (b t) c h w") + x = self.norm(x) + q, k, v = self.to_qkv(x).reshape(b * t, 1, c * 3, -1).permute(0, 1, 3, 2).contiguous().chunk(3, dim=-1) + x = F.scaled_dot_product_attention(q, k, v) + x = x.squeeze(1).permute(0, 2, 1).reshape(b * t, c, h, w) + x = self.proj(x) + x = rearrange(x, "(b t) c h w-> b c t h w", t=t) + return x + identity + + +# ── Encoder / Decoder ── + + +class Encoder3d(nn.Module): + def __init__( + self, + dim=128, + z_dim=4, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_downsample=[True, True, False], + dropout=0.0, + ): + super().__init__() + dims = [dim * u for u in [1] + dim_mult] + scale = 1.0 + + self.conv1 = CausalConv3d(3, dims[0], 3, padding=1) + + downsamples = [] + for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): + for _ in range(num_res_blocks): + downsamples.append(ResidualBlock(in_dim, out_dim, dropout)) + if scale in attn_scales: + downsamples.append(AttentionBlock(out_dim)) + in_dim = out_dim + if i != len(dim_mult) - 1: + mode = "downsample3d" if temperal_downsample[i] else "downsample2d" + downsamples.append(Resample(out_dim, mode=mode)) + scale /= 2.0 + self.downsamples = nn.Sequential(*downsamples) + + self.middle = nn.Sequential( + ResidualBlock(out_dim, out_dim, dropout), AttentionBlock(out_dim), ResidualBlock(out_dim, out_dim, dropout) + ) + + self.head = nn.Sequential( + RMS_norm(out_dim, images=False), nn.SiLU(), CausalConv3d(out_dim, z_dim, 3, padding=1) + ) + + def forward(self, x, feat_cache=None, feat_idx=None): + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2) + x = self.conv1(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = self.conv1(x) + + for layer in self.downsamples: + if feat_cache is not None: + x, feat_cache, feat_idx = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + for layer in self.middle: + if check_is_instance(layer, ResidualBlock) and feat_cache is not None: + x, feat_cache, feat_idx = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + for layer in self.head: + if check_is_instance(layer, CausalConv3d) and feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat( + [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2 + ) + x = layer(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = layer(x) + return x, feat_cache, feat_idx + + +class Decoder3d(nn.Module): + def __init__( + self, + dim=128, + z_dim=4, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_upsample=[False, True, True], + dropout=0.0, + use_nested_grad_checkpoint=True, + ): + super().__init__() + self.use_nested_grad_checkpoint = use_nested_grad_checkpoint + + dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]] + scale = 1.0 / 2 ** (len(dim_mult) - 2) + + self.conv1 = CausalConv3d(z_dim, dims[0], 3, padding=1) + setattr(self.conv1, "decoder", True) + + res_block_1 = ResidualBlock(dims[0], dims[0], dropout) + attn_block_2 = AttentionBlock(dims[0]) + res_block_3 = ResidualBlock(dims[0], dims[0], dropout) + for m in (res_block_1, attn_block_2, res_block_3): + setattr(m, "decoder", True) + self.middle = nn.Sequential(res_block_1, attn_block_2, res_block_3) + + upsamples = [] + for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): + if i in (1, 2, 3): + in_dim = in_dim // 2 + for _ in range(num_res_blocks + 1): + res_block = ResidualBlock(in_dim, out_dim, dropout) + setattr(res_block, "decoder", True) + upsamples.append(res_block) + if scale in attn_scales: + attn_block = AttentionBlock(out_dim) + setattr(attn_block, "decoder", True) + upsamples.append(attn_block) + in_dim = out_dim + if i != len(dim_mult) - 1: + mode = "upsample3d" if temperal_upsample[i] else "upsample2d" + resample = Resample(out_dim, mode=mode) + setattr(resample, "decoder", True) + upsamples.append(resample) + scale *= 2.0 + self.upsamples = nn.Sequential(*upsamples) + + causal_conv3d = CausalConv3d(out_dim, 3, 3, padding=1) + setattr(causal_conv3d, "decoder", True) + self.head = nn.Sequential(RMS_norm(out_dim, images=False), nn.SiLU(), causal_conv3d) + + self.unsample_splits = [ + len(self.upsamples) // 15 * 10, + len(self.upsamples) // 15 * 13, + ] + + def forward(self, x): + def custom_forward1(x): + x = self.conv1(x) + for layer in self.middle: + x = layer(x) + for layer in self.upsamples[: self.unsample_splits[0]]: + x = layer(x) + return x + + def custom_forward2(x): + for layer in self.upsamples[self.unsample_splits[0] : self.unsample_splits[1]]: + x = layer(x) + return x + + def custom_forward3(x): + for layer in self.upsamples[self.unsample_splits[1] :]: + x = layer(x) + for layer in self.head: + x = layer(x) + return x + + if self.use_nested_grad_checkpoint: + x = torch.utils.checkpoint.checkpoint(custom_forward1, x, use_reentrant=True) + x = torch.utils.checkpoint.checkpoint(custom_forward2, x, use_reentrant=True) + x = torch.utils.checkpoint.checkpoint(custom_forward3, x, use_reentrant=True) + else: + x = custom_forward1(x) + x = custom_forward2(x) + x = custom_forward3(x) + return x + + +def count_conv3d(model): + return sum(1 for m in model.modules() if isinstance(m, CausalConv3d)) + + +# ── VideoVAE_ (core model) ── + + +class VideoVAE_(nn.Module): + def __init__( + self, + dim=96, + z_dim=16, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_downsample=[False, True, True], + dropout=0.0, + use_nested_grad_checkpoint=True, + ): + super().__init__() + self.dim = dim + self.z_dim = z_dim + self.dim_mult = dim_mult + self.num_res_blocks = num_res_blocks + self.temperal_downsample = temperal_downsample + self.temperal_upsample = temperal_downsample[::-1] + self.use_nested_grad_checkpoint = use_nested_grad_checkpoint + + self.encoder = Encoder3d( + dim, z_dim * 2, dim_mult, num_res_blocks, attn_scales, self.temperal_downsample, dropout + ) + self.conv1 = CausalConv3d(z_dim * 2, z_dim * 2, 1) + self.conv2 = CausalConv3d(z_dim, z_dim, 1) + self.decoder = Decoder3d( + dim, + z_dim, + dim_mult, + num_res_blocks, + attn_scales, + self.temperal_upsample, + dropout, + use_nested_grad_checkpoint, + ) + + def encode(self, x, scale): + self.clear_cache() + t = x.shape[2] + iter_ = 1 + (t - 1) // 4 + for i in range(iter_): + self._enc_conv_idx = [0] + if i == 0: + out, self._enc_feat_map, self._enc_conv_idx = self.encoder( + x[:, :, :1, :, :], feat_cache=self._enc_feat_map, feat_idx=self._enc_conv_idx + ) + else: + out_, self._enc_feat_map, self._enc_conv_idx = self.encoder( + x[:, :, 1 + 4 * (i - 1) : 1 + 4 * i, :, :], + feat_cache=self._enc_feat_map, + feat_idx=self._enc_conv_idx, + ) + out = torch.cat([out, out_], 2) + + mu, log_var = self.conv1(out).chunk(2, dim=1) + if isinstance(scale[0], torch.Tensor): + scale = [s.to(dtype=mu.dtype, device=mu.device) for s in scale] + mu = (mu - scale[0].view(1, self.z_dim, 1, 1, 1)) * scale[1].view(1, self.z_dim, 1, 1, 1) + else: + scale = scale.to(dtype=mu.dtype, device=mu.device) + mu = (mu - scale[0]) * scale[1] + return mu + + def decode(self, z, scale): + if isinstance(scale[0], torch.Tensor): + scale = [s.to(dtype=z.dtype, device=z.device) for s in scale] + z = z / scale[1].view(1, self.z_dim, 1, 1, 1) + scale[0].view(1, self.z_dim, 1, 1, 1) + else: + scale = scale.to(dtype=z.dtype, device=z.device) + z = z / scale[1] + scale[0] + + def create_custom_forward(module): + def custom_forward(*inputs): + return module(*inputs) + + return custom_forward + + x = torch.utils.checkpoint.checkpoint(create_custom_forward(self.conv2), z, use_reentrant=False) + out = torch.utils.checkpoint.checkpoint(create_custom_forward(self.decoder), x, use_reentrant=True) + return out + + def clear_cache(self): + self._conv_num = count_conv3d(self.decoder) + self._conv_idx = [0] + self._feat_map = [None] * self._conv_num + self._enc_conv_num = count_conv3d(self.encoder) + self._enc_conv_idx = [0] + self._enc_feat_map = [None] * self._enc_conv_num + + +# ── WanVideoVAE (top-level wrapper) ── + + +def _replace_conv_with_act_grad_only(model): + """Replace Conv layers with act-grad-only versions (post-construction). + + Avoids global monkey-patching of nn.Conv2d / CausalConv3d. + """ + for name, module in model.named_modules(): + if isinstance(module, CausalConv3d) and not isinstance(module, CausalConv3dActGradOnly): + parent_name, child_name = name.rsplit(".", 1) if "." in name else ("", name) + parent = model if parent_name == "" else dict(model.named_modules())[parent_name] + new_module = CausalConv3dActGradOnly.__new__(CausalConv3dActGradOnly) + nn.Conv3d.__init__( + new_module, + module.in_channels, + module.out_channels, + module.kernel_size, + module.stride, + (0, 0, 0), + module.dilation, + module.groups, + module.bias is not None, + ) + new_module._padding = module._padding + new_module.padding = module.padding + new_module.time_kernel_size = module.time_kernel_size + new_module.weight = module.weight + new_module.weight.requires_grad = False + if module.bias is not None: + new_module.bias = module.bias + new_module.bias.requires_grad = False + setattr(parent, child_name, new_module) + + elif isinstance(module, nn.Conv2d) and not isinstance(module, Conv2dActGradOnly): + parent_name, child_name = name.rsplit(".", 1) if "." in name else ("", name) + parent = model if parent_name == "" else dict(model.named_modules())[parent_name] + new_module = Conv2dActGradOnly( + module.in_channels, + module.out_channels, + module.kernel_size, + module.stride, + module.padding, + module.dilation, + module.groups, + module.bias is not None, + ) + new_module.weight = module.weight + new_module.weight.requires_grad = False + if module.bias is not None: + new_module.bias = module.bias + new_module.bias.requires_grad = False + setattr(parent, child_name, new_module) + + +def _get_task_range(total, world_size, rank): + """Evenly divide tasks across ranks.""" + per_rank = (total + world_size - 1) // world_size + start = rank * per_rank + end = min(start + per_rank, total) + return start, end + + +class _SimpleConfig: + """Minimal config object for diffusers pipeline compatibility. + + WanPipeline.__init__ reads vae.config.scale_factor_temporal/spatial. + """ + + def __init__(self, **kwargs): + for k, v in kwargs.items(): + setattr(self, k, v) + + +class _DeterministicLatentDist: + """Drop-in replacement for diffusers' ``DiagonalGaussianDistribution``. + + UniRL's I2V image-condition encode stage calls + ``vae.encode(x).latent_dist.mode()`` (see ``image_encode.py:100``) + to get the deterministic mode of the encoded Gaussian. ROLL's + ``VideoVAE_.encode`` only returns the (already-normalized) mean + directly, so ``mode()`` and ``sample()`` here are identical. + """ + + def __init__(self, latent: torch.Tensor) -> None: + self._latent = latent + + def mode(self) -> torch.Tensor: + return self._latent + + def sample(self, *_args, **_kwargs) -> torch.Tensor: + return self._latent + + +class _EncoderOutput: + """Drop-in replacement for diffusers' ``AutoencoderKLOutput``.""" + + def __init__(self, latent: torch.Tensor) -> None: + self.latent_dist = _DeterministicLatentDist(latent) + + +class WanVideoVAE(nn.Module): + """Wan Video VAE with training optimizations. + + Wraps VideoVAE_ with: + - Built-in latent normalization (mean/std) + - Tiled encode/decode for large resolutions + - Tiled parallel decode across SP ranks + - Optional act-grad-only conv optimization + """ + + def __init__(self, z_dim=16, use_nested_grad_checkpoint=True, use_act_grad_only_conv=True): + super().__init__() + + mean = [ + -0.7571, + -0.7089, + -0.9113, + 0.1075, + -0.1745, + 0.9653, + -0.1517, + 1.5508, + 0.4134, + -0.0715, + 0.5517, + -0.3632, + -0.1922, + -0.9497, + 0.2503, + -0.2921, + ] + std = [ + 2.8184, + 1.4541, + 2.3275, + 2.6558, + 1.2196, + 1.7708, + 2.6052, + 2.0743, + 3.2687, + 2.1526, + 2.8652, + 1.5579, + 1.6382, + 1.1253, + 2.8251, + 1.9160, + ] + self.mean = torch.tensor(mean) + self.std = torch.tensor(std) + self.scale = [self.mean, 1.0 / self.std] + self.upsampling_factor = 8 + self.z_dim = z_dim + + # diffusers-compatible config for WanPipeline.__init__ compatibility. + # WanPipeline reads vae.config.scale_factor_temporal/spatial in __init__. + # ``scaling_factor`` is exposed only as a legacy fallback — the real + # un-normalization always uses ``latents_mean`` / ``latents_std``. + self.config = _SimpleConfig( + scale_factor_temporal=4, + scale_factor_spatial=8, + latents_mean=mean, + latents_std=std, + z_dim=z_dim, + scaling_factor=1.0, + ) + + # Build model with standard conv, then optionally replace with act-grad-only + self.model = ( + VideoVAE_( + z_dim=z_dim, + use_nested_grad_checkpoint=use_nested_grad_checkpoint, + ) + .eval() + .requires_grad_(False) + ) + + if use_act_grad_only_conv: + _replace_conv_with_act_grad_only(self.model) + + # ── Properties / dtype shim ── + + @property + def dtype(self) -> torch.dtype: + """Parameter dtype, matching the diffusers ``vae.dtype`` convention. + + UniRL's ``image_encode.py:94`` reads ``vae.dtype`` to cast the + VAE input. Expose it as a property over the first parameter so it + stays in sync if the module is later ``.to(dtype=...)``-cast. + """ + return next(self.parameters()).dtype + + # ── Encode / Decode ── + + def single_encode(self, video, device): + video = video.to(device) + return self.model.encode(video, self.scale) + + def single_decode(self, hidden_state, device): + hidden_state = hidden_state.to(device) + video = self.model.decode(hidden_state, self.scale) + return video.clamp_(-1, 1) + + def encode(self, videos, device=None, tiled=False, tile_size=(34, 34), tile_stride=(18, 16)): + """Encode videos to latent space. + + Two calling conventions are supported (so this class can drop into + both UniRL's batched-encode path and diffusers' single-tensor + ``vae.encode(x).latent_dist.mode()`` contract used by the I2V + image-condition encode stage): + + 1. **Batched / list**: ``videos`` is a list of ``(C, T, H, W)`` + tensors or a 5D ``(B, C, T, H, W)`` tensor, ``device`` must be + passed; returns a stacked ``(B, z_dim, T_lat, H_lat, W_lat)`` + latent tensor directly. + 2. **Diffusers-compatible**: ``videos`` is a 5D + ``(B, C, T, H, W)`` tensor, ``device`` is omitted; returns an + ``_EncoderOutput`` whose ``.latent_dist.mode()`` / + ``.latent_dist.sample()`` both return the (deterministic) + normalized latent. The VAE is non-stochastic in this port + , so ``mode()`` and ``sample()`` agree. + """ + # Diffusers-compatible path: single 5D tensor, infer device. + if device is None and isinstance(videos, torch.Tensor) and videos.dim() == 5: + target_device = videos.device + latents = self._encode_batched(videos, target_device, tiled, tile_size, tile_stride) + return _EncoderOutput(latents) + if device is None: + raise ValueError("WanVideoVAE.encode: ``device`` is required for the batched/list calling convention.") + return self._encode_batched(videos, device, tiled, tile_size, tile_stride) + + def _encode_batched(self, videos, device, tiled, tile_size, tile_stride): + if isinstance(videos, torch.Tensor) and videos.dim() == 5: + videos = [videos[i] for i in range(videos.shape[0])] + + hidden_states = [] + for video in videos: + video = video.unsqueeze(0) + if tiled: + tile_size_px = (tile_size[0] * self.upsampling_factor, tile_size[1] * self.upsampling_factor) + tile_stride_px = (tile_stride[0] * self.upsampling_factor, tile_stride[1] * self.upsampling_factor) + hidden_state = self.tiled_encode(video, device, tile_size_px, tile_stride_px) + else: + hidden_state = self.single_encode(video, device) + hidden_states.append(hidden_state.squeeze(0)) + return torch.stack(hidden_states) + + def decode(self, hidden_states, device=None, tiled=True, sp_group=None, tile_size=(34, 34), tile_stride=(18, 16)): + """Decode latents to video. + + ``device`` defaults to the latent's own device (so callers + upstream of UniRL's decode stage can ``vae.decode(z)`` without + threading the device through). ``tiled=True`` is the default + because uniform-tile + per-tile checkpointing is the only path + that survives BPTT memory pressure at 480×832×81f resolution. + """ + if device is None: + device = hidden_states.device + if tiled: + if sp_group is not None: + video = self.tiled_parallel_decode(hidden_states, device, tile_size, tile_stride, sp_group) + else: + video = self.tiled_decode(hidden_states, device, tile_size, tile_stride) + else: + video = self.single_decode(hidden_states, device) + return video + + # ── Tiled operations ── + + def build_1d_mask(self, length, left_bound, right_bound, border_width): + x = torch.ones((length,)) + if not left_bound: + x[:border_width] = (torch.arange(border_width) + 1) / border_width + if not right_bound: + x[-border_width:] = torch.flip((torch.arange(border_width) + 1) / border_width, dims=(0,)) + return x + + def build_mask(self, data, is_bound, border_width): + _, _, _, H, W = data.shape + h = self.build_1d_mask(H, is_bound[0], is_bound[1], border_width[0]) + w = self.build_1d_mask(W, is_bound[2], is_bound[3], border_width[1]) + h = repeat(h, "H -> H W", H=H, W=W) + w = repeat(w, "W -> H W", H=H, W=W) + mask = torch.stack([h, w]).min(dim=0).values + return rearrange(mask, "H W -> 1 1 1 H W") + + def _make_tile_tasks(self, H, W, size_h, size_w, stride_h, stride_w): + tasks = [] + for h in range(0, H, stride_h): + if h - stride_h >= 0 and h - stride_h + size_h >= H: + continue + for w in range(0, W, stride_w): + if w - stride_w >= 0 and w - stride_w + size_w >= W: + continue + tasks.append((h, h + size_h, w, w + size_w)) + return tasks + + def tiled_decode(self, hidden_states, device, tile_size, tile_stride): + B, _, T, H, W = hidden_states.shape + size_h, size_w = tile_size + stride_h, stride_w = tile_stride + tasks = self._make_tile_tasks(H, W, size_h, size_w, stride_h, stride_w) + + out_T = T * 4 - 3 + weight = torch.zeros( + (B, 1, out_T, H * self.upsampling_factor, W * self.upsampling_factor), + dtype=hidden_states.dtype, + device=device, + ) + values = torch.zeros( + (B, 3, out_T, H * self.upsampling_factor, W * self.upsampling_factor), + dtype=hidden_states.dtype, + device=device, + ) + + for h, h_, w, w_ in tasks: + batch = hidden_states[:, :, :, h:h_, w:w_].to(device) + batch = self.model.decode(batch, self.scale).to(device) + mask = self.build_mask( + batch, + is_bound=(h == 0, h_ >= H, w == 0, w_ >= W), + border_width=( + (size_h - stride_h) * self.upsampling_factor, + (size_w - stride_w) * self.upsampling_factor, + ), + ).to(dtype=hidden_states.dtype, device=device) + + th = h * self.upsampling_factor + tw = w * self.upsampling_factor + value_slice = values[:, :, :, th : th + batch.shape[3], tw : tw + batch.shape[4]] + weight_slice = weight[:, :, :, th : th + batch.shape[3], tw : tw + batch.shape[4]] + value_slice += batch * mask + weight_slice += mask.expand_as(weight_slice) + + values = values / weight + return values.clamp_(-1, 1) + + def tiled_parallel_decode(self, hidden_states, device, tile_size, tile_stride, sp_group): + """Tiled decode with SP parallelism — each rank decodes a subset of tiles. + + Args: + sp_group: torch.distributed ProcessGroup for sequence parallelism. + """ + B, _, T, H, W = hidden_states.shape + size_h, size_w = tile_size + stride_h, stride_w = tile_stride + tile_img_h = size_h * self.upsampling_factor + tile_img_w = size_w * self.upsampling_factor + tasks = self._make_tile_tasks(H, W, size_h, size_w, stride_h, stride_w) + + out_T = T * 4 - 3 + weight = torch.zeros( + (B, 1, out_T, H * self.upsampling_factor, W * self.upsampling_factor), + dtype=hidden_states.dtype, + device=device, + ) + values = torch.zeros( + (B, 3, out_T, H * self.upsampling_factor, W * self.upsampling_factor), + dtype=hidden_states.dtype, + device=device, + ) + + world_size = dist.get_world_size(sp_group) + sp_rank = dist.get_rank(sp_group) + task_start, task_end = _get_task_range(len(tasks), world_size, sp_rank) + rank_tasks = tasks[task_start:task_end] + + all_decoded = [] + for h, h_, w, w_ in rank_tasks: + batch = hidden_states[:, :, :, h:h_, w:w_].to(device) + batch = self.model.decode(batch, self.scale).to(device) + # Pad to uniform tile size for all_gather + padding = (0, tile_img_w - batch.shape[-1], 0, tile_img_h - batch.shape[-2]) + if padding != (0, 0, 0, 0): + batch = F.pad(batch, padding) + all_decoded.append(batch) + + # Stack and all_gather across SP ranks + local_stack = torch.stack(all_decoded, dim=0) # (num_local_tasks, B, C, T, H, W) + gathered = [torch.empty_like(local_stack) for _ in range(world_size)] + # Pad to max local tasks across ranks + max_tasks = max( + (task_end - task_start) + for s, e in [_get_task_range(len(tasks), world_size, r) for r in range(world_size)] + for task_start, task_end in [(s, e)] + ) + if local_stack.shape[0] < max_tasks: + pad_n = max_tasks - local_stack.shape[0] + local_stack = torch.cat( + [ + local_stack, + torch.zeros(pad_n, *local_stack.shape[1:], dtype=local_stack.dtype, device=local_stack.device), + ], + dim=0, + ) + + gathered = [torch.empty_like(local_stack) for _ in range(world_size)] + dist.all_gather(gathered, local_stack.contiguous(), group=sp_group) + all_decoded_global = torch.cat(gathered, dim=0) + + # Reassemble tiles (using global task order) + for i, (h, h_, w, w_) in enumerate(tasks): + if i >= all_decoded_global.shape[0]: + break + latent_h = hidden_states[:, :, :, h:h_, w:w_].shape[-2] + latent_w = hidden_states[:, :, :, h:h_, w:w_].shape[-1] + img_h = latent_h * self.upsampling_factor + img_w = latent_w * self.upsampling_factor + decoded_tile = all_decoded_global[i][:, :, :, :img_h, :img_w] + + mask = self.build_mask( + decoded_tile, + is_bound=(h == 0, h_ >= H, w == 0, w_ >= W), + border_width=( + (size_h - stride_h) * self.upsampling_factor, + (size_w - stride_w) * self.upsampling_factor, + ), + ).to(dtype=hidden_states.dtype, device=device) + + th = h * self.upsampling_factor + tw = w * self.upsampling_factor + value_slice = values[:, :, :, th : th + decoded_tile.shape[3], tw : tw + decoded_tile.shape[4]] + weight_slice = weight[:, :, :, th : th + decoded_tile.shape[3], tw : tw + decoded_tile.shape[4]] + value_slice += decoded_tile * mask + weight_slice += mask.expand_as(weight_slice) + + values = values / weight + return values.clamp_(-1, 1) + + def tiled_encode(self, video, device, tile_size, tile_stride): + _, _, T, H, W = video.shape + size_h, size_w = tile_size + stride_h, stride_w = tile_stride + tasks = self._make_tile_tasks(H, W, size_h, size_w, stride_h, stride_w) + + out_T = (T + 3) // 4 + data_device = "cpu" + weight = torch.zeros( + (1, 1, out_T, H // self.upsampling_factor, W // self.upsampling_factor), + dtype=video.dtype, + device=data_device, + ) + values = torch.zeros( + (1, self.z_dim, out_T, H // self.upsampling_factor, W // self.upsampling_factor), + dtype=video.dtype, + device=data_device, + ) + + for h, h_, w, w_ in tasks: + batch = video[:, :, :, h:h_, w:w_].to(device) + batch = self.model.encode(batch, self.scale).to(data_device) + mask = self.build_mask( + batch, + is_bound=(h == 0, h_ >= H, w == 0, w_ >= W), + border_width=( + (size_h - stride_h) // self.upsampling_factor, + (size_w - stride_w) // self.upsampling_factor, + ), + ).to(dtype=video.dtype, device=data_device) + + th = h // self.upsampling_factor + tw = w // self.upsampling_factor + values[:, :, :, th : th + batch.shape[3], tw : tw + batch.shape[4]] += batch * mask + weight[:, :, :, th : th + batch.shape[3], tw : tw + batch.shape[4]] += mask + + return values / weight + + # ── HF Diffusers weight loading ── + + @classmethod + def load_from_diffusers(cls, pretrained_path, **kwargs): + """Load WanVideoVAE from HuggingFace diffusers format. + + Args: + pretrained_path: Path to model directory containing vae/ subfolder, + or direct path to vae directory. + """ + vae_dir = pretrained_path + if os.path.isdir(os.path.join(pretrained_path, "vae")): + vae_dir = os.path.join(pretrained_path, "vae") + + safetensors_path = os.path.join(vae_dir, "diffusion_pytorch_model.safetensors") + if os.path.exists(safetensors_path): + from safetensors.torch import load_file + + hf_sd = load_file(safetensors_path) + else: + bin_path = os.path.join(vae_dir, "diffusion_pytorch_model.bin") + hf_sd = torch.load(bin_path, map_location="cpu") + + vae = cls(**kwargs) + converted_sd = convert_diffusers_state_dict(hf_sd) + vae.load_state_dict(converted_sd, strict=True) + return vae + + +# ── State dict converter ── + + +def convert_diffusers_state_dict(hf_sd: dict) -> OrderedDict: + """Convert HuggingFace AutoencoderKLWan state dict to WanVideoVAE format. + + The two models have identical architecture (194 params, same shapes) but + different key naming. This function handles the mapping including the + mid_block reorder: HF [attn, res0, res1] → ROLL [res0, attn, res1]. + """ + new_sd = OrderedDict() + + # ResidualBlock internal mapping + # HF: norm1.gamma, conv1.{w,b}, norm2.gamma, conv2.{w,b}, conv_shortcut.{w,b} + # ROLL: residual.{0.gamma, 2.{w,b}, 3.gamma, 6.{w,b}}, shortcut.{w,b} + resblock_map = { + "norm1.gamma": "residual.0.gamma", + "conv1.weight": "residual.2.weight", + "conv1.bias": "residual.2.bias", + "norm2.gamma": "residual.3.gamma", + "conv2.weight": "residual.6.weight", + "conv2.bias": "residual.6.bias", + "conv_shortcut.weight": "shortcut.weight", + "conv_shortcut.bias": "shortcut.bias", + } + + # AttentionBlock internal mapping (keys match directly) + # HF: norm.gamma, to_qkv.{w,b}, proj.{w,b} + # ROLL: norm.gamma, to_qkv.{w,b}, proj.{w,b} + + for hf_key, tensor in hf_sd.items(): + roll_key = _convert_single_key(hf_key, resblock_map) + new_sd[roll_key] = tensor + + return new_sd + + +def _convert_single_key(hf_key: str, resblock_map: dict) -> str: + """Convert a single HF key to ROLL key.""" + + # quant_conv / post_quant_conv + if hf_key.startswith("quant_conv."): + return "model.conv1." + hf_key[len("quant_conv.") :] + if hf_key.startswith("post_quant_conv."): + return "model.conv2." + hf_key[len("post_quant_conv.") :] + + # Encoder + if hf_key.startswith("encoder."): + return "model.encoder." + _convert_encoder_key(hf_key[len("encoder.") :], resblock_map) + + # Decoder + if hf_key.startswith("decoder."): + return "model.decoder." + _convert_decoder_key(hf_key[len("decoder.") :], resblock_map) + + raise ValueError(f"Unknown HF key: {hf_key}") + + +def _convert_encoder_key(key: str, rb_map: dict) -> str: + # conv_in → conv1 + if key.startswith("conv_in."): + return "conv1." + key[len("conv_in.") :] + + # norm_out → head.0 + if key.startswith("norm_out."): + return "head.0." + key[len("norm_out.") :] + + # conv_out → head.2 + if key.startswith("conv_out."): + return "head.2." + key[len("conv_out.") :] + + # down_blocks.{i}.{suffix} + if key.startswith("down_blocks."): + rest = key[len("down_blocks.") :] + dot = rest.index(".") + block_idx = rest[:dot] + suffix = rest[dot + 1 :] + return "downsamples." + block_idx + "." + _convert_resblock_suffix(suffix, rb_map) + + # mid_block: HF [attn, res0, res1] → ROLL [res0(middle.0), attn(middle.1), res1(middle.2)] + if key.startswith("mid_block."): + rest = key[len("mid_block.") :] + if rest.startswith("resnets.0."): + suffix = rest[len("resnets.0.") :] + return "middle.0." + _convert_resblock_suffix(suffix, rb_map) + if rest.startswith("attentions.0."): + suffix = rest[len("attentions.0.") :] + return "middle.1." + suffix + if rest.startswith("resnets.1."): + suffix = rest[len("resnets.1.") :] + return "middle.2." + _convert_resblock_suffix(suffix, rb_map) + + raise ValueError(f"Unknown encoder key: {key}") + + +def _convert_decoder_key(key: str, rb_map: dict) -> str: + # conv_in → conv1 + if key.startswith("conv_in."): + return "conv1." + key[len("conv_in.") :] + + # norm_out → head.0 + if key.startswith("norm_out."): + return "head.0." + key[len("norm_out.") :] + + # conv_out → head.2 + if key.startswith("conv_out."): + return "head.2." + key[len("conv_out.") :] + + # mid_block: same reorder as encoder + if key.startswith("mid_block."): + rest = key[len("mid_block.") :] + if rest.startswith("resnets.0."): + suffix = rest[len("resnets.0.") :] + return "middle.0." + _convert_resblock_suffix(suffix, rb_map) + if rest.startswith("attentions.0."): + suffix = rest[len("attentions.0.") :] + return "middle.1." + suffix + if rest.startswith("resnets.1."): + suffix = rest[len("resnets.1.") :] + return "middle.2." + _convert_resblock_suffix(suffix, rb_map) + + # up_blocks.{block_i}.resnets.{res_j}.{suffix} → upsamples.{flat_idx}.{suffix} + # up_blocks.{block_i}.upsamplers.0.{suffix} → upsamples.{flat_idx}.{suffix} + if key.startswith("up_blocks."): + return _convert_upblock_key(key[len("up_blocks.") :], rb_map) + + raise ValueError(f"Unknown decoder key: {key}") + + +def _convert_upblock_key(key: str, rb_map: dict) -> str: + """Convert decoder up_blocks flat key to ROLL upsamples flat index. + + HF decoder structure (Wan 2.1 VideoVAE_): + - up_blocks.0: 3 resnets + 1 upsampler → upsamples indices 0,1,2,3 + - up_blocks.1: 3 resnets + 1 upsampler → upsamples indices 4,5,6,7 + - up_blocks.2: 3 resnets + 1 upsampler → upsamples indices 8,9,10,11 + - up_blocks.3: 3 resnets (no upsampler) → upsamples indices 12,13,14 + + ROLL's Decoder3d has 15 sequential upsamples: + - indices 0-2: ResBlocks (up_blocks.0.resnets.0-2) + - index 3: Resample (up_blocks.0.upsamplers.0) + - indices 4-6: ResBlocks (up_blocks.1.resnets.0-2) + - index 7: Resample (up_blocks.1.upsamplers.0) + - indices 8-10: ResBlocks (up_blocks.2.resnets.0-2) + - index 11: Resample (up_blocks.2.upsamplers.0) + - indices 12-14: ResBlocks (up_blocks.3.resnets.0-2) + """ + dot = key.index(".") + block_idx = int(key[:dot]) + rest = key[dot + 1 :] + + # Base offset: each block contributes 4 items (3 resnets + 1 upsampler) + # except the last block which has 3 items (3 resnets, no upsampler) + base = block_idx * 4 # works for blocks 0,1,2; block 3 won't have upsampler + + if rest.startswith("resnets."): + rest2 = rest[len("resnets.") :] + dot2 = rest2.index(".") + res_idx = int(rest2[:dot2]) + suffix = rest2[dot2 + 1 :] + flat_idx = base + res_idx + return f"upsamples.{flat_idx}." + _convert_resblock_suffix(suffix, rb_map) + + if rest.startswith("upsamplers.0."): + suffix = rest[len("upsamplers.0.") :] + flat_idx = base + 3 # upsampler comes after 3 resnets + return f"upsamples.{flat_idx}." + suffix + + raise ValueError(f"Unknown up_block key: {key}") + + +def _convert_resblock_suffix(suffix: str, rb_map: dict) -> str: + """Convert a ResidualBlock suffix from HF to ROLL format.""" + for hf_pattern, roll_pattern in rb_map.items(): + if suffix == hf_pattern or suffix.startswith(hf_pattern): + return suffix.replace(hf_pattern, roll_pattern, 1) + # Pass through (e.g. resample.*, time_conv.*) + return suffix diff --git a/unirl/reward/base.py b/unirl/reward/base.py index a6c8f0148..1daca48f4 100644 --- a/unirl/reward/base.py +++ b/unirl/reward/base.py @@ -81,12 +81,16 @@ class DifferentiableReward(Protocol): def compute_rewards_differentiable( self, - images_tensor: "torch.Tensor", + media_tensor: "torch.Tensor", prompts: List[str], - records: Optional[List[dict]] = None, + records: Optional[List[dict[str, object]]] = None, ) -> "torch.Tensor": - """Score a grad-carrying image tensor ``[B, C, H, W]`` in ``[0, 1]`` → - ``[B]`` reward tensor with ``grad_fn`` intact (no ``no_grad``/``.item()``).""" + """Score grad-carrying image ``[B,C,H,W]`` or video ``[B,C,T,H,W]`` media. + + Returns a ``[B]`` reward tensor with ``grad_fn`` intact. ``records`` may + carry per-sample metadata such as ``ref_video_path`` for recipe-local + video rewards. + """ ... diff --git a/unirl/reward/service.py b/unirl/reward/service.py index ae3b58670..7893293de 100644 --- a/unirl/reward/service.py +++ b/unirl/reward/service.py @@ -10,13 +10,13 @@ from __future__ import annotations import logging -from typing import Dict, Optional +from typing import Dict, List, Optional import torch from unirl.distributed.group.dispatch import Dispatch, distributed from unirl.distributed.group.remote import Remote -from unirl.types.primitives import Images, Texts, primitive_modality_key +from unirl.types.primitives import primitive_modality_key from unirl.types.reward import RewardRequest, RewardResponse from unirl.types.sample import Primitive, Sample, _part_with_field from unirl.types.sampling import ARSamplingParams @@ -111,16 +111,23 @@ def compute_rewards(self, request: RewardRequest) -> RewardResponse: return self.backend.compute_rewards(request) @distributed(dispatch_mode=Dispatch.DP_SCATTER) - def score_differentiable(self, *, images: Images, prompts: Texts) -> torch.Tensor: - """ReFL scoring: score grad-carrying ``images`` (pixels ``[B, C, H, W]`` in - ``[0, 1]``) against ``prompts`` and return a ``[B]`` reward tensor with - ``grad_fn`` intact. - - Deliberately bypasses :meth:`score_and_attach` / ``RewardRequest.images`` - (those go through ``tensor_frame_to_pil`` + ``torch.tensor(...)``, which - detach). Under ``enable_grad()`` the framework marks ``images.pixels`` as a - grad leaf and chains the returned reward's grad back to it. The backend must - satisfy the :class:`~unirl.reward.base.DifferentiableReward` Protocol. + def score_differentiable( + self, + media_tensor: torch.Tensor, + prompts: List[str], + records: Optional[List[dict]] = None, + ) -> torch.Tensor: + """ReFL scoring: score grad-carrying ``media_tensor`` (image ``[B, C, H, W]`` + or video ``[B, C, T, H, W]``) against ``prompts`` and return a ``[B]`` reward + tensor with ``grad_fn`` intact. + + Deliberately bypasses :meth:`score_and_attach` / ``RewardRequest`` media + conversion (those go through ``tensor_frame_to_pil`` + ``torch.tensor(...)``, + which detach). Under ``enable_grad()`` the framework marks ``media_tensor`` as + a grad leaf and chains the returned reward's grad back to it. ``records`` + carries optional per-sample metadata (e.g. ``ref_video_path`` for recipe-local + video rewards). The backend must satisfy the + :class:`~unirl.reward.base.DifferentiableReward` Protocol. """ if not isinstance(self.backend, DifferentiableReward): raise TypeError( @@ -128,7 +135,7 @@ def score_differentiable(self, *, images: Images, prompts: Texts) -> torch.Tenso f"{type(self.backend).__name__} is not a DifferentiableReward — ReFL " f"needs a differentiable in-process reward (e.g. pickscore/clip/hpsv2)." ) - return self.backend.compute_rewards_differentiable(images.pixels, list(prompts.texts)) + return self.backend.compute_rewards_differentiable(media_tensor, list(prompts), records=records) @distributed(dispatch_mode=Dispatch.DP_SCATTER) def score_and_attach(self, sample: Sample) -> Sample: diff --git a/unirl/train/backend/base_backend.py b/unirl/train/backend/base_backend.py index 4dee2c641..108d49929 100644 --- a/unirl/train/backend/base_backend.py +++ b/unirl/train/backend/base_backend.py @@ -50,7 +50,7 @@ ) from unirl.train.configs import EmaFullConfig, EmaLoraConfig, FSDPConfig, LoraConfig from unirl.train.ema import EMA, Shadow, inject_mirror, inject_nft, make_decay_fn -from unirl.train.lora import inject_lora +from unirl.train.lora import inject_lora, resolve_target_modules_pattern from unirl.train.optim import build_lr_scheduler, build_optimizer if TYPE_CHECKING: @@ -208,6 +208,7 @@ def _inject_structural( rank=lora_cfg.rank, alpha=lora_cfg.alpha, target_modules=lora_cfg.target_modules, + module_prefix=lora_cfg.module_prefix, exclude_modules=lora_cfg.exclude_modules, dropout=lora_cfg.dropout, bias=lora_cfg.bias, @@ -273,11 +274,22 @@ def _finalize_construction( # Checkpointed for export tooling: the LoRA fold needs scaling = # alpha / rank, and alpha is not derivable from the weights. active_lora = lora_cfg or ema_lora_cfg + recorded_target_modules = None + if active_lora is not None: + recorded_target_modules = active_lora.target_modules + if lora_cfg is not None: + recorded_target_modules, _ = resolve_target_modules_pattern( + target_modules=lora_cfg.target_modules, + module_prefix=lora_cfg.module_prefix, + ) self._lora_meta = ( { "rank": active_lora.rank, "alpha": active_lora.alpha, - "target_modules": active_lora.target_modules, + # Store the selector actually passed to PEFT. For prefixed LoRA + # this is a regex, so exported adapters keep the same subtree + # restriction instead of matching equivalent suffixes elsewhere. + "target_modules": recorded_target_modules, "exclude_modules": active_lora.exclude_modules, "dropout": active_lora.dropout, "bias": active_lora.bias, diff --git a/unirl/train/configs.py b/unirl/train/configs.py index e2a719b70..57f44ab14 100644 --- a/unirl/train/configs.py +++ b/unirl/train/configs.py @@ -19,6 +19,9 @@ class LoraConfig: # Like ``target_modules``, tuples match exact names/suffixes and strings are # regular expressions. This lets a regex exclude an entire frozen sub-tower. exclude_modules: Any = None + # Restrict sequence-style target suffixes to a named model subtree. + # This cannot be combined with regex or ``all-linear`` target strings. + module_prefix: str = "" dropout: float = 0.0 bias: str = "none" task_type: str = "FEATURE_EXTRACTION" diff --git a/unirl/train/lora.py b/unirl/train/lora.py index a7e6796e6..7b1579b9e 100644 --- a/unirl/train/lora.py +++ b/unirl/train/lora.py @@ -9,6 +9,7 @@ from __future__ import annotations import logging +import re from contextlib import contextmanager from functools import partial from typing import Iterator, Optional, Sequence, Union @@ -43,12 +44,46 @@ def normalize_optional_module_selection( return None if modules is None else normalize_module_selection(modules) +def resolve_target_modules_pattern( + *, + target_modules: ModuleSelection, + module_prefix: str = "", +) -> tuple[PeftModuleSelection, str]: + """Resolve PEFT targets, optionally restricting suffixes to one subtree.""" + normalized = normalize_module_selection(target_modules) + if not module_prefix: + logged = normalized if isinstance(normalized, str) else tuple(normalized) + return normalized, str(logged) + + if isinstance(normalized, str): + raise ValueError( + "resolve_target_modules_pattern: module_prefix cannot be combined " + "with a regex or 'all-linear' target_modules string; provide an " + "explicit sequence of target-module suffixes." + ) + if not normalized: + raise ValueError( + "resolve_target_modules_pattern: module_prefix is set but " + "target_modules is empty; provide at least one target-module suffix." + ) + normalized_prefix = str(module_prefix).strip(".") + if not normalized_prefix: + raise ValueError( + "resolve_target_modules_pattern: module_prefix must contain a model subtree name, not only dots." + ) + prefix_re = re.escape(normalized_prefix) + leaves_re = "|".join(re.escape(str(target)) for target in normalized) + pattern = rf"^{prefix_re}\.(?:.*\.)?(?:{leaves_re})$" + return pattern, pattern + + def inject_lora( model: nn.Module, *, rank: int, alpha: int, target_modules: ModuleSelection, + module_prefix: str = "", exclude_modules: Optional[ModuleSelection] = None, dropout: float = 0.0, bias: str = "none", @@ -58,11 +93,16 @@ def inject_lora( """Inject a single LoRA adapter. No Shadow, no EMA.""" from peft import LoraConfig, inject_adapter_in_model + peft_target_modules, log_target = resolve_target_modules_pattern( + target_modules=target_modules, + module_prefix=module_prefix, + ) + peft_cfg = LoraConfig( r=int(rank), lora_alpha=int(alpha), lora_dropout=float(dropout), - target_modules=normalize_module_selection(target_modules), + target_modules=peft_target_modules, exclude_modules=normalize_optional_module_selection(exclude_modules), bias=str(bias), task_type=str(task_type), @@ -71,7 +111,6 @@ def inject_lora( if _current_rank() == 0: n_trainable = sum(1 for p in model.parameters() if p.requires_grad) - logged_targets = target_modules if isinstance(target_modules, str) else tuple(target_modules) logged_exclusions = ( exclude_modules if isinstance(exclude_modules, str) or exclude_modules is None else tuple(exclude_modules) ) @@ -80,7 +119,7 @@ def inject_lora( adapter_name, rank, alpha, - logged_targets, + log_target, logged_exclusions, n_trainable, ) @@ -136,4 +175,5 @@ def _current_rank() -> int: "inject_lora", "normalize_module_selection", "normalize_optional_module_selection", + "resolve_target_modules_pattern", ] diff --git a/unirl/train/optim.py b/unirl/train/optim.py index 3f20310de..5438cf653 100644 --- a/unirl/train/optim.py +++ b/unirl/train/optim.py @@ -131,8 +131,8 @@ def build_lr_scheduler( Supports the same backend-override path as :func:`build_optimizer`. Returns ``None`` if ``config.type`` is not one of the supported values - (``constant`` / ``linear`` / ``cosine``) and the backend did not provide - an override. + (``constant`` / ``linear`` / ``linear_warmup`` / ``cosine``) and the backend + did not provide an override. """ del actor if backend is not None: @@ -167,6 +167,27 @@ def lr_lambda(step: int) -> float: return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda) + if scheduler_type == "linear_warmup": + constant = torch.optim.lr_scheduler.LinearLR( + optimizer, + start_factor=1.0, + end_factor=1.0, + total_iters=1, + ) + if warmup_steps <= 0: + return constant + warmup = torch.optim.lr_scheduler.LinearLR( + optimizer, + start_factor=1e-8, + end_factor=1.0, + total_iters=warmup_steps, + ) + return torch.optim.lr_scheduler.SequentialLR( + optimizer, + schedulers=[warmup, constant], + milestones=[warmup_steps], + ) + if scheduler_type == "cosine": def lr_lambda(step: int) -> float: diff --git a/unirl/trainer/refl.py b/unirl/trainer/refl.py index 53312cd7e..84742a2e7 100644 --- a/unirl/trainer/refl.py +++ b/unirl/trainer/refl.py @@ -130,7 +130,7 @@ def train_step(self, prompts: Texts, *, rollout_id: int) -> Tuple[float, float, t0 = time.perf_counter() with enable_grad(): images = self.policy.sample_and_decode(prompts=prompts, rollout_id=rollout_id) - rewards = self.reward.score_differentiable(images=images, prompts=prompts) + rewards = self.reward.score_differentiable(images.pixels, list(prompts.texts)) # Detached value for logging (does not disturb the worker-side graph). mean_reward = float(hydrate(rewards).float().mean().item()) self.policy.loss_backward(rewards=rewards) @@ -197,7 +197,7 @@ def _eval_pass( sub = Texts(texts=texts[start : start + chunk]) images = self.policy.eval_sample(prompts=sub, rollout_id=step, guidance_scale=self.eval_cfg_text_scale) for name, reward in scorers: - r = hydrate(reward.score_differentiable(images=images, prompts=sub)).float() + r = hydrate(reward.score_differentiable(images.pixels, list(sub.texts))).float() sums[name] += float(r.sum().item()) counts[name] += int(r.numel()) return {name: sums[name] / max(1, counts[name]) for name, _ in scorers}