from __future__ import annotations import inspect import math import time import warnings from typing import Any, Callable, Dict, List, Optional, Tuple, Union from dataclasses import dataclass from einops import rearrange, repeat import PIL.Image import numpy as np import torch from torch import nn import torch.nn.functional as F from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer from diffusers.pipelines.controlnet.pipeline_controlnet import ( StableDiffusionSafetyChecker, EXAMPLE_DOC_STRING, ) from diffusers.pipelines.controlnet.pipeline_controlnet_img2img import ( StableDiffusionControlNetImg2ImgPipeline as DiffusersStableDiffusionControlNetImg2ImgPipeline, ) from diffusers.configuration_utils import FrozenDict from diffusers.models import AutoencoderKL, ControlNetModel from diffusers.pipelines.controlnet.multicontrolnet import MultiControlNetModel from diffusers.pipelines.stable_diffusion.safety_checker import ( StableDiffusionSafetyChecker, ) from diffusers.schedulers import KarrasDiffusionSchedulers from diffusers.utils import ( deprecate, logging, BaseOutput, replace_example_docstring, ) from diffusers.utils.torch_utils import is_compiled_module from diffusers.loaders import TextualInversionLoaderMixin from diffusers.models.attention import ( BasicTransformerBlock as DiffusersBasicTransformerBlock, ) from mmcm.vision.process.correct_color import ( hist_match_color_video_batch, hist_match_video_bcthw, ) from ..models.attention import BasicTransformerBlock from ..models.unet_3d_condition import UNet3DConditionModel from ..utils.noise_util import random_noise, video_fusion_noise from ..data.data_util import ( adaptive_instance_normalization, align_repeat_tensor_single_dim, batch_adain_conditioned_tensor, batch_concat_two_tensor_with_index, batch_index_select, fuse_part_tensor, ) from ..utils.text_emb_util import encode_weighted_prompt from ..utils.tensor_util import his_match from ..utils.timesteps_util import generate_parameters_with_timesteps from .context import get_context_scheduler, prepare_global_context logger = logging.get_logger(__name__) # pylint: disable=invalid-name @dataclass class VideoPipelineOutput(BaseOutput): videos: Union[torch.Tensor, np.ndarray] latents: Union[torch.Tensor, np.ndarray] videos_mid: Union[torch.Tensor, np.ndarray] down_block_res_samples: Tuple[torch.FloatTensor] = None mid_block_res_samples: torch.FloatTensor = None up_block_res_samples: torch.FloatTensor = None mid_video_latents: List[torch.FloatTensor] = None mid_video_noises: List[torch.FloatTensor] = None def torch_dfs(model: torch.nn.Module): result = [model] for child in model.children(): result += torch_dfs(child) return result def prepare_image( image, # b c t h w batch_size, device, dtype, image_processor: Callable, num_images_per_prompt: int = 1, width=None, height=None, ): if isinstance(image, List) and isinstance(image[0], str): raise NotImplementedError if isinstance(image, List) and isinstance(image[0], np.ndarray): image = np.concatenate(image, axis=0) if isinstance(image, np.ndarray): image = torch.from_numpy(image) if image.ndim == 5: image = rearrange(image, "b c t h w-> (b t) c h w") if height is None: height = image.shape[-2] if width is None: width = image.shape[-1] width, height = (x - x % image_processor.vae_scale_factor for x in (width, height)) if height != image.shape[-2] or width != image.shape[-1]: image = torch.nn.functional.interpolate( image, size=(height, width), mode="bilinear" ) image = image.to(dtype=torch.float32) / 255.0 do_normalize = image_processor.config.do_normalize if image.min() < 0: warnings.warn( "Passing `image` as torch tensor with value range in [-1,1] is deprecated. The expected value range for image tensor is [0,1] " f"when passing as pytorch tensor or numpy Array. You passed `image` with value range [{image.min()},{image.max()}]", FutureWarning, ) do_normalize = False if do_normalize: image = image_processor.normalize(image) 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) return image class MusevControlNetPipeline( DiffusersStableDiffusionControlNetImg2ImgPipeline, TextualInversionLoaderMixin ): """ a union diffusers pipeline, support 1. text2image model only, or text2video model, by setting skip_temporal_layer 2. text2video, image2video, video2video; 3. multi controlnet 4. IPAdapter 5. referencenet 6. IPAdapterFaceID """ _optional_components = [ "safety_checker", "feature_extractor", ] print_idx = 0 def __init__( self, vae: AutoencoderKL, unet: UNet3DConditionModel, scheduler: KarrasDiffusionSchedulers, controlnet: ControlNetModel | List[ControlNetModel] | Tuple[ControlNetModel] | MultiControlNetModel, text_encoder: CLIPTextModel, tokenizer: CLIPTokenizer, safety_checker: StableDiffusionSafetyChecker, feature_extractor: CLIPImageProcessor, # | MultiControlNetModel = None, # text_encoder: CLIPTextModel = None, # tokenizer: CLIPTokenizer = None, # safety_checker: StableDiffusionSafetyChecker = None, # feature_extractor: CLIPImageProcessor = None, requires_safety_checker: bool = False, referencenet: nn.Module = None, vision_clip_extractor: nn.Module = None, ip_adapter_image_proj: nn.Module = None, face_emb_extractor: nn.Module = None, facein_image_proj: nn.Module = None, ip_adapter_face_emb_extractor: nn.Module = None, ip_adapter_face_image_proj: nn.Module = None, pose_guider: nn.Module = None, ): super().__init__( vae, text_encoder, tokenizer, unet, controlnet, scheduler, safety_checker, feature_extractor, requires_safety_checker, ) self.referencenet = referencenet # ip_adapter if isinstance(vision_clip_extractor, nn.Module): vision_clip_extractor.to(dtype=self.unet.dtype, device=self.unet.device) self.vision_clip_extractor = vision_clip_extractor if isinstance(ip_adapter_image_proj, nn.Module): ip_adapter_image_proj.to(dtype=self.unet.dtype, device=self.unet.device) self.ip_adapter_image_proj = ip_adapter_image_proj # facein if isinstance(face_emb_extractor, nn.Module): face_emb_extractor.to(dtype=self.unet.dtype, device=self.unet.device) self.face_emb_extractor = face_emb_extractor if isinstance(facein_image_proj, nn.Module): facein_image_proj.to(dtype=self.unet.dtype, device=self.unet.device) self.facein_image_proj = facein_image_proj # ip_adapter_face if isinstance(ip_adapter_face_emb_extractor, nn.Module): ip_adapter_face_emb_extractor.to( dtype=self.unet.dtype, device=self.unet.device ) self.ip_adapter_face_emb_extractor = ip_adapter_face_emb_extractor if isinstance(ip_adapter_face_image_proj, nn.Module): ip_adapter_face_image_proj.to( dtype=self.unet.dtype, device=self.unet.device ) self.ip_adapter_face_image_proj = ip_adapter_face_image_proj if isinstance(pose_guider, nn.Module): pose_guider.to(dtype=self.unet.dtype, device=self.unet.device) self.pose_guider = pose_guider def decode_latents(self, latents): batch_size = latents.shape[0] latents = rearrange(latents, "b c f h w -> (b f) c h w") video = super().decode_latents(latents=latents) video = rearrange(video, "(b f) h w c -> b c f h w", b=batch_size) return video def prepare_latents( self, batch_size: int, num_channels_latents: int, video_length: int, height: int, width: int, dtype: torch.dtype, device: torch.device, generator: torch.Generator, latents: torch.Tensor = None, w_ind_noise: float = 0.5, image: torch.Tensor = None, timestep: int = None, initial_common_latent: torch.Tensor = None, noise_type: str = "random", add_latents_noise: bool = False, need_img_based_video_noise: bool = False, condition_latents: torch.Tensor = None, img_weight=1e-3, ) -> torch.Tensor: """ 支持多种情况下的latens： img_based_latents: 当Image t=1，latents=None时，使用image赋值到shape，然后加噪；适用于text2video、middle2video。 video_based_latents：image =shape或Latents!=None时，加噪，适用于video2video； noise_latents：当image 和latents都为None时，生成随机噪声，适用于text2video support multi latents condition: img_based_latents: when Image t=1, latents=None, use image to assign to shape, then add noise; suitable for text2video, middle2video. video_based_latents: image =shape or Latents!=None, add noise, suitable for video2video; noise_laten: when image and latents are both None, generate random noise, suitable for text2video Args: batch_size (int): _description_ num_channels_latents (int): _description_ video_length (int): _description_ height (int): _description_ width (int): _description_ dtype (torch.dtype): _description_ device (torch.device): _description_ generator (torch.Generator): _description_ latents (torch.Tensor, optional): _description_. Defaults to None. w_ind_noise (float, optional): _description_. Defaults to 0.5. image (torch.Tensor, optional): _description_. Defaults to None. timestep (int, optional): _description_. Defaults to None. initial_common_latent (torch.Tensor, optional): _description_. Defaults to None. noise_type (str, optional): _description_. Defaults to "random". add_latents_noise (bool, optional): _description_. Defaults to False. need_img_based_video_noise (bool, optional): _description_. Defaults to False. condition_latents (torch.Tensor, optional): _description_. Defaults to None. img_weight (_type_, optional): _description_. Defaults to 1e-3. Raises: ValueError: _description_ ValueError: _description_ ValueError: _description_ Returns: torch.Tensor: latents """ # ref https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/controlnet/pipeline_controlnet_img2img.py#L691 # ref https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/controlnet/pipeline_controlnet.py#L659 shape = ( batch_size, num_channels_latents, video_length, height // self.vae_scale_factor, width // self.vae_scale_factor, ) if latents is None or (latents is not None and add_latents_noise): if noise_type == "random": noise = random_noise( shape=shape, dtype=dtype, device=device, generator=generator ) elif noise_type == "video_fusion": noise = video_fusion_noise( shape=shape, dtype=dtype, device=device, generator=generator, w_ind_noise=w_ind_noise, initial_common_noise=initial_common_latent, ) if ( need_img_based_video_noise and condition_latents is not None and image is None and latents is None ): if self.print_idx == 0: logger.debug( ( f"need_img_based_video_noise, condition_latents={condition_latents.shape}," f"batch_size={batch_size}, noise={noise.shape}, video_length={video_length}" ) ) condition_latents = condition_latents.mean(dim=2, keepdim=True) condition_latents = repeat( condition_latents, "b c t h w->b c (t x) h w", x=video_length ) noise = ( img_weight**0.5 * condition_latents + (1 - img_weight) ** 0.5 * noise ) if self.print_idx == 0: logger.debug(f"noise={noise.shape}") if image is not None: if image.ndim == 5: image = rearrange(image, "b c t h w->(b t) c h w") image = image.to(device=device, dtype=dtype) 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 isinstance(generator, list): init_latents = [ # self.vae.encode(image[i : i + 1]).latent_dist.sample(generator[i]) self.vae.encode(image[i : i + 1]).latent_dist.mean for i in range(batch_size) ] init_latents = torch.cat(init_latents, dim=0) else: # init_latents = self.vae.encode(image).latent_dist.sample(generator) init_latents = self.vae.encode(image).latent_dist.mean init_latents = self.vae.config.scaling_factor * init_latents # scale the initial noise by the standard deviation required by the scheduler if ( batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] == 0 ): # expand init_latents for batch_size deprecation_message = ( f"You have passed {batch_size} text prompts (`prompt`), but only {init_latents.shape[0]} initial" " images (`image`). Initial images are now duplicating to match the number of text prompts. Note" " that this behavior is deprecated and will be removed in a version 1.0.0. Please make sure to update" " your script to pass as many initial images as text prompts to suppress this warning." ) deprecate( "len(prompt) != len(image)", "1.0.0", deprecation_message, standard_warn=False, ) additional_image_per_prompt = batch_size // init_latents.shape[0] init_latents = torch.cat( [init_latents] * additional_image_per_prompt, dim=0 ) elif ( batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] != 0 ): raise ValueError( f"Cannot duplicate `image` of batch size {init_latents.shape[0]} to {batch_size} text prompts." ) else: init_latents = torch.cat([init_latents], dim=0) if init_latents.shape[2] != shape[3] and init_latents.shape[3] != shape[4]: init_latents = torch.nn.functional.interpolate( init_latents, size=(shape[3], shape[4]), mode="bilinear", ) init_latents = rearrange( init_latents, "(b t) c h w-> b c t h w", t=video_length ) if self.print_idx == 0: logger.debug(f"init_latensts={init_latents.shape}") if latents is None: if image is None: latents = noise * self.scheduler.init_noise_sigma else: if self.print_idx == 0: logger.debug(f"prepare latents, image is not None") latents = self.scheduler.add_noise(init_latents, noise, timestep) else: if isinstance(latents, np.ndarray): latents = torch.from_numpy(latents) latents = latents.to(device=device, dtype=dtype) if add_latents_noise: latents = self.scheduler.add_noise(latents, noise, timestep) else: latents = latents * self.scheduler.init_noise_sigma if latents.shape != shape: raise ValueError( f"Unexpected latents shape, got {latents.shape}, expected {shape}" ) latents = latents.to(device, dtype=dtype) return latents def prepare_image( self, image, # b c t h w batch_size, num_images_per_prompt, device, dtype, width=None, height=None, ): return prepare_image( image=image, batch_size=batch_size, num_images_per_prompt=num_images_per_prompt, device=device, dtype=dtype, width=width, height=height, image_processor=self.image_processor, ) def prepare_control_image( self, image, # b c t h w width, height, batch_size, num_images_per_prompt, device, dtype, do_classifier_free_guidance=False, guess_mode=False, ): image = prepare_image( image=image, batch_size=batch_size, num_images_per_prompt=num_images_per_prompt, device=device, dtype=dtype, width=width, height=height, image_processor=self.control_image_processor, ) if do_classifier_free_guidance and not guess_mode: image = torch.cat([image] * 2) return image def check_inputs( self, prompt, image, callback_steps, negative_prompt=None, prompt_embeds=None, negative_prompt_embeds=None, controlnet_conditioning_scale=1, control_guidance_start=0, control_guidance_end=1, ): # TODO: to implement if image is not None: return super().check_inputs( prompt, image, callback_steps, negative_prompt, prompt_embeds, negative_prompt_embeds, controlnet_conditioning_scale, control_guidance_start, control_guidance_end, ) def hist_match_with_vis_cond( self, video: np.ndarray, target: np.ndarray ) -> np.ndarray: """ video: b c t1 h w target: b c t2(=1) h w """ video = hist_match_video_bcthw(video, target, value=255.0) return video def get_facein_image_emb( self, refer_face_image, device, dtype, batch_size, do_classifier_free_guidance ): # refer_face_image and its face_emb if self.print_idx == 0: logger.debug( f"face_emb_extractor={type(self.face_emb_extractor)}, facein_image_proj={type(self.facein_image_proj)}, refer_face_image={type(refer_face_image)}, " ) if ( self.face_emb_extractor is not None and self.facein_image_proj is not None and refer_face_image is not None ): if self.print_idx == 0: logger.debug(f"refer_face_image={refer_face_image.shape}") if isinstance(refer_face_image, np.ndarray): refer_face_image = torch.from_numpy(refer_face_image) refer_face_image_facein = refer_face_image n_refer_face_image = refer_face_image_facein.shape[2] refer_face_image_facein = rearrange( refer_face_image, "b c t h w-> (b t) h w c" ) # refer_face_image_emb： bt d或者 bt h w d ( refer_face_image_emb, refer_align_face_image, ) = self.face_emb_extractor.extract_images( refer_face_image_facein, return_type="torch" ) refer_face_image_emb = refer_face_image_emb.to(device=device, dtype=dtype) if self.print_idx == 0: logger.debug(f"refer_face_image_emb={refer_face_image_emb.shape}") if refer_face_image_emb.shape == 2: refer_face_image_emb = rearrange(refer_face_image_emb, "bt d-> bt 1 d") elif refer_face_image_emb.shape == 4: refer_face_image_emb = rearrange( refer_face_image_emb, "bt h w d-> bt (h w) d" ) refer_face_image_emb_bk = refer_face_image_emb refer_face_image_emb = self.facein_image_proj(refer_face_image_emb) # Todo:当前不支持 IPAdapterPlus的vision_clip的输出 refer_face_image_emb = rearrange( refer_face_image_emb, "(b t) n q-> b (t n) q", t=n_refer_face_image, ) refer_face_image_emb = align_repeat_tensor_single_dim( refer_face_image_emb, target_length=batch_size, dim=0 ) if do_classifier_free_guidance: # TODO：固定特征，有优化空间 # TODO: fix the feature, there is optimization space uncond_refer_face_image_emb = self.facein_image_proj( torch.zeros_like(refer_face_image_emb_bk).to( device=device, dtype=dtype ) ) # Todo:当前可能不支持 IPAdapterPlus的vision_clip的输出 # TODO: do not support IPAdapterPlus's vision_clip's output uncond_refer_face_image_emb = rearrange( uncond_refer_face_image_emb, "(b t) n q-> b (t n) q", t=n_refer_face_image, ) uncond_refer_face_image_emb = align_repeat_tensor_single_dim( uncond_refer_face_image_emb, target_length=batch_size, dim=0 ) if self.print_idx == 0: logger.debug( f"uncond_refer_face_image_emb, {uncond_refer_face_image_emb.shape}" ) logger.debug(f"refer_face_image_emb, {refer_face_image_emb.shape}") refer_face_image_emb = torch.concat( [ uncond_refer_face_image_emb, refer_face_image_emb, ], ) else: refer_face_image_emb = None if self.print_idx == 0: logger.debug(f"refer_face_image_emb={type(refer_face_image_emb)}") return refer_face_image_emb def get_ip_adapter_face_emb( self, refer_face_image, device, dtype, batch_size, do_classifier_free_guidance ): # refer_face_image and its ip_adapter_face_emb if self.print_idx == 0: logger.debug( f"face_emb_extractor={type(self.face_emb_extractor)}, ip_adapter__image_proj={type(self.facein_image_proj)}, refer_face_image={type(refer_face_image)}, " ) if ( self.ip_adapter_face_emb_extractor is not None and self.ip_adapter_face_image_proj is not None and refer_face_image is not None ): if self.print_idx == 0: logger.debug(f"refer_face_image={refer_face_image.shape}") if isinstance(refer_face_image, np.ndarray): refer_face_image = torch.from_numpy(refer_face_image) refer_ip_adapter_face_image = refer_face_image n_refer_face_image = refer_ip_adapter_face_image.shape[2] refer_ip_adapter_face_image = rearrange( refer_ip_adapter_face_image, "b c t h w-> (b t) h w c" ) # refer_face_image_emb： bt d or bt h w d ( refer_face_image_emb, refer_align_face_image, ) = self.ip_adapter_face_emb_extractor.extract_images( refer_ip_adapter_face_image, return_type="torch" ) refer_face_image_emb = refer_face_image_emb.to(device=device, dtype=dtype) if self.print_idx == 0: logger.debug(f"refer_face_image_emb={refer_face_image_emb.shape}") if refer_face_image_emb.shape == 2: refer_face_image_emb = rearrange(refer_face_image_emb, "bt d-> bt 1 d") elif refer_face_image_emb.shape == 4: refer_face_image_emb = rearrange( refer_face_image_emb, "bt h w d-> bt (h w) d" ) refer_face_image_emb_bk = refer_face_image_emb refer_face_image_emb = self.ip_adapter_face_image_proj(refer_face_image_emb) refer_face_image_emb = rearrange( refer_face_image_emb, "(b t) n q-> b (t n) q", t=n_refer_face_image, ) refer_face_image_emb = align_repeat_tensor_single_dim( refer_face_image_emb, target_length=batch_size, dim=0 ) if do_classifier_free_guidance: # TODO：固定特征，有优化空间 # TODO: fix the feature, there is optimization space uncond_refer_face_image_emb = self.ip_adapter_face_image_proj( torch.zeros_like(refer_face_image_emb_bk).to( device=device, dtype=dtype ) ) # TODO: 当前可能不支持 IPAdapterPlus的vision_clip的输出 # TODO: do not support IPAdapterPlus's vision_clip's output uncond_refer_face_image_emb = rearrange( uncond_refer_face_image_emb, "(b t) n q-> b (t n) q", t=n_refer_face_image, ) uncond_refer_face_image_emb = align_repeat_tensor_single_dim( uncond_refer_face_image_emb, target_length=batch_size, dim=0 ) if self.print_idx == 0: logger.debug( f"uncond_refer_face_image_emb, {uncond_refer_face_image_emb.shape}" ) logger.debug(f"refer_face_image_emb, {refer_face_image_emb.shape}") refer_face_image_emb = torch.concat( [ uncond_refer_face_image_emb, refer_face_image_emb, ], ) else: refer_face_image_emb = None if self.print_idx == 0: logger.debug(f"ip_adapter_face_emb={type(refer_face_image_emb)}") return refer_face_image_emb def get_ip_adapter_image_emb( self, ip_adapter_image, device, dtype, batch_size, do_classifier_free_guidance, height, width, ): # refer_image vision_clip and its ipadapter_emb if self.print_idx == 0: logger.debug( f"vision_clip_extractor={type(self.vision_clip_extractor)}," f"ip_adapter_image_proj={type(self.ip_adapter_image_proj)}," f"ip_adapter_image={type(ip_adapter_image)}," ) if self.vision_clip_extractor is not None and ip_adapter_image is not None: if self.print_idx == 0: logger.debug(f"ip_adapter_image={ip_adapter_image.shape}") if isinstance(ip_adapter_image, np.ndarray): ip_adapter_image = torch.from_numpy(ip_adapter_image) # ip_adapter_image = ip_adapter_image.to(device=device, dtype=dtype) n_ip_adapter_image = ip_adapter_image.shape[2] ip_adapter_image = rearrange(ip_adapter_image, "b c t h w-> (b t) h w c") ip_adapter_image_emb = self.vision_clip_extractor.extract_images( ip_adapter_image, target_height=height, target_width=width, return_type="torch", ) if ip_adapter_image_emb.ndim == 2: ip_adapter_image_emb = rearrange(ip_adapter_image_emb, "b q-> b 1 q") ip_adapter_image_emb_bk = ip_adapter_image_emb # 存在只需要image_prompt、但不需要 proj的场景，如使用image_prompt替代text_prompt # There are scenarios where only image_prompt is needed, but proj is not needed, such as using image_prompt instead of text_prompt if self.ip_adapter_image_proj is not None: logger.debug(f"ip_adapter_image_proj is None, ") ip_adapter_image_emb = self.ip_adapter_image_proj(ip_adapter_image_emb) # TODO: 当前不支持 IPAdapterPlus的vision_clip的输出 # TODO: do not support IPAdapterPlus's vision_clip's output ip_adapter_image_emb = rearrange( ip_adapter_image_emb, "(b t) n q-> b (t n) q", t=n_ip_adapter_image, ) ip_adapter_image_emb = align_repeat_tensor_single_dim( ip_adapter_image_emb, target_length=batch_size, dim=0 ) if do_classifier_free_guidance: # TODO：固定特征，有优化空间 # TODO: fix the feature, there is optimization space if self.ip_adapter_image_proj is not None: uncond_ip_adapter_image_emb = self.ip_adapter_image_proj( torch.zeros_like(ip_adapter_image_emb_bk).to( device=device, dtype=dtype ) ) if self.print_idx == 0: logger.debug( f"uncond_ip_adapter_image_emb use ip_adapter_image_proj(zero_like)" ) else: uncond_ip_adapter_image_emb = torch.zeros_like(ip_adapter_image_emb) if self.print_idx == 0: logger.debug(f"uncond_ip_adapter_image_emb use zero_like") # TODO:当前可能不支持 IPAdapterPlus的vision_clip的输出 # TODO: do not support IPAdapterPlus's vision_clip's output uncond_ip_adapter_image_emb = rearrange( uncond_ip_adapter_image_emb, "(b t) n q-> b (t n) q", t=n_ip_adapter_image, ) uncond_ip_adapter_image_emb = align_repeat_tensor_single_dim( uncond_ip_adapter_image_emb, target_length=batch_size, dim=0 ) if self.print_idx == 0: logger.debug( f"uncond_ip_adapter_image_emb, {uncond_ip_adapter_image_emb.shape}" ) logger.debug(f"ip_adapter_image_emb, {ip_adapter_image_emb.shape}") # uncond_ip_adapter_image_emb = torch.zeros_like(ip_adapter_image_emb) ip_adapter_image_emb = torch.concat( [ uncond_ip_adapter_image_emb, ip_adapter_image_emb, ], ) else: ip_adapter_image_emb = None if self.print_idx == 0: logger.debug(f"ip_adapter_image_emb={type(ip_adapter_image_emb)}") return ip_adapter_image_emb def get_referencenet_image_vae_emb( self, refer_image, batch_size, num_videos_per_prompt, device, dtype, do_classifier_free_guidance, width: int = None, height: int = None, ): # prepare_referencenet_emb if self.print_idx == 0: logger.debug( f"referencenet={type(self.referencenet)}, refer_image={type(refer_image)}" ) if self.referencenet is not None and refer_image is not None: n_refer_image = refer_image.shape[2] refer_image_vae = self.prepare_image( refer_image, batch_size=batch_size * num_videos_per_prompt, num_images_per_prompt=num_videos_per_prompt, device=device, dtype=dtype, width=width, height=height, ) # ref_hidden_states = self.vae.encode(refer_image_vae).latent_dist.sample() refer_image_vae_emb = self.vae.encode(refer_image_vae).latent_dist.mean refer_image_vae_emb = self.vae.config.scaling_factor * refer_image_vae_emb logger.debug(f"refer_image_vae_emb={refer_image_vae_emb.shape}") if do_classifier_free_guidance: # 1. zeros_like image # uncond_refer_image_vae_emb = self.vae.encode( # torch.zeros_like(refer_image_vae) # ).latent_dist.mean # uncond_refer_image_vae_emb = ( # self.vae.config.scaling_factor * uncond_refer_image_vae_emb # ) # 2. zeros_like image vae emb # uncond_refer_image_vae_emb = torch.zeros_like(refer_image_vae_emb) # uncond_refer_image_vae_emb = rearrange( # uncond_refer_image_vae_emb, # "(b t) c h w-> b c t h w", # t=n_refer_image, # ) # refer_image_vae_emb = rearrange( # refer_image_vae_emb, "(b t) c h w-> b c t h w", t=n_refer_image # ) # refer_image_vae_emb = torch.concat( # [uncond_refer_image_vae_emb, refer_image_vae_emb], dim=0 # ) # refer_image_vae_emb = rearrange( # refer_image_vae_emb, "b c t h w-> (b t) c h w" # ) # logger.debug(f"refer_image_vae_emb={refer_image_vae_emb.shape}") # 3. uncond_refer_image_vae_emb = refer_image_vae_emb uncond_refer_image_vae_emb = refer_image_vae_emb uncond_refer_image_vae_emb = rearrange( uncond_refer_image_vae_emb, "(b t) c h w-> b c t h w", t=n_refer_image, ) refer_image_vae_emb = rearrange( refer_image_vae_emb, "(b t) c h w-> b c t h w", t=n_refer_image ) refer_image_vae_emb = torch.concat( [uncond_refer_image_vae_emb, refer_image_vae_emb], dim=0 ) refer_image_vae_emb = rearrange( refer_image_vae_emb, "b c t h w-> (b t) c h w" ) logger.debug(f"refer_image_vae_emb={refer_image_vae_emb.shape}") else: refer_image_vae_emb = None return refer_image_vae_emb def get_referencenet_emb( self, refer_image_vae_emb, refer_image, batch_size, num_videos_per_prompt, device, dtype, ip_adapter_image_emb, do_classifier_free_guidance, prompt_embeds, ref_timestep_int: int = 0, ): # prepare_referencenet_emb if self.print_idx == 0: logger.debug( f"referencenet={type(self.referencenet)}, refer_image={type(refer_image)}" ) if ( self.referencenet is not None and refer_image_vae_emb is not None and refer_image is not None ): n_refer_image = refer_image.shape[2] # ref_timestep = ( # torch.ones((refer_image_vae_emb.shape[0],), device=device) # * ref_timestep_int # ) ref_timestep = torch.zeros_like(ref_timestep_int) # referencenet 优先使用 ip_adapter 中图像提取到的 clip_vision_emb if ip_adapter_image_emb is not None: refer_prompt_embeds = ip_adapter_image_emb else: refer_prompt_embeds = prompt_embeds if self.print_idx == 0: logger.debug( f"use referencenet: n_refer_image={n_refer_image}, refer_image_vae_emb={refer_image_vae_emb.shape}, ref_timestep={ref_timestep.shape}" ) if prompt_embeds is not None: logger.debug(f"prompt_embeds={prompt_embeds.shape},") # refer_image_vae_emb = self.scheduler.scale_model_input( # refer_image_vae_emb, ref_timestep # ) # self.scheduler._step_index = None # self.scheduler.is_scale_input_called = False referencenet_params = { "sample": refer_image_vae_emb, "encoder_hidden_states": refer_prompt_embeds, "timestep": ref_timestep, "num_frames": n_refer_image, "return_ndim": 5, } ( down_block_refer_embs, mid_block_refer_emb, refer_self_attn_emb, ) = self.referencenet(**referencenet_params) # many ways to prepare negative referencenet emb # mode 1 # zero shape like ref_image # if do_classifier_free_guidance: # # mode 2: # # if down_block_refer_embs is not None: # # down_block_refer_embs = [ # # torch.cat([x] * 2) for x in down_block_refer_embs # # ] # # if mid_block_refer_emb is not None: # # mid_block_refer_emb = torch.cat([mid_block_refer_emb] * 2) # # if refer_self_attn_emb is not None: # # refer_self_attn_emb = [ # # torch.cat([x] * 2) for x in refer_self_attn_emb # # ] # # mode 3 # if down_block_refer_embs is not None: # down_block_refer_embs = [ # torch.cat([torch.zeros_like(x), x]) # for x in down_block_refer_embs # ] # if mid_block_refer_emb is not None: # mid_block_refer_emb = torch.cat( # [torch.zeros_like(mid_block_refer_emb), mid_block_refer_emb] * 2 # ) # if refer_self_attn_emb is not None: # refer_self_attn_emb = [ # torch.cat([torch.zeros_like(x), x]) for x in refer_self_attn_emb # ] else: down_block_refer_embs = None mid_block_refer_emb = None refer_self_attn_emb = None if self.print_idx == 0: logger.debug(f"down_block_refer_embs={type(down_block_refer_embs)}") logger.debug(f"mid_block_refer_emb={type(mid_block_refer_emb)}") logger.debug(f"refer_self_attn_emb={type(refer_self_attn_emb)}") return down_block_refer_embs, mid_block_refer_emb, refer_self_attn_emb def prepare_condition_latents_and_index( self, condition_images, condition_latents, video_length, batch_size, dtype, device, latent_index, vision_condition_latent_index, ): # prepare condition_latents if condition_images is not None and condition_latents is None: # condition_latents = self.vae.encode(condition_images).latent_dist.sample() condition_latents = self.vae.encode(condition_images).latent_dist.mean condition_latents = self.vae.config.scaling_factor * condition_latents condition_latents = rearrange( condition_latents, "(b t) c h w-> b c t h w", b=batch_size ) if self.print_idx == 0: logger.debug( f"condition_latents from condition_images, shape is condition_latents={condition_latents.shape}", ) if condition_latents is not None: total_frames = condition_latents.shape[2] + video_length if isinstance(condition_latents, np.ndarray): condition_latents = torch.from_numpy(condition_latents) condition_latents = condition_latents.to(dtype=dtype, device=device) # if condition is None, mean condition_latents head, generated video is tail if vision_condition_latent_index is not None: # vision_condition_latent_index should be list, whose length is condition_latents.shape[2] # -1 -> will be converted to condition_latents.shape[2]+video_length vision_condition_latent_index_lst = [ i_v if i_v != -1 else total_frames - 1 for i_v in vision_condition_latent_index ] vision_condition_latent_index = torch.LongTensor( vision_condition_latent_index_lst, ).to(device=device) if self.print_idx == 0: logger.debug( f"vision_condition_latent_index {type(vision_condition_latent_index)}, {vision_condition_latent_index}" ) else: # [0, condition_latents.shape[2]] vision_condition_latent_index = torch.arange( condition_latents.shape[2], dtype=torch.long, device=device ) vision_condition_latent_index_lst = ( vision_condition_latent_index.tolist() ) if latent_index is None: # [condition_latents.shape[2], condition_latents.shape[2]+video_length] latent_index_lst = sorted( list( set(range(total_frames)) - set(vision_condition_latent_index_lst) ) ) latent_index = torch.LongTensor( latent_index_lst, ).to(device=device) if vision_condition_latent_index is not None: vision_condition_latent_index = vision_condition_latent_index.to( device=device ) if self.print_idx == 0: logger.debug( f"pipeline vision_condition_latent_index ={vision_condition_latent_index.shape}, {vision_condition_latent_index}" ) if latent_index is not None: latent_index = latent_index.to(device=device) if self.print_idx == 0: logger.debug( f"pipeline latent_index ={latent_index.shape}, {latent_index}" ) logger.debug(f"condition_latents={type(condition_latents)}") logger.debug(f"latent_index={type(latent_index)}") logger.debug( f"vision_condition_latent_index={type(vision_condition_latent_index)}" ) return condition_latents, latent_index, vision_condition_latent_index def prepare_controlnet_and_guidance_parameter( self, control_guidance_start, control_guidance_end ): controlnet = ( self.controlnet._orig_mod if is_compiled_module(self.controlnet) else self.controlnet ) # align format for control guidance if not isinstance(control_guidance_start, list) and isinstance( control_guidance_end, list ): control_guidance_start = len(control_guidance_end) * [ control_guidance_start ] elif not isinstance(control_guidance_end, list) and isinstance( control_guidance_start, list ): control_guidance_end = len(control_guidance_start) * [control_guidance_end] elif not isinstance(control_guidance_start, list) and not isinstance( control_guidance_end, list ): mult = ( len(controlnet.nets) if isinstance(controlnet, MultiControlNetModel) else 1 ) control_guidance_start, control_guidance_end = mult * [ control_guidance_start ], mult * [control_guidance_end] return controlnet, control_guidance_start, control_guidance_end def prepare_controlnet_guess_mode(self, controlnet, guess_mode): global_pool_conditions = ( controlnet.config.global_pool_conditions if isinstance(controlnet, ControlNetModel) else controlnet.nets[0].config.global_pool_conditions ) guess_mode = guess_mode or global_pool_conditions return guess_mode def prepare_controlnet_image_and_latents( self, controlnet, width, height, batch_size, num_videos_per_prompt, device, dtype, controlnet_latents=None, controlnet_condition_latents=None, control_image=None, controlnet_condition_images=None, guess_mode=False, do_classifier_free_guidance=False, ): if isinstance(controlnet, ControlNetModel): if controlnet_latents is not None: if isinstance(controlnet_latents, np.ndarray): controlnet_latents = torch.from_numpy(controlnet_latents) if controlnet_condition_latents is not None: if isinstance(controlnet_condition_latents, np.ndarray): controlnet_condition_latents = torch.from_numpy( controlnet_condition_latents ) # TODO：使用index进行concat controlnet_latents = torch.concat( [controlnet_condition_latents, controlnet_latents], dim=2 ) if not guess_mode and do_classifier_free_guidance: controlnet_latents = torch.concat([controlnet_latents] * 2, dim=0) controlnet_latents = rearrange( controlnet_latents, "b c t h w->(b t) c h w" ) controlnet_latents = controlnet_latents.to(device=device, dtype=dtype) if self.print_idx == 0: logger.debug( f"call, controlnet_latents.shape, f{controlnet_latents.shape}" ) else: # TODO: concat with index if isinstance(control_image, np.ndarray): control_image = torch.from_numpy(control_image) if controlnet_condition_images is not None: if isinstance(controlnet_condition_images, np.ndarray): controlnet_condition_images = torch.from_numpy( controlnet_condition_images ) control_image = torch.concatenate( [controlnet_condition_images, control_image], dim=2 ) control_image = self.prepare_control_image( image=control_image, width=width, height=height, batch_size=batch_size * num_videos_per_prompt, num_images_per_prompt=num_videos_per_prompt, device=device, dtype=controlnet.dtype, do_classifier_free_guidance=do_classifier_free_guidance, guess_mode=guess_mode, ) height, width = control_image.shape[-2:] if self.print_idx == 0: logger.debug(f"call, control_image.shape , {control_image.shape}") elif isinstance(controlnet, MultiControlNetModel): control_images = [] # TODO: directly support contronet_latent instead of frames if ( controlnet_latents is not None and controlnet_condition_latents is not None ): raise NotImplementedError for i, control_image_ in enumerate(control_image): if controlnet_condition_images is not None and isinstance( controlnet_condition_images, list ): if isinstance(controlnet_condition_images[i], np.ndarray): control_image_ = np.concatenate( [controlnet_condition_images[i], control_image_], axis=2 ) control_image_ = self.prepare_control_image( image=control_image_, width=width, height=height, batch_size=batch_size * num_videos_per_prompt, num_images_per_prompt=num_videos_per_prompt, device=device, dtype=controlnet.dtype, do_classifier_free_guidance=do_classifier_free_guidance, guess_mode=guess_mode, ) control_images.append(control_image_) control_image = control_images height, width = control_image[0].shape[-2:] else: assert False if control_image is not None: if not isinstance(control_image, list): if self.print_idx == 0: logger.debug(f"control_image shape is {control_image.shape}") else: if self.print_idx == 0: logger.debug(f"control_image shape is {control_image[0].shape}") return control_image, controlnet_latents def get_controlnet_emb( self, run_controlnet, guess_mode, do_classifier_free_guidance, latents, prompt_embeds, latent_model_input, controlnet_keep, controlnet_conditioning_scale, control_image, controlnet_latents, i, t, ): if run_controlnet and self.pose_guider is None: # controlnet(s) inference if guess_mode and do_classifier_free_guidance: # Infer ControlNet only for the conditional batch. control_model_input = latents control_model_input = self.scheduler.scale_model_input( control_model_input, t ) controlnet_prompt_embeds = prompt_embeds.chunk(2)[1] else: control_model_input = latent_model_input controlnet_prompt_embeds = prompt_embeds if isinstance(controlnet_keep[i], list): cond_scale = [ c * s for c, s in zip(controlnet_conditioning_scale, controlnet_keep[i]) ] else: cond_scale = controlnet_conditioning_scale * controlnet_keep[i] control_model_input_reshape = rearrange( control_model_input, "b c t h w -> (b t) c h w" ) logger.debug( f"control_model_input_reshape={control_model_input_reshape.shape}, controlnet_prompt_embeds={controlnet_prompt_embeds.shape}" ) encoder_hidden_states_repeat = align_repeat_tensor_single_dim( controlnet_prompt_embeds, target_length=control_model_input_reshape.shape[0], dim=0, ) if self.print_idx == 0: logger.debug( f"control_model_input_reshape={control_model_input_reshape.shape}, " f"encoder_hidden_states_repeat={encoder_hidden_states_repeat.shape}, " ) down_block_res_samples, mid_block_res_sample = self.controlnet( control_model_input_reshape, t, encoder_hidden_states_repeat, controlnet_cond=control_image, controlnet_cond_latents=controlnet_latents, conditioning_scale=cond_scale, guess_mode=guess_mode, return_dict=False, ) if self.print_idx == 0: logger.debug( f"controlnet, len(down_block_res_samples, {len(down_block_res_samples)}", ) for i_tmp, tmp in enumerate(down_block_res_samples): logger.debug( f"controlnet down_block_res_samples i={i_tmp}, down_block_res_sample={tmp.shape}" ) logger.debug( f"controlnet mid_block_res_sample, {mid_block_res_sample.shape}" ) if guess_mode and do_classifier_free_guidance: # Infered ControlNet only for the conditional batch. # To apply the output of ControlNet to both the unconditional and conditional batches, # add 0 to the unconditional batch to keep it unchanged. down_block_res_samples = [ torch.cat([torch.zeros_like(d), d]) for d in down_block_res_samples ] mid_block_res_sample = torch.cat( [ torch.zeros_like(mid_block_res_sample), mid_block_res_sample, ] ) else: down_block_res_samples = None mid_block_res_sample = None return down_block_res_samples, mid_block_res_sample @torch.no_grad() @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, video_length: Optional[int], prompt: Union[str, List[str]] = None, # b c t h w image: Union[ torch.FloatTensor, PIL.Image.Image, np.ndarray, List[torch.FloatTensor], List[PIL.Image.Image], List[np.ndarray], ] = None, control_image: Union[ torch.FloatTensor, PIL.Image.Image, np.ndarray, List[torch.FloatTensor], List[PIL.Image.Image], List[np.ndarray], ] = None, # b c t(1) ho wo condition_images: Optional[torch.FloatTensor] = None, condition_latents: Optional[torch.FloatTensor] = None, latents: Optional[torch.FloatTensor] = None, add_latents_noise: bool = False, height: Optional[int] = None, width: Optional[int] = None, strength: float = 0.8, num_inference_steps: int = 50, guidance_scale: float = 7.5, guidance_scale_end: float = None, guidance_scale_method: str = "linear", negative_prompt: Optional[Union[str, List[str]]] = None, num_videos_per_prompt: Optional[int] = 1, eta: float = 0.0, generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, # b c t(1) hi wi controlnet_condition_images: Optional[torch.FloatTensor] = None, # b c t(1) ho wo controlnet_condition_latents: Optional[torch.FloatTensor] = None, controlnet_latents: Union[torch.FloatTensor, np.ndarray] = None, prompt_embeds: Optional[torch.FloatTensor] = None, negative_prompt_embeds: Optional[torch.FloatTensor] = None, output_type: Optional[str] = "tensor", return_dict: bool = True, callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None, callback_steps: int = 1, cross_attention_kwargs: Optional[Dict[str, Any]] = None, controlnet_conditioning_scale: Union[float, List[float]] = 1.0, guess_mode: bool = False, control_guidance_start: Union[float, List[float]] = 0.0, control_guidance_end: Union[float, List[float]] = 1.0, need_middle_latents: bool = False, w_ind_noise: float = 0.5, initial_common_latent: Optional[torch.FloatTensor] = None, latent_index: torch.LongTensor = None, vision_condition_latent_index: torch.LongTensor = None, # noise parameters noise_type: str = "random", need_img_based_video_noise: bool = False, skip_temporal_layer: bool = False, img_weight: float = 1e-3, need_hist_match: bool = False, motion_speed: float = 8.0, refer_image: Optional[Tuple[torch.Tensor, np.array]] = None, ip_adapter_image: Optional[Tuple[torch.Tensor, np.array]] = None, refer_face_image: Optional[Tuple[torch.Tensor, np.array]] = None, ip_adapter_scale: float = 1.0, facein_scale: float = 1.0, ip_adapter_face_scale: float = 1.0, ip_adapter_face_image: Optional[Tuple[torch.Tensor, np.array]] = None, prompt_only_use_image_prompt: bool = False, # serial_denoise parameter start record_mid_video_noises: bool = False, last_mid_video_noises: List[torch.Tensor] = None, record_mid_video_latents: bool = False, last_mid_video_latents: List[torch.TensorType] = None, video_overlap: int = 1, # serial_denoise parameter end # parallel_denoise parameter start # refer to https://github.com/MooreThreads/Moore-AnimateAnyone/blob/master/src/pipelines/pipeline_pose2vid_long.py#L354 context_schedule="uniform", context_frames=12, context_stride=1, context_overlap=4, context_batch_size=1, interpolation_factor=1, # parallel_denoise parameter end ): r""" 旨在兼容text2video、text2image、img2img、video2video、是否有controlnet等的通用pipeline。目前仅不支持img2img、video2video。 支持多片段同时denoise，交叉部分加权平均 当 skip_temporal_layer 为 False 时, unet 起 video 生成作用；skip_temporal_layer为True时，unet起原image作用。 当controlnet的所有入参为None，等价于走的是text2video pipeline； 当 condition_latents、controlnet_condition_images、controlnet_condition_latents为None时，表示不走首帧条件生成的时序condition pipeline 现在没有考虑对 `num_videos_per_prompt` 的兼容性，不是1可能报错； if skip_temporal_layer is False, unet motion layer works, else unet only run text2image layers. if parameters about controlnet are None, means text2video pipeline; if ondition_latents、controlnet_condition_images、controlnet_condition_latents are None, means only run text2video without vision condition images. By now, code works well with `num_videos_per_prpmpt=1`, !=1 may be wrong. Args: prompt (`str` or `List[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. image (`torch.FloatTensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.FloatTensor]`, `List[PIL.Image.Image]`, `List[np.ndarray]`,: `List[List[torch.FloatTensor]]`, `List[List[np.ndarray]]` or `List[List[PIL.Image.Image]]`): The ControlNet input condition. ControlNet uses this input condition to generate guidance to Unet. If the type is specified as `Torch.FloatTensor`, it is passed to ControlNet as is. `PIL.Image.Image` can also be accepted as an image. The dimensions of the output image defaults to `image`'s dimensions. If height and/or width are passed, `image` is resized according to them. If multiple ControlNets are specified in init, images must be passed as a list such that each element of the list can be correctly batched for input to a single controlnet. condition_latents： 与latents相对应，是Latents的时序condition，一般为首帧，b c t(1) ho wo be corresponding to latents, vision condtion latents, usually first frame, should be b c t(1) ho wo. controlnet_latents: 与image二选一，image会被转化成controlnet_latents Choose either image or controlnet_latents. If image is chosen, it will be converted to controlnet_latents. controlnet_condition_images: Optional[torch.FloatTensor]# b c t(1) ho wo，与image相对应，会和image在t通道concat一起，然后转化成 controlnet_latents b c t(1) ho wo, corresponding to image, will be concatenated along the t channel with image and then converted to controlnet_latents. controlnet_condition_latents: Optional[torch.FloatTensor]:# b c t(1) ho wo，会和 controlnet_latents 在t 通道concat一起，转化成 controlnet_latents b c t(1) ho wo will be concatenated along the t channel with controlnet_latents and converted to controlnet_latents. height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): The height in pixels of the generated image. width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): The width in pixels of the generated image. num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. guidance_scale (`float`, *optional*, defaults to 7.5): Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598). `guidance_scale` is defined as `w` of equation 2. of [Imagen Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, usually at the expense of lower image quality. negative_prompt (`str` or `List[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). strength (`float`, *optional*, defaults to 0.8): Indicates extent to transform the reference `image`. Must be between 0 and 1. `image` is used as a starting point and more noise is added the higher the `strength`. The number of denoising steps depends on the amount of noise initially added. When `strength` is 1, added noise is maximum and the denoising process runs for the full number of iterations specified in `num_inference_steps`. A value of 1 essentially ignores `image`. num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to [`schedulers.DDIMScheduler`], will be ignored for others. generator (`torch.Generator` or `List[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.FloatTensor`, *optional*): Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image generation. Can be used to tweak the same generation with different prompts. If not provided, a latents tensor will ge generated by sampling using the supplied random `generator`. prompt_embeds (`torch.FloatTensor`, *optional*): Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not provided, text embeddings will be generated from `prompt` input argument. negative_prompt_embeds (`torch.FloatTensor`, *optional*): Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input argument. output_type (`str`, *optional*, defaults to `"pil"`): The output format of the generate image. Choose between [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`. return_dict (`bool`, *optional*, defaults to `True`): Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a plain tuple. callback (`Callable`, *optional*): A function that will be called every `callback_steps` steps during inference. The function will be called with the following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`. callback_steps (`int`, *optional*, defaults to 1): The frequency at which the `callback` function will be called. If not specified, the callback will be called at every step. cross_attention_kwargs (`dict`, *optional*): A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under `self.processor` in [diffusers.cross_attention](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/cross_attention.py). controlnet_conditioning_scale (`float` or `List[float]`, *optional*, defaults to 1.0): The outputs of the controlnet are multiplied by `controlnet_conditioning_scale` before they are added to the residual in the original unet. If multiple ControlNets are specified in init, you can set the corresponding scale as a list. guess_mode (`bool`, *optional*, defaults to `False`): In this mode, the ControlNet encoder will try best to recognize the content of the input image even if you remove all prompts. The `guidance_scale` between 3.0 and 5.0 is recommended. control_guidance_start (`float` or `List[float]`, *optional*, defaults to 0.0): The percentage of total steps at which the controlnet starts applying. control_guidance_end (`float` or `List[float]`, *optional*, defaults to 1.0): The percentage of total steps at which the controlnet stops applying. skip_temporal_layer (`bool`: default to False) 为False时，unet起video生成作用,会运行时序生成的block；skip_temporal_layer为True时，unet起原image作用，跳过时序生成的block。 need_img_based_video_noise: bool = False, 当只有首帧latents时，是否需要扩展为video noise; num_videos_per_prompt: now only support 1. Examples: Returns: [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`: [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] if `return_dict` is True, otherwise a `tuple. When returning a tuple, the first element is a list with the generated images, and the second element is a list of `bool`s denoting whether the corresponding generated image likely represents "not-safe-for-work" (nsfw) content, according to the `safety_checker`. """ run_controlnet = control_image is not None or controlnet_latents is not None if run_controlnet: ( controlnet, control_guidance_start, control_guidance_end, ) = self.prepare_controlnet_and_guidance_parameter( control_guidance_start=control_guidance_start, control_guidance_end=control_guidance_end, ) # 1. Check inputs. Raise error if not correct self.check_inputs( prompt, control_image, callback_steps, negative_prompt, prompt_embeds, negative_prompt_embeds, controlnet_conditioning_scale, control_guidance_start, control_guidance_end, ) # 2. Define call parameters if prompt is not None and isinstance(prompt, str): batch_size = 1 elif prompt is not None and isinstance(prompt, list): batch_size = len(prompt) else: batch_size = prompt_embeds.shape[0] device = self._execution_device dtype = self.unet.dtype # print("pipeline unet dtype", dtype) # 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 run_controlnet: if isinstance(controlnet, MultiControlNetModel) and isinstance( controlnet_conditioning_scale, float ): controlnet_conditioning_scale = [controlnet_conditioning_scale] * len( controlnet.nets ) guess_mode = self.prepare_controlnet_guess_mode( controlnet=controlnet, guess_mode=guess_mode, ) # 3. Encode input prompt text_encoder_lora_scale = ( cross_attention_kwargs.get("scale", None) if cross_attention_kwargs is not None else None ) if self.text_encoder is not None: prompt_embeds = encode_weighted_prompt( self, prompt, device, num_videos_per_prompt, do_classifier_free_guidance, negative_prompt, prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_prompt_embeds, # lora_scale=text_encoder_lora_scale, ) logger.debug(f"use text_encoder prepare prompt_emb={prompt_embeds.shape}") else: prompt_embeds = None if image is not None: image = self.prepare_image( image, width=width, height=height, batch_size=batch_size * num_videos_per_prompt, num_images_per_prompt=num_videos_per_prompt, device=device, dtype=dtype, ) if self.print_idx == 0: logger.debug(f"image={image.shape}") if condition_images is not None: condition_images = self.prepare_image( condition_images, width=width, height=height, batch_size=batch_size * num_videos_per_prompt, num_images_per_prompt=num_videos_per_prompt, device=device, dtype=dtype, ) if self.print_idx == 0: logger.debug(f"condition_images={condition_images.shape}") # 4. Prepare image if run_controlnet: ( control_image, controlnet_latents, ) = self.prepare_controlnet_image_and_latents( controlnet=controlnet, width=width, height=height, batch_size=batch_size, num_videos_per_prompt=num_videos_per_prompt, device=device, dtype=dtype, controlnet_condition_latents=controlnet_condition_latents, control_image=control_image, controlnet_condition_images=controlnet_condition_images, guess_mode=guess_mode, do_classifier_free_guidance=do_classifier_free_guidance, controlnet_latents=controlnet_latents, ) # 5. Prepare timesteps self.scheduler.set_timesteps(num_inference_steps, device=device) if strength and (image is not None and latents is not None): if self.print_idx == 0: logger.debug( f"prepare timesteps, with get_timesteps strength={strength}, num_inference_steps={num_inference_steps}" ) timesteps, num_inference_steps = self.get_timesteps( num_inference_steps, strength, device ) else: if self.print_idx == 0: logger.debug(f"prepare timesteps, without get_timesteps") timesteps = self.scheduler.timesteps latent_timestep = timesteps[:1].repeat( batch_size * num_videos_per_prompt ) # 6. Prepare latent variables ( condition_latents, latent_index, vision_condition_latent_index, ) = self.prepare_condition_latents_and_index( condition_images=condition_images, condition_latents=condition_latents, video_length=video_length, batch_size=batch_size, dtype=dtype, device=device, latent_index=latent_index, vision_condition_latent_index=vision_condition_latent_index, ) if vision_condition_latent_index is None: n_vision_cond = 0 else: n_vision_cond = vision_condition_latent_index.shape[0] num_channels_latents = self.unet.config.in_channels if self.print_idx == 0: logger.debug(f"pipeline controlnet, start prepare latents") latents = self.prepare_latents( batch_size=batch_size * num_videos_per_prompt, num_channels_latents=num_channels_latents, video_length=video_length, height=height, width=width, dtype=dtype, device=device, generator=generator, latents=latents, image=image, timestep=latent_timestep, w_ind_noise=w_ind_noise, initial_common_latent=initial_common_latent, noise_type=noise_type, add_latents_noise=add_latents_noise, need_img_based_video_noise=need_img_based_video_noise, condition_latents=condition_latents, img_weight=img_weight, ) if self.print_idx == 0: logger.debug(f"pipeline controlnet, finish prepare latents={latents.shape}") # 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) if noise_type == "video_fusion" and "noise_type" in set( inspect.signature(self.scheduler.step).parameters.keys() ): extra_step_kwargs["w_ind_noise"] = w_ind_noise extra_step_kwargs["noise_type"] = noise_type # extra_step_kwargs["noise_offset"] = noise_offset # 7.1 Create tensor stating which controlnets to keep if run_controlnet: controlnet_keep = [] for i in range(len(timesteps)): keeps = [ 1.0 - float(i / len(timesteps) < s or (i + 1) / len(timesteps) > e) for s, e in zip(control_guidance_start, control_guidance_end) ] controlnet_keep.append( keeps[0] if isinstance(controlnet, ControlNetModel) else keeps ) else: controlnet_keep = None # 8. Denoising loop num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order if skip_temporal_layer: self.unet.set_skip_temporal_layers(True) n_timesteps = len(timesteps) guidance_scale_lst = generate_parameters_with_timesteps( start=guidance_scale, stop=guidance_scale_end, num=n_timesteps, method=guidance_scale_method, ) if self.print_idx == 0: logger.debug( f"guidance_scale_lst, {guidance_scale_method}, {guidance_scale}, {guidance_scale_end}, {guidance_scale_lst}" ) ip_adapter_image_emb = self.get_ip_adapter_image_emb( ip_adapter_image=ip_adapter_image, batch_size=batch_size, device=device, dtype=dtype, do_classifier_free_guidance=do_classifier_free_guidance, height=height, width=width, ) # 当前仅当没有ip_adapter时，按照参数 prompt_only_use_image_prompt 要求是否完全替换 image_prompt_emb # only if ip_adapter is None and prompt_only_use_image_prompt is True, use image_prompt_emb replace text_prompt if ( ip_adapter_image_emb is not None and prompt_only_use_image_prompt and not self.unet.ip_adapter_cross_attn ): prompt_embeds = ip_adapter_image_emb logger.debug(f"use ip_adapter_image_emb replace prompt_embeds") refer_face_image_emb = self.get_facein_image_emb( refer_face_image=refer_face_image, batch_size=batch_size, device=device, dtype=dtype, do_classifier_free_guidance=do_classifier_free_guidance, ) ip_adapter_face_emb = self.get_ip_adapter_face_emb( refer_face_image=ip_adapter_face_image, batch_size=batch_size, device=device, dtype=dtype, do_classifier_free_guidance=do_classifier_free_guidance, ) refer_image_vae_emb = self.get_referencenet_image_vae_emb( refer_image=refer_image, device=device, dtype=dtype, do_classifier_free_guidance=do_classifier_free_guidance, num_videos_per_prompt=num_videos_per_prompt, batch_size=batch_size, width=width, height=height, ) if self.pose_guider is not None and control_image is not None: if self.print_idx == 0: logger.debug(f"pose_guider, controlnet_image={control_image.shape}") control_image = rearrange( control_image, " (b t) c h w->b c t h w", t=video_length ) pose_guider_emb = self.pose_guider(control_image) pose_guider_emb = rearrange(pose_guider_emb, "b c t h w-> (b t) c h w") else: pose_guider_emb = None logger.debug(f"prompt_embeds={prompt_embeds.shape}") if control_image is not None: if isinstance(control_image, list): logger.debug(f"control_imageis list, num={len(control_image)}") control_image = [ rearrange( control_image_tmp, " (b t) c h w->b c t h w", b=(int(do_classifier_free_guidance) * 1 + 1) * batch_size, ) for control_image_tmp in control_image ] else: logger.debug(f"control_image={control_image.shape}, before") control_image = rearrange( control_image, " (b t) c h w->b c t h w", b=(int(do_classifier_free_guidance) * 1 + 1) * batch_size, ) logger.debug(f"control_image={control_image.shape}, after") if controlnet_latents is not None: if isinstance(controlnet_latents, list): logger.debug( f"controlnet_latents is list, num={len(controlnet_latents)}" ) controlnet_latents = [ rearrange( controlnet_latents_tmp, " (b t) c h w->b c t h w", b=(int(do_classifier_free_guidance) * 1 + 1) * batch_size, ) for controlnet_latents_tmp in controlnet_latents ] else: logger.debug(f"controlnet_latents={controlnet_latents.shape}, before") controlnet_latents = rearrange( controlnet_latents, " (b t) c h w->b c t h w", b=(int(do_classifier_free_guidance) * 1 + 1) * batch_size, ) logger.debug(f"controlnet_latents={controlnet_latents.shape}, after") videos_mid = [] mid_video_noises = [] if record_mid_video_noises else None mid_video_latents = [] if record_mid_video_latents else None global_context = prepare_global_context( context_schedule=context_schedule, num_inference_steps=num_inference_steps, time_size=latents.shape[2], context_frames=context_frames, context_stride=context_stride, context_overlap=context_overlap, context_batch_size=context_batch_size, ) logger.debug( f"context_schedule={context_schedule}, time_size={latents.shape[2]}, context_frames={context_frames}, context_stride={context_stride}, context_overlap={context_overlap}, context_batch_size={context_batch_size}" ) logger.debug(f"global_context={global_context}") # iterative denoise with self.progress_bar(total=num_inference_steps) as progress_bar: for i, t in enumerate(timesteps): # 使用 last_mid_video_latents 来影响初始化latent，该部分效果较差，暂留代码 # use last_mide_video_latents to affect initial latent. works bad, Temporarily reserved if i == 0: if record_mid_video_latents: mid_video_latents.append(latents[:, :, -video_overlap:]) if record_mid_video_noises: mid_video_noises.append(None) if ( last_mid_video_latents is not None and len(last_mid_video_latents) > 0 ): if self.print_idx == 1: logger.debug( f"{i}, last_mid_video_latents={last_mid_video_latents[i].shape}" ) latents = fuse_part_tensor( last_mid_video_latents[0], latents, video_overlap, weight=0.1, skip_step=0, ) noise_pred = torch.zeros( ( latents.shape[0] * (2 if do_classifier_free_guidance else 1), *latents.shape[1:], ), device=latents.device, dtype=latents.dtype, ) counter = torch.zeros( (1, 1, latents.shape[2], 1, 1), device=latents.device, dtype=latents.dtype, ) if i == 0: ( down_block_refer_embs, mid_block_refer_emb, refer_self_attn_emb, ) = self.get_referencenet_emb( refer_image_vae_emb=refer_image_vae_emb, refer_image=refer_image, device=device, dtype=dtype, do_classifier_free_guidance=do_classifier_free_guidance, num_videos_per_prompt=num_videos_per_prompt, prompt_embeds=prompt_embeds, ip_adapter_image_emb=ip_adapter_image_emb, batch_size=batch_size, ref_timestep_int=t, ) for context in global_context: # expand the latents if we are doing classifier free guidance latents_c = torch.cat([latents[:, :, c] for c in context]) latent_index_c = ( torch.cat([latent_index[c] for c in context]) if latent_index is not None else None ) latent_model_input = latents_c.to(device).repeat( 2 if do_classifier_free_guidance else 1, 1, 1, 1, 1 ) latent_model_input = self.scheduler.scale_model_input( latent_model_input, t ) sub_latent_index_c = ( torch.LongTensor( torch.arange(latent_index_c.shape[-1]) + n_vision_cond ).to(device=latents_c.device) if latent_index is not None else None ) if condition_latents is not None: latent_model_condition = ( torch.cat([condition_latents] * 2) if do_classifier_free_guidance else latents ) if self.print_idx == 0: logger.debug( f"vision_condition_latent_index, {vision_condition_latent_index.shape}, vision_condition_latent_index" ) logger.debug( f"latent_model_condition, {latent_model_condition.shape}" ) logger.debug(f"latent_index, {latent_index_c.shape}") logger.debug( f"latent_model_input, {latent_model_input.shape}" ) logger.debug(f"sub_latent_index_c, {sub_latent_index_c}") latent_model_input = batch_concat_two_tensor_with_index( data1=latent_model_condition, data1_index=vision_condition_latent_index, data2=latent_model_input, data2_index=sub_latent_index_c, dim=2, ) if control_image is not None: if vision_condition_latent_index is not None: # 获取 vision_condition 对应的 control_imgae/control_latent 部分 # generate control_image/control_latent corresponding to vision_condition controlnet_condtion_latent_index = ( vision_condition_latent_index.clone().cpu().tolist() ) if self.print_idx == 0: logger.debug( f"context={context}, controlnet_condtion_latent_index={controlnet_condtion_latent_index}" ) controlnet_context = [ controlnet_condtion_latent_index + [c_i + n_vision_cond for c_i in c] for c in context ] else: controlnet_context = context if self.print_idx == 0: logger.debug( f"controlnet_context={controlnet_context}, latent_model_input={latent_model_input.shape}" ) if isinstance(control_image, list): control_image_c = [ torch.cat( [ control_image_tmp[:, :, c] for c in controlnet_context ] ) for control_image_tmp in control_image ] control_image_c = [ rearrange(control_image_tmp, " b c t h w-> (b t) c h w") for control_image_tmp in control_image_c ] else: control_image_c = torch.cat( [control_image[:, :, c] for c in controlnet_context] ) control_image_c = rearrange( control_image_c, " b c t h w-> (b t) c h w" ) else: control_image_c = None if controlnet_latents is not None: if vision_condition_latent_index is not None: # 获取 vision_condition 对应的 control_imgae/control_latent 部分 # generate control_image/control_latent corresponding to vision_condition controlnet_condtion_latent_index = ( vision_condition_latent_index.clone().cpu().tolist() ) if self.print_idx == 0: logger.debug( f"context={context}, controlnet_condtion_latent_index={controlnet_condtion_latent_index}" ) controlnet_context = [ controlnet_condtion_latent_index + [c_i + n_vision_cond for c_i in c] for c in context ] else: controlnet_context = context if self.print_idx == 0: logger.debug( f"controlnet_context={controlnet_context}, controlnet_latents={controlnet_latents.shape}, latent_model_input={latent_model_input.shape}," ) controlnet_latents_c = torch.cat( [controlnet_latents[:, :, c] for c in controlnet_context] ) controlnet_latents_c = rearrange( controlnet_latents_c, " b c t h w-> (b t) c h w" ) else: controlnet_latents_c = None ( down_block_res_samples, mid_block_res_sample, ) = self.get_controlnet_emb( run_controlnet=run_controlnet, guess_mode=guess_mode, do_classifier_free_guidance=do_classifier_free_guidance, latents=latents_c, prompt_embeds=prompt_embeds, latent_model_input=latent_model_input, control_image=control_image_c, controlnet_latents=controlnet_latents_c, controlnet_keep=controlnet_keep, t=t, i=i, controlnet_conditioning_scale=controlnet_conditioning_scale, ) if self.print_idx == 0: logger.debug( f"{i}, latent_model_input={latent_model_input.shape}, sub_latent_index_c={sub_latent_index_c}" f"{vision_condition_latent_index}" ) # time.sleep(10) noise_pred_c = self.unet( latent_model_input, t, encoder_hidden_states=prompt_embeds, cross_attention_kwargs=cross_attention_kwargs, down_block_additional_residuals=down_block_res_samples, mid_block_additional_residual=mid_block_res_sample, return_dict=False, sample_index=sub_latent_index_c, vision_conditon_frames_sample_index=vision_condition_latent_index, sample_frame_rate=motion_speed, down_block_refer_embs=down_block_refer_embs, mid_block_refer_emb=mid_block_refer_emb, refer_self_attn_emb=refer_self_attn_emb, vision_clip_emb=ip_adapter_image_emb, face_emb=refer_face_image_emb, ip_adapter_scale=ip_adapter_scale, facein_scale=facein_scale, ip_adapter_face_emb=ip_adapter_face_emb, ip_adapter_face_scale=ip_adapter_face_scale, do_classifier_free_guidance=do_classifier_free_guidance, pose_guider_emb=pose_guider_emb, )[0] if condition_latents is not None: noise_pred_c = batch_index_select( noise_pred_c, dim=2, index=sub_latent_index_c ).contiguous() if self.print_idx == 0: logger.debug( f"{i}, latent_model_input={latent_model_input.shape}, noise_pred_c={noise_pred_c.shape}, {len(context)}, {len(context[0])}" ) for j, c in enumerate(context): noise_pred[:, :, c] = noise_pred[:, :, c] + noise_pred_c counter[:, :, c] = counter[:, :, c] + 1 noise_pred = noise_pred / counter if ( last_mid_video_noises is not None and len(last_mid_video_noises) > 0 and i <= num_inference_steps // 2 # 是个超参数 super paramter ): if self.print_idx == 1: logger.debug( f"{i}, last_mid_video_noises={last_mid_video_noises[i].shape}" ) noise_pred = fuse_part_tensor( last_mid_video_noises[i + 1], noise_pred, video_overlap, weight=0.01, skip_step=1, ) if record_mid_video_noises: mid_video_noises.append(noise_pred[:, :, -video_overlap:]) # perform guidance if do_classifier_free_guidance: noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) noise_pred = noise_pred_uncond + guidance_scale_lst[i] * ( noise_pred_text - noise_pred_uncond ) if self.print_idx == 0: logger.debug( f"before step, noise_pred={noise_pred.shape}, {noise_pred.device}, latents={latents.shape}, {latents.device}, t={t}" ) # compute the previous noisy sample x_t -> x_t-1 latents = self.scheduler.step( noise_pred, t, latents, **extra_step_kwargs, ).prev_sample if ( last_mid_video_latents is not None and len(last_mid_video_latents) > 0 and i <= 1 # 超参数, super parameter ): if self.print_idx == 1: logger.debug( f"{i}, last_mid_video_latents={last_mid_video_latents[i].shape}" ) latents = fuse_part_tensor( last_mid_video_latents[i + 1], latents, video_overlap, weight=0.1, skip_step=0, ) if record_mid_video_latents: mid_video_latents.append(latents[:, :, -video_overlap:]) if need_middle_latents is True: videos_mid.append(self.decode_latents(latents)) # 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) self.print_idx += 1 if condition_latents is not None: latents = batch_concat_two_tensor_with_index( data1=condition_latents, data1_index=vision_condition_latent_index, data2=latents, data2_index=latent_index, dim=2, ) video = self.decode_latents(latents) if skip_temporal_layer: self.unet.set_skip_temporal_layers(False) if need_hist_match: video[:, :, latent_index, :, :] = self.hist_match_with_vis_cond( batch_index_select(video, index=latent_index, dim=2), batch_index_select(video, index=vision_condition_latent_index, dim=2), ) # Convert to tensor if output_type == "tensor": videos_mid = [torch.from_numpy(x) for x in videos_mid] video = torch.from_numpy(video) else: latents = latents.cpu().numpy() if not return_dict: return ( video, latents, videos_mid, mid_video_latents, mid_video_noises, ) return VideoPipelineOutput( videos=video, latents=latents, videos_mid=videos_mid, mid_video_latents=mid_video_latents, mid_video_noises=mid_video_noises, )