import time import torch from omegaconf import DictConfig from einops import rearrange from pipeline import ( CausalInferencePipeline, ) from utils.misc import set_seed from demo_utils.memory import gpu device = torch.device("cuda") def get_video(video): video = rearrange(video, "b t c h w -> b t h w c").cpu() video = 255.0 * video assert video.shape[0] == 1 video = video[0] return video class StreamInferenceWrapper: def __init__( self, stream_model_config: DictConfig, checkpoint_path: str, total_generate_block_number: int, use_ema: bool = True, seed: int = 0, ): torch.set_grad_enabled(False) # Initialize pipeline assert hasattr(stream_model_config, "denoising_step_list") self.pipeline = CausalInferencePipeline( stream_model_config, device=device, ) state_dict = torch.load(checkpoint_path, map_location="cpu") self.pipeline.generator.load_state_dict( state_dict[("generator" if not use_ema else "generator_ema")] ) self.pipeline = self.pipeline.to(dtype=torch.bfloat16) self.pipeline.text_encoder.to(device=gpu) self.pipeline.generator.to(device=gpu) if not stream_model_config.vae_offload_cpu: self.pipeline.vae.to(device=gpu) else: self.pipeline.vae.to(device="cpu") # Offload VAE to CPU set_seed(seed) self.seed = seed self.initial_noise = torch.randn( [ 1, total_generate_block_number * self.pipeline.num_frame_per_block, 16, 60, 104, ], device=device, dtype=torch.bfloat16, ) self.recorded_latents = None self.video = None self.stream_model_config = stream_model_config print(f""" {self.__class__.__name__}.__init__(): {self.initial_noise.shape = } {self.stream_model_config = } """) def block_to_latent_index(self, block_index: int) -> int: return block_index * self.pipeline.num_frame_per_block def latent_to_video_index(self, latent_frame_index: int) -> int: if latent_frame_index <= 0: return 0 return (latent_frame_index - 1) * 4 + 1 def block_to_video_index(self, block_index: int) -> int: return self.latent_to_video_index(self.block_to_latent_index(block_index)) def get_sampled_noise( self, start_block_index: int, end_block_index: int, ): current_start_latent_frame_index = self.block_to_latent_index(start_block_index) current_end_latent_frame_index = self.block_to_latent_index(end_block_index) print(f"{current_start_latent_frame_index = } | {current_end_latent_frame_index = }") assert current_start_latent_frame_index < self.initial_noise.shape[1] assert current_end_latent_frame_index <= self.initial_noise.shape[1] sampled_noise = self.initial_noise[ :, current_start_latent_frame_index:current_end_latent_frame_index, ..., ] return sampled_noise def get_initial_latents( self, start_block_index: int, ): if self.recorded_latents is None: return None print(f"{start_block_index = }") return self.recorded_latents[ :, : self.block_to_latent_index(start_block_index), ] def decode_to_pixel( self, latents: torch.Tensor, ): start_decode_time = time.time() # Move VAE to GPU if offloaded if self.stream_model_config.vae_offload_cpu: self.pipeline.vae.to(device=gpu) video = self.pipeline.vae.decode_to_pixel(latents, use_cache=False) # Optionally move VAE back to CPU after decoding if self.stream_model_config.vae_offload_cpu: self.pipeline.vae.to(device="cpu") video = (video * 0.5 + 0.5).clamp(0, 1) print( f"{self.__class__.__name__}.decode_to_pixel() VAE decode time: {time.time() - start_decode_time:.2f} seconds" ) return video def update_video( self, video: torch.Tensor, start_latent_frame_index: int, ): video = get_video(video) # t, h, w, c start_video_frame_index = self.latent_to_video_index(start_latent_frame_index) if self.video is None: self.video = video else: self.video = self.video[:start_video_frame_index] self.video = torch.cat([self.video, video], dim=0) def decode_and_update_video( self, start_block_index: int, end_block_index: int, ): if start_block_index == 0: current_chunk_latent = self.recorded_latents[ :, self.block_to_latent_index(start_block_index) : self.block_to_latent_index( end_block_index ), ] current_chunk_video = self.decode_to_pixel(current_chunk_latent) else: current_chunk_latent = self.recorded_latents[ :, self.block_to_latent_index(start_block_index - 1) : self.block_to_latent_index( end_block_index ), ] current_chunk_video = self.decode_to_pixel(current_chunk_latent) current_chunk_video = current_chunk_video[:, 9:] self.update_video(current_chunk_video, self.block_to_latent_index(start_block_index)) def inference( self, start_block_index: int, end_block_index: int, prompt: str, ): assert start_block_index >= 0 assert end_block_index > start_block_index print(f""" {self.__class__.__name__}.inference(): {start_block_index = } | {end_block_index = } """) sampled_noise = self.get_sampled_noise(start_block_index, end_block_index) prompts = [prompt] initial_latents = self.get_initial_latents( start_block_index, ) if initial_latents is not None: print(f"{initial_latents.shape = }") latents_result = self.pipeline.inference( noise=sampled_noise, text_prompts=prompts, return_latents=True, initial_latent=initial_latents, do_not_decode_video=True, do_not_recompute_initial_latents=True, ) latents = latents_result print(f"{latents.shape = }") if self.recorded_latents is None: self.recorded_latents = latents else: self.recorded_latents = torch.concat( [ self.recorded_latents[:, :0], latents, ], dim=1, ) print(f"{self.recorded_latents.shape = }") self.decode_and_update_video(start_block_index, end_block_index) return ( self.video, self.video[self.block_to_video_index(start_block_index) :], ) def reset( self, ): self.recorded_latents = None self.video = None # Clear VAE cache self.pipeline.vae.model.clear_cache() # Optionally move VAE back to CPU after reset if offloading if self.stream_model_config.vae_offload_cpu: self.pipeline.vae.to(device="cpu") if self.pipeline.is_kv_cache_initialized(): self.pipeline._reset_kv_cache() if self.pipeline.is_crossattn_cache_initialized(): self.pipeline._reset_crossattn_cache()