File size: 5,470 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 | """Shared generator setup and few-step rollout utilities.
This module contains only generator behavior used by distribution matching. Loss
construction lives in :mod:`model.distribution_matching`.
"""
from __future__ import annotations
from typing import Optional
import torch
import torch.distributed as dist
from model.base import SelfForcingModel
class FewStepGenerationModel(SelfForcingModel):
"""Initialize the causal generator and expose deterministic replay helpers."""
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.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_steps = self.denoising_step_list.to(device=device)
if denoising_steps.numel() > 0 and denoising_steps[-1].item() == 0:
denoising_steps = denoising_steps[:-1]
self.generation_denoising_step_list = denoising_steps
@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.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.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:
"""Run detached rollout steps followed by one differentiated 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
|