diff --git a/fastvideo/pipelines/stages/causal_denoising.py b/fastvideo/pipelines/stages/causal_denoising.py index 0afb07406..804d22f3f 100644 --- a/fastvideo/pipelines/stages/causal_denoising.py +++ b/fastvideo/pipelines/stages/causal_denoising.py @@ -34,7 +34,11 @@ class CausalDMDDenosingStage(DenoisingStage): Denoising stage for causal diffusion. """ - def __init__(self, transformer, scheduler, transformer_2=None, vae=None) -> None: + def __init__(self, + transformer, + scheduler, + transformer_2=None, + vae=None) -> None: super().__init__(transformer, scheduler, transformer_2) # KV and cross-attention cache state (initialized on first forward) self.transformer = transformer @@ -110,14 +114,14 @@ class CausalDMDDenosingStage(DenoisingStage): # Initialize or reset caches kv_cache1 = self._initialize_kv_cache(batch_size=latents.shape[0], - dtype=target_dtype, - device=latents.device) + dtype=target_dtype, + device=latents.device) kv_cache2 = None if boundary_timestep is not None: # Initialize the low noise kv cache kv_cache2 = self._initialize_kv_cache(batch_size=latents.shape[0], - dtype=target_dtype, - device=latents.device) + dtype=target_dtype, + device=latents.device) def _get_kv_cache(timestep: float) -> dict: if boundary_timestep is not None: @@ -128,11 +132,11 @@ class CausalDMDDenosingStage(DenoisingStage): return kv_cache1 crossattn_cache = self._initialize_crossattn_cache( - batch_size=latents.shape[0], - max_text_len=fastvideo_args.pipeline_config. - text_encoder_configs[0].arch_config.text_len, - dtype=target_dtype, - device=latents.device) + batch_size=latents.shape[0], + max_text_len=fastvideo_args.pipeline_config.text_encoder_configs[0]. + arch_config.text_len, + dtype=target_dtype, + device=latents.device) pos_start_base = 0 @@ -151,15 +155,17 @@ class CausalDMDDenosingStage(DenoisingStage): if (hasattr(self.vae, "shift_factor") and self.vae.shift_factor is not None): if isinstance(self.vae.shift_factor, torch.Tensor): - first_frame_latent -= self.vae.shift_factor.to(first_frame_latent.device, first_frame_latent.dtype) + first_frame_latent -= self.vae.shift_factor.to( + first_frame_latent.device, first_frame_latent.dtype) else: first_frame_latent -= self.vae.shift_factor if isinstance(self.vae.scaling_factor, torch.Tensor): - first_frame_latent = first_frame_latent * self.vae.scaling_factor.to(first_frame_latent.device, first_frame_latent.dtype) + first_frame_latent = first_frame_latent * self.vae.scaling_factor.to( + first_frame_latent.device, first_frame_latent.dtype) else: first_frame_latent = first_frame_latent * self.vae.scaling_factor - + if fastvideo_args.vae_cpu_offload: self.vae = self.vae.to("cpu") @@ -173,11 +179,24 @@ class CausalDMDDenosingStage(DenoisingStage): set_forward_context(current_timestep=0, attn_metadata=None, forward_batch=batch): - self.transformer( + self.transformer( + first_frame_latent.to(target_dtype), + prompt_embeds, + t_zero, + kv_cache=kv_cache1, + crossattn_cache=crossattn_cache, + current_start=(pos_start_base + start_index) * + self.frame_seq_length, + start_frame=start_index, + **image_kwargs, + **pos_cond_kwargs, + ) + if boundary_timestep is not None: + self.transformer_2( first_frame_latent.to(target_dtype), prompt_embeds, t_zero, - kv_cache=kv_cache1, + kv_cache=kv_cache2, crossattn_cache=crossattn_cache, current_start=(pos_start_base + start_index) * self.frame_seq_length, @@ -185,24 +204,11 @@ class CausalDMDDenosingStage(DenoisingStage): **image_kwargs, **pos_cond_kwargs, ) - if boundary_timestep is not None: - self.transformer_2( - first_frame_latent.to(target_dtype), - prompt_embeds, - t_zero, - kv_cache=kv_cache2, - crossattn_cache=crossattn_cache, - current_start=(pos_start_base + start_index) * - self.frame_seq_length, - start_frame=start_index, - **image_kwargs, - **pos_cond_kwargs, - ) - + start_index += 1 block_sizes.pop(0) latents[:, :, :1, :, :] = first_frame_latent - + # DMD loop in causal blocks with self.progress_bar(total=len(block_sizes) * len(timesteps)) as progress_bar: @@ -287,8 +293,10 @@ class CausalDMDDenosingStage(DenoisingStage): pred_noise=pred_noise_btchw.flatten(0, 1), noise_input_latent=noise_latents.flatten(0, 1), timestep=t_expand, - boundary_timestep=torch.ones_like(t_expand) * boundary_timestep, - scheduler=self.scheduler).unflatten(0, pred_noise_btchw.shape[:2]) + boundary_timestep=torch.ones_like(t_expand) * + boundary_timestep, + scheduler=self.scheduler).unflatten( + 0, pred_noise_btchw.shape[:2]) else: pred_video_btchw = pred_noise_to_pred_video( pred_noise=pred_noise_btchw.flatten(0, 1), @@ -309,25 +317,28 @@ class CausalDMDDenosingStage(DenoisingStage): batch.generator, list) else batch.generator)).to(self.device) noise_btchw = noise - if boundary_timestep is not None and i < len(high_noise_timesteps) - 1: + if boundary_timestep is not None and i < len( + high_noise_timesteps) - 1: noise_latents_btchw = self.scheduler.add_noise_high( pred_video_btchw.flatten(0, 1), - noise_btchw.flatten(0, 1), - next_timestep, - torch.ones_like(next_timestep) * boundary_timestep).unflatten(0, pred_video_btchw.shape[:2]) - elif boundary_timestep is not None and i == len(high_noise_timesteps) - 1: + noise_btchw.flatten(0, 1), next_timestep, + torch.ones_like(next_timestep) * + boundary_timestep).unflatten( + 0, pred_video_btchw.shape[:2]) + elif boundary_timestep is not None and i == len( + high_noise_timesteps) - 1: noise_latents_btchw = pred_video_btchw else: noise_latents_btchw = self.scheduler.add_noise( pred_video_btchw.flatten(0, 1), noise_btchw.flatten(0, 1), - next_timestep).unflatten(0, - pred_video_btchw.shape[:2]) + next_timestep).unflatten( + 0, pred_video_btchw.shape[:2]) current_latents = noise_latents_btchw.permute( 0, 2, 1, 3, 4) else: current_latents = pred_video_btchw.permute( - 0, 2, 1, 3, 4) + 0, 2, 1, 3, 4) if progress_bar is not None: progress_bar.update() @@ -364,7 +375,7 @@ class CausalDMDDenosingStage(DenoisingStage): **image_kwargs, **pos_cond_kwargs, ) - + self.transformer( context_bcthw, prompt_embeds, @@ -377,7 +388,7 @@ class CausalDMDDenosingStage(DenoisingStage): **image_kwargs, **pos_cond_kwargs, ) - + start_index += current_num_frames batch.latents = latents diff --git a/fastvideo/pipelines/stages/input_validation.py b/fastvideo/pipelines/stages/input_validation.py index b1f2b4628..3d263673a 100644 --- a/fastvideo/pipelines/stages/input_validation.py +++ b/fastvideo/pipelines/stages/input_validation.py @@ -5,7 +5,6 @@ Input validation stage for diffusion pipelines. import torch import torchvision.transforms.functional as TF -from PIL import Image from fastvideo.fastvideo_args import FastVideoArgs from fastvideo.logger import init_logger @@ -14,7 +13,6 @@ from fastvideo.pipelines.pipeline_batch_info import ForwardBatch from fastvideo.pipelines.stages.base import PipelineStage from fastvideo.pipelines.stages.validators import (StageValidators, VerificationResult) -from fastvideo.utils import best_output_size logger = init_logger(__name__) @@ -106,7 +104,9 @@ class InputValidationStage(PipelineStage): batch.pil_image = image # further processing for ti2v task - if (fastvideo_args.pipeline_config.ti2v_task or fastvideo_args.pipeline_config.is_causal) and batch.pil_image is not None: + if (fastvideo_args.pipeline_config.ti2v_task + or fastvideo_args.pipeline_config.is_causal + ) and batch.pil_image is not None: img = batch.pil_image # ih, iw = img.height, img.width # patch_size = fastvideo_args.pipeline_config.dit_config.arch_config.patch_size @@ -126,8 +126,7 @@ class InputValidationStage(PipelineStage): # y1 = (img.height - oh) // 2 # img = img.crop((x1, y1, x1 + ow, y1 + oh)) # assert img.width == ow and img.height == oh - logger.info("img height: %s, img width: %s", img.height, - img.width) + logger.info("img height: %s, img width: %s", img.height, img.width) oh = img.height ow = img.width # to tensor