MuseV/musev/pipelines/pipeline_controlnet.py

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,
 )