File size: 5,657 Bytes
2ea2835 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 | """Shared generator runtime used by distribution matching objectives.
This module contains rollout mechanics only. Distribution objectives belong in
their concrete model classes; frozen-encoder losses and statistics are not part
of the generator runtime.
"""
from typing import Optional
import torch
import torch.distributed as dist
from model.base import SelfForcingModel
class GeneratorRuntimeBase(SelfForcingModel):
"""Initialize the generator and provide common differentiable rollouts."""
uses_critic = False
uses_vae_for_loss = True
def __init__(self, args, device):
super().__init__(args, device)
self.causal_generator = bool(getattr(args, "causal", True))
self.num_frame_per_block = getattr(args, "num_frame_per_block", 1)
self.same_step_across_blocks = getattr(
args, "same_step_across_blocks", True)
self.num_training_frames = getattr(args, "num_training_frames", 21)
if self.causal_generator and self.num_frame_per_block > 1:
self.generator.model.num_frame_per_block = self.num_frame_per_block
self.independent_first_frame = getattr(
args, "independent_first_frame", False)
if self.causal_generator and self.independent_first_frame:
self.generator.model.independent_first_frame = True
if args.gradient_checkpointing:
self.generator.enable_gradient_checkpointing()
self.inference_pipeline = None
self.generation_guidance_scale = float(
getattr(args, "guidance_scale", 1.0))
if not hasattr(self, "denoising_step_list"):
raise ValueError(
"few-step generation requires denoising_step_list")
denoising_step_list = self.denoising_step_list.to(device=device)
if (
denoising_step_list.numel() > 0
and denoising_step_list[-1].item() == 0
):
denoising_step_list = denoising_step_list[:-1]
self.generation_denoising_step_list = denoising_step_list
@staticmethod
def _apply_gradient_mask(
latents: torch.Tensor,
gradient_mask: Optional[torch.Tensor],
) -> torch.Tensor:
if gradient_mask is None:
return latents
frame_mask = gradient_mask.flatten(2).any(dim=2).view(
latents.shape[0], latents.shape[1], 1, 1, 1)
return torch.where(frame_mask, latents, latents.detach())
def _predict_x0_with_cfg(
self,
noisy_latent: torch.Tensor,
conditional_dict: dict,
unconditional_dict: dict,
timestep: torch.Tensor,
) -> torch.Tensor:
_, x0_cond = self.generator(
noisy_image_or_video=noisy_latent,
conditional_dict=conditional_dict,
timestep=timestep,
)
if self.generation_guidance_scale == 1.0:
return x0_cond
_, x0_uncond = self.generator(
noisy_image_or_video=noisy_latent,
conditional_dict=unconditional_dict,
timestep=timestep,
)
return x0_uncond + self.generation_guidance_scale * (
x0_cond - x0_uncond)
def _run_few_step_generator(
self,
image_or_video_shape,
conditional_dict: dict,
unconditional_dict: dict,
exit_index: Optional[int] = None,
) -> torch.Tensor:
"""Roll out detached steps, then differentiate one selected exit step."""
latent = torch.randn(
image_or_video_shape,
device=self.device,
dtype=self.dtype,
)
batch_size, num_frames = latent.shape[:2]
denoising_steps = self.generation_denoising_step_list.to(
device=self.device)
num_steps = len(denoising_steps)
if exit_index is None:
if num_steps <= 1:
exit_index = 0
else:
index_tensor = torch.randint(
0, num_steps, (1,), device=self.device)
if dist.is_available() and dist.is_initialized():
dist.broadcast(index_tensor, src=0)
exit_index = int(index_tensor.item())
for index in range(num_steps):
current_timestep = denoising_steps[index]
timestep = torch.ones(
[batch_size, num_frames],
device=self.device,
dtype=current_timestep.dtype,
) * current_timestep
if index < exit_index:
with torch.no_grad():
x0_pred = self._predict_x0_with_cfg(
noisy_latent=latent,
conditional_dict=conditional_dict,
unconditional_dict=unconditional_dict,
timestep=timestep,
)
next_timestep = denoising_steps[index + 1]
latent = self.scheduler.add_noise(
x0_pred.flatten(0, 1),
torch.randn_like(x0_pred.flatten(0, 1)),
next_timestep * torch.ones(
[batch_size * num_frames],
device=self.device,
dtype=next_timestep.dtype,
),
).unflatten(0, x0_pred.shape[:2]).to(self.dtype)
else:
return self._predict_x0_with_cfg(
noisy_latent=latent,
conditional_dict=conditional_dict,
unconditional_dict=unconditional_dict,
timestep=timestep,
)
return latent
|