114 lines
2.8 KiB
Python
114 lines
2.8 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Single diffusion step for PPO training phase."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import torch
|
|
|
|
from fastvideo.train.methods.rl.utils.sde import (
|
|
sde_step_with_logprob,
|
|
)
|
|
|
|
|
|
def compute_log_prob(
|
|
model,
|
|
scheduler,
|
|
sample: dict[str, torch.Tensor],
|
|
j: int,
|
|
embeds: torch.Tensor,
|
|
negative_embeds: torch.Tensor | None,
|
|
guidance_scale: float,
|
|
use_cfg: bool,
|
|
noise_level: float,
|
|
sde_type: str,
|
|
diffusion_clip: bool = False,
|
|
diffusion_clip_value: float = 0.45,
|
|
) -> tuple[
|
|
torch.Tensor,
|
|
torch.Tensor,
|
|
torch.Tensor,
|
|
torch.Tensor,
|
|
torch.Tensor,
|
|
torch.Tensor,
|
|
float,
|
|
]:
|
|
"""Run one diffusion step and return log-probability.
|
|
|
|
Uses model.forward_transformer_raw() for the forward
|
|
pass.
|
|
|
|
Args:
|
|
model: WanModel instance.
|
|
scheduler: Noise scheduler.
|
|
sample: Dict with latents, next_latents, timesteps.
|
|
j: Timestep index within the trajectory.
|
|
embeds: Conditional text embeddings.
|
|
negative_embeds: Unconditional embeddings (or None).
|
|
guidance_scale: CFG scale.
|
|
use_cfg: Whether to use classifier-free guidance.
|
|
noise_level: SDE noise level.
|
|
sde_type: 'flow_sde' or 'flow_cps'.
|
|
diffusion_clip: Clip SDE variance.
|
|
diffusion_clip_value: Clip threshold.
|
|
|
|
Returns:
|
|
(prev_sample, log_prob, prev_sample_mean,
|
|
std_dev_t, dt_sqrt, sigma, sigma_max)
|
|
"""
|
|
dtype = embeds.dtype
|
|
latents_j = sample["latents"][:, j]
|
|
timestep_j = sample["timesteps"][:, j]
|
|
|
|
if use_cfg and negative_embeds is not None:
|
|
noise_pred_text = model.forward_transformer_raw(
|
|
latents_j.to(dtype),
|
|
timestep_j,
|
|
embeds,
|
|
)
|
|
noise_pred_uncond = model.forward_transformer_raw(
|
|
latents_j.to(dtype),
|
|
timestep_j,
|
|
negative_embeds,
|
|
)
|
|
noise_pred = (
|
|
noise_pred_uncond
|
|
+ guidance_scale
|
|
* (noise_pred_text - noise_pred_uncond)
|
|
)
|
|
else:
|
|
noise_pred = model.forward_transformer_raw(
|
|
latents_j.to(dtype),
|
|
timestep_j,
|
|
embeds,
|
|
)
|
|
|
|
(
|
|
prev_sample,
|
|
log_prob,
|
|
prev_sample_mean,
|
|
std_dev_t,
|
|
dt_sqrt,
|
|
sigma,
|
|
sigma_max,
|
|
) = sde_step_with_logprob(
|
|
scheduler,
|
|
noise_pred.float(),
|
|
timestep_j,
|
|
latents_j.float(),
|
|
noise_level=noise_level,
|
|
prev_sample=sample["next_latents"][:, j].float(),
|
|
sde_type=sde_type,
|
|
diffusion_clip=diffusion_clip,
|
|
diffusion_clip_value=diffusion_clip_value,
|
|
return_sqrt_dt_and_std_dev_t=True,
|
|
)
|
|
return (
|
|
prev_sample,
|
|
log_prob,
|
|
prev_sample_mean,
|
|
std_dev_t,
|
|
dt_sqrt,
|
|
sigma,
|
|
sigma_max,
|
|
)
|