# Adapted from https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py import cv2 import os import re import inspect from typing import Callable, List, Optional, Union, Tuple from dataclasses import dataclass import numpy as np import PIL.Image import torch from torch import nn from diffusers.utils import is_accelerate_available, PIL_INTERPOLATION from packaging import version from transformers import CLIPTextModel, CLIPTokenizer from diffusers.modeling_utils import ModelMixin from diffusers.configuration_utils import FrozenDict from diffusers.models import AutoencoderKL from diffusers.pipeline_utils import DiffusionPipeline from ..models.controlnet import ControlNetModel, ControlNetOutput from diffusers.schedulers import ( DDIMScheduler, DPMSolverMultistepScheduler, EulerAncestralDiscreteScheduler, EulerDiscreteScheduler, LMSDiscreteScheduler, PNDMScheduler, ) from diffusers.utils import deprecate, logging, BaseOutput from einops import rearrange from ..data.dataset import TuneAVideoDataset from ..models.unet import UNet3DConditionModel from ..loaders import TextualInversionLoaderMixin from ..util import ddim_inversion, controlnet_image_preprocessing logger = logging.get_logger(__name__) # pylint: disable=invalid-name @dataclass class TuneAVideoPipelineOutput(BaseOutput): videos: Union[torch.Tensor, np.ndarray] class MultiControlNetModel(ModelMixin): r""" Multiple `ControlNetModel` wrapper class for Multi-ControlNet This module is a wrapper for multiple instances of the `ControlNetModel`. The `forward()` API is designed to be compatible with `ControlNetModel`. Args: controlnets (`List[ControlNetModel]`): Provides additional conditioning to the unet during the denoising process. You must set multiple `ControlNetModel` as a list. """ def __init__(self, controlnets: Union[List[ControlNetModel], Tuple[ControlNetModel]]): super().__init__() self.nets = nn.ModuleList(controlnets) def forward( self, sample: torch.FloatTensor, timestep: Union[torch.Tensor, float, int], encoder_hidden_states: torch.Tensor, controlnet_cond: List[torch.tensor], conditioning_scale: List[float], class_labels: Optional[torch.Tensor] = None, timestep_cond: Optional[torch.Tensor] = None, attention_mask: Optional[torch.Tensor] = None, return_dict: bool = True, ) -> Union[ControlNetOutput, Tuple]: for i, (image, scale, controlnet) in enumerate(zip(controlnet_cond, conditioning_scale, self.nets)): down_samples, mid_sample = controlnet( sample, timestep, encoder_hidden_states, image, scale, class_labels, timestep_cond, attention_mask, return_dict, ) # merge samples if i == 0: down_block_res_samples, mid_block_res_sample = down_samples, mid_sample else: down_block_res_samples = [ samples_prev + samples_curr for samples_prev, samples_curr in zip(down_block_res_samples, down_samples) ] mid_block_res_sample += mid_sample return down_block_res_samples, mid_block_res_sample class TuneAVideoPipeline(DiffusionPipeline, TextualInversionLoaderMixin): _optional_components = [] def __init__( self, vae: AutoencoderKL, text_encoder: CLIPTextModel, tokenizer: CLIPTokenizer, unet: UNet3DConditionModel, scheduler: Union[ DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler, EulerDiscreteScheduler, EulerAncestralDiscreteScheduler, DPMSolverMultistepScheduler, ], controlnet: Union[List[ControlNetModel], Tuple[ControlNetModel], MultiControlNetModel] = None, ): super().__init__() if hasattr(scheduler.config, "steps_offset") and scheduler.config.steps_offset != 1: deprecation_message = ( f"The configuration file of this scheduler: {scheduler} is outdated. `steps_offset`" f" should be set to 1 instead of {scheduler.config.steps_offset}. Please make sure " "to update the config accordingly as leaving `steps_offset` might led to incorrect results" " in future versions. If you have downloaded this checkpoint from the Hugging Face Hub," " it would be very nice if you could open a Pull request for the `scheduler/scheduler_config.json`" " file" ) deprecate("steps_offset!=1", "1.0.0", deprecation_message, standard_warn=False) new_config = dict(scheduler.config) new_config["steps_offset"] = 1 scheduler._internal_dict = FrozenDict(new_config) if hasattr(scheduler.config, "clip_sample") and scheduler.config.clip_sample is True: deprecation_message = ( f"The configuration file of this scheduler: {scheduler} has not set the configuration `clip_sample`." " `clip_sample` should be set to False in the configuration file. Please make sure to update the" " config accordingly as not setting `clip_sample` in the config might lead to incorrect results in" " future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it would be very" " nice if you could open a Pull request for the `scheduler/scheduler_config.json` file" ) deprecate("clip_sample not set", "1.0.0", deprecation_message, standard_warn=False) new_config = dict(scheduler.config) new_config["clip_sample"] = False scheduler._internal_dict = FrozenDict(new_config) is_unet_version_less_0_9_0 = hasattr(unet.config, "_diffusers_version") and version.parse( version.parse(unet.config._diffusers_version).base_version ) < version.parse("0.9.0.dev0") is_unet_sample_size_less_64 = hasattr(unet.config, "sample_size") and unet.config.sample_size < 64 if is_unet_version_less_0_9_0 and is_unet_sample_size_less_64: deprecation_message = ( "The configuration file of the unet has set the default `sample_size` to smaller than" " 64 which seems highly unlikely. If your checkpoint is a fine-tuned version of any of the" " following: \n- CompVis/stable-diffusion-v1-4 \n- CompVis/stable-diffusion-v1-3 \n-" " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- runwayml/stable-diffusion-v1-5" " \n- runwayml/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the" " configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`" " in the config might lead to incorrect results in future versions. If you have downloaded this" " checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for" " the `unet/config.json` file" ) deprecate("sample_size<64", "1.0.0", deprecation_message, standard_warn=False) new_config = dict(unet.config) new_config["sample_size"] = 64 unet._internal_dict = FrozenDict(new_config) if isinstance(controlnet, (list, tuple)): controlnet = MultiControlNetModel(controlnet) self.register_modules( vae=vae, text_encoder=text_encoder, tokenizer=tokenizer, unet=unet, controlnet=controlnet, scheduler=scheduler, ) self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) def enable_vae_slicing(self): self.vae.enable_slicing() def disable_vae_slicing(self): self.vae.disable_slicing() def enable_sequential_cpu_offload(self, gpu_id=0): if is_accelerate_available(): from accelerate import cpu_offload else: raise ImportError("Please install accelerate via `pip install accelerate`") device = torch.device(f"cuda:{gpu_id}") for cpu_offloaded_model in [self.unet, self.text_encoder, self.vae, self.controlnet]: if cpu_offloaded_model is not None: cpu_offload(cpu_offloaded_model, device) @property def _execution_device(self): if self.device != torch.device("meta") or not hasattr(self.unet, "_hf_hook"): return self.device for module in self.unet.modules(): if ( hasattr(module, "_hf_hook") and hasattr(module._hf_hook, "execution_device") and module._hf_hook.execution_device is not None ): return torch.device(module._hf_hook.execution_device) return self.device def _encode_prompt(self, prompt, device, num_videos_per_prompt, do_classifier_free_guidance, negative_prompt=None, prompt_embeds: Optional[torch.FloatTensor] = None, negative_prompt_embeds: Optional[torch.FloatTensor] = None): if prompt is not None and isinstance(prompt, str): batch_size = 1 else: batch_size = prompt_embeds.shape[0] if prompt_embeds is None: text_inputs = self.tokenizer( prompt, padding="max_length", max_length=self.tokenizer.model_max_length, truncation=True, return_tensors="pt", ) text_input_ids = text_inputs.input_ids untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal( text_input_ids, untruncated_ids ): removed_text = self.tokenizer.batch_decode( untruncated_ids[:, self.tokenizer.model_max_length - 1: -1] ) logger.warning( "The following part of your input was truncated because CLIP can only handle sequences up to" f" {self.tokenizer.model_max_length} tokens: {removed_text}" ) if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask: attention_mask = text_inputs.attention_mask.to(device) else: attention_mask = None prompt_embeds = self.text_encoder( text_input_ids.to(device), attention_mask=attention_mask, ) prompt_embeds = prompt_embeds[0] prompt_embeds = prompt_embeds.to(dtype=self.text_encoder.dtype, device=device) bs_embed, seq_len, _ = prompt_embeds.shape # duplicate text embeddings for each generation per prompt, using mps friendly method prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1) prompt_embeds = prompt_embeds.view(bs_embed * num_videos_per_prompt, seq_len, -1) # get unconditional embeddings for classifier free guidance if do_classifier_free_guidance and negative_prompt_embeds is None: uncond_tokens: List[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif type(prompt) is not type(negative_prompt): raise TypeError( f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=" f" {type(prompt)}." ) elif isinstance(negative_prompt, str): uncond_tokens = [negative_prompt] elif batch_size != len(negative_prompt): raise ValueError( f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:" f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches" " the batch size of `prompt`." ) else: uncond_tokens = negative_prompt max_length = prompt_embeds.shape[1] uncond_input = self.tokenizer( uncond_tokens, padding="max_length", max_length=max_length, truncation=True, return_tensors="pt", ) if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask: attention_mask = uncond_input.attention_mask.to(device) else: attention_mask = None negative_prompt_embeds = self.text_encoder( uncond_input.input_ids.to(device), attention_mask=attention_mask, ) negative_prompt_embeds = negative_prompt_embeds[0] if do_classifier_free_guidance: # duplicate unconditional embeddings for each generation per prompt, using mps friendly method seq_len = negative_prompt_embeds.shape[1] negative_prompt_embeds = negative_prompt_embeds.to(dtype=self.text_encoder.dtype, device=device) negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_videos_per_prompt, 1) negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1) # For classifier free guidance, we need to do two forward passes. # Here we concatenate the unconditional and text embeddings into a single batch # to avoid doing two forward passes prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) return prompt_embeds def decode_latents(self, latents): video_length = latents.shape[2] latents = 1 / 0.18215 * latents latents = rearrange(latents, "b c f h w -> (b f) c h w") video = self.vae.decode(latents).sample video = rearrange(video, "(b f) c h w -> b c f h w", f=video_length) video = (video / 2 + 0.5).clamp(0, 1) # we always cast to float32 as this does not cause significant overhead and is compatible with bfloa16 video = video.cpu().float().numpy() return video def prepare_extra_step_kwargs(self, generator, eta): # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502 # and should be between [0, 1] accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) extra_step_kwargs = {} if accepts_eta: extra_step_kwargs["eta"] = eta # check if the scheduler accepts generator accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) if accepts_generator: extra_step_kwargs["generator"] = generator return extra_step_kwargs def check_inputs(self, prompt, height, width, callback_steps): if not isinstance(prompt, str) and prompt is not None: raise ValueError(f"`prompt` has to be of type `str` but is {type(prompt)}") if height % 8 != 0 or width % 8 != 0: raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.") if (callback_steps is None) or ( callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0) ): raise ValueError( f"`callback_steps` has to be a positive integer but is {callback_steps} of type" f" {type(callback_steps)}." ) def prepare_image(self, image, width, height, batch_size, num_images_per_prompt, device, dtype, do_classifier_free_guidance): if not isinstance(image, torch.Tensor): if isinstance(image, PIL.Image.Image): image = [image] if isinstance(image[0], PIL.Image.Image): image = [ np.array(i.resize((width, height), resample=PIL_INTERPOLATION["lanczos"]))[None, :] for i in image ] image = np.concatenate(image, axis=0) image = np.array(image).astype(np.float32) / 255.0 image = image.transpose(0, 3, 1, 2) image = torch.from_numpy(image) elif isinstance(image[0], torch.Tensor): image = torch.cat(image, dim=0) image_batch_size = image.shape[0] if image_batch_size == 1: repeat_by = batch_size else: # image batch size is the same as prompt batch size repeat_by = num_images_per_prompt image = image.repeat_interleave(repeat_by, dim=0) image = image.to(device=device, dtype=dtype) if do_classifier_free_guidance: image = torch.cat([image] * 2) return image def prepare_latents(self, ddim_prompt, scheduler_path, video_input_dataloader, latent_timestep, use_vid2vid, use_inv_latent, num_inv_steps, batch_size, num_channels_latents, video_length, height, width, dtype, device, generator, latents=None): shape = (batch_size, num_channels_latents, video_length, height // self.vae_scale_factor, width // self.vae_scale_factor) if isinstance(generator, list) and len(generator) != batch_size: raise ValueError( f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" f" size of {batch_size}. Make sure the batch size matches the length of the generators." ) if latents is None: rand_device = "cpu" if device.type == "mps" else device if use_vid2vid or use_inv_latent: for _, batch_all in enumerate(video_input_dataloader): batch = batch_all pixel_values = batch["pixel_values"].to(self.vae.dtype) video_length_in = pixel_values.shape[1] pixel_values = rearrange(pixel_values, "b f c h w -> (b f) c h w") init_latents = self.vae.encode(pixel_values.to(device)).latent_dist.sample() init_latents = rearrange(init_latents, "(b f) c h w -> b c f h w", f=video_length_in) init_latents = init_latents * 0.18215 shape = init_latents.shape if use_inv_latent: ddim_inv_scheduler = DDIMScheduler.from_pretrained(scheduler_path, subfolder='scheduler') ddim_inv_scheduler.set_timesteps(num_inv_steps) else: init_latents = None if isinstance(generator, list): shape = (1,) + shape[1:] latents = [] if use_vid2vid and use_inv_latent: for i in range(batch_size): noise = ddim_inversion(self, ddim_inv_scheduler, init_latents, num_inv_steps=num_inv_steps, prompt=ddim_prompt)[-1].to(self.vae.dtype) latents.append(self.scheduler.add_noise(init_latents.to(rand_device), noise, latent_timestep)) elif use_vid2vid: for i in range(batch_size): noise = torch.randn(shape, generator=generator[i], device=rand_device, dtype=dtype) latents.append(self.scheduler.add_noise(init_latents.to(rand_device), noise, latent_timestep)) elif use_inv_latent: for i in range(batch_size): latents.append(ddim_inversion(self, ddim_inv_scheduler, init_latents, num_inv_steps=num_inv_steps, prompt=ddim_prompt)[-1].to(self.vae.dtype)) else: for i in range(batch_size): latents.append(torch.randn(shape, generator=generator[i], device=rand_device, dtype=dtype)) latents = torch.cat(latents, dim=0).to(device) else: if use_vid2vid and use_inv_latent: noise = ddim_inversion(self, ddim_inv_scheduler, init_latents, num_inv_steps=num_inv_steps, prompt=ddim_prompt)[-1].to(self.vae.dtype) latents = self.scheduler.add_noise(init_latents.to(rand_device), noise, latent_timestep).to(device) elif use_vid2vid: noise = torch.randn(shape, generator=generator, device=rand_device, dtype=dtype) latents = self.scheduler.add_noise(init_latents.to(rand_device), noise, latent_timestep).to(device) elif use_inv_latent: latents = ddim_inversion(self, ddim_inv_scheduler, init_latents, num_inv_steps=num_inv_steps, prompt=ddim_prompt)[-1].to(self.vae.dtype).to(device) else: latents = torch.randn(shape, generator=generator, device=rand_device, dtype=dtype).to(device) else: if latents.shape != shape: raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {shape}") latents = latents.to(device) # scale the initial noise by the standard deviation required by the scheduler latents = latents * self.scheduler.init_noise_sigma return latents def get_timesteps(self, num_inference_steps, strength): # get the original timestep using init_timestep init_timestep = min(int(num_inference_steps * strength), num_inference_steps) t_start = max(num_inference_steps - init_timestep, 0) timesteps = self.scheduler.timesteps[t_start:] return timesteps, num_inference_steps - t_start def PIL_load_video(self, videoFileName, sample_frame_rate, sample_start_idx, video_length): file_list = sorted(os.listdir(videoFileName), key=lambda s: sum(((s, int(n)) for s, n in re.findall(r'(\D+)(\d+)', 'a%s0' % s)), ())) sample_frames = [] sample_index = list(range(sample_start_idx, len(file_list), sample_frame_rate))[:video_length] for i in sample_index: if i >= len(file_list): raise ValueError(f"Unexpected sample index, got {i}, expected less than {len(file_list)}") img = cv2.imread(os.path.join(videoFileName, file_list[i])) img_array = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) sample_frames.append(PIL.Image.fromarray(img_array)) return sample_frames @torch.no_grad() def __call__( self, prompt: Optional[str] = None, video_length: Optional[int] = None, height: Optional[int] = None, width: Optional[int] = None, num_inference_steps: int = 50, guidance_scale: float = 7.5, negative_prompt: Optional[str] = None, prompt_embeds: Optional[torch.FloatTensor] = None, negative_prompt_embeds: Optional[torch.FloatTensor] = None, num_videos_per_prompt: Optional[int] = 1, eta: float = 0.0, generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, latents: Optional[torch.FloatTensor] = None, output_type: Optional[str] = "tensor", return_dict: bool = True, callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None, callback_steps: Optional[int] = 1, use_vid2vid: bool = False, video_input_path: str = None, video_strength: float = 0.0, sample_frame_rate: int = 0, sample_start_idx: int = 0, use_inv_latent: bool = False, num_inv_steps: int = 0, scheduler_path: str = None, ddim_prompt: str = "", use_controlnet: bool = False, video_prepare_type_list: List[str] = None, controlnet_conditioning_scale: List[float] = None, controlnet_video_path: List[str] = None, **kwargs, ): # Default height and width to unet height = height or self.unet.config.sample_size * self.vae_scale_factor width = width or self.unet.config.sample_size * self.vae_scale_factor # Check inputs. Raise error if not correct self.check_inputs(prompt, height, width, callback_steps) # Define call parameters if prompt is not None and isinstance(prompt, str): batch_size = 1 else: batch_size = prompt_embeds.shape[0] device = self._execution_device # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2) # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1` # corresponds to doing no classifier free guidance. do_classifier_free_guidance = guidance_scale > 1.0 if isinstance(self.controlnet, MultiControlNetModel) and isinstance(controlnet_conditioning_scale, float): controlnet_conditioning_scale = [controlnet_conditioning_scale] * len(self.controlnet.nets) # Encode input prompt text_embeddings = self._encode_prompt( prompt, device, num_videos_per_prompt, do_classifier_free_guidance, negative_prompt, prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_prompt_embeds ) # Load controlnet if use_controlnet: prepare_images = [] images_raw = self.PIL_load_video(video_input_path, sample_frame_rate, sample_start_idx, video_length) for controlnet_num in range(len(video_prepare_type_list)): prepare_images_ = [] if controlnet_video_path[controlnet_num] is False: images = controlnet_image_preprocessing(images_raw, video_prepare_type_list[controlnet_num]) else: images = self.PIL_load_video(controlnet_video_path[controlnet_num], sample_frame_rate, sample_start_idx, video_length) images = controlnet_image_preprocessing(images, video_prepare_type_list[controlnet_num] + "_input") for image in images: # Prepare image image = self.prepare_image(image, width, height, batch_size * num_videos_per_prompt, num_videos_per_prompt, device, self.controlnet.dtype, do_classifier_free_guidance) prepare_images_.append(image) prepare_images.append(prepare_images_) # Prepare timesteps self.scheduler.set_timesteps(num_inference_steps, device=device) # Prepare timesteps latent_timestep = None video_input_dataloader = None timesteps = self.scheduler.timesteps if use_vid2vid or use_inv_latent: if use_vid2vid: # Prepare timesteps timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, video_strength) latent_timestep = timesteps[:1].repeat(batch_size * num_videos_per_prompt) # Prepare video input video_input_dataset = TuneAVideoDataset(prompt="", video_path=video_input_path, n_sample_frames=video_length, width=width, height=height, sample_frame_rate=sample_frame_rate, sample_start_idx=sample_start_idx) video_input_dataset.prompt_ids = self.tokenizer(video_input_dataset.prompt, max_length=self.tokenizer.model_max_length, padding="max_length", truncation=True, return_tensors="pt").input_ids[0] # Preprocessing the dataset video_input_dataloader = torch.utils.data.DataLoader(video_input_dataset, batch_size=1) # Prepare latent variables num_channels_latents = self.unet.in_channels latents = self.prepare_latents( ddim_prompt, scheduler_path, video_input_dataloader, latent_timestep, use_vid2vid, use_inv_latent, num_inv_steps, batch_size * num_videos_per_prompt, num_channels_latents, video_length, height, width, text_embeddings.dtype, device, generator, latents, ) latents_dtype = latents.dtype # Prepare extra step kwargs. extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) # Denoising loop num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order with self.progress_bar(total=num_inference_steps) as progress_bar: for i, t in enumerate(timesteps): # expand the latents if we are doing classifier free guidance latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) # use Controlnet if use_controlnet: latent_model_input_list = torch.unbind(latent_model_input, dim=2) down_block_res_samples_list = [] mid_block_res_sample_list = [] for latent_model_input_num in range(len(latent_model_input_list)): down_block_res_samples_temp, mid_block_res_sample = self.controlnet( latent_model_input_list[latent_model_input_num], t, encoder_hidden_states=text_embeddings, controlnet_cond=[prepare_image[latent_model_input_num] for prepare_image in prepare_images], conditioning_scale=controlnet_conditioning_scale, return_dict=False, ) down_block_res_samples_list.append([ torch.unsqueeze(down_block_res_sample, dim=2) for down_block_res_sample in down_block_res_samples_temp ]) mid_block_res_sample_list.append(torch.unsqueeze(mid_block_res_sample, dim=2)) down_block_res_samples = [ torch.cat(tuple(down_block_res_sample_list), dim=2) for down_block_res_sample_list in [list(items) for items in zip(*down_block_res_samples_list)]] mid_block_res_sample = torch.cat(tuple(mid_block_res_sample_list), dim=2) # predict the noise residual noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=text_embeddings, down_block_additional_residuals=down_block_res_samples, mid_block_additional_residual=mid_block_res_sample).sample.to(dtype=latents_dtype) else: # predict the noise residual noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample.to(dtype=latents_dtype) # perform guidance if do_classifier_free_guidance: noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) # compute the previous noisy sample x_t -> x_t-1 latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample # call the callback, if provided if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): progress_bar.update() if callback is not None and i % callback_steps == 0: callback(i, t, latents) # Post-processing video = self.decode_latents(latents) # Convert to tensor if output_type == "tensor": video = torch.from_numpy(video) if not return_dict: return video return TuneAVideoPipelineOutput(videos=video)