| """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 |
|
|
