stages ¶

def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
    """Denoise latents while keeping scheduler samples and predictions in fp32."""
    if batch.latents is None or batch.timesteps is None:
        raise ValueError("LingBot-Video denoising requires latents and timesteps")
    device = get_local_torch_device()
    transformer_dtype = next(self.transformer.parameters()).dtype
    condition, condition_mask = self._prepare_conditions(batch, transformer_dtype, device)
    latents = batch.latents.to(device=device, dtype=torch.float32)
    do_cfg = self._uses_cfg(batch)
    negative = negative_mask = None
    if do_cfg and not batch.batch_cfg:
        negative, negative_mask = self._negative_condition(batch, condition, condition_mask, transformer_dtype,
                                                           device)
    trajectory: list[torch.Tensor] = []
    trajectory_timesteps: list[torch.Tensor] = []
    for timestep in batch.timesteps:
        timestep_batch = self._transformer_timestep(timestep, transformer_dtype).expand(1).to(device)
        latent_input = latents
        if do_cfg and batch.batch_cfg:
            latent_input = torch.cat((latents, latents), dim=0)
            timestep_batch = torch.cat((timestep_batch, timestep_batch), dim=0)
        autocast_enabled = device.type == "cuda" and transformer_dtype != torch.float32
        with torch.autocast(device_type=device.type, dtype=transformer_dtype, enabled=autocast_enabled):
            prediction = self.transformer(
                latent_input,
                timestep_batch,
                condition,
                encoder_attention_mask=condition_mask,
                return_dict=False,
            )[0].float()
        if do_cfg:
            if batch.batch_cfg:
                conditional, unconditional = prediction.chunk(2, dim=0)
            else:
                with torch.autocast(
                        device_type=device.type,
                        dtype=transformer_dtype,
                        enabled=autocast_enabled,
                ):
                    unconditional = self.transformer(
                        latents,
                        timestep_batch,
                        negative,
                        encoder_attention_mask=negative_mask,
                        return_dict=False,
                    )[0].float()
                conditional = prediction
            guidance_scale = batch.guidance_scale_2 if self.refiner else batch.guidance_scale
            if guidance_scale is None:
                raise ValueError("LingBot-Video CFG requires a guidance scale")
            prediction = unconditional + guidance_scale * (conditional - unconditional)
        latents = self.scheduler.step(
            prediction,
            timestep,
            latents,
            return_dict=False,
            generator=batch.generator,
        )[0].float()
        if batch.return_trajectory_latents:
            trajectory.append(latents.detach().cpu())
            trajectory_timesteps.append(timestep.detach().cpu())
    batch.latents = latents
    if trajectory:
        batch.trajectory_latents = torch.stack(trajectory, dim=1)
        batch.trajectory_timesteps = trajectory_timesteps
    return batch

def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
    """Run shared validation, then enforce LingBot temporal and spatial geometry."""
    batch = super().forward(batch, fastvideo_args)
    if not isinstance(batch.num_frames, int):
        raise TypeError("LingBot-Video num_frames must be an integer")
    if batch.num_frames != 1 and (batch.num_frames - 1) % 4 != 0:
        raise ValueError(f"num_frames must be 1 or 4n+1, got {batch.num_frames}")
    if not isinstance(batch.height, int) or not isinstance(batch.width, int):
        raise TypeError("LingBot-Video height and width must be integers")
    if batch.height % 16 != 0 or batch.width % 16 != 0:
        raise ValueError(f"height and width must be divisible by 16, got {batch.height}x{batch.width}")
    if isinstance(batch.prompt, list) and len(batch.prompt) != 1:
        raise ValueError("LingBot-Video currently supports prompt batch size one")
    if self.refiner_enabled and fastvideo_args.output_type == "latent":
        raise ValueError("LingBot-Video refinement requires decoded pixel output")
    return batch

def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
    """Generate or validate one normalized fp32 latent video."""
    del fastvideo_args
    if not all(isinstance(value, int) for value in (batch.num_frames, batch.height, batch.width)):
        raise TypeError("latent geometry must contain integer frames, height, and width")
    shape = (
        1,
        self.transformer.num_channels_latents,
        (batch.num_frames - 1) // 4 + 1,
        batch.height // 8,
        batch.width // 8,
    )
    device = get_local_torch_device()
    if batch.latents is None:
        batch.latents = torch.randn(
            shape,
            generator=batch.generator,
            device=device,
            dtype=torch.float32,
        )
    else:
        if tuple(batch.latents.shape) != shape:
            raise ValueError(f"supplied latent shape {tuple(batch.latents.shape)} does not match {shape}")
        batch.latents = batch.latents.to(device=device, dtype=torch.float32)
    batch.raw_latent_shape = shape
    return batch

def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
    """Prepare high-resolution refiner latents and its exact truncated sigma schedule."""
    if batch.output is None or batch.output.ndim != 5:
        raise ValueError("LingBot-Video refinement requires a decoded base video")
    if not isinstance(batch.height_sr, int) or not isinstance(batch.width_sr, int):
        raise TypeError("LingBot-Video refinement requires integer height_sr and width_sr")
    if batch.height_sr % 16 != 0 or batch.width_sr % 16 != 0:
        raise ValueError("LingBot-Video refiner height_sr and width_sr must be divisible by 16")
    if batch.seed is None:
        raise ValueError("LingBot-Video refinement requires a seed")
    device = get_local_torch_device()
    if isinstance(self.vae, torch.nn.Module):
        self.vae.to(device)
    generator = torch.Generator(device=device).manual_seed(batch.seed)
    resized = self._resize_video(batch.output, batch.height_sr, batch.width_sr)
    encoded = self._encode_video(resized, generator, device)
    noise = torch.randn(encoded.shape, generator=generator, device=device, dtype=encoded.dtype)
    batch.latents = ((1.0 - batch.t_thresh) * encoded + batch.t_thresh * noise).float()
    batch.generator = generator
    batch.height = batch.height_sr
    batch.width = batch.width_sr
    batch.raw_latent_shape = tuple(batch.latents.shape)
    batch.extra["lingbot_video_base_shape"] = tuple(batch.output.shape)
    batch.output = None

    shift = fastvideo_args.pipeline_config.flow_shift
    if shift is None:
        raise ValueError("LingBot-Video refinement requires a flow shift")
    sigmas = _compute_refiner_sigmas(
        float(self.scheduler.sigma_max),
        float(self.scheduler.sigma_min),
        batch.num_inference_steps_sr,
        float(shift),
        float(batch.t_thresh),
    )
    self.scheduler.set_timesteps(len(sigmas), device=device, sigmas=sigmas, shift=1.0)
    batch.timesteps = self.scheduler.timesteps
    if getattr(fastvideo_args, "vae_cpu_offload", False):
        self.vae.to("cpu")
    return batch