import copy from typing import Any, Callable, Dict, Iterable, Union import PIL import cv2 import torch import argparse import datetime import logging import inspect import math import os import shutil from typing import Dict, List, Optional, Tuple from pprint import pformat, pprint from collections import OrderedDict from dataclasses import dataclass import gc import time import numpy as np from omegaconf import OmegaConf from omegaconf import SCMode import torch from torch import nn import torch.nn.functional as F import torch.utils.checkpoint from einops import rearrange, repeat import pandas as pd import h5py from diffusers.models.autoencoder_kl import AutoencoderKL from diffusers.models.modeling_utils import load_state_dict from diffusers.utils import ( logging, BaseOutput, logging, ) from diffusers.utils.dummy_pt_objects import ConsistencyDecoderVAE from diffusers.utils.import_utils import is_xformers_available from mmcm.utils.seed_util import set_all_seed from mmcm.vision.data.video_dataset import DecordVideoDataset from mmcm.vision.process.correct_color import hist_match_video_bcthw from mmcm.vision.process.image_process import ( batch_dynamic_crop_resize_images, batch_dynamic_crop_resize_images_v2, ) from mmcm.vision.utils.data_type_util import is_video from mmcm.vision.feature_extractor.controlnet import load_controlnet_model from ..schedulers import ( EulerDiscreteScheduler, LCMScheduler, DDIMScheduler, DDPMScheduler, ) from ..models.unet_3d_condition import UNet3DConditionModel from .pipeline_controlnet import ( MusevControlNetPipeline, VideoPipelineOutput as PipelineVideoPipelineOutput, ) from ..utils.util import save_videos_grid_with_opencv from ..utils.model_util import ( update_pipeline_basemodel, update_pipeline_lora_model, update_pipeline_lora_models, update_pipeline_model_parameters, ) 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] controlnet_cond: Union[torch.Tensor, np.ndarray] generated_videos: Union[torch.Tensor, np.ndarray] def update_controlnet_processor_params( src: Union[Dict, List[Dict]], dst: Union[Dict, List[Dict]] ): """merge dst into src""" if isinstance(src, list) and not isinstance(dst, List): dst = [dst] * len(src) if isinstance(src, list) and isinstance(dst, list): return [ update_controlnet_processor_params(src[i], dst[i]) for i in range(len(src)) ] if src is None: dct = {} else: dct = copy.deepcopy(src) if dst is None: dst = {} dct.update(dst) return dct class DiffusersPipelinePredictor(object): """wraper of diffusers pipeline, support generation function interface. support 1. text2video: inputs include text, image(optional), refer_image(optional) 2. video2video: 1. use controlnet to control spatial 2. or use video fuse noise to denoise """ def __init__( self, sd_model_path: str, unet: nn.Module, controlnet_name: Union[str, List[str]] = None, controlnet: nn.Module = None, lora_dict: Dict[str, Dict] = None, requires_safety_checker: bool = False, device: str = "cuda", dtype: torch.dtype = torch.float16, # controlnet parameters start need_controlnet_processor: bool = True, need_controlnet: bool = True, image_resolution: int = 512, detect_resolution: int = 512, include_body: bool = True, hand_and_face: bool = None, include_face: bool = False, include_hand: bool = True, negative_embedding: List = None, # controlnet parameters end enable_xformers_memory_efficient_attention: bool = True, lcm_lora_dct: Dict = None, referencenet: nn.Module = None, ip_adapter_image_proj: nn.Module = None, vision_clip_extractor: 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, vae_model: Optional[Tuple[nn.Module, str]] = None, pose_guider: Optional[nn.Module] = None, enable_zero_snr: bool = False, ) -> None: self.sd_model_path = sd_model_path self.unet = unet self.controlnet_name = controlnet_name self.controlnet = controlnet self.requires_safety_checker = requires_safety_checker self.device = device self.dtype = dtype self.need_controlnet_processor = need_controlnet_processor self.need_controlnet = need_controlnet self.need_controlnet_processor = need_controlnet_processor self.image_resolution = image_resolution self.detect_resolution = detect_resolution self.include_body = include_body self.hand_and_face = hand_and_face self.include_face = include_face self.include_hand = include_hand self.negative_embedding = negative_embedding self.device = device self.dtype = dtype self.lcm_lora_dct = lcm_lora_dct if controlnet is None and controlnet_name is not None: controlnet, controlnet_processor, processor_params = load_controlnet_model( controlnet_name, device=device, dtype=dtype, need_controlnet_processor=need_controlnet_processor, need_controlnet=need_controlnet, image_resolution=image_resolution, detect_resolution=detect_resolution, include_body=include_body, include_face=include_face, hand_and_face=hand_and_face, include_hand=include_hand, ) self.controlnet_processor = controlnet_processor self.controlnet_processor_params = processor_params logger.debug(f"init controlnet controlnet_name={controlnet_name}") if controlnet is not None: controlnet = controlnet.to(device=device, dtype=dtype) controlnet.eval() if pose_guider is not None: pose_guider = pose_guider.to(device=device, dtype=dtype) pose_guider.eval() unet.to(device=device, dtype=dtype) unet.eval() if referencenet is not None: referencenet.to(device=device, dtype=dtype) referencenet.eval() if ip_adapter_image_proj is not None: ip_adapter_image_proj.to(device=device, dtype=dtype) ip_adapter_image_proj.eval() if vision_clip_extractor is not None: vision_clip_extractor.to(device=device, dtype=dtype) vision_clip_extractor.eval() if face_emb_extractor is not None: face_emb_extractor.to(device=device, dtype=dtype) face_emb_extractor.eval() if facein_image_proj is not None: facein_image_proj.to(device=device, dtype=dtype) facein_image_proj.eval() if isinstance(vae_model, str): # TODO: poor implementation, to improve if "consistency" in vae_model: vae = ConsistencyDecoderVAE.from_pretrained(vae_model) else: vae = AutoencoderKL.from_pretrained(vae_model) elif isinstance(vae_model, nn.Module): vae = vae_model else: vae = None if vae is not None: vae.to(device=device, dtype=dtype) vae.eval() if ip_adapter_face_emb_extractor is not None: ip_adapter_face_emb_extractor.to(device=device, dtype=dtype) ip_adapter_face_emb_extractor.eval() if ip_adapter_face_image_proj is not None: ip_adapter_face_image_proj.to(device=device, dtype=dtype) ip_adapter_face_image_proj.eval() params = { "pretrained_model_name_or_path": sd_model_path, "controlnet": controlnet, "unet": unet, "requires_safety_checker": requires_safety_checker, "torch_dtype": dtype, "torch_device": device, "referencenet": referencenet, "ip_adapter_image_proj": ip_adapter_image_proj, "vision_clip_extractor": vision_clip_extractor, "facein_image_proj": facein_image_proj, "face_emb_extractor": face_emb_extractor, "ip_adapter_face_emb_extractor": ip_adapter_face_emb_extractor, "ip_adapter_face_image_proj": ip_adapter_face_image_proj, "pose_guider": pose_guider, } if vae is not None: params["vae"] = vae pipeline = MusevControlNetPipeline.from_pretrained(**params) pipeline = pipeline.to(torch_device=device, torch_dtype=dtype) logger.debug( f"init pipeline from sd_model_path={sd_model_path}, device={device}, dtype={dtype}" ) if ( negative_embedding is not None and pipeline.text_encoder is not None and pipeline.tokenizer is not None ): for neg_emb_path, neg_token in negative_embedding: pipeline.load_textual_inversion(neg_emb_path, token=neg_token) # pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config) # pipe.enable_model_cpu_offload() if not enable_zero_snr: pipeline.scheduler = EulerDiscreteScheduler.from_config( pipeline.scheduler.config ) # pipeline.scheduler = DDIMScheduler.from_config( # pipeline.scheduler.config, # 该部分会影响生成视频的亮度，不适用于首帧给定的视频生成 # this part will change brightness of video, not suitable for image2video mode # rescale_betas_zero_snr affect the brightness of the generated video, not suitable for vision condition images mode # # rescale_betas_zero_snr=True, # ) # pipeline.scheduler = DDPMScheduler.from_config(pipeline.scheduler.config) else: # moore scheduler, just for codetest pipeline.scheduler = DDIMScheduler( beta_start=0.00085, beta_end=0.012, beta_schedule="linear", clip_sample=False, steps_offset=1, ### Zero-SNR params prediction_type="v_prediction", rescale_betas_zero_snr=True, timestep_spacing="trailing", ) pipeline.enable_vae_slicing() self.enable_xformers_memory_efficient_attention = ( enable_xformers_memory_efficient_attention ) if enable_xformers_memory_efficient_attention: if is_xformers_available(): pipeline.enable_xformers_memory_efficient_attention() else: raise ValueError( "xformers is not available. Make sure it is installed correctly" ) self.pipeline = pipeline self.unload_dict = [] # keep lora state if lora_dict is not None: self.load_lora(lora_dict=lora_dict) logger.debug("load lora {}".format(" ".join(list(lora_dict.keys())))) if lcm_lora_dct is not None: self.pipeline.scheduler = LCMScheduler.from_config( self.pipeline.scheduler.config ) self.load_lora(lora_dict=lcm_lora_dct) logger.debug("load lcm lora {}".format(" ".join(list(lcm_lora_dct.keys())))) # logger.debug("Unet3Model Parameters") # logger.debug(pformat(self.__dict__)) def load_lora( self, lora_dict: Dict[str, Dict], ): self.pipeline, unload_dict = update_pipeline_lora_models( self.pipeline, lora_dict, device=self.device ) self.unload_dict += unload_dict def unload_lora(self): for layer_data in self.unload_dict: layer = layer_data["layer"] added_weight = layer_data["added_weight"] layer.weight.data -= added_weight self.unload_dict = [] gc.collect() torch.cuda.empty_cache() def update_unet(self, unet: nn.Module): self.pipeline.unet = unet.to(device=self.device, dtype=self.dtype) def update_sd_model(self, model_path: str, text_model_path: str): self.pipeline = update_pipeline_basemodel( self.pipeline, model_path, text_sd_model_path=text_model_path, device=self.device, ) def update_sd_model_and_unet( self, lora_sd_path: str, lora_path: str, sd_model_path: str = None ): self.pipeline = update_pipeline_model_parameters( self.pipeline, model_path=lora_sd_path, lora_path=lora_path, text_model_path=sd_model_path, device=self.device, ) def update_controlnet(self, controlnet_name=Union[str, List[str]]): self.pipeline.controlnet = load_controlnet_model(controlnet_name).to( device=self.device, dtype=self.dtype ) def run_pipe_text2video( self, video_length: int, prompt: Union[str, List[str]] = None, # b c t h w height: Optional[int] = None, width: Optional[int] = None, video_num_inference_steps: int = 50, video_guidance_scale: float = 7.5, video_guidance_scale_end: float = 3.5, video_guidance_scale_method: str = "linear", strength: float = 0.8, video_negative_prompt: Optional[Union[str, List[str]]] = None, 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, same_seed: Optional[Union[int, List[int]]] = None, # b c t(1) ho wo condition_latents: Optional[torch.FloatTensor] = None, latents: Optional[torch.FloatTensor] = None, prompt_embeds: Optional[torch.FloatTensor] = None, negative_prompt_embeds: Optional[torch.FloatTensor] = None, guidance_scale: float = 7.5, num_inference_steps: int = 50, 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, 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, n_vision_condition: int = 1, noise_type: str = "random", max_batch_num: int = 30, need_img_based_video_noise: bool = False, condition_images: torch.Tensor = None, fix_condition_images: bool = False, redraw_condition_image: bool = False, img_weight: float = 1e-3, motion_speed: float = 8.0, need_hist_match: bool = False, refer_image: Optional[ Tuple[np.ndarray, torch.Tensor, List[str], List[np.ndarray]] ] = None, ip_adapter_image: Optional[Tuple[torch.Tensor, np.array]] = None, fixed_refer_image: bool = True, fixed_ip_adapter_image: bool = True, redraw_condition_image_with_ipdapter: bool = True, redraw_condition_image_with_referencenet: bool = True, refer_face_image: Optional[Tuple[torch.Tensor, np.array]] = None, fixed_refer_face_image: bool = True, redraw_condition_image_with_facein: bool = True, ip_adapter_scale: float = 1.0, redraw_condition_image_with_ip_adapter_face: bool = True, facein_scale: float = 1.0, ip_adapter_face_scale: float = 1.0, prompt_only_use_image_prompt: bool = False, # serial_denoise parameter start record_mid_video_noises: bool = False, record_mid_video_latents: bool = False, video_overlap: int = 1, # serial_denoise parameter end # parallel_denoise parameter start context_schedule="uniform", context_frames=12, context_stride=1, context_overlap=4, context_batch_size=1, interpolation_factor=1, # parallel_denoise parameter end ): """ generate long video with end2end mode 1. prepare vision condition image by assingning, redraw, or generation with text2image module with skip_temporal_layer=True; 2. use image or latest of vision condition image to generate first shot; 3. use last n (1) image or last latent of last shot as new vision condition latent to generate next shot 4. repeat n_batch times between 2 and 3 类似img2img pipeline refer_image和ip_adapter_image的来源： 1. 输入给定； 2. 当未输入时，纯text2video生成首帧，并赋值更新refer_image和ip_adapter_image; 3. 当有输入，但是因为redraw更新了首帧时，也需要赋值更新refer_image和ip_adapter_image; refer_image和ip_adapter_image的作用： 1. 当无首帧图像时，用于生成首帧； 2. 用于生成视频。 similar to diffusers img2img pipeline. three ways to prepare refer_image and ip_adapter_image 1. from input parameter 2. when input paramter is None, use text2video to generate vis cond image, and use as refer_image and ip_adapter_image too. 3. given from input paramter, but still redraw, update with redrawn vis cond image. """ # crop resize images if condition_images is not None: logger.debug( f"center crop resize condition_images={condition_images.shape}, to height={height}, width={width}" ) condition_images = batch_dynamic_crop_resize_images_v2( condition_images, target_height=height, target_width=width, ) if refer_image is not None: logger.debug( f"center crop resize refer_image to height={height}, width={width}" ) refer_image = batch_dynamic_crop_resize_images_v2( refer_image, target_height=height, target_width=width, ) if ip_adapter_image is not None: logger.debug( f"center crop resize ip_adapter_image to height={height}, width={width}" ) ip_adapter_image = batch_dynamic_crop_resize_images_v2( ip_adapter_image, target_height=height, target_width=width, ) if refer_face_image is not None: logger.debug( f"center crop resize refer_face_image to height={height}, width={width}" ) refer_face_image = batch_dynamic_crop_resize_images_v2( refer_face_image, target_height=height, target_width=width, ) run_video_length = video_length # generate vision condition frame start # if condition_images is None, generate with refer_image, ip_adapter_image # if condition_images not None and need redraw, according to redraw_condition_image_with_ipdapter, redraw_condition_image_with_referencenet, refer_image, ip_adapter_image if n_vision_condition > 0: if condition_images is None and condition_latents is None: logger.debug("run_pipe_text2video, generate first_image") ( condition_images, condition_latents, _, _, _, ) = self.pipeline( prompt=prompt, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale, negative_prompt=negative_prompt, video_length=1, height=height, width=width, return_dict=False, skip_temporal_layer=True, output_type="np", generator=generator, w_ind_noise=w_ind_noise, need_img_based_video_noise=need_img_based_video_noise, refer_image=refer_image if redraw_condition_image_with_referencenet else None, ip_adapter_image=ip_adapter_image if redraw_condition_image_with_ipdapter else None, refer_face_image=refer_face_image if redraw_condition_image_with_facein else None, ip_adapter_scale=ip_adapter_scale, facein_scale=facein_scale, ip_adapter_face_scale=ip_adapter_face_scale, ip_adapter_face_image=refer_face_image if redraw_condition_image_with_ip_adapter_face else None, prompt_only_use_image_prompt=prompt_only_use_image_prompt, ) run_video_length = video_length - 1 elif ( condition_images is not None and redraw_condition_image and condition_latents is None ): logger.debug("run_pipe_text2video, redraw first_image") ( condition_images, condition_latents, _, _, _, ) = self.pipeline( prompt=prompt, image=condition_images, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale, negative_prompt=negative_prompt, strength=strength, video_length=condition_images.shape[2], height=height, width=width, return_dict=False, skip_temporal_layer=True, output_type="np", generator=generator, w_ind_noise=w_ind_noise, need_img_based_video_noise=need_img_based_video_noise, refer_image=refer_image if redraw_condition_image_with_referencenet else None, ip_adapter_image=ip_adapter_image if redraw_condition_image_with_ipdapter else None, refer_face_image=refer_face_image if redraw_condition_image_with_facein else None, ip_adapter_scale=ip_adapter_scale, facein_scale=facein_scale, ip_adapter_face_scale=ip_adapter_face_scale, ip_adapter_face_image=refer_face_image if redraw_condition_image_with_ip_adapter_face else None, prompt_only_use_image_prompt=prompt_only_use_image_prompt, ) else: condition_images = None condition_latents = None # generate vision condition frame end # refer_image and ip_adapter_image, update mode from 2 and 3 as mentioned above start if ( refer_image is not None and redraw_condition_image and condition_images is not None ): refer_image = condition_images * 255.0 logger.debug(f"update refer_image because of redraw_condition_image") elif ( refer_image is None and self.pipeline.referencenet is not None and condition_images is not None ): refer_image = condition_images * 255.0 logger.debug(f"update refer_image because of generate first_image") # ipadapter_image if ( ip_adapter_image is not None and redraw_condition_image and condition_images is not None ): ip_adapter_image = condition_images * 255.0 logger.debug(f"update ip_adapter_image because of redraw_condition_image") elif ( ip_adapter_image is None and self.pipeline.ip_adapter_image_proj is not None and condition_images is not None ): ip_adapter_image = condition_images * 255.0 logger.debug(f"update ip_adapter_image because of generate first_image") # refer_image and ip_adapter_image, update mode from 2 and 3 as mentioned above end # refer_face_image, update mode from 2 and 3 as mentioned above start if ( refer_face_image is not None and redraw_condition_image and condition_images is not None ): refer_face_image = condition_images * 255.0 logger.debug(f"update refer_face_image because of redraw_condition_image") elif ( refer_face_image is None and self.pipeline.facein_image_proj is not None and condition_images is not None ): refer_face_image = condition_images * 255.0 logger.debug(f"update face_image because of generate first_image") # refer_face_image, update mode from 2 and 3 as mentioned above end last_mid_video_noises = None last_mid_video_latents = None initial_common_latent = None out_videos = [] for i_batch in range(max_batch_num): logger.debug(f"sd_pipeline_predictor, run_pipe_text2video: {i_batch}") if max_batch_num is not None and i_batch == max_batch_num: break if i_batch == 0: result_overlap = 0 else: if n_vision_condition > 0: # ignore condition_images if condition_latents is not None in pipeline if not fix_condition_images: logger.debug(f"{i_batch}, update condition_latents") condition_latents = out_latents_batch[ :, :, -n_vision_condition:, :, : ] else: logger.debug(f"{i_batch}, do not update condition_latents") result_overlap = n_vision_condition if not fixed_refer_image and n_vision_condition > 0: logger.debug("ref_image use last frame of last generated out video") refer_image = out_batch[:, :, -n_vision_condition:, :, :] * 255.0 else: logger.debug("use given fixed ref_image") if not fixed_ip_adapter_image and n_vision_condition > 0: logger.debug( "ip_adapter_image use last frame of last generated out video" ) ip_adapter_image = ( out_batch[:, :, -n_vision_condition:, :, :] * 255.0 ) else: logger.debug("use given fixed ip_adapter_image") if not fixed_refer_face_image and n_vision_condition > 0: logger.debug( "refer_face_image use last frame of last generated out video" ) refer_face_image = ( out_batch[:, :, -n_vision_condition:, :, :] * 255.0 ) else: logger.debug("use given fixed ip_adapter_image") run_video_length = video_length if same_seed is not None: _, generator = set_all_seed(same_seed) out = self.pipeline( video_length=run_video_length, # int prompt=prompt, num_inference_steps=video_num_inference_steps, height=height, width=width, generator=generator, condition_images=condition_images, condition_latents=condition_latents, # b co t(1) ho wo skip_temporal_layer=False, output_type="np", noise_type=noise_type, negative_prompt=video_negative_prompt, guidance_scale=video_guidance_scale, guidance_scale_end=video_guidance_scale_end, guidance_scale_method=video_guidance_scale_method, w_ind_noise=w_ind_noise, need_img_based_video_noise=need_img_based_video_noise, img_weight=img_weight, motion_speed=motion_speed, vision_condition_latent_index=vision_condition_latent_index, refer_image=refer_image, ip_adapter_image=ip_adapter_image, refer_face_image=refer_face_image, ip_adapter_scale=ip_adapter_scale, facein_scale=facein_scale, ip_adapter_face_scale=ip_adapter_face_scale, ip_adapter_face_image=refer_face_image, prompt_only_use_image_prompt=prompt_only_use_image_prompt, initial_common_latent=initial_common_latent, # serial_denoise parameter start record_mid_video_noises=record_mid_video_noises, last_mid_video_noises=last_mid_video_noises, record_mid_video_latents=record_mid_video_latents, last_mid_video_latents=last_mid_video_latents, video_overlap=video_overlap, # serial_denoise parameter end # parallel_denoise parameter start context_schedule=context_schedule, context_frames=context_frames, context_stride=context_stride, context_overlap=context_overlap, context_batch_size=context_batch_size, interpolation_factor=interpolation_factor, # parallel_denoise parameter end ) logger.debug( f"run_pipe_text2video, out.videos.shape, i_batch={i_batch}, videos={out.videos.shape}, result_overlap={result_overlap}" ) out_batch = out.videos[:, :, result_overlap:, :, :] out_latents_batch = out.latents[:, :, result_overlap:, :, :] out_videos.append(out_batch) out_videos = np.concatenate(out_videos, axis=2) if need_hist_match: out_videos[:, :, 1:, :, :] = hist_match_video_bcthw( out_videos[:, :, 1:, :, :], out_videos[:, :, :1, :, :], value=255.0 ) return out_videos def run_pipe_with_latent_input( self, ): pass def run_pipe_middle2video_with_middle(self, middle: Tuple[str, Iterable]): pass def run_pipe_video2video( self, video: Tuple[str, Iterable], time_size: int = None, sample_rate: int = None, overlap: int = None, step: int = None, prompt: Union[str, List[str]] = None, # b c t h w height: Optional[int] = None, width: Optional[int] = None, num_inference_steps: int = 50, video_num_inference_steps: int = 50, guidance_scale: float = 7.5, video_guidance_scale: float = 7.5, video_guidance_scale_end: float = 3.5, video_guidance_scale_method: str = "linear", video_negative_prompt: Optional[Union[str, List[str]]] = None, negative_prompt: Optional[Union[str, List[str]]] = None, num_videos_per_prompt: Optional[int] = 1, negative_prompt_embeds: Optional[torch.FloatTensor] = None, eta: float = 0.0, generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, controlnet_latents: Union[torch.FloatTensor, np.ndarray] = 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, # b c t(1) ho wo condition_latents: Optional[torch.FloatTensor] = None, condition_images: Optional[torch.FloatTensor] = None, fix_condition_images: bool = False, latents: Optional[torch.FloatTensor] = None, 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, img_weight: float = 0.001, initial_common_latent: Optional[torch.FloatTensor] = None, latent_index: torch.LongTensor = None, vision_condition_latent_index: torch.LongTensor = None, noise_type: str = "random", controlnet_processor_params: Dict = None, need_return_videos: bool = False, need_return_condition: bool = False, max_batch_num: int = 30, strength: float = 0.8, video_strength: float = 0.8, need_video2video: bool = False, need_img_based_video_noise: bool = False, need_hist_match: bool = False, end_to_end: bool = True, refer_image: Optional[ Tuple[np.ndarray, torch.Tensor, List[str], List[np.ndarray]] ] = None, ip_adapter_image: Optional[Tuple[torch.Tensor, np.array]] = None, fixed_refer_image: bool = True, fixed_ip_adapter_image: bool = True, redraw_condition_image: bool = False, redraw_condition_image_with_ipdapter: bool = True, redraw_condition_image_with_referencenet: bool = True, refer_face_image: Optional[Tuple[torch.Tensor, np.array]] = None, fixed_refer_face_image: bool = True, redraw_condition_image_with_facein: bool = True, ip_adapter_scale: float = 1.0, facein_scale: float = 1.0, ip_adapter_face_scale: float = 1.0, redraw_condition_image_with_ip_adapter_face: bool = False, n_vision_condition: int = 1, prompt_only_use_image_prompt: bool = False, motion_speed: float = 8.0, # serial_denoise parameter start record_mid_video_noises: bool = False, record_mid_video_latents: bool = False, video_overlap: int = 1, # serial_denoise parameter end # parallel_denoise parameter start context_schedule="uniform", context_frames=12, context_stride=1, context_overlap=4, context_batch_size=1, interpolation_factor=1, # parallel_denoise parameter end # 支持 video_path 时多种输入 # TODO:// video_has_condition =False，当且仅支持 video_is_middle=True, 待后续重构 # TODO:// when video_has_condition =False, video_is_middle should be True. video_is_middle: bool = False, video_has_condition: bool = True, ): """ 类似controlnet text2img pipeline。 输入视频，用视频得到controlnet condition。 目前仅支持time_size == step，overlap=0 输出视频长度=输入视频长度 similar to controlnet text2image pipeline, generate video with controlnet condition from given video. By now, sliding window only support time_size == step, overlap = 0. """ if isinstance(video, str): video_reader = DecordVideoDataset( video, time_size=time_size, step=step, overlap=overlap, sample_rate=sample_rate, device="cpu", data_type="rgb", channels_order="c t h w", drop_last=True, ) else: video_reader = video videos = [] if need_return_videos else None out_videos = [] out_condition = ( [] if need_return_condition and self.pipeline.controlnet is not None else None ) # crop resize images if condition_images is not None: logger.debug( f"center crop resize condition_images={condition_images.shape}, to height={height}, width={width}" ) condition_images = batch_dynamic_crop_resize_images_v2( condition_images, target_height=height, target_width=width, ) if refer_image is not None: logger.debug( f"center crop resize refer_image to height={height}, width={width}" ) refer_image = batch_dynamic_crop_resize_images_v2( refer_image, target_height=height, target_width=width, ) if ip_adapter_image is not None: logger.debug( f"center crop resize ip_adapter_image to height={height}, width={width}" ) ip_adapter_image = batch_dynamic_crop_resize_images_v2( ip_adapter_image, target_height=height, target_width=width, ) if refer_face_image is not None: logger.debug( f"center crop resize refer_face_image to height={height}, width={width}" ) refer_face_image = batch_dynamic_crop_resize_images_v2( refer_face_image, target_height=height, target_width=width, ) first_image = None last_mid_video_noises = None last_mid_video_latents = None initial_common_latent = None # initial_common_latent = torch.randn((1, 4, 1, 112, 64)).to( # device=self.device, dtype=self.dtype # ) for i_batch, item in enumerate(video_reader): logger.debug(f"\n sd_pipeline_predictor, run_pipe_video2video: {i_batch}") if max_batch_num is not None and i_batch == max_batch_num: break # read and prepare video batch batch = item.data batch = batch_dynamic_crop_resize_images( batch, target_height=height, target_width=width, ) batch = batch[np.newaxis, ...] batch_size, channel, video_length, video_height, video_width = batch.shape # extract controlnet middle if self.pipeline.controlnet is not None: batch = rearrange(batch, "b c t h w-> (b t) h w c") controlnet_processor_params = update_controlnet_processor_params( src=self.controlnet_processor_params, dst=controlnet_processor_params, ) if not video_is_middle: batch_condition = self.controlnet_processor( data=batch, data_channel_order="b h w c", target_height=height, target_width=width, return_type="np", return_data_channel_order="b c h w", input_rgb_order="rgb", processor_params=controlnet_processor_params, ) else: # TODO: 临时用于可视化输入的 controlnet middle 序列，后续待拆到 middl2video中，也可以增加参数支持 # TODO: only use video_path is controlnet middle output, to improved batch_condition = rearrange( copy.deepcopy(batch), " b h w c-> b c h w" ) # 当前仅当 输入是 middle、condition_image的pose在middle首帧之前，需要重新生成condition_images的pose并绑定到middle_batch上 # when video_path is middle seq and condition_image is not aligned with middle seq, # regenerate codntion_images pose, and then concat into middle_batch, if ( i_batch == 0 and not video_has_condition and video_is_middle and condition_images is not None ): condition_images_reshape = rearrange( condition_images, "b c t h w-> (b t) h w c" ) condition_images_condition = self.controlnet_processor( data=condition_images_reshape, data_channel_order="b h w c", target_height=height, target_width=width, return_type="np", return_data_channel_order="b c h w", input_rgb_order="rgb", processor_params=controlnet_processor_params, ) condition_images_condition = rearrange( condition_images_condition, "(b t) c h w-> b c t h w", b=batch_size, ) else: condition_images_condition = None if not isinstance(batch_condition, list): batch_condition = rearrange( batch_condition, "(b t) c h w-> b c t h w", b=batch_size ) if condition_images_condition is not None: batch_condition = np.concatenate( [ condition_images_condition, batch_condition, ], axis=2, ) # 此时 batch_condition 比 batch 多了一帧，为了最终视频能 concat 存储，替换下 # 当前仅适用于 condition_images_condition 不为None # when condition_images_condition is not None, batch_condition has more frames than batch batch = rearrange(batch_condition, "b c t h w ->(b t) h w c") else: batch_condition = [ rearrange(x, "(b t) c h w-> b c t h w", b=batch_size) for x in batch_condition ] if condition_images_condition is not None: batch_condition = [ np.concatenate( [condition_images_condition, batch_condition_tmp], axis=2, ) for batch_condition_tmp in batch_condition ] batch = rearrange(batch, "(b t) h w c -> b c t h w", b=batch_size) else: batch_condition = None # condition [0,255] # latent: [0,1] # 按需求生成多个片段， # generate multi video_shot # 第一个片段 会特殊处理，需要生成首帧 # first shot is special because of first frame. # 后续片段根据拿前一个片段结果，首尾相连的方式生成。 # use last frame of last shot as the first frame of the current shot # TODO: 当前独立拆开实现，待后续融合到一起实现 # TODO: to optimize implementation way if n_vision_condition == 0: actual_video_length = video_length control_image = batch_condition first_image_controlnet_condition = None first_image_latents = None if need_video2video: video = batch else: video = None result_overlap = 0 else: if i_batch == 0: if self.pipeline.controlnet is not None: if not isinstance(batch_condition, list): first_image_controlnet_condition = batch_condition[ :, :, :1, :, : ] else: first_image_controlnet_condition = [ x[:, :, :1, :, :] for x in batch_condition ] else: first_image_controlnet_condition = None if need_video2video: if condition_images is None: video = batch[:, :, :1, :, :] else: video = condition_images else: video = None if condition_images is not None and not redraw_condition_image: first_image = condition_images first_image_latents = None else: ( first_image, first_image_latents, _, _, _, ) = self.pipeline( prompt=prompt, image=video, control_image=first_image_controlnet_condition, num_inference_steps=num_inference_steps, video_length=1, height=height, width=width, return_dict=False, skip_temporal_layer=True, output_type="np", generator=generator, negative_prompt=negative_prompt, controlnet_conditioning_scale=controlnet_conditioning_scale, control_guidance_start=control_guidance_start, control_guidance_end=control_guidance_end, w_ind_noise=w_ind_noise, strength=strength, refer_image=refer_image if redraw_condition_image_with_referencenet else None, ip_adapter_image=ip_adapter_image if redraw_condition_image_with_ipdapter else None, refer_face_image=refer_face_image if redraw_condition_image_with_facein else None, ip_adapter_scale=ip_adapter_scale, facein_scale=facein_scale, ip_adapter_face_scale=ip_adapter_face_scale, ip_adapter_face_image=refer_face_image if redraw_condition_image_with_ip_adapter_face else None, prompt_only_use_image_prompt=prompt_only_use_image_prompt, ) if refer_image is not None: refer_image = first_image * 255.0 if ip_adapter_image is not None: ip_adapter_image = first_image * 255.0 # 首帧用于后续推断可以直接用first_image_latent不需要 first_image了 first_image = None if self.pipeline.controlnet is not None: if not isinstance(batch_condition, list): control_image = batch_condition[:, :, 1:, :, :] logger.debug(f"control_image={control_image.shape}") else: control_image = [x[:, :, 1:, :, :] for x in batch_condition] else: control_image = None actual_video_length = time_size - int(video_has_condition) if need_video2video: video = batch[:, :, 1:, :, :] else: video = None result_overlap = 0 else: actual_video_length = time_size if self.pipeline.controlnet is not None: if not fix_condition_images: logger.debug( f"{i_batch}, update first_image_controlnet_condition" ) if not isinstance(last_batch_condition, list): first_image_controlnet_condition = last_batch_condition[ :, :, -1:, :, : ] else: first_image_controlnet_condition = [ x[:, :, -1:, :, :] for x in last_batch_condition ] else: logger.debug( f"{i_batch}, do not update first_image_controlnet_condition" ) control_image = batch_condition else: control_image = None first_image_controlnet_condition = None if not fix_condition_images: logger.debug(f"{i_batch}, update condition_images") first_image_latents = out_latents_batch[:, :, -1:, :, :] else: logger.debug(f"{i_batch}, do not update condition_images") if need_video2video: video = batch else: video = None result_overlap = 1 # 更新 ref_image和 ipadapter_image if not fixed_refer_image: logger.debug( "ref_image use last frame of last generated out video" ) refer_image = ( out_batch[:, :, -n_vision_condition:, :, :] * 255.0 ) else: logger.debug("use given fixed ref_image") if not fixed_ip_adapter_image: logger.debug( "ip_adapter_image use last frame of last generated out video" ) ip_adapter_image = ( out_batch[:, :, -n_vision_condition:, :, :] * 255.0 ) else: logger.debug("use given fixed ip_adapter_image") # face image if not fixed_ip_adapter_image: logger.debug( "refer_face_image use last frame of last generated out video" ) refer_face_image = ( out_batch[:, :, -n_vision_condition:, :, :] * 255.0 ) else: logger.debug("use given fixed ip_adapter_image") out = self.pipeline( video_length=actual_video_length, # int prompt=prompt, num_inference_steps=video_num_inference_steps, height=height, width=width, generator=generator, image=video, control_image=control_image, # b ci(3) t hi wi controlnet_condition_images=first_image_controlnet_condition, # b ci(3) t(1) hi wi # controlnet_condition_images=np.zeros_like( # first_image_controlnet_condition # ), # b ci(3) t(1) hi wi condition_images=first_image, condition_latents=first_image_latents, # b co t(1) ho wo skip_temporal_layer=False, output_type="np", noise_type=noise_type, negative_prompt=video_negative_prompt, need_img_based_video_noise=need_img_based_video_noise, controlnet_conditioning_scale=controlnet_conditioning_scale, control_guidance_start=control_guidance_start, control_guidance_end=control_guidance_end, w_ind_noise=w_ind_noise, img_weight=img_weight, motion_speed=video_reader.sample_rate, guidance_scale=video_guidance_scale, guidance_scale_end=video_guidance_scale_end, guidance_scale_method=video_guidance_scale_method, strength=video_strength, refer_image=refer_image, ip_adapter_image=ip_adapter_image, refer_face_image=refer_face_image, ip_adapter_scale=ip_adapter_scale, facein_scale=facein_scale, ip_adapter_face_scale=ip_adapter_face_scale, ip_adapter_face_image=refer_face_image, prompt_only_use_image_prompt=prompt_only_use_image_prompt, initial_common_latent=initial_common_latent, # serial_denoise parameter start record_mid_video_noises=record_mid_video_noises, last_mid_video_noises=last_mid_video_noises, record_mid_video_latents=record_mid_video_latents, last_mid_video_latents=last_mid_video_latents, video_overlap=video_overlap, # serial_denoise parameter end # parallel_denoise parameter start context_schedule=context_schedule, context_frames=context_frames, context_stride=context_stride, context_overlap=context_overlap, context_batch_size=context_batch_size, interpolation_factor=interpolation_factor, # parallel_denoise parameter end ) last_batch = batch last_batch_condition = batch_condition last_mid_video_latents = out.mid_video_latents last_mid_video_noises = out.mid_video_noises out_batch = out.videos[:, :, result_overlap:, :, :] out_latents_batch = out.latents[:, :, result_overlap:, :, :] out_videos.append(out_batch) if need_return_videos: videos.append(batch) if out_condition is not None: out_condition.append(batch_condition) out_videos = np.concatenate(out_videos, axis=2) if need_return_videos: videos = np.concatenate(videos, axis=2) if out_condition is not None: if not isinstance(out_condition[0], list): out_condition = np.concatenate(out_condition, axis=2) else: out_condition = [ [out_condition[j][i] for j in range(len(out_condition))] for i in range(len(out_condition[0])) ] out_condition = [np.concatenate(x, axis=2) for x in out_condition] if need_hist_match: videos[:, :, 1:, :, :] = hist_match_video_bcthw( videos[:, :, 1:, :, :], videos[:, :, :1, :, :], value=255.0 ) return out_videos, out_condition, videos