wehub-resource-sync
546 lines
25 KiB
Python
546 lines
25 KiB
Python
# Adopted from https://github.com/guandeh17/Self-Forcing
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
from typing import Tuple
|
|
from einops import rearrange
|
|
from torch import nn
|
|
import torch.distributed as dist
|
|
import torch
|
|
import math
|
|
|
|
from pipeline import SelfForcingTrainingPipeline
|
|
from utils.config import section_get
|
|
from utils.loss import get_denoising_loss
|
|
from utils.wan_5b_wrapper import WanDiffusionWrapper, WanTextEncoder, WanVAEWrapper
|
|
|
|
|
|
def build_default_denoising_step_list(sampling_steps, num_train_timesteps=1000, shift=1.0, include_zero=True):
|
|
sigmas = torch.linspace(1.0, 0.0, int(sampling_steps) + 1, dtype=torch.float32)[:-1]
|
|
sigmas = shift * sigmas / (1 + (shift - 1) * sigmas)
|
|
timesteps = (sigmas * num_train_timesteps).to(torch.long)
|
|
if include_zero:
|
|
timesteps = torch.cat([timesteps, torch.zeros(1, dtype=torch.long)])
|
|
return timesteps
|
|
|
|
|
|
class BaseModel(nn.Module):
|
|
def __init__(self, args, device):
|
|
super().__init__()
|
|
print("args.model_kwargs.model_name", args.model_kwargs.model_name)
|
|
self._initialize_models(args, device)
|
|
|
|
self.device = device
|
|
self.args = args
|
|
self.independent_first_frame = getattr(args, "independent_first_frame", False)
|
|
self.dtype = torch.bfloat16 if args.mixed_precision else torch.float32
|
|
self.denoising_step_list = None
|
|
if getattr(args, "denoising_step_list", None) is not None:
|
|
self.denoising_step_list = torch.tensor(args.denoising_step_list, dtype=torch.long)
|
|
if getattr(args, "warp_denoising_step", False):
|
|
timesteps = torch.cat((self.scheduler.timesteps.cpu(), torch.tensor([0], dtype=torch.float32)))
|
|
self.denoising_step_list = timesteps[1000 - self.denoising_step_list]
|
|
elif getattr(args, "sampling_steps", None) is not None:
|
|
self.denoising_step_list = build_default_denoising_step_list(
|
|
sampling_steps=args.sampling_steps,
|
|
num_train_timesteps=getattr(args, "num_train_timestep", self.scheduler.num_train_timesteps),
|
|
shift=getattr(args, "timestep_shift", self.scheduler.shift),
|
|
include_zero=True,
|
|
)
|
|
|
|
def _initialize_models(self, args, device):
|
|
self.real_model_name = getattr(args, "real_name", "Wan2.2-TI2V-5B")
|
|
self.fake_model_name = getattr(args, "fake_name", "Wan2.2-TI2V-5B")
|
|
self.local_attn_size = section_get(
|
|
args,
|
|
"inference",
|
|
"local_attn_size",
|
|
getattr(args, "model_kwargs", {}).get("local_attn_size", -1),
|
|
aliases=("inference_local_attn_size",),
|
|
)
|
|
all_causal = getattr(args, "all_causal", False)
|
|
score_is_causal = all_causal
|
|
|
|
model_name = args.model_kwargs.get("model_name", "Wan2.2-TI2V-5B")
|
|
if "5B" not in model_name:
|
|
raise ValueError(f"Only Wan2.2-TI2V-5B is supported in this release, got {model_name}")
|
|
if not dist.is_initialized() or dist.get_rank() == 0:
|
|
tag = "all-causal 5B mode" if all_causal else "Wan2.2-TI2V-5B"
|
|
print(f"Using {tag}")
|
|
|
|
# Generator
|
|
generator_is_causal = getattr(args, "generator_is_causal", True)
|
|
self.generator = WanDiffusionWrapper(**getattr(args, "model_kwargs", {}), is_causal=generator_is_causal)
|
|
self.generator.model.requires_grad_(True)
|
|
|
|
# Real Score
|
|
real_kwargs = getattr(args, "real_model_kwargs", {"model_name": self.real_model_name})
|
|
self.real_score = WanDiffusionWrapper(**real_kwargs, is_causal=score_is_causal)
|
|
self.real_score.model.requires_grad_(False)
|
|
|
|
# Fake Score
|
|
fake_kwargs = getattr(args, "fake_model_kwargs", {"model_name": self.fake_model_name})
|
|
self.fake_score = WanDiffusionWrapper(**fake_kwargs, is_causal=score_is_causal)
|
|
self.fake_score.model.requires_grad_(True)
|
|
|
|
# Text Encoder & VAE
|
|
self.text_encoder = WanTextEncoder()
|
|
self.text_encoder.requires_grad_(False)
|
|
|
|
self.vae = WanVAEWrapper()
|
|
self.vae.requires_grad_(False)
|
|
|
|
self.scheduler = self.generator.get_scheduler()
|
|
self.scheduler.timesteps = self.scheduler.timesteps.to(device)
|
|
|
|
def _get_timestep(
|
|
self,
|
|
min_timestep: int,
|
|
max_timestep: int,
|
|
batch_size: int,
|
|
num_frame: int,
|
|
num_frame_per_block: int,
|
|
uniform_timestep: bool = False
|
|
) -> torch.Tensor:
|
|
"""
|
|
Randomly generate a timestep tensor based on the generator's task type. It uniformly samples a timestep
|
|
from the range [min_timestep, max_timestep], and returns a tensor of shape [batch_size, num_frame].
|
|
- If uniform_timestep, it will use the same timestep for all frames.
|
|
- If not uniform_timestep, it will use a different timestep for each block.
|
|
"""
|
|
if uniform_timestep:
|
|
timestep = torch.randint(
|
|
min_timestep,
|
|
max_timestep,
|
|
[batch_size, 1],
|
|
device=self.device,
|
|
dtype=torch.long
|
|
).repeat(1, num_frame)
|
|
return timestep
|
|
else:
|
|
timestep = torch.randint(
|
|
min_timestep,
|
|
max_timestep,
|
|
[batch_size, num_frame],
|
|
device=self.device,
|
|
dtype=torch.long
|
|
)
|
|
# make the noise level the same within every block
|
|
if self.independent_first_frame and not getattr(self.args, "i2v", False):
|
|
# the first frame is always kept the same
|
|
timestep_from_second = timestep[:, 1:]
|
|
timestep_from_second = timestep_from_second.reshape(
|
|
timestep_from_second.shape[0], -1, num_frame_per_block)
|
|
timestep_from_second[:, :, 1:] = timestep_from_second[:, :, 0:1]
|
|
timestep_from_second = timestep_from_second.reshape(
|
|
timestep_from_second.shape[0], -1)
|
|
timestep = torch.cat([timestep[:, 0:1], timestep_from_second], dim=1)
|
|
else:
|
|
timestep = timestep.reshape(
|
|
timestep.shape[0], -1, num_frame_per_block)
|
|
timestep[:, :, 1:] = timestep[:, :, 0:1]
|
|
timestep = timestep.reshape(timestep.shape[0], -1)
|
|
return timestep
|
|
|
|
|
|
class SelfForcingModel(BaseModel):
|
|
def __init__(self, args, device):
|
|
super().__init__(args, device)
|
|
self.denoising_loss_func = get_denoising_loss(getattr(args, "denoising_loss_type", "flow"))()
|
|
|
|
def _run_generator(
|
|
self,
|
|
image_or_video_shape,
|
|
conditional_dict: dict,
|
|
initial_latent: torch.tensor = None,
|
|
slice_last_frames: int = 21,
|
|
noise=None,
|
|
clean_latent: torch.Tensor = None,
|
|
) -> Tuple[torch.Tensor, torch.Tensor]:
|
|
"""
|
|
Optionally simulate the generator's input from noise using backward simulation
|
|
and then run the generator for one-step.
|
|
Input:
|
|
- image_or_video_shape: a list containing the shape of the image or video [B, F, C, H, W].
|
|
- conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).
|
|
- initial_latent: a tensor containing the initial latents [B, F, C, H, W].
|
|
- slice_last_frames: number of frames to keep from the end.
|
|
- noise: optional pre-sampled noise.
|
|
- clean_latent: a tensor [B, F, C, H, W] for off-policy mode (backward_simulation=False).
|
|
Output:
|
|
- pred_image: a tensor with shape [B, F, C, H, W].
|
|
- gradient_mask: boolean tensor or None.
|
|
- denoised_timestep_from: int or None.
|
|
- denoised_timestep_to: int or None.
|
|
"""
|
|
use_backward_simulation = getattr(self.args, "backward_simulation", True)
|
|
|
|
if use_backward_simulation:
|
|
return self._run_generator_backward_simulation(
|
|
image_or_video_shape=image_or_video_shape,
|
|
conditional_dict=conditional_dict,
|
|
initial_latent=initial_latent,
|
|
slice_last_frames=slice_last_frames,
|
|
noise=noise,
|
|
)
|
|
else:
|
|
assert clean_latent is not None, "clean_latent is required when backward_simulation=False"
|
|
return self._run_generator_off_policy(
|
|
image_or_video_shape=image_or_video_shape,
|
|
conditional_dict=conditional_dict,
|
|
clean_latent=clean_latent,
|
|
initial_latent=initial_latent,
|
|
)
|
|
|
|
def _run_generator_off_policy(
|
|
self,
|
|
image_or_video_shape,
|
|
conditional_dict: dict,
|
|
clean_latent: torch.Tensor,
|
|
initial_latent: torch.Tensor = None,
|
|
) -> Tuple[torch.Tensor, torch.Tensor, int, int]:
|
|
"""
|
|
Off-policy generator: add noise to clean_latent at each timestep in
|
|
denoising_step_list, randomly pick one, and run the generator for a
|
|
single forward pass. Returns (pred_x0, gradient_mask,
|
|
denoised_timestep_from, denoised_timestep_to).
|
|
"""
|
|
batch_size, num_frame = image_or_video_shape[:2]
|
|
denoising_step_list = self.denoising_step_list.to(self.device)
|
|
|
|
# Build noisy versions at every timestep in the schedule
|
|
simulated_noisy_input = []
|
|
for ts in denoising_step_list:
|
|
rand_noise = torch.randn_like(clean_latent)
|
|
noisy_timestep = ts * torch.ones(
|
|
[batch_size, num_frame], device=self.device, dtype=torch.long)
|
|
|
|
if ts.item() != 0:
|
|
noisy_image = self.scheduler.add_noise(
|
|
clean_latent.flatten(0, 1),
|
|
rand_noise.flatten(0, 1),
|
|
noisy_timestep.flatten(0, 1),
|
|
).unflatten(0, (batch_size, num_frame))
|
|
else:
|
|
noisy_image = clean_latent
|
|
simulated_noisy_input.append(noisy_image)
|
|
|
|
simulated_noisy_input = torch.stack(simulated_noisy_input, dim=1) # [B, T, F, C, H, W]
|
|
|
|
# Randomly sample a step index [B, F], uniform within each block
|
|
num_steps = len(denoising_step_list)
|
|
generator_is_causal = getattr(self.args, "generator_is_causal", True)
|
|
if not generator_is_causal:
|
|
# Bidirectional generator: all frames must share the same timestep
|
|
index = torch.randint(0, num_steps, [batch_size, 1],
|
|
device=self.device, dtype=torch.long).expand(-1, num_frame).contiguous()
|
|
else:
|
|
index = torch.randint(0, num_steps, [batch_size, num_frame],
|
|
device=self.device, dtype=torch.long)
|
|
# Make the index the same within every block
|
|
if self.independent_first_frame and not getattr(self.args, "i2v", False):
|
|
idx_rest = index[:, 1:]
|
|
idx_rest = idx_rest.reshape(batch_size, -1, self.num_frame_per_block)
|
|
idx_rest[:, :, 1:] = idx_rest[:, :, 0:1]
|
|
index = torch.cat([index[:, :1], idx_rest.reshape(batch_size, -1)], dim=1)
|
|
else:
|
|
index = index.reshape(batch_size, -1, self.num_frame_per_block)
|
|
index[:, :, 1:] = index[:, :, 0:1]
|
|
index = index.reshape(batch_size, -1)
|
|
|
|
# Gather the noisy input corresponding to the sampled index
|
|
noisy_input = torch.gather(
|
|
simulated_noisy_input, dim=1,
|
|
index=index.reshape(batch_size, 1, num_frame, 1, 1, 1).expand(
|
|
-1, -1, -1, *image_or_video_shape[2:])
|
|
).squeeze(1) # [B, F, C, H, W]
|
|
|
|
timestep = denoising_step_list[index] # [B, F]
|
|
context_frames = int(initial_latent.shape[1]) if initial_latent is not None else 0
|
|
if context_frames > 0:
|
|
if context_frames >= num_frame:
|
|
raise ValueError(
|
|
f"initial_latent has {context_frames} frames but training clip has {num_frame}."
|
|
)
|
|
noisy_input[:, :context_frames] = initial_latent.to(
|
|
device=noisy_input.device,
|
|
dtype=noisy_input.dtype,
|
|
)
|
|
timestep[:, :context_frames] = 0
|
|
|
|
# Single forward pass through the generator
|
|
_, pred_x0 = self.generator(
|
|
noisy_image_or_video=noisy_input,
|
|
conditional_dict=conditional_dict,
|
|
timestep=timestep.float(),
|
|
clean_x=clean_latent if getattr(self.args, "teacher_forcing", False) else None,
|
|
)
|
|
pred_x0 = pred_x0.to(self.dtype)
|
|
|
|
# Derive denoised_timestep_from / to from the sampled index for ts_schedule
|
|
# Use the first batch element's first block index as the representative scalar
|
|
rep_idx = index[0, 0].item()
|
|
denoised_timestep_from = denoising_step_list[rep_idx].item()
|
|
if rep_idx + 1 < num_steps:
|
|
denoised_timestep_to = denoising_step_list[rep_idx + 1].item()
|
|
else:
|
|
denoised_timestep_to = 0
|
|
|
|
gradient_mask = None
|
|
if context_frames > 0:
|
|
pred_x0[:, :context_frames] = initial_latent.to(
|
|
device=pred_x0.device,
|
|
dtype=pred_x0.dtype,
|
|
)
|
|
gradient_mask = torch.ones_like(pred_x0, dtype=torch.bool)
|
|
gradient_mask[:, :context_frames] = False
|
|
return pred_x0, gradient_mask, denoised_timestep_from, denoised_timestep_to
|
|
|
|
def _run_generator_backward_simulation(
|
|
self,
|
|
image_or_video_shape,
|
|
conditional_dict: dict,
|
|
initial_latent: torch.tensor = None,
|
|
slice_last_frames: int = 21,
|
|
noise=None,
|
|
) -> Tuple[torch.Tensor, torch.Tensor]:
|
|
"""
|
|
On-policy generator via backward simulation (original path).
|
|
"""
|
|
if initial_latent is not None:
|
|
conditional_dict["initial_latent"] = initial_latent
|
|
noise_shape = image_or_video_shape.copy()
|
|
|
|
separate_first_frame = self.independent_first_frame and not getattr(self.args, "i2v", False)
|
|
min_num_frames = (self.min_num_training_frames - 1) if separate_first_frame else self.min_num_training_frames
|
|
max_num_frames = self.num_training_frames - 1 if separate_first_frame else self.num_training_frames
|
|
assert max_num_frames % self.num_frame_per_block == 0
|
|
assert min_num_frames % self.num_frame_per_block == 0
|
|
max_num_blocks = max_num_frames // self.num_frame_per_block
|
|
min_num_blocks = min_num_frames // self.num_frame_per_block
|
|
num_generated_blocks = torch.randint(min_num_blocks, max_num_blocks + 1, (1,), device=self.device)
|
|
dist.broadcast(num_generated_blocks, src=0)
|
|
num_generated_blocks = num_generated_blocks.item()
|
|
num_generated_frames = num_generated_blocks * self.num_frame_per_block
|
|
if separate_first_frame and initial_latent is None:
|
|
num_generated_frames += 1
|
|
min_num_frames += 1
|
|
noise_shape[1] = num_generated_frames
|
|
if noise is not None:
|
|
noise = noise[:, :num_generated_frames]
|
|
else:
|
|
noise = torch.randn(noise_shape, device=self.device, dtype=self.dtype)
|
|
|
|
pred_image_or_video, denoised_timestep_from, denoised_timestep_to = self._consistency_backward_simulation(
|
|
noise=noise,
|
|
slice_last_frames=slice_last_frames,
|
|
**conditional_dict,
|
|
)
|
|
|
|
if slice_last_frames != -1 and pred_image_or_video.shape[1] > slice_last_frames:
|
|
with torch.no_grad():
|
|
if slice_last_frames > 1:
|
|
latent_to_decode = pred_image_or_video[:, :-(slice_last_frames - 1), ...]
|
|
else:
|
|
latent_to_decode = pred_image_or_video
|
|
pixels = self.vae.decode_to_pixel(latent_to_decode)
|
|
frame = pixels[:, -1:, ...].to(self.dtype)
|
|
frame = rearrange(frame, "b t c h w -> b c t h w")
|
|
image_latent = self.vae.encode_to_latent(frame).to(self.dtype)
|
|
if slice_last_frames > 1:
|
|
last_frames = pred_image_or_video[:, -(slice_last_frames - 1):, ...]
|
|
pred_image_or_video_sliced = torch.cat([image_latent, last_frames], dim=1)
|
|
else:
|
|
pred_image_or_video_sliced = image_latent
|
|
if num_generated_frames != min_num_frames:
|
|
gradient_mask = torch.ones_like(pred_image_or_video_sliced, dtype=torch.bool)
|
|
if self.independent_first_frame:
|
|
gradient_mask[:, :1] = False
|
|
else:
|
|
gradient_mask[:, :self.num_frame_per_block] = False
|
|
else:
|
|
gradient_mask = None
|
|
else:
|
|
pred_image_or_video_sliced = pred_image_or_video
|
|
if num_generated_frames != min_num_frames:
|
|
gradient_mask = torch.ones_like(pred_image_or_video_sliced, dtype=torch.bool)
|
|
if self.independent_first_frame:
|
|
gradient_mask[:, :1] = False
|
|
else:
|
|
gradient_mask[:, :self.num_frame_per_block] = False
|
|
else:
|
|
gradient_mask = None
|
|
|
|
pred_image_or_video_sliced = pred_image_or_video_sliced.to(self.dtype)
|
|
return pred_image_or_video_sliced, gradient_mask, denoised_timestep_from, denoised_timestep_to
|
|
|
|
def _consistency_backward_simulation(
|
|
self,
|
|
noise: torch.Tensor,
|
|
slice_last_frames: int = 21,
|
|
**conditional_dict: dict
|
|
) -> torch.Tensor:
|
|
"""
|
|
Simulate the generator's input from noise to avoid training/inference mismatch.
|
|
See Sec 4.5 of the DMD2 paper (https://arxiv.org/abs/2405.14867) for details.
|
|
Here we use the consistency sampler (https://arxiv.org/abs/2303.01469)
|
|
Input:
|
|
- noise: a tensor sampled from N(0, 1) with shape [B, F, C, H, W] where the number of frame is 1 for images.
|
|
- conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).
|
|
Output:
|
|
- output: a tensor with shape [B, T, F, C, H, W].
|
|
T is the total number of timesteps. output[0] is a pure noise and output[i] and i>0
|
|
represents the x0 prediction at each timestep.
|
|
"""
|
|
generator_is_causal = getattr(self.args, "generator_is_causal", True)
|
|
if not generator_is_causal:
|
|
return self._bidirectional_backward_simulation(
|
|
noise=noise,
|
|
slice_last_frames=slice_last_frames,
|
|
**conditional_dict,
|
|
)
|
|
|
|
if self.inference_pipeline is None:
|
|
self._initialize_inference_pipeline()
|
|
|
|
return self.inference_pipeline.inference_with_trajectory(
|
|
noise=noise, **conditional_dict, slice_last_frames=slice_last_frames
|
|
)
|
|
|
|
def _bidirectional_backward_simulation(
|
|
self,
|
|
noise: torch.Tensor,
|
|
slice_last_frames: int = 21,
|
|
**conditional_dict: dict
|
|
) -> Tuple[torch.Tensor, int, int]:
|
|
"""
|
|
Backward simulation for bidirectional (non-causal) generator.
|
|
All frames are processed at once at each denoising step — no KV cache,
|
|
no block-by-block processing.
|
|
"""
|
|
from wan_5b.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
|
|
|
|
batch_size, num_frames = noise.shape[:2]
|
|
|
|
# Resolve single-segment prompt for bidirectional model
|
|
prompt_embeds = conditional_dict["prompt_embeds"]
|
|
num_segments = prompt_embeds.shape[0] // batch_size
|
|
if num_segments > 1:
|
|
prompt_embeds = prompt_embeds.reshape(
|
|
batch_size, num_segments, *prompt_embeds.shape[1:])[:, 0]
|
|
cond = {**conditional_dict, "prompt_embeds": prompt_embeds}
|
|
|
|
# Setup UniPC scheduler
|
|
sampling_steps = getattr(self.args, "sampling_steps", None) or len(self.denoising_step_list)
|
|
shift = self.scheduler.shift
|
|
num_train_timesteps = self.scheduler.num_train_timesteps
|
|
sample_scheduler = FlowUniPCMultistepScheduler(
|
|
num_train_timesteps=num_train_timesteps, shift=1, use_dynamic_shifting=False)
|
|
sample_scheduler.set_timesteps(sampling_steps, device=noise.device, shift=shift)
|
|
unipc_timesteps = sample_scheduler.timesteps
|
|
num_denoising_steps = len(unipc_timesteps)
|
|
|
|
# Pick a random exit step (synchronized across ranks)
|
|
last_step_only = getattr(self.args, "last_step_only", False)
|
|
if last_step_only:
|
|
exit_step = num_denoising_steps - 1
|
|
else:
|
|
exit_step_t = torch.randint(0, num_denoising_steps, (1,), device=self.device)
|
|
dist.broadcast(exit_step_t, src=0)
|
|
exit_step = exit_step_t.item()
|
|
|
|
# Multi-step denoising loop (full-sequence bidirectional forward each step)
|
|
latents = noise
|
|
for index, t in enumerate(unipc_timesteps):
|
|
timestep = t * torch.ones(
|
|
[batch_size, num_frames], device=noise.device, dtype=torch.float32)
|
|
|
|
if index < exit_step:
|
|
with torch.no_grad():
|
|
flow_pred, _ = self.generator(
|
|
noisy_image_or_video=latents,
|
|
conditional_dict=cond,
|
|
timestep=timestep,
|
|
)
|
|
latents = sample_scheduler.step(
|
|
flow_pred, t, latents, return_dict=False)[0]
|
|
else:
|
|
# Exit step: forward with gradient
|
|
flow_pred, denoised_pred = self.generator(
|
|
noisy_image_or_video=latents,
|
|
conditional_dict=cond,
|
|
timestep=timestep,
|
|
)
|
|
break
|
|
|
|
# Compute denoised_timestep_from / to
|
|
if exit_step == num_denoising_steps - 1:
|
|
denoised_timestep_to = 0
|
|
denoised_timestep_from = 1000 - torch.argmin(
|
|
(self.scheduler.timesteps.cuda() - unipc_timesteps[exit_step].cuda()).abs(), dim=0).item()
|
|
else:
|
|
denoised_timestep_to = 1000 - torch.argmin(
|
|
(self.scheduler.timesteps.cuda() - unipc_timesteps[exit_step + 1].cuda()).abs(), dim=0).item()
|
|
denoised_timestep_from = 1000 - torch.argmin(
|
|
(self.scheduler.timesteps.cuda() - unipc_timesteps[exit_step].cuda()).abs(), dim=0).item()
|
|
|
|
return denoised_pred, denoised_timestep_from, denoised_timestep_to
|
|
|
|
def _initialize_inference_pipeline(self):
|
|
"""
|
|
Lazy initialize the inference pipeline during the first backward simulation run.
|
|
Here we encapsulate the inference code with a model-dependent outside function.
|
|
We pass our FSDP-wrapped modules into the pipeline to save memory.
|
|
"""
|
|
local_attn_size = section_get(
|
|
self.args,
|
|
"inference",
|
|
"local_attn_size",
|
|
getattr(self.args, "model_kwargs", {}).get("local_attn_size", -1),
|
|
aliases=("inference_local_attn_size",),
|
|
)
|
|
sink_size = section_get(
|
|
self.args,
|
|
"inference",
|
|
"sink_size",
|
|
getattr(self.args, "model_kwargs", {}).get("sink_size", 0),
|
|
aliases=("inference_sink_size",),
|
|
)
|
|
multi_shot_sink = section_get(self.args, "inference", "multi_shot_sink", False)
|
|
multi_shot_rope_offset = section_get(
|
|
self.args,
|
|
"inference",
|
|
"multi_shot_rope_offset",
|
|
0.0,
|
|
)
|
|
scene_cut_prefix = section_get(self.args, "inference", "scene_cut_prefix", "[SCENE_CUT]")
|
|
slice_last_frames = getattr(self.args, "slice_last_frames", 21)
|
|
# do not use self.num_training_frames, because it is changed by generator_loss and critic_loss
|
|
num_training_frames = getattr(self.args, "num_training_frames")
|
|
if local_attn_size == -1:
|
|
kv_cache_size = num_training_frames
|
|
else:
|
|
kv_cache_size = min(local_attn_size + slice_last_frames, num_training_frames)
|
|
frame_seq_length = math.prod(self.args.image_or_video_shape[-2:]) // 4
|
|
self.inference_pipeline = SelfForcingTrainingPipeline(
|
|
denoising_step_list=self.denoising_step_list,
|
|
scheduler=self.scheduler,
|
|
generator=self.generator,
|
|
num_frame_per_block=self.num_frame_per_block,
|
|
independent_first_frame=self.independent_first_frame,
|
|
same_step_across_blocks=getattr(
|
|
self.args, "same_step_across_blocks", getattr(self, "same_step_across_blocks", False)
|
|
),
|
|
last_step_only=getattr(self.args, "last_step_only", False),
|
|
num_max_frames=kv_cache_size,
|
|
context_noise=getattr(self.args, "context_noise", 0),
|
|
sampling_steps=getattr(self.args, "sampling_steps", None),
|
|
local_attn_size=local_attn_size,
|
|
sink_size=sink_size,
|
|
multi_shot_sink=multi_shot_sink,
|
|
scene_cut_prefix=scene_cut_prefix,
|
|
multi_shot_rope_offset=multi_shot_rope_offset,
|
|
slice_last_frames=slice_last_frames,
|
|
num_training_frames=num_training_frames,
|
|
model_name=getattr(self.args, "model_kwargs", {}).get("model_name", None),
|
|
frame_seq_length=frame_seq_length,
|
|
)
|