init
This commit is contained in:
smthem
+46
-21
@@ -6,6 +6,8 @@ import gc
|
||||
import numpy as np
|
||||
from diffusers import StableDiffusionPipeline
|
||||
from .lav_relight import load_ic_light_model,infer_relight
|
||||
from .lav_wan_relight import load_ic_light_wan,infer_relight_wan
|
||||
from .lav_cog_relight import load_ic_light_cog,infer_relight_cog
|
||||
from .node_utils import load_images,tensor2pil_list
|
||||
import folder_paths
|
||||
from .src.ic_light import BGSource
|
||||
@@ -31,6 +33,7 @@ class Light_A_Video_Loader:
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"repo": ("STRING", {"default":"F:/test/ComfyUI/models/diffusers/Wan-AI/Wan2.1-T2V-1.3B-Diffusers"},),
|
||||
"model": (folder_paths.get_filename_list("checkpoints"),),
|
||||
"motion_adapter_model": (["none"] + folder_paths.get_filename_list("controlnet"),),
|
||||
"ic_light_model": (["none"] + folder_paths.get_filename_list("controlnet"),),
|
||||
@@ -43,35 +46,47 @@ class Light_A_Video_Loader:
|
||||
FUNCTION = "loader_main"
|
||||
CATEGORY = "Light_A_Video"
|
||||
|
||||
def loader_main(self, model,motion_adapter_model, ic_light_model,mode):
|
||||
def loader_main(self, repo,model,motion_adapter_model, ic_light_model,mode):
|
||||
|
||||
adopted_dtype = torch.float16
|
||||
|
||||
original_config_file=os.path.join(folder_paths.models_dir,"configs","v1-inference.yaml")
|
||||
# ic light
|
||||
sd_repo = os.path.join(current_node_path, "sd_repo")
|
||||
motion_repo=os.path.join(current_node_path,"animate_repo")
|
||||
original_config_file=os.path.join(sd_repo,"v1-inference.yaml") #fix for desktop comfyUI
|
||||
if model!="none":
|
||||
ckpt_path=folder_paths.get_full_path("checkpoints",model)
|
||||
ckpt_path=folder_paths.get_full_path("checkpoints",model)
|
||||
else:
|
||||
raise "no checkpoint"
|
||||
raise "no sd1.5 checkpoint"
|
||||
try:
|
||||
pipeline = StableDiffusionPipeline.from_single_file(
|
||||
sd_pipe = StableDiffusionPipeline.from_single_file(
|
||||
ckpt_path,config=sd_repo, original_config=original_config_file)
|
||||
except:
|
||||
pipeline = StableDiffusionPipeline.from_single_file(
|
||||
sd_pipe = StableDiffusionPipeline.from_single_file(
|
||||
ckpt_path, config=sd_repo,original_config_file=original_config_file)
|
||||
|
||||
# load model
|
||||
print("***********Load model ***********")
|
||||
motion_adapter_model=folder_paths.get_full_path("controlnet",motion_adapter_model)
|
||||
ic_light_model=folder_paths.get_full_path("controlnet",ic_light_model)
|
||||
|
||||
pipe,ic_light_pipe=load_ic_light_model(pipeline,ic_light_model,ckpt_path,sd_repo,motion_repo,motion_adapter_model,device,adopted_dtype,mode)
|
||||
print("***********Load model done ***********")
|
||||
|
||||
# video model
|
||||
if repo:
|
||||
if "wan" in repo.lower():
|
||||
print("***********Load wan diffuser ***********")
|
||||
pipe,ic_light_pipe=load_ic_light_wan(repo,sd_pipe,sd_repo,ckpt_path,ic_light_model,device,adopted_dtype)
|
||||
video_mode="wan"
|
||||
elif "cog" in repo.lower():
|
||||
print("***********Load cogvideox diffuser ***********")
|
||||
pipe,ic_light_pipe=load_ic_light_cog(repo,sd_pipe,sd_repo,ckpt_path,ic_light_model,device,adopted_dtype)
|
||||
video_mode="cog"
|
||||
else:
|
||||
raise "no string match wan or cog,check your repo name"
|
||||
else:
|
||||
# load animatediff model
|
||||
motion_repo=os.path.join(current_node_path,"animate_repo")
|
||||
print("***********Load animatediff model ***********")
|
||||
motion_adapter_model=folder_paths.get_full_path("controlnet",motion_adapter_model)
|
||||
pipe,ic_light_pipe=load_ic_light_model(sd_pipe,ic_light_model,ckpt_path,sd_repo,motion_repo,motion_adapter_model,device,adopted_dtype,mode)
|
||||
video_mode="animate"
|
||||
print("***********Load model done ***********")
|
||||
gc.collect()
|
||||
torch.cuda.empty_cache()
|
||||
return ({"model":pipe,"ic_light_pipe":ic_light_pipe,"mode":mode,"adopted_dtype":adopted_dtype},)
|
||||
return ({"model":pipe,"ic_light_pipe":ic_light_pipe,"mode":mode,"adopted_dtype":adopted_dtype,"video_mode":video_mode},)
|
||||
|
||||
|
||||
class Light_A_Video_Sampler:
|
||||
@@ -84,7 +99,8 @@ class Light_A_Video_Sampler:
|
||||
"required": {
|
||||
"model": ("MODEL_Light_A_Video",),
|
||||
"images": ("IMAGE",),
|
||||
"relight_prompt": ("STRING", {"default": "a car driving on the street, neon light", "multiline": True}),
|
||||
"relight_prompt": ("STRING", {"default": "a bear walking on the rock, nature lighting, soft light", "multiline": True}),
|
||||
"vdm_prompt": ("STRING", {"default": "a bear walking on the rock", "multiline": True}),
|
||||
"inpaint_prompt": ("STRING", {"default": "a car driving on the beach, sunset over sea", "multiline": True}),
|
||||
"n_prompt": ("STRING", {"default": "bad quality, worse quality", "multiline": True}),
|
||||
"seed": ("INT", {"default": 0, "min": 0, "max": MAX_SEED}),
|
||||
@@ -93,7 +109,8 @@ class Light_A_Video_Sampler:
|
||||
"text_guide_scale": ("INT", {"default": 2, "min": 1, "max": 20, "step": 1, "display": "number"}),
|
||||
"width": ("INT", {"default": 512, "min": 64, "max": 2048, "step": 64, "display": "number"}),
|
||||
"height": ("INT", {"default": 512, "min": 64, "max": 2048, "step": 64, "display": "number"}),
|
||||
"bg_target": (["LEFT", "NONE", "RIGHT", "TOP", "BOTTOM"],),
|
||||
"num_frames": ("INT", {"default": 49, "min": 10, "max": 1024, "step": 1, "display": "number"}),
|
||||
"bg_target": (["LEFT", "RIGHT", "TOP", "BOTTOM","NONE",],),
|
||||
"mask_repo": ("STRING", {"default": "ZhengPeng7/BiRefNet"},),},
|
||||
"optional": {"mask_img": ("IMAGE",),
|
||||
"fps": ("FLOAT", {"default": 8.0, "min": 8.0, "max": 100.0, "step": 0.1}),
|
||||
@@ -105,7 +122,7 @@ class Light_A_Video_Sampler:
|
||||
FUNCTION = "sampler_main"
|
||||
CATEGORY = "Light_A_Video"
|
||||
|
||||
def sampler_main(self, model,images,relight_prompt,inpaint_prompt, n_prompt,seed, num_step, strength,text_guide_scale,width, height,bg_target,mask_repo,**kwargs):
|
||||
def sampler_main(self, model,images,relight_prompt,vdm_prompt,inpaint_prompt, n_prompt,seed, num_step, strength,text_guide_scale,width, height,num_frames,bg_target,mask_repo,**kwargs):
|
||||
set_all_seed(42)
|
||||
|
||||
local_sam=os.path.join(Light_A_Video_weigths_path,"sam2_b.pt")
|
||||
@@ -128,11 +145,14 @@ class Light_A_Video_Sampler:
|
||||
fps=kwargs.get("fps",8.0)
|
||||
ic_light_pipe=model.get("ic_light_pipe")
|
||||
pipe=model.get("model")
|
||||
video_mode=model.get("video_mode")
|
||||
mode=model.get("mode")
|
||||
adopted_dtype=model.get("adopted_dtype")
|
||||
|
||||
if isinstance(mask_img,torch.Tensor):
|
||||
mask_list=tensor2pil_list(mask_img,width,height)
|
||||
# for i,mask in enumerate(mask_list):
|
||||
# mask.save(f"mask{i}.png")
|
||||
if len(mask_list)==1:
|
||||
mask_list=mask_list*len(ref_image_list)
|
||||
print("not enough mask, repeat mask to match the number of images")
|
||||
@@ -140,8 +160,13 @@ class Light_A_Video_Sampler:
|
||||
mask_list=None
|
||||
|
||||
print("***********Start infer ***********")
|
||||
iamge = infer_relight(ic_light_pipe,pipe,strength,num_step,text_guide_scale,seed,width,height,n_prompt,relight_prompt,inpaint_prompt,ref_image_list,bg_source,mode,mask_list,device,adopted_dtype,mask_repo,fps,local_sam)
|
||||
gc.collect()
|
||||
if video_mode=="wan":
|
||||
iamge = infer_relight_wan(ic_light_pipe,pipe,strength,num_step,text_guide_scale,seed,width,height,n_prompt,relight_prompt,vdm_prompt,ref_image_list,bg_source,num_frames,mode,mask_list,device,adopted_dtype,mask_repo,fps,local_sam)
|
||||
elif video_mode=="cog":
|
||||
iamge = infer_relight_cog(ic_light_pipe,pipe,strength,num_step,text_guide_scale,seed,width,height,n_prompt,relight_prompt,vdm_prompt,ref_image_list,bg_source,num_frames,mode,mask_list,device,adopted_dtype,mask_repo,fps,local_sam)
|
||||
else:
|
||||
iamge = infer_relight(ic_light_pipe,pipe,strength,num_step,text_guide_scale,seed,width,height,n_prompt,relight_prompt,inpaint_prompt,ref_image_list,bg_source,mode,mask_list,device,adopted_dtype,mask_repo,fps,local_sam)
|
||||
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
return (load_images(iamge),)
|
||||
|
||||
|
Before Width: | Height: | Size: 460 KiB After Width: | Height: | Size: 460 KiB |
|
Before Width: | Height: | Size: 390 KiB After Width: | Height: | Size: 390 KiB |
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,14 @@
|
||||
n_prompt: "bad quality, worse quality"
|
||||
vdm_prompt: "a bear walking on the rock"
|
||||
relight_prompt: "a bear walking on the rock, nature lighting, soft light"
|
||||
video_path: "input_cog/bear.mp4"
|
||||
bg_source: "TOP" ## LEFT, RIGHT, BOTTOM, TOP
|
||||
save_path: "output"
|
||||
|
||||
width: 720
|
||||
height: 480
|
||||
strength: 0.4
|
||||
gamma: 0.7
|
||||
num_step: 25
|
||||
text_guide_scale: 2
|
||||
seed: 42
|
||||
@@ -0,0 +1,15 @@
|
||||
n_prompt: "bad quality, worse quality"
|
||||
vdm_prompt: "a bear walking on the rock"
|
||||
relight_prompt: "a bear walking on the rock, nature lighting, soft light"
|
||||
video_path: "input_wan/bear.mp4"
|
||||
bg_source: "TOP" ## LEFT, RIGHT, BOTTOM, TOP
|
||||
save_path: "output"
|
||||
|
||||
width: 832
|
||||
height: 480
|
||||
num_frames: 49
|
||||
strength: 0.4
|
||||
gamma: 0.7
|
||||
num_step: 25
|
||||
text_guide_scale: 2
|
||||
seed: 42
|
||||
@@ -0,0 +1,15 @@
|
||||
n_prompt: "bad quality, worse quality"
|
||||
vdm_prompt: "a man walking in the factory"
|
||||
relight_prompt: "a man walking in the factory, red and blue neon light"
|
||||
video_path: "input_wan/man.mp4"
|
||||
bg_source: "LEFT" ## LEFT, RIGHT, BOTTOM, TOP
|
||||
save_path: "output"
|
||||
|
||||
width: 512
|
||||
height: 512
|
||||
num_frames: 49
|
||||
strength: 0.4
|
||||
gamma: 0.7
|
||||
num_step: 25
|
||||
text_guide_scale: 2
|
||||
seed: 42
|
||||
@@ -0,0 +1,236 @@
|
||||
import os
|
||||
import torch
|
||||
import imageio
|
||||
import argparse
|
||||
import safetensors.torch as sf
|
||||
from omegaconf import OmegaConf
|
||||
import torch.nn.functional as F
|
||||
from torch.hub import download_url_to_file
|
||||
|
||||
from diffusers import CogVideoXDDIMScheduler
|
||||
from transformers import CLIPTextModel, CLIPTokenizer
|
||||
from diffusers.models.attention_processor import AttnProcessor2_0
|
||||
from diffusers import AutoencoderKL, UNet2DConditionModel, DPMSolverMultistepScheduler
|
||||
|
||||
from .src.ic_light import BGSource
|
||||
from .src.ic_light_pipe import StableDiffusionImg2ImgPipeline
|
||||
from .src.cogvideo_pipe import CogVideoXVideoToVideoPipeline
|
||||
from .utils.tools import set_all_seed, read_video
|
||||
|
||||
|
||||
|
||||
def load_ic_light_cog(repo,sd_pipe,sd_repo,ckpt_path,ic_light_model,device,adopted_dtype):
|
||||
|
||||
|
||||
# config = OmegaConf.load(args.config)
|
||||
# device = torch.device('cuda')
|
||||
# adopted_dtype = torch.float16
|
||||
# set_all_seed(42)
|
||||
|
||||
## vdm model
|
||||
pipe = CogVideoXVideoToVideoPipeline.from_pretrained(repo, torch_dtype=adopted_dtype)
|
||||
pipe.scheduler = CogVideoXDDIMScheduler.from_config(pipe.scheduler.config)
|
||||
pipe = pipe.to(device=device, dtype=adopted_dtype)
|
||||
pipe.vae.requires_grad_(False)
|
||||
pipe.transformer.requires_grad_(False)
|
||||
|
||||
## module
|
||||
tokenizer = CLIPTokenizer.from_pretrained(sd_repo, subfolder="tokenizer")
|
||||
text_encoder = sd_pipe.text_encoder
|
||||
vae = sd_pipe.vae
|
||||
unet =sd_pipe.unet
|
||||
|
||||
with torch.no_grad():
|
||||
new_conv_in = torch.nn.Conv2d(8, unet.conv_in.out_channels, unet.conv_in.kernel_size, unet.conv_in.stride, unet.conv_in.padding)
|
||||
new_conv_in.weight.zero_()
|
||||
new_conv_in.weight[:, :4, :, :].copy_(unet.conv_in.weight)
|
||||
new_conv_in.bias = unet.conv_in.bias
|
||||
unet.conv_in = new_conv_in
|
||||
unet_original_forward = unet.forward
|
||||
|
||||
def hooked_unet_forward(sample, timestep, encoder_hidden_states, **kwargs):
|
||||
c_concat = kwargs['cross_attention_kwargs']['concat_conds'].to(sample)
|
||||
c_concat = torch.cat([c_concat] * (sample.shape[0] // c_concat.shape[0]), dim=0)
|
||||
new_sample = torch.cat([sample, c_concat], dim=1)
|
||||
kwargs['cross_attention_kwargs'] = {}
|
||||
return unet_original_forward(new_sample, timestep, encoder_hidden_states, **kwargs)
|
||||
unet.forward = hooked_unet_forward
|
||||
|
||||
## ic-light model loader
|
||||
if not os.path.exists(ic_light_model):
|
||||
download_url_to_file(url='https://huggingface.co/lllyasviel/ic-light/resolve/main/iclight_sd15_fc.safetensors',
|
||||
dst=ic_light_model) #TO DO: change to ic-light model
|
||||
sd_offset = sf.load_file(ic_light_model)
|
||||
sd_origin = unet.state_dict()
|
||||
sd_merged = {k: sd_origin[k] + sd_offset[k] for k in sd_origin.keys()}
|
||||
unet.load_state_dict(sd_merged, strict=True)
|
||||
del sd_offset, sd_origin, sd_merged
|
||||
text_encoder = text_encoder.to(device=device, dtype=adopted_dtype)
|
||||
vae = vae.to(device=device, dtype=adopted_dtype)
|
||||
unet = unet.to(device=device, dtype=adopted_dtype)
|
||||
|
||||
unet.set_attn_processor(AttnProcessor2_0())
|
||||
vae.set_attn_processor(AttnProcessor2_0())
|
||||
|
||||
# Consistent light attention
|
||||
@torch.inference_mode()
|
||||
def custom_forward_CLA(self,
|
||||
hidden_states,
|
||||
gamma=0.7,
|
||||
encoder_hidden_states=None,
|
||||
attention_mask=None,
|
||||
cross_attention_kwargs=None
|
||||
):
|
||||
|
||||
batch_size, sequence_length, channel = hidden_states.shape
|
||||
|
||||
residual = hidden_states
|
||||
input_ndim = hidden_states.ndim
|
||||
if input_ndim == 4:
|
||||
batch_size, channel, height, width = hidden_states.shape
|
||||
hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
|
||||
|
||||
if attention_mask is not None:
|
||||
if attention_mask.shape[-1] != query.shape[1]:
|
||||
target_length = query.shape[1]
|
||||
attention_mask = F.pad(attention_mask, (0, target_length), value=0.0)
|
||||
attention_mask = attention_mask.repeat_interleave(self.heads, dim=0)
|
||||
if self.group_norm is not None:
|
||||
hidden_states = self.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
|
||||
if encoder_hidden_states is None:
|
||||
encoder_hidden_states = hidden_states
|
||||
|
||||
query = self.to_q(hidden_states)
|
||||
key = self.to_k(encoder_hidden_states)
|
||||
value = self.to_v(encoder_hidden_states)
|
||||
inner_dim = key.shape[-1]
|
||||
head_dim = inner_dim // self.heads
|
||||
query = query.view(batch_size, -1, self.heads, head_dim).transpose(1, 2)
|
||||
key = key.view(batch_size, -1, self.heads, head_dim).transpose(1, 2)
|
||||
value = value.view(batch_size, -1, self.heads, head_dim).transpose(1, 2)
|
||||
|
||||
hidden_states = F.scaled_dot_product_attention(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False)
|
||||
shape = query.shape
|
||||
mean_key = key.reshape(2,-1,shape[1],shape[2],shape[3]).mean(dim=1,keepdim=True).expand(-1,shape[0]//2,-1,-1,-1).reshape(shape[0],shape[1],shape[2],shape[3])
|
||||
mean_value = value.reshape(2,-1,shape[1],shape[2],shape[3]).mean(dim=1,keepdim=True).expand(-1,shape[0]//2,-1,-1,-1).reshape(shape[0],shape[1],shape[2],shape[3])
|
||||
hidden_states_mean = F.scaled_dot_product_attention(query, mean_key, mean_value, attn_mask=None, dropout_p=0.0, is_causal=False)
|
||||
|
||||
hidden_states = (1-gamma)*hidden_states + gamma*hidden_states_mean
|
||||
hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, self.heads * head_dim)
|
||||
hidden_states = hidden_states.to(query.dtype)
|
||||
hidden_states = self.to_out[0](hidden_states)
|
||||
hidden_states = self.to_out[1](hidden_states)
|
||||
|
||||
if input_ndim == 4:
|
||||
hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
|
||||
|
||||
if self.residual_connection:
|
||||
hidden_states = hidden_states + residual
|
||||
|
||||
hidden_states = hidden_states / self.rescale_output_factor
|
||||
return hidden_states
|
||||
|
||||
### attention
|
||||
from types import MethodType
|
||||
@torch.inference_mode()
|
||||
def prep_unet_self_attention(unet):
|
||||
|
||||
for name, module in unet.named_modules():
|
||||
module_name = type(module).__name__
|
||||
|
||||
name_split_list = name.split(".")
|
||||
cond_1 = name_split_list[0] in "up_blocks"
|
||||
cond_2 = name_split_list[-1] in ('attn1')
|
||||
|
||||
if "Attention" in module_name and cond_1 and cond_2:
|
||||
cond_3 = name_split_list[1]
|
||||
if cond_3 not in "3":
|
||||
module.forward = MethodType(custom_forward_CLA, module)
|
||||
|
||||
return unet
|
||||
|
||||
## attn module
|
||||
unet = prep_unet_self_attention(unet)
|
||||
|
||||
## ic-light-scheduler
|
||||
ic_light_scheduler = DPMSolverMultistepScheduler(
|
||||
num_train_timesteps=1000,
|
||||
beta_start=0.00085,
|
||||
beta_end=0.012,
|
||||
algorithm_type="sde-dpmsolver++",
|
||||
use_karras_sigmas=True,
|
||||
steps_offset=1
|
||||
)
|
||||
ic_light_pipe = StableDiffusionImg2ImgPipeline(
|
||||
vae=vae,
|
||||
text_encoder=text_encoder,
|
||||
tokenizer=tokenizer,
|
||||
unet=unet,
|
||||
scheduler=ic_light_scheduler,
|
||||
safety_checker=None,
|
||||
requires_safety_checker=False,
|
||||
feature_extractor=None,
|
||||
image_encoder=None
|
||||
)
|
||||
ic_light_pipe = ic_light_pipe.to(device=device, dtype=adopted_dtype)
|
||||
ic_light_pipe.vae.requires_grad_(False)
|
||||
ic_light_pipe.unet.requires_grad_(False)
|
||||
|
||||
return pipe,ic_light_pipe
|
||||
|
||||
def infer_relight_cog(ic_light_pipe,pipe,strength,num_step,text_guide_scale,seed,image_width,image_height,n_prompt,relight_prompt,vdm_prompt,video_list,bg_source,num_frames,mode,mask_list,device,adopted_dtype,repo,fps,local_sam):
|
||||
############################# params ######################################
|
||||
# strength = config.get("strength", 0.4)
|
||||
# num_step = config.get("num_step", 25)
|
||||
# text_guide_scale = config.get("text_guide_scale", 2)
|
||||
# seed = config.get("seed")
|
||||
# image_width = config.get("width", 720)
|
||||
# image_height = config.get("height", 480)
|
||||
# negative_prompt = config.get("n_prompt", "")
|
||||
# vdm_prompt = config.get("vdm_prompt", "")
|
||||
# relight_prompt = config.get("relight_prompt", "")
|
||||
# video_path = config.get("video_path", "")
|
||||
# bg_source = BGSource[config.get("bg_source")]
|
||||
# save_path = config.get("save_path")
|
||||
|
||||
############################## infer #####################################
|
||||
generator = torch.manual_seed(seed)
|
||||
# video_name = os.path.basename(video_path)
|
||||
# video_list, video_name = read_video(video_path, image_width, image_height)
|
||||
|
||||
print("################## begin ##################")
|
||||
with torch.no_grad():
|
||||
num_inference_steps = int(round(num_step / strength))
|
||||
|
||||
output = pipe(
|
||||
ic_light_pipe=ic_light_pipe,
|
||||
relight_prompt=relight_prompt,
|
||||
bg_source=bg_source,
|
||||
video=video_list,
|
||||
prompt=vdm_prompt,
|
||||
strength=strength,
|
||||
negative_prompt=n_prompt,
|
||||
guidance_scale=text_guide_scale,
|
||||
num_inference_steps=num_inference_steps,
|
||||
height=image_height,
|
||||
width=image_width,
|
||||
generator=generator,
|
||||
)
|
||||
|
||||
frames = output.frames[0]
|
||||
# results_path = f"{save_path}/relight_{video_name}"
|
||||
# imageio.mimwrite(results_path, frames, fps=14)
|
||||
# print(f"relight! \n prompt:{relight_prompt}, light:{bg_source.value}, save in {results_path}.")
|
||||
return frames
|
||||
|
||||
# if __name__ == "__main__":
|
||||
# parser = argparse.ArgumentParser()
|
||||
|
||||
# parser.add_argument("--sd_model", type=str, default="stablediffusionapi/realistic-vision-v51")
|
||||
# parser.add_argument("--vdm_model", type=str, default="THUDM/CogVideoX-2b")
|
||||
# parser.add_argument("--ic_light_model", type=str, default="./models/iclight_sd15_fc.safetensors")
|
||||
|
||||
# parser.add_argument("--config", type=str, default="configs/cog_relight/bear.yaml", help="the config file for each sample.")
|
||||
|
||||
# args = parser.parse_args()
|
||||
# main(args)
|
||||
@@ -0,0 +1,440 @@
|
||||
import os
|
||||
import torch
|
||||
import imageio
|
||||
import argparse
|
||||
import numpy as np
|
||||
import safetensors.torch as sf
|
||||
|
||||
from omegaconf import OmegaConf
|
||||
import torch.nn.functional as F
|
||||
from torch.hub import download_url_to_file
|
||||
import gc
|
||||
from transformers import CLIPTextModel, CLIPTokenizer
|
||||
from diffusers import AutoencoderKL, UNet2DConditionModel, DPMSolverMultistepScheduler
|
||||
from diffusers.models.attention_processor import AttnProcessor2_0
|
||||
|
||||
from .src.ic_light_pipe import StableDiffusionImg2ImgPipeline
|
||||
from .src.wan_pipe import WanVideoToVideoPipeline
|
||||
from .src.ic_light import BGSource
|
||||
from .utils.tools import set_all_seed, read_video
|
||||
|
||||
|
||||
def load_ic_light_wan(repo,sd_pipe,sd_repo,ckpt_path,ic_light_model,device,adopted_dtype):
|
||||
# config = OmegaConf.load(args.config)
|
||||
# device = torch.device('cuda')
|
||||
# adopted_dtype = torch.float16
|
||||
#set_all_seed(42)
|
||||
## vdm model
|
||||
from diffusers import AutoencoderKLWan
|
||||
vae = AutoencoderKLWan.from_pretrained(repo, subfolder="vae", torch_dtype=adopted_dtype)
|
||||
pipe = WanVideoToVideoPipeline.from_pretrained(repo, vae=vae, torch_dtype=adopted_dtype)
|
||||
|
||||
pipe = pipe.to(device=device, dtype=adopted_dtype)
|
||||
pipe.vae.requires_grad_(False)
|
||||
pipe.transformer.requires_grad_(False)
|
||||
|
||||
## module
|
||||
tokenizer = CLIPTokenizer.from_pretrained(sd_repo, subfolder="tokenizer")
|
||||
unet=sd_pipe.unet #UNet2DConditionModel.from_pretrained(sd_repo, subfolder="unet") 确认是否有效
|
||||
text_encoder = sd_pipe.text_encoder
|
||||
vae = sd_pipe.vae
|
||||
del sd_pipe
|
||||
gc.collect()
|
||||
|
||||
|
||||
with torch.no_grad():
|
||||
new_conv_in = torch.nn.Conv2d(8, unet.conv_in.out_channels, unet.conv_in.kernel_size, unet.conv_in.stride, unet.conv_in.padding)
|
||||
new_conv_in.weight.zero_() #torch.Size([320, 8, 3, 3])
|
||||
new_conv_in.weight[:, :4, :, :].copy_(unet.conv_in.weight)
|
||||
new_conv_in.bias = unet.conv_in.bias
|
||||
unet.conv_in = new_conv_in
|
||||
unet_original_forward = unet.forward
|
||||
def hooked_unet_forward(sample, timestep, encoder_hidden_states, **kwargs):
|
||||
c_concat = kwargs['cross_attention_kwargs']['concat_conds'].to(sample)
|
||||
c_concat = torch.cat([c_concat] * (sample.shape[0] // c_concat.shape[0]), dim=0)
|
||||
new_sample = torch.cat([sample, c_concat], dim=1)
|
||||
kwargs['cross_attention_kwargs'] = {}
|
||||
return unet_original_forward(new_sample, timestep, encoder_hidden_states, **kwargs)
|
||||
unet.forward = hooked_unet_forward
|
||||
|
||||
## ic-light model loader
|
||||
if not os.path.exists(ic_light_model):
|
||||
download_url_to_file(url='https://huggingface.co/lllyasviel/ic-light/resolve/main/iclight_sd15_fc.safetensors',
|
||||
dst=ic_light_model) #TO DO: change to ic-light model
|
||||
sd_offset = sf.load_file(ic_light_model)
|
||||
sd_origin = unet.state_dict()
|
||||
sd_merged = {k: sd_origin[k] + sd_offset[k] for k in sd_origin.keys()}
|
||||
unet.load_state_dict(sd_merged, strict=True)
|
||||
del sd_offset, sd_origin, sd_merged
|
||||
text_encoder = text_encoder.to(device=device, dtype=adopted_dtype)
|
||||
vae = vae.to(device=device, dtype=adopted_dtype)
|
||||
unet = unet.to(device=device, dtype=adopted_dtype)
|
||||
|
||||
unet.set_attn_processor(AttnProcessor2_0())
|
||||
vae.set_attn_processor(AttnProcessor2_0())
|
||||
@torch.inference_mode()
|
||||
def custom_forward_CLA(self,
|
||||
hidden_states,
|
||||
gamma=0.7,
|
||||
encoder_hidden_states=None,
|
||||
attention_mask=None,
|
||||
cross_attention_kwargs=None
|
||||
):
|
||||
|
||||
batch_size, sequence_length, channel = hidden_states.shape
|
||||
|
||||
residual = hidden_states
|
||||
input_ndim = hidden_states.ndim
|
||||
if input_ndim == 4:
|
||||
batch_size, channel, height, width = hidden_states.shape
|
||||
hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
|
||||
|
||||
if attention_mask is not None:
|
||||
if attention_mask.shape[-1] != query.shape[1]:
|
||||
target_length = query.shape[1]
|
||||
attention_mask = F.pad(attention_mask, (0, target_length), value=0.0)
|
||||
attention_mask = attention_mask.repeat_interleave(self.heads, dim=0)
|
||||
if self.group_norm is not None:
|
||||
hidden_states = self.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
|
||||
if encoder_hidden_states is None:
|
||||
encoder_hidden_states = hidden_states
|
||||
|
||||
query = self.to_q(hidden_states)
|
||||
key = self.to_k(encoder_hidden_states)
|
||||
value = self.to_v(encoder_hidden_states)
|
||||
inner_dim = key.shape[-1]
|
||||
head_dim = inner_dim // self.heads
|
||||
query = query.view(batch_size, -1, self.heads, head_dim).transpose(1, 2)
|
||||
key = key.view(batch_size, -1, self.heads, head_dim).transpose(1, 2)
|
||||
value = value.view(batch_size, -1, self.heads, head_dim).transpose(1, 2)
|
||||
|
||||
hidden_states = F.scaled_dot_product_attention(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False)
|
||||
shape = query.shape
|
||||
mean_key = key.reshape(2,-1,shape[1],shape[2],shape[3]).mean(dim=1,keepdim=True).expand(-1,shape[0]//2,-1,-1,-1).reshape(shape[0],shape[1],shape[2],shape[3])
|
||||
mean_value = value.reshape(2,-1,shape[1],shape[2],shape[3]).mean(dim=1,keepdim=True).expand(-1,shape[0]//2,-1,-1,-1).reshape(shape[0],shape[1],shape[2],shape[3])
|
||||
hidden_states_mean = F.scaled_dot_product_attention(query, mean_key, mean_value, attn_mask=None, dropout_p=0.0, is_causal=False)
|
||||
|
||||
hidden_states = (1-gamma)*hidden_states + gamma*hidden_states_mean
|
||||
hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, self.heads * head_dim)
|
||||
hidden_states = hidden_states.to(query.dtype)
|
||||
|
||||
hidden_states = self.to_out[0](hidden_states)
|
||||
hidden_states = self.to_out[1](hidden_states)
|
||||
|
||||
if input_ndim == 4:
|
||||
hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
|
||||
|
||||
if self.residual_connection:
|
||||
hidden_states = hidden_states + residual
|
||||
|
||||
hidden_states = hidden_states / self.rescale_output_factor
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
return hidden_states
|
||||
|
||||
### attention
|
||||
from types import MethodType
|
||||
@torch.inference_mode()
|
||||
def prep_unet_self_attention(unet):
|
||||
|
||||
for name, module in unet.named_modules():
|
||||
module_name = type(module).__name__
|
||||
|
||||
name_split_list = name.split(".")
|
||||
cond_1 = name_split_list[0] in "up_blocks"
|
||||
cond_2 = name_split_list[-1] in ('attn1')
|
||||
|
||||
if "Attention" in module_name and cond_1 and cond_2:
|
||||
cond_3 = name_split_list[1]
|
||||
if cond_3 not in "3":
|
||||
module.forward = MethodType(custom_forward_CLA, module)
|
||||
|
||||
return unet
|
||||
|
||||
## attn module
|
||||
unet = prep_unet_self_attention(unet)
|
||||
|
||||
## ic-light-scheduler
|
||||
ic_light_scheduler = DPMSolverMultistepScheduler(
|
||||
num_train_timesteps=1000,
|
||||
beta_start=0.00085,
|
||||
beta_end=0.012,
|
||||
algorithm_type="sde-dpmsolver++",
|
||||
use_karras_sigmas=True,
|
||||
steps_offset=1
|
||||
)
|
||||
ic_light_pipe = StableDiffusionImg2ImgPipeline(
|
||||
vae=vae,
|
||||
text_encoder=text_encoder,
|
||||
tokenizer=tokenizer,
|
||||
unet=unet,
|
||||
scheduler=ic_light_scheduler,
|
||||
safety_checker=None,
|
||||
requires_safety_checker=False,
|
||||
feature_extractor=None,
|
||||
image_encoder=None
|
||||
)
|
||||
ic_light_pipe = ic_light_pipe.to(device=device, dtype=adopted_dtype)
|
||||
ic_light_pipe.vae.requires_grad_(False)
|
||||
ic_light_pipe.unet.requires_grad_(False)
|
||||
|
||||
pipeline = {
|
||||
"transformer": pipe.transformer,
|
||||
"text_encoder": pipe.text_encoder,
|
||||
"vae": pipe.vae
|
||||
}
|
||||
|
||||
|
||||
return pipe,ic_light_pipe
|
||||
|
||||
|
||||
def infer_relight_wan(ic_light_pipe,pipe,strength,num_step,text_guide_scale,seed,image_width,image_height,n_prompt,relight_prompt,vdm_prompt,video_list,bg_source,num_frames,mode,mask_list,device,adopted_dtype,repo,fps,local_sam):
|
||||
############################## infer #####################################
|
||||
generator = torch.manual_seed(seed)
|
||||
# video_name = os.path.basename(video_path)
|
||||
# video_list, video_name = read_video(video_path, image_width, image_height)
|
||||
print("################## begin ##################")
|
||||
with torch.no_grad():
|
||||
num_inference_steps = int(round(num_step / strength))
|
||||
output = pipe(
|
||||
ic_light_pipe=ic_light_pipe,
|
||||
relight_prompt=relight_prompt,
|
||||
bg_source=bg_source,
|
||||
video=video_list,
|
||||
prompt=vdm_prompt,
|
||||
negative_prompt=n_prompt,
|
||||
strength=strength,
|
||||
guidance_scale=text_guide_scale,
|
||||
num_inference_steps=num_inference_steps,
|
||||
height=image_height,
|
||||
num_frames=num_frames,
|
||||
width=image_width,
|
||||
generator=generator,
|
||||
)
|
||||
|
||||
frames = output.frames[0]
|
||||
# frames = (frames * 255).astype(np.uint8)
|
||||
# results_path = f"{save_path}/relight_{video_name}"
|
||||
# imageio.mimwrite(results_path, frames, fps=14)
|
||||
# print(f"relight! prompt:{relight_prompt}, light:{bg_source.value}, save in {results_path}.")
|
||||
return frames
|
||||
|
||||
|
||||
#def main(args):
|
||||
|
||||
# config = OmegaConf.load(args.config)
|
||||
# device = torch.device('cuda')
|
||||
# adopted_dtype = torch.float16
|
||||
# set_all_seed(42)
|
||||
|
||||
# ## vdm model
|
||||
# vae = AutoencoderKLWan.from_pretrained(args.vdm_model, subfolder="vae", torch_dtype=adopted_dtype)
|
||||
# pipe = WanVideoToVideoPipeline.from_pretrained(args.vdm_model, vae=vae, torch_dtype=adopted_dtype)
|
||||
|
||||
# pipe = pipe.to(device=device, dtype=adopted_dtype)
|
||||
# pipe.vae.requires_grad_(False)
|
||||
# pipe.transformer.requires_grad_(False)
|
||||
|
||||
# ## module
|
||||
# tokenizer = CLIPTokenizer.from_pretrained(args.sd_model, subfolder="tokenizer")
|
||||
# text_encoder = CLIPTextModel.from_pretrained(args.sd_model, subfolder="text_encoder")
|
||||
# vae = AutoencoderKL.from_pretrained(args.sd_model, subfolder="vae")
|
||||
# unet = UNet2DConditionModel.from_pretrained(args.sd_model, subfolder="unet")
|
||||
# with torch.no_grad():
|
||||
# new_conv_in = torch.nn.Conv2d(8, unet.conv_in.out_channels, unet.conv_in.kernel_size, unet.conv_in.stride, unet.conv_in.padding)
|
||||
# new_conv_in.weight.zero_() #torch.Size([320, 8, 3, 3])
|
||||
# new_conv_in.weight[:, :4, :, :].copy_(unet.conv_in.weight)
|
||||
# new_conv_in.bias = unet.conv_in.bias
|
||||
# unet.conv_in = new_conv_in
|
||||
# unet_original_forward = unet.forward
|
||||
|
||||
# def hooked_unet_forward(sample, timestep, encoder_hidden_states, **kwargs):
|
||||
# c_concat = kwargs['cross_attention_kwargs']['concat_conds'].to(sample)
|
||||
# c_concat = torch.cat([c_concat] * (sample.shape[0] // c_concat.shape[0]), dim=0)
|
||||
# new_sample = torch.cat([sample, c_concat], dim=1)
|
||||
# kwargs['cross_attention_kwargs'] = {}
|
||||
# return unet_original_forward(new_sample, timestep, encoder_hidden_states, **kwargs)
|
||||
# unet.forward = hooked_unet_forward
|
||||
|
||||
# ## ic-light model loader
|
||||
# if not os.path.exists(args.ic_light_model):
|
||||
# download_url_to_file(url='https://huggingface.co/lllyasviel/ic-light/resolve/main/iclight_sd15_fc.safetensors',
|
||||
# dst=args.ic_light_model)
|
||||
# sd_offset = sf.load_file(args.ic_light_model)
|
||||
# sd_origin = unet.state_dict()
|
||||
# sd_merged = {k: sd_origin[k] + sd_offset[k] for k in sd_origin.keys()}
|
||||
# unet.load_state_dict(sd_merged, strict=True)
|
||||
# del sd_offset, sd_origin, sd_merged
|
||||
# text_encoder = text_encoder.to(device=device, dtype=adopted_dtype)
|
||||
# vae = vae.to(device=device, dtype=adopted_dtype)
|
||||
# unet = unet.to(device=device, dtype=adopted_dtype)
|
||||
# unet.set_attn_processor(AttnProcessor2_0())
|
||||
# vae.set_attn_processor(AttnProcessor2_0())
|
||||
|
||||
# Consistent light attention
|
||||
# @torch.inference_mode()
|
||||
# def custom_forward_CLA(self,
|
||||
# hidden_states,
|
||||
# gamma=0.7,
|
||||
# encoder_hidden_states=None,
|
||||
# attention_mask=None,
|
||||
# cross_attention_kwargs=None
|
||||
# ):
|
||||
|
||||
# batch_size, sequence_length, channel = hidden_states.shape
|
||||
|
||||
# residual = hidden_states
|
||||
# input_ndim = hidden_states.ndim
|
||||
# if input_ndim == 4:
|
||||
# batch_size, channel, height, width = hidden_states.shape
|
||||
# hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
|
||||
|
||||
# if attention_mask is not None:
|
||||
# if attention_mask.shape[-1] != query.shape[1]:
|
||||
# target_length = query.shape[1]
|
||||
# attention_mask = F.pad(attention_mask, (0, target_length), value=0.0)
|
||||
# attention_mask = attention_mask.repeat_interleave(self.heads, dim=0)
|
||||
# if self.group_norm is not None:
|
||||
# hidden_states = self.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
|
||||
# if encoder_hidden_states is None:
|
||||
# encoder_hidden_states = hidden_states
|
||||
|
||||
# query = self.to_q(hidden_states)
|
||||
# key = self.to_k(encoder_hidden_states)
|
||||
# value = self.to_v(encoder_hidden_states)
|
||||
# inner_dim = key.shape[-1]
|
||||
# head_dim = inner_dim // self.heads
|
||||
# query = query.view(batch_size, -1, self.heads, head_dim).transpose(1, 2)
|
||||
# key = key.view(batch_size, -1, self.heads, head_dim).transpose(1, 2)
|
||||
# value = value.view(batch_size, -1, self.heads, head_dim).transpose(1, 2)
|
||||
|
||||
# hidden_states = F.scaled_dot_product_attention(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False)
|
||||
# shape = query.shape
|
||||
# mean_key = key.reshape(2,-1,shape[1],shape[2],shape[3]).mean(dim=1,keepdim=True).expand(-1,shape[0]//2,-1,-1,-1).reshape(shape[0],shape[1],shape[2],shape[3])
|
||||
# mean_value = value.reshape(2,-1,shape[1],shape[2],shape[3]).mean(dim=1,keepdim=True).expand(-1,shape[0]//2,-1,-1,-1).reshape(shape[0],shape[1],shape[2],shape[3])
|
||||
# hidden_states_mean = F.scaled_dot_product_attention(query, mean_key, mean_value, attn_mask=None, dropout_p=0.0, is_causal=False)
|
||||
|
||||
# hidden_states = (1-gamma)*hidden_states + gamma*hidden_states_mean
|
||||
# hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, self.heads * head_dim)
|
||||
# hidden_states = hidden_states.to(query.dtype)
|
||||
|
||||
# hidden_states = self.to_out[0](hidden_states)
|
||||
# hidden_states = self.to_out[1](hidden_states)
|
||||
|
||||
# if input_ndim == 4:
|
||||
# hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
|
||||
|
||||
# if self.residual_connection:
|
||||
# hidden_states = hidden_states + residual
|
||||
|
||||
# hidden_states = hidden_states / self.rescale_output_factor
|
||||
|
||||
|
||||
# return hidden_states
|
||||
|
||||
# ### attention
|
||||
# from types import MethodType
|
||||
# @torch.inference_mode()
|
||||
# def prep_unet_self_attention(unet):
|
||||
|
||||
# for name, module in unet.named_modules():
|
||||
# module_name = type(module).__name__
|
||||
|
||||
# name_split_list = name.split(".")
|
||||
# cond_1 = name_split_list[0] in "up_blocks"
|
||||
# cond_2 = name_split_list[-1] in ('attn1')
|
||||
|
||||
# if "Attention" in module_name and cond_1 and cond_2:
|
||||
# cond_3 = name_split_list[1]
|
||||
# if cond_3 not in "3":
|
||||
# module.forward = MethodType(custom_forward_CLA, module)
|
||||
|
||||
# return unet
|
||||
|
||||
# ## attn module
|
||||
# unet = prep_unet_self_attention(unet)
|
||||
|
||||
# ## ic-light-scheduler
|
||||
# ic_light_scheduler = DPMSolverMultistepScheduler(
|
||||
# num_train_timesteps=1000,
|
||||
# beta_start=0.00085,
|
||||
# beta_end=0.012,
|
||||
# algorithm_type="sde-dpmsolver++",
|
||||
# use_karras_sigmas=True,
|
||||
# steps_offset=1
|
||||
# )
|
||||
# ic_light_pipe = StableDiffusionImg2ImgPipeline(
|
||||
# vae=vae,
|
||||
# text_encoder=text_encoder,
|
||||
# tokenizer=tokenizer,
|
||||
# unet=unet,
|
||||
# scheduler=ic_light_scheduler,
|
||||
# safety_checker=None,
|
||||
# requires_safety_checker=False,
|
||||
# feature_extractor=None,
|
||||
# image_encoder=None
|
||||
# )
|
||||
# ic_light_pipe = ic_light_pipe.to(device=device, dtype=adopted_dtype)
|
||||
# ic_light_pipe.vae.requires_grad_(False)
|
||||
# ic_light_pipe.unet.requires_grad_(False)
|
||||
|
||||
############################# params ######################################
|
||||
# strength = config.get("strength", 0.4)
|
||||
# num_step = config.get("num_step", 25)
|
||||
# text_guide_scale = config.get("text_guide_scale", 2)
|
||||
# seed = config.get("seed")
|
||||
# image_width = config.get("width", 512)
|
||||
# image_height = config.get("height", 512)
|
||||
# negative_prompt = config.get("n_prompt", "")
|
||||
# vdm_prompt = config.get("vdm_prompt", "")
|
||||
# relight_prompt = config.get("relight_prompt", "")
|
||||
# video_path = config.get("video_path", "")
|
||||
# bg_source = BGSource[config.get("bg_source")]
|
||||
# save_path = config.get("save_path")
|
||||
# num_frames = config.get("num_frames", 49)
|
||||
|
||||
# ############################## infer #####################################
|
||||
# generator = torch.manual_seed(seed)
|
||||
# video_name = os.path.basename(video_path)
|
||||
# video_list, video_name = read_video(video_path, image_width, image_height)
|
||||
|
||||
# print("################## begin ##################")
|
||||
# with torch.no_grad():
|
||||
# num_inference_steps = int(round(num_step / strength))
|
||||
|
||||
# output = pipe(
|
||||
# ic_light_pipe=ic_light_pipe,
|
||||
# relight_prompt=relight_prompt,
|
||||
# bg_source=bg_source,
|
||||
# video=video_list,
|
||||
# prompt=vdm_prompt,
|
||||
# negative_prompt=negative_prompt,
|
||||
# strength=strength,
|
||||
# guidance_scale=text_guide_scale,
|
||||
# num_inference_steps=num_inference_steps,
|
||||
# height=image_height,
|
||||
# num_frames=num_frames,
|
||||
# width=image_width,
|
||||
# generator=generator,
|
||||
# )
|
||||
|
||||
# frames = output.frames[0]
|
||||
# frames = (frames * 255).astype(np.uint8)
|
||||
# results_path = f"{save_path}/relight_{video_name}"
|
||||
# imageio.mimwrite(results_path, frames, fps=14)
|
||||
# print(f"relight! prompt:{relight_prompt}, light:{bg_source.value}, save in {results_path}.")
|
||||
|
||||
# if __name__ == "__main__":
|
||||
# parser = argparse.ArgumentParser()
|
||||
|
||||
# parser.add_argument("--sd_model", type=str, default="stablediffusionapi/realistic-vision-v51")
|
||||
# parser.add_argument("--vdm_model", type=str, default="Wan-AI/Wan2.1-T2V-1.3B-Diffusers")
|
||||
# parser.add_argument("--ic_light_model", type=str, default="./models/iclight_sd15_fc.safetensors")
|
||||
|
||||
# parser.add_argument("--config", type=str, default="configs/wan_relight/man.yaml", help="the config file for each sample.")
|
||||
|
||||
# args = parser.parse_args()
|
||||
# main(args)
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
+21
-4
@@ -38,14 +38,14 @@ from diffusers.pipelines.free_init_utils import FreeInitMixin
|
||||
from diffusers.pipelines.free_noise_utils import AnimateDiffFreeNoiseMixin
|
||||
from diffusers.pipelines.pipeline_utils import DiffusionPipeline, StableDiffusionMixin
|
||||
from diffusers.pipelines.animatediff.pipeline_output import AnimateDiffPipelineOutput
|
||||
|
||||
from packaging import version
|
||||
from .ic_light import Relighter
|
||||
from einops import rearrange
|
||||
from diffusers.utils import export_to_gif
|
||||
from .animatediff_eul import eul_step
|
||||
import math
|
||||
from .tools import vis_video
|
||||
|
||||
import diffusers
|
||||
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
|
||||
|
||||
EXAMPLE_DOC_STRING = """
|
||||
@@ -977,6 +977,9 @@ class AnimateDiffVideoToVideoPipeline(
|
||||
else None
|
||||
)
|
||||
|
||||
|
||||
required_version = "0.32.1"
|
||||
|
||||
num_free_init_iters = self._free_init_num_iters if self.free_init_enabled else 1
|
||||
for free_init_iter in range(num_free_init_iters):
|
||||
if self.free_init_enabled:
|
||||
@@ -989,7 +992,8 @@ class AnimateDiffVideoToVideoPipeline(
|
||||
|
||||
self._num_timesteps = len(timesteps)
|
||||
num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
|
||||
|
||||
batch_size_ = 2 # 低版本的diffusers仅支持2个批次输入
|
||||
j=0
|
||||
# 8. Denoising loop
|
||||
with self.progress_bar(total=self._num_timesteps) as progress_bar:
|
||||
for i, t in enumerate(timesteps):
|
||||
@@ -997,11 +1001,24 @@ class AnimateDiffVideoToVideoPipeline(
|
||||
latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents
|
||||
latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
|
||||
|
||||
if version.parse(diffusers.__version__) < version.parse(required_version):
|
||||
start_idx = (j * batch_size_)
|
||||
end_idx = start_idx + batch_size_
|
||||
if end_idx > prompt_embeds.shape[0]:
|
||||
# 如果end_idx超出边界，从头开始截取
|
||||
prompt_embeds_batch = prompt_embeds[start_idx:]
|
||||
j = 0
|
||||
else:
|
||||
prompt_embeds_batch = prompt_embeds[start_idx:end_idx]
|
||||
j += 1
|
||||
else:
|
||||
prompt_embeds_batch = prompt_embeds
|
||||
|
||||
# predict the noise residual
|
||||
noise_pred = self.unet(
|
||||
latent_model_input,
|
||||
t,
|
||||
encoder_hidden_states=prompt_embeds,
|
||||
encoder_hidden_states=prompt_embeds_batch,
|
||||
cross_attention_kwargs=self.cross_attention_kwargs,
|
||||
added_cond_kwargs=added_cond_kwargs,
|
||||
).sample
|
||||
|
||||
@@ -0,0 +1,152 @@
|
||||
# Copyright 2024 The CogVideoX team, Tsinghua University & ZhipuAI and The HuggingFace Team.
|
||||
# All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
# DISCLAIMER: This code is strongly influenced by https://github.com/pesser/pytorch_diffusion
|
||||
# and https://github.com/hojonathanho/diffusion
|
||||
|
||||
import math
|
||||
from dataclasses import dataclass
|
||||
from typing import List, Optional, Tuple, Union
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from diffusers.configuration_utils import ConfigMixin, register_to_config
|
||||
from diffusers.utils import BaseOutput
|
||||
from diffusers.schedulers.scheduling_utils import KarrasDiffusionSchedulers, SchedulerMixin
|
||||
|
||||
|
||||
@dataclass
|
||||
# Copied from diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput with DDPM->DDIM
|
||||
class DDIMSchedulerOutput(BaseOutput):
|
||||
"""
|
||||
Output class for the scheduler's `step` function output.
|
||||
|
||||
Args:
|
||||
prev_sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` for images):
|
||||
Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the
|
||||
denoising loop.
|
||||
pred_original_sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` for images):
|
||||
The predicted denoised sample `(x_{0})` based on the model output from the current timestep.
|
||||
`pred_original_sample` can be used to preview progress or for guidance.
|
||||
"""
|
||||
|
||||
prev_sample: torch.Tensor
|
||||
pred_original_sample: Optional[torch.Tensor] = None
|
||||
|
||||
|
||||
# Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar
|
||||
def cog_ddim_step(
|
||||
self,
|
||||
model_output: torch.Tensor,
|
||||
timestep: int,
|
||||
sample: torch.Tensor,
|
||||
fusion_target,
|
||||
eta: float = 0.0,
|
||||
use_clipped_model_output: bool = False,
|
||||
generator=None,
|
||||
variance_noise: Optional[torch.Tensor] = None,
|
||||
return_dict: bool = True,
|
||||
) -> Union[DDIMSchedulerOutput, Tuple]:
|
||||
"""
|
||||
Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
|
||||
process from the learned model outputs (most often the predicted noise).
|
||||
|
||||
Args:
|
||||
model_output (`torch.Tensor`):
|
||||
The direct output from learned diffusion model.
|
||||
timestep (`float`):
|
||||
The current discrete timestep in the diffusion chain.
|
||||
sample (`torch.Tensor`):
|
||||
A current instance of a sample created by the diffusion process.
|
||||
eta (`float`):
|
||||
The weight of noise for added noise in diffusion step.
|
||||
use_clipped_model_output (`bool`, defaults to `False`):
|
||||
If `True`, computes "corrected" `model_output` from the clipped predicted original sample. Necessary
|
||||
because predicted original sample is clipped to [-1, 1] when `self.config.clip_sample` is `True`. If no
|
||||
clipping has happened, "corrected" `model_output` would coincide with the one provided as input and
|
||||
`use_clipped_model_output` has no effect.
|
||||
generator (`torch.Generator`, *optional*):
|
||||
A random number generator.
|
||||
variance_noise (`torch.Tensor`):
|
||||
Alternative to generating noise with `generator` by directly providing the noise for the variance
|
||||
itself. Useful for methods such as [`CycleDiffusion`].
|
||||
return_dict (`bool`, *optional*, defaults to `True`):
|
||||
Whether or not to return a [`~schedulers.scheduling_ddim.DDIMSchedulerOutput`] or `tuple`.
|
||||
|
||||
Returns:
|
||||
[`~schedulers.scheduling_ddim.DDIMSchedulerOutput`] or `tuple`:
|
||||
If return_dict is `True`, [`~schedulers.scheduling_ddim.DDIMSchedulerOutput`] is returned, otherwise a
|
||||
tuple is returned where the first element is the sample tensor.
|
||||
|
||||
"""
|
||||
if self.num_inference_steps is None:
|
||||
raise ValueError(
|
||||
"Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler"
|
||||
)
|
||||
|
||||
# See formulas (12) and (16) of DDIM paper https://arxiv.org/pdf/2010.02502.pdf
|
||||
# Ideally, read DDIM paper in-detail understanding
|
||||
|
||||
# Notation (<variable name> -> <name in paper>
|
||||
# - pred_noise_t -> e_theta(x_t, t)
|
||||
# - pred_original_sample -> f_theta(x_t, t) or x_0
|
||||
# - std_dev_t -> sigma_t
|
||||
# - eta -> η
|
||||
# - pred_sample_direction -> "direction pointing to x_t"
|
||||
# - pred_prev_sample -> "x_t-1"
|
||||
|
||||
# 1. get previous step value (=t-1)
|
||||
prev_timestep = timestep - self.config.num_train_timesteps // self.num_inference_steps
|
||||
|
||||
# 2. compute alphas, betas
|
||||
alpha_prod_t = self.alphas_cumprod[timestep]
|
||||
alpha_prod_t_prev = self.alphas_cumprod[prev_timestep] if prev_timestep >= 0 else self.final_alpha_cumprod
|
||||
|
||||
beta_prod_t = 1 - alpha_prod_t
|
||||
|
||||
# 3. compute predicted original sample from predicted noise also called
|
||||
# "predicted x_0" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
|
||||
# To make style tests pass, commented out `pred_epsilon` as it is an unused variable
|
||||
if self.config.prediction_type == "epsilon":
|
||||
pred_original_sample = (sample - beta_prod_t ** (0.5) * model_output) / alpha_prod_t ** (0.5)
|
||||
# pred_epsilon = model_output
|
||||
elif self.config.prediction_type == "sample":
|
||||
pred_original_sample = model_output
|
||||
# pred_epsilon = (sample - alpha_prod_t ** (0.5) * pred_original_sample) / beta_prod_t ** (0.5)
|
||||
elif self.config.prediction_type == "v_prediction":
|
||||
pred_original_sample = (alpha_prod_t**0.5) * sample - (beta_prod_t**0.5) * model_output
|
||||
# pred_epsilon = (alpha_prod_t**0.5) * model_output + (beta_prod_t**0.5) * sample
|
||||
else:
|
||||
raise ValueError(
|
||||
f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`, or"
|
||||
" `v_prediction`"
|
||||
)
|
||||
|
||||
## insert fusion target
|
||||
pred_original_sample = fusion_target
|
||||
|
||||
a_t = ((1 - alpha_prod_t_prev) / (1 - alpha_prod_t)) ** 0.5
|
||||
b_t = alpha_prod_t_prev**0.5 - alpha_prod_t**0.5 * a_t
|
||||
|
||||
prev_sample = a_t * sample + b_t * pred_original_sample
|
||||
|
||||
if not return_dict:
|
||||
return (
|
||||
prev_sample,
|
||||
pred_original_sample,
|
||||
)
|
||||
|
||||
return DDIMSchedulerOutput(prev_sample=prev_sample, pred_original_sample=pred_original_sample)
|
||||
@@ -0,0 +1,846 @@
|
||||
# Copyright 2024 The CogVideoX team, Tsinghua University & ZhipuAI and The HuggingFace Team.
|
||||
# All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import inspect
|
||||
import math
|
||||
from typing import Any, Callable, Dict, List, Optional, Tuple, Union
|
||||
|
||||
import torch
|
||||
from PIL import Image
|
||||
from transformers import T5EncoderModel, T5Tokenizer
|
||||
|
||||
from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback
|
||||
from diffusers.loaders import CogVideoXLoraLoaderMixin
|
||||
from diffusers.models import AutoencoderKLCogVideoX, CogVideoXTransformer3DModel
|
||||
from diffusers.models.embeddings import get_3d_rotary_pos_embed
|
||||
from diffusers.pipelines.pipeline_utils import DiffusionPipeline
|
||||
from diffusers.schedulers import CogVideoXDDIMScheduler, CogVideoXDPMScheduler
|
||||
from diffusers.utils import logging, replace_example_docstring
|
||||
from diffusers.utils.torch_utils import randn_tensor
|
||||
from diffusers.video_processor import VideoProcessor
|
||||
from diffusers.pipelines.cogvideo.pipeline_output import CogVideoXPipelineOutput
|
||||
|
||||
from .cogvideo_ddim import cog_ddim_step
|
||||
from .ic_light import Relighter
|
||||
from einops import rearrange
|
||||
from diffusers.utils import export_to_gif
|
||||
import math
|
||||
|
||||
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
|
||||
|
||||
|
||||
EXAMPLE_DOC_STRING = """
|
||||
Examples:
|
||||
```python
|
||||
>>> import torch
|
||||
>>> from diffusers import CogVideoXDPMScheduler, CogVideoXVideoToVideoPipeline
|
||||
>>> from diffusers.utils import export_to_video, load_video
|
||||
|
||||
>>> # Models: "THUDM/CogVideoX-2b" or "THUDM/CogVideoX-5b"
|
||||
>>> pipe = CogVideoXVideoToVideoPipeline.from_pretrained("THUDM/CogVideoX-5b", torch_dtype=torch.bfloat16)
|
||||
>>> pipe.to("cuda")
|
||||
>>> pipe.scheduler = CogVideoXDPMScheduler.from_config(pipe.scheduler.config)
|
||||
|
||||
>>> input_video = load_video(
|
||||
... "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/hiker.mp4"
|
||||
... )
|
||||
>>> prompt = (
|
||||
... "An astronaut stands triumphantly at the peak of a towering mountain. Panorama of rugged peaks and "
|
||||
... "valleys. Very futuristic vibe and animated aesthetic. Highlights of purple and golden colors in "
|
||||
... "the scene. The sky is looks like an animated/cartoonish dream of galaxies, nebulae, stars, planets, "
|
||||
... "moons, but the remainder of the scene is mostly realistic."
|
||||
... )
|
||||
|
||||
>>> video = pipe(
|
||||
... video=input_video, prompt=prompt, strength=0.8, guidance_scale=6, num_inference_steps=50
|
||||
... ).frames[0]
|
||||
>>> export_to_video(video, "output.mp4", fps=8)
|
||||
```
|
||||
"""
|
||||
|
||||
|
||||
# Similar to diffusers.pipelines.hunyuandit.pipeline_hunyuandit.get_resize_crop_region_for_grid
|
||||
def get_resize_crop_region_for_grid(src, tgt_width, tgt_height):
|
||||
tw = tgt_width
|
||||
th = tgt_height
|
||||
h, w = src
|
||||
r = h / w
|
||||
if r > (th / tw):
|
||||
resize_height = th
|
||||
resize_width = int(round(th / h * w))
|
||||
else:
|
||||
resize_width = tw
|
||||
resize_height = int(round(tw / w * h))
|
||||
|
||||
crop_top = int(round((th - resize_height) / 2.0))
|
||||
crop_left = int(round((tw - resize_width) / 2.0))
|
||||
|
||||
return (crop_top, crop_left), (crop_top + resize_height, crop_left + resize_width)
|
||||
|
||||
|
||||
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
|
||||
def retrieve_timesteps(
|
||||
scheduler,
|
||||
num_inference_steps: Optional[int] = None,
|
||||
device: Optional[Union[str, torch.device]] = None,
|
||||
timesteps: Optional[List[int]] = None,
|
||||
sigmas: Optional[List[float]] = None,
|
||||
**kwargs,
|
||||
):
|
||||
r"""
|
||||
Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
|
||||
custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
|
||||
|
||||
Args:
|
||||
scheduler (`SchedulerMixin`):
|
||||
The scheduler to get timesteps from.
|
||||
num_inference_steps (`int`):
|
||||
The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
|
||||
must be `None`.
|
||||
device (`str` or `torch.device`, *optional*):
|
||||
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
|
||||
timesteps (`List[int]`, *optional*):
|
||||
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
|
||||
`num_inference_steps` and `sigmas` must be `None`.
|
||||
sigmas (`List[float]`, *optional*):
|
||||
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
|
||||
`num_inference_steps` and `timesteps` must be `None`.
|
||||
|
||||
Returns:
|
||||
`Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
|
||||
second element is the number of inference steps.
|
||||
"""
|
||||
if timesteps is not None and sigmas is not None:
|
||||
raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
|
||||
if timesteps is not None:
|
||||
accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
|
||||
if not accepts_timesteps:
|
||||
raise ValueError(
|
||||
f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
|
||||
f" timestep schedules. Please check whether you are using the correct scheduler."
|
||||
)
|
||||
scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
|
||||
timesteps = scheduler.timesteps
|
||||
num_inference_steps = len(timesteps)
|
||||
elif sigmas is not None:
|
||||
accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
|
||||
if not accept_sigmas:
|
||||
raise ValueError(
|
||||
f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
|
||||
f" sigmas schedules. Please check whether you are using the correct scheduler."
|
||||
)
|
||||
scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
|
||||
timesteps = scheduler.timesteps
|
||||
num_inference_steps = len(timesteps)
|
||||
else:
|
||||
scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
|
||||
timesteps = scheduler.timesteps
|
||||
return timesteps, num_inference_steps
|
||||
|
||||
|
||||
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
|
||||
def retrieve_latents(
|
||||
encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
|
||||
):
|
||||
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
|
||||
return encoder_output.latent_dist.sample(generator)
|
||||
elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax":
|
||||
return encoder_output.latent_dist.mode()
|
||||
elif hasattr(encoder_output, "latents"):
|
||||
return encoder_output.latents
|
||||
else:
|
||||
raise AttributeError("Could not access latents of provided encoder_output")
|
||||
|
||||
|
||||
class CogVideoXVideoToVideoPipeline(DiffusionPipeline, CogVideoXLoraLoaderMixin):
|
||||
r"""
|
||||
Pipeline for video-to-video generation using CogVideoX.
|
||||
|
||||
This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the
|
||||
library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)
|
||||
|
||||
Args:
|
||||
vae ([`AutoencoderKL`]):
|
||||
Variational Auto-Encoder (VAE) Model to encode and decode videos to and from latent representations.
|
||||
text_encoder ([`T5EncoderModel`]):
|
||||
Frozen text-encoder. CogVideoX uses
|
||||
[T5](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5EncoderModel); specifically the
|
||||
[t5-v1_1-xxl](https://huggingface.co/PixArt-alpha/PixArt-alpha/tree/main/t5-v1_1-xxl) variant.
|
||||
tokenizer (`T5Tokenizer`):
|
||||
Tokenizer of class
|
||||
[T5Tokenizer](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5Tokenizer).
|
||||
transformer ([`CogVideoXTransformer3DModel`]):
|
||||
A text conditioned `CogVideoXTransformer3DModel` to denoise the encoded video latents.
|
||||
scheduler ([`SchedulerMixin`]):
|
||||
A scheduler to be used in combination with `transformer` to denoise the encoded video latents.
|
||||
"""
|
||||
|
||||
_optional_components = []
|
||||
model_cpu_offload_seq = "text_encoder->transformer->vae"
|
||||
|
||||
_callback_tensor_inputs = [
|
||||
"latents",
|
||||
"prompt_embeds",
|
||||
"negative_prompt_embeds",
|
||||
]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
tokenizer: T5Tokenizer,
|
||||
text_encoder: T5EncoderModel,
|
||||
vae: AutoencoderKLCogVideoX,
|
||||
transformer: CogVideoXTransformer3DModel,
|
||||
scheduler: Union[CogVideoXDDIMScheduler, CogVideoXDPMScheduler],
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.register_modules(
|
||||
tokenizer=tokenizer, text_encoder=text_encoder, vae=vae, transformer=transformer, scheduler=scheduler
|
||||
)
|
||||
|
||||
self.vae_scale_factor_spatial = (
|
||||
2 ** (len(self.vae.config.block_out_channels) - 1) if hasattr(self, "vae") and self.vae is not None else 8
|
||||
)
|
||||
self.vae_scale_factor_temporal = (
|
||||
self.vae.config.temporal_compression_ratio if hasattr(self, "vae") and self.vae is not None else 4
|
||||
)
|
||||
self.vae_scaling_factor_image = (
|
||||
self.vae.config.scaling_factor if hasattr(self, "vae") and self.vae is not None else 0.7
|
||||
)
|
||||
|
||||
self.video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial)
|
||||
|
||||
# Copied from diffusers.pipelines.cogvideo.pipeline_cogvideox.CogVideoXPipeline._get_t5_prompt_embeds
|
||||
def _get_t5_prompt_embeds(
|
||||
self,
|
||||
prompt: Union[str, List[str]] = None,
|
||||
num_videos_per_prompt: int = 1,
|
||||
max_sequence_length: int = 226,
|
||||
device: Optional[torch.device] = None,
|
||||
dtype: Optional[torch.dtype] = None,
|
||||
):
|
||||
device = device or self._execution_device
|
||||
dtype = dtype or self.text_encoder.dtype
|
||||
|
||||
prompt = [prompt] if isinstance(prompt, str) else prompt
|
||||
batch_size = len(prompt)
|
||||
|
||||
text_inputs = self.tokenizer(
|
||||
prompt,
|
||||
padding="max_length",
|
||||
max_length=max_sequence_length,
|
||||
truncation=True,
|
||||
add_special_tokens=True,
|
||||
return_tensors="pt",
|
||||
)
|
||||
text_input_ids = text_inputs.input_ids
|
||||
untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids
|
||||
|
||||
if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids):
|
||||
removed_text = self.tokenizer.batch_decode(untruncated_ids[:, max_sequence_length - 1 : -1])
|
||||
logger.warning(
|
||||
"The following part of your input was truncated because `max_sequence_length` is set to "
|
||||
f" {max_sequence_length} tokens: {removed_text}"
|
||||
)
|
||||
|
||||
prompt_embeds = self.text_encoder(text_input_ids.to(device))[0]
|
||||
prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
|
||||
|
||||
# duplicate text embeddings for each generation per prompt, using mps friendly method
|
||||
_, seq_len, _ = prompt_embeds.shape
|
||||
prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1)
|
||||
prompt_embeds = prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1)
|
||||
|
||||
return prompt_embeds
|
||||
|
||||
# Copied from diffusers.pipelines.cogvideo.pipeline_cogvideox.CogVideoXPipeline.encode_prompt
|
||||
def encode_prompt(
|
||||
self,
|
||||
prompt: Union[str, List[str]],
|
||||
negative_prompt: Optional[Union[str, List[str]]] = None,
|
||||
do_classifier_free_guidance: bool = True,
|
||||
num_videos_per_prompt: int = 1,
|
||||
prompt_embeds: Optional[torch.Tensor] = None,
|
||||
negative_prompt_embeds: Optional[torch.Tensor] = None,
|
||||
max_sequence_length: int = 226,
|
||||
device: Optional[torch.device] = None,
|
||||
dtype: Optional[torch.dtype] = None,
|
||||
):
|
||||
r"""
|
||||
Encodes the prompt into text encoder hidden states.
|
||||
|
||||
Args:
|
||||
prompt (`str` or `List[str]`, *optional*):
|
||||
prompt to be encoded
|
||||
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`).
|
||||
do_classifier_free_guidance (`bool`, *optional*, defaults to `True`):
|
||||
Whether to use classifier free guidance or not.
|
||||
num_videos_per_prompt (`int`, *optional*, defaults to 1):
|
||||
Number of videos that should be generated per prompt. torch device to place the resulting embeddings on
|
||||
prompt_embeds (`torch.Tensor`, *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.Tensor`, *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.
|
||||
device: (`torch.device`, *optional*):
|
||||
torch device
|
||||
dtype: (`torch.dtype`, *optional*):
|
||||
torch dtype
|
||||
"""
|
||||
device = device or self._execution_device
|
||||
|
||||
prompt = [prompt] if isinstance(prompt, str) else prompt
|
||||
if prompt is not None:
|
||||
batch_size = len(prompt)
|
||||
else:
|
||||
batch_size = prompt_embeds.shape[0]
|
||||
|
||||
if prompt_embeds is None:
|
||||
prompt_embeds = self._get_t5_prompt_embeds(
|
||||
prompt=prompt,
|
||||
num_videos_per_prompt=num_videos_per_prompt,
|
||||
max_sequence_length=max_sequence_length,
|
||||
device=device,
|
||||
dtype=dtype,
|
||||
)
|
||||
|
||||
if do_classifier_free_guidance and negative_prompt_embeds is None:
|
||||
negative_prompt = negative_prompt or ""
|
||||
negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt
|
||||
|
||||
if prompt is not None and type(prompt) is not type(negative_prompt):
|
||||
raise TypeError(
|
||||
f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
|
||||
f" {type(prompt)}."
|
||||
)
|
||||
elif batch_size != len(negative_prompt):
|
||||
raise ValueError(
|
||||
f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
|
||||
f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
|
||||
" the batch size of `prompt`."
|
||||
)
|
||||
|
||||
negative_prompt_embeds = self._get_t5_prompt_embeds(
|
||||
prompt=negative_prompt,
|
||||
num_videos_per_prompt=num_videos_per_prompt,
|
||||
max_sequence_length=max_sequence_length,
|
||||
device=device,
|
||||
dtype=dtype,
|
||||
)
|
||||
|
||||
return prompt_embeds, negative_prompt_embeds
|
||||
|
||||
def prepare_latents(
|
||||
self,
|
||||
video: Optional[torch.Tensor] = None, ## 输入 torch.Size([1, 3, 49, 480, 720])
|
||||
batch_size: int = 1,
|
||||
num_channels_latents: int = 16,
|
||||
height: int = 60,
|
||||
width: int = 90,
|
||||
dtype: Optional[torch.dtype] = None,
|
||||
device: Optional[torch.device] = None,
|
||||
generator: Optional[torch.Generator] = None,
|
||||
latents: Optional[torch.Tensor] = None,
|
||||
timestep: Optional[torch.Tensor] = None,
|
||||
):
|
||||
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."
|
||||
)
|
||||
|
||||
num_frames = (video.size(2) - 1) // self.vae_scale_factor_temporal + 1 if latents is None else latents.size(1)
|
||||
|
||||
shape = ( ## (1, 13, 16, 60, 90)
|
||||
batch_size,
|
||||
num_frames,
|
||||
num_channels_latents,
|
||||
height // self.vae_scale_factor_spatial, ## 8倍降采样
|
||||
width // self.vae_scale_factor_spatial,
|
||||
)
|
||||
|
||||
## true
|
||||
if latents is None:
|
||||
if isinstance(generator, list):
|
||||
if 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."
|
||||
)
|
||||
|
||||
init_latents = [
|
||||
retrieve_latents(self.vae.encode(video[i].unsqueeze(0)), generator[i]) for i in range(batch_size)
|
||||
]
|
||||
else:
|
||||
init_latents = [retrieve_latents(self.vae.encode(vid.unsqueeze(0)), generator) for vid in video]
|
||||
|
||||
init_latents = torch.cat(init_latents, dim=0).to(dtype).permute(0, 2, 1, 3, 4)
|
||||
init_latents = self.vae_scaling_factor_image * init_latents
|
||||
noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
|
||||
latents = self.scheduler.add_noise(init_latents, noise, timestep)
|
||||
else:
|
||||
latents = latents.to(device)
|
||||
|
||||
# scale the initial noise by the standard deviation required by the scheduler
|
||||
latents = latents * self.scheduler.init_noise_sigma
|
||||
return latents
|
||||
|
||||
# Copied from diffusers.pipelines.cogvideo.pipeline_cogvideox.CogVideoXPipeline.decode_latents
|
||||
def decode_latents(self, latents: torch.Tensor) -> torch.Tensor:
|
||||
latents = latents.permute(0, 2, 1, 3, 4) # [batch_size, num_channels, num_frames, height, width]
|
||||
latents = 1 / self.vae_scaling_factor_image * latents
|
||||
|
||||
frames = self.vae.decode(latents).sample
|
||||
return frames
|
||||
|
||||
# Copied from diffusers.pipelines.animatediff.pipeline_animatediff_video2video.AnimateDiffVideoToVideoPipeline.get_timesteps
|
||||
def get_timesteps(self, num_inference_steps, timesteps, strength, device):
|
||||
# get the original timestep using init_timestep
|
||||
init_timestep = min(int(num_inference_steps * strength), num_inference_steps)
|
||||
|
||||
t_start = max(num_inference_steps - init_timestep, 0)
|
||||
timesteps = timesteps[t_start * self.scheduler.order :]
|
||||
|
||||
return timesteps, num_inference_steps - t_start
|
||||
|
||||
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs
|
||||
def prepare_extra_step_kwargs(self, generator, eta):
|
||||
# prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
|
||||
# eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
|
||||
# eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
|
||||
# and should be between [0, 1]
|
||||
|
||||
accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())
|
||||
extra_step_kwargs = {}
|
||||
if accepts_eta:
|
||||
extra_step_kwargs["eta"] = eta
|
||||
|
||||
# check if the scheduler accepts generator
|
||||
accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())
|
||||
if accepts_generator:
|
||||
extra_step_kwargs["generator"] = generator
|
||||
return extra_step_kwargs
|
||||
|
||||
def check_inputs(
|
||||
self,
|
||||
prompt,
|
||||
height,
|
||||
width,
|
||||
strength,
|
||||
negative_prompt,
|
||||
callback_on_step_end_tensor_inputs,
|
||||
video=None,
|
||||
latents=None,
|
||||
prompt_embeds=None,
|
||||
negative_prompt_embeds=None,
|
||||
):
|
||||
if height % 8 != 0 or width % 8 != 0:
|
||||
raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")
|
||||
|
||||
if strength < 0 or strength > 1:
|
||||
raise ValueError(f"The value of strength should in [0.0, 1.0] but is {strength}")
|
||||
|
||||
if callback_on_step_end_tensor_inputs is not None and not all(
|
||||
k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs
|
||||
):
|
||||
raise ValueError(
|
||||
f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}"
|
||||
)
|
||||
if prompt is not None and prompt_embeds is not None:
|
||||
raise ValueError(
|
||||
f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
|
||||
" only forward one of the two."
|
||||
)
|
||||
elif prompt is None and prompt_embeds is None:
|
||||
raise ValueError(
|
||||
"Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
|
||||
)
|
||||
elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
|
||||
raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
|
||||
|
||||
if prompt is not None and negative_prompt_embeds is not None:
|
||||
raise ValueError(
|
||||
f"Cannot forward both `prompt`: {prompt} and `negative_prompt_embeds`:"
|
||||
f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
|
||||
)
|
||||
|
||||
if negative_prompt is not None and negative_prompt_embeds is not None:
|
||||
raise ValueError(
|
||||
f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"
|
||||
f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
|
||||
)
|
||||
|
||||
if prompt_embeds is not None and negative_prompt_embeds is not None:
|
||||
if prompt_embeds.shape != negative_prompt_embeds.shape:
|
||||
raise ValueError(
|
||||
"`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"
|
||||
f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"
|
||||
f" {negative_prompt_embeds.shape}."
|
||||
)
|
||||
|
||||
if video is not None and latents is not None:
|
||||
raise ValueError("Only one of `video` or `latents` should be provided")
|
||||
|
||||
# Copied from diffusers.pipelines.cogvideo.pipeline_cogvideox.CogVideoXPipeline.fuse_qkv_projections
|
||||
def fuse_qkv_projections(self) -> None:
|
||||
r"""Enables fused QKV projections."""
|
||||
self.fusing_transformer = True
|
||||
self.transformer.fuse_qkv_projections()
|
||||
|
||||
# Copied from diffusers.pipelines.cogvideo.pipeline_cogvideox.CogVideoXPipeline.unfuse_qkv_projections
|
||||
def unfuse_qkv_projections(self) -> None:
|
||||
r"""Disable QKV projection fusion if enabled."""
|
||||
if not self.fusing_transformer:
|
||||
logger.warning("The Transformer was not initially fused for QKV projections. Doing nothing.")
|
||||
else:
|
||||
self.transformer.unfuse_qkv_projections()
|
||||
self.fusing_transformer = False
|
||||
|
||||
# Copied from diffusers.pipelines.cogvideo.pipeline_cogvideox.CogVideoXPipeline._prepare_rotary_positional_embeddings
|
||||
def _prepare_rotary_positional_embeddings(
|
||||
self,
|
||||
height: int,
|
||||
width: int,
|
||||
num_frames: int,
|
||||
device: torch.device,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
grid_height = height // (self.vae_scale_factor_spatial * self.transformer.config.patch_size)
|
||||
grid_width = width // (self.vae_scale_factor_spatial * self.transformer.config.patch_size)
|
||||
base_size_width = 720 // (self.vae_scale_factor_spatial * self.transformer.config.patch_size)
|
||||
base_size_height = 480 // (self.vae_scale_factor_spatial * self.transformer.config.patch_size)
|
||||
|
||||
grid_crops_coords = get_resize_crop_region_for_grid(
|
||||
(grid_height, grid_width), base_size_width, base_size_height
|
||||
)
|
||||
freqs_cos, freqs_sin = get_3d_rotary_pos_embed(
|
||||
embed_dim=self.transformer.config.attention_head_dim,
|
||||
crops_coords=grid_crops_coords,
|
||||
grid_size=(grid_height, grid_width),
|
||||
temporal_size=num_frames,
|
||||
)
|
||||
|
||||
freqs_cos = freqs_cos.to(device=device)
|
||||
freqs_sin = freqs_sin.to(device=device)
|
||||
return freqs_cos, freqs_sin
|
||||
|
||||
@property
|
||||
def guidance_scale(self):
|
||||
return self._guidance_scale
|
||||
|
||||
@property
|
||||
def num_timesteps(self):
|
||||
return self._num_timesteps
|
||||
|
||||
@property
|
||||
def attention_kwargs(self):
|
||||
return self._attention_kwargs
|
||||
|
||||
@property
|
||||
def interrupt(self):
|
||||
return self._interrupt
|
||||
|
||||
@torch.no_grad()
|
||||
@replace_example_docstring(EXAMPLE_DOC_STRING)
|
||||
def __call__(
|
||||
self,
|
||||
ic_light_pipe=None,
|
||||
relight_prompt=None,
|
||||
bg_source=None,
|
||||
video: List[Image.Image] = None,
|
||||
prompt: Optional[Union[str, List[str]]] = None,
|
||||
negative_prompt: Optional[Union[str, List[str]]] = None,
|
||||
height: int = 480,
|
||||
width: int = 720,
|
||||
num_inference_steps: int = 50,
|
||||
timesteps: Optional[List[int]] = None,
|
||||
strength: float = 0.8,
|
||||
guidance_scale: float = 6,
|
||||
use_dynamic_cfg: bool = False,
|
||||
num_videos_per_prompt: int = 1,
|
||||
eta: float = 0.0,
|
||||
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
|
||||
latents: Optional[torch.FloatTensor] = None,
|
||||
prompt_embeds: Optional[torch.FloatTensor] = None,
|
||||
negative_prompt_embeds: Optional[torch.FloatTensor] = None,
|
||||
output_type: str = "pil",
|
||||
return_dict: bool = True,
|
||||
attention_kwargs: Optional[Dict[str, Any]] = None,
|
||||
callback_on_step_end: Optional[
|
||||
Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
|
||||
] = None,
|
||||
callback_on_step_end_tensor_inputs: List[str] = ["latents"],
|
||||
max_sequence_length: int = 226,
|
||||
) -> Union[CogVideoXPipelineOutput, Tuple]:
|
||||
"""
|
||||
Function invoked when calling the pipeline for generation.
|
||||
|
||||
Args:
|
||||
video (`List[PIL.Image.Image]`):
|
||||
The input video to condition the generation on. Must be a list of images/frames of the video.
|
||||
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.
|
||||
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`).
|
||||
height (`int`, *optional*, defaults to self.transformer.config.sample_height * self.vae_scale_factor_spatial):
|
||||
The height in pixels of the generated image. This is set to 480 by default for the best results.
|
||||
width (`int`, *optional*, defaults to self.transformer.config.sample_height * self.vae_scale_factor_spatial):
|
||||
The width in pixels of the generated image. This is set to 720 by default for the best results.
|
||||
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.
|
||||
timesteps (`List[int]`, *optional*):
|
||||
Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
|
||||
in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
|
||||
passed will be used. Must be in descending order.
|
||||
strength (`float`, *optional*, defaults to 0.8):
|
||||
Higher strength leads to more differences between original video and generated video.
|
||||
guidance_scale (`float`, *optional*, defaults to 7.0):
|
||||
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.
|
||||
num_videos_per_prompt (`int`, *optional*, defaults to 1):
|
||||
The number of videos to generate per prompt.
|
||||
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_xl.StableDiffusionXLPipelineOutput`] instead
|
||||
of a plain tuple.
|
||||
attention_kwargs (`dict`, *optional*):
|
||||
A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
|
||||
`self.processor` in
|
||||
[diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
|
||||
callback_on_step_end (`Callable`, *optional*):
|
||||
A function that calls at the end of each denoising steps during the inference. The function is called
|
||||
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
|
||||
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
|
||||
`callback_on_step_end_tensor_inputs`.
|
||||
callback_on_step_end_tensor_inputs (`List`, *optional*):
|
||||
The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
|
||||
will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
|
||||
`._callback_tensor_inputs` attribute of your pipeline class.
|
||||
max_sequence_length (`int`, defaults to `226`):
|
||||
Maximum sequence length in encoded prompt. Must be consistent with
|
||||
`self.transformer.config.max_text_seq_length` otherwise may lead to poor results.
|
||||
|
||||
Examples:
|
||||
|
||||
Returns:
|
||||
[`~pipelines.cogvideo.pipeline_output.CogVideoXPipelineOutput`] or `tuple`:
|
||||
[`~pipelines.cogvideo.pipeline_output.CogVideoXPipelineOutput`] if `return_dict` is True, otherwise a
|
||||
`tuple`. When returning a tuple, the first element is a list with the generated images.
|
||||
"""
|
||||
|
||||
if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)):
|
||||
callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs
|
||||
|
||||
num_videos_per_prompt = 1
|
||||
|
||||
self._guidance_scale = guidance_scale
|
||||
self._attention_kwargs = attention_kwargs
|
||||
self._interrupt = False
|
||||
|
||||
# 2. Default 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
|
||||
|
||||
# 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
|
||||
|
||||
# 3. Encode input prompt
|
||||
prompt_embeds, negative_prompt_embeds = self.encode_prompt(
|
||||
prompt,
|
||||
negative_prompt,
|
||||
do_classifier_free_guidance,
|
||||
num_videos_per_prompt=num_videos_per_prompt,
|
||||
prompt_embeds=prompt_embeds,
|
||||
negative_prompt_embeds=negative_prompt_embeds,
|
||||
max_sequence_length=max_sequence_length,
|
||||
device=device,
|
||||
)
|
||||
if do_classifier_free_guidance:
|
||||
prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
|
||||
|
||||
# 4. Prepare timesteps
|
||||
timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps)
|
||||
timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, timesteps, strength, device)
|
||||
latent_timestep = timesteps[:1].repeat(batch_size * num_videos_per_prompt)
|
||||
self._num_timesteps = len(timesteps)
|
||||
|
||||
# 5. Prepare latents
|
||||
if latents is None:
|
||||
video = self.video_processor.preprocess_video(video, height=height, width=width)
|
||||
video = video.to(device=device, dtype=prompt_embeds.dtype)
|
||||
org_target = video
|
||||
|
||||
latent_channels = self.transformer.config.in_channels
|
||||
latents = self.prepare_latents(
|
||||
video,
|
||||
batch_size * num_videos_per_prompt,
|
||||
latent_channels,
|
||||
height,
|
||||
width,
|
||||
prompt_embeds.dtype,
|
||||
device,
|
||||
generator,
|
||||
latents,
|
||||
latent_timestep,
|
||||
)
|
||||
|
||||
## Init Relighter
|
||||
num_frames = video.shape[2]
|
||||
relighter = Relighter(
|
||||
pipeline=ic_light_pipe,
|
||||
relight_prompt=relight_prompt,
|
||||
bg_source=bg_source,
|
||||
generator=generator,
|
||||
num_frames=num_frames,
|
||||
image_width=width,
|
||||
image_height=height,
|
||||
)
|
||||
|
||||
# 6. 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)
|
||||
|
||||
# 7. Create rotary embeds if required
|
||||
image_rotary_emb = (
|
||||
self._prepare_rotary_positional_embeddings(height, width, latents.size(1), device)
|
||||
if self.transformer.config.use_rotary_positional_embeddings
|
||||
else None
|
||||
)
|
||||
|
||||
# 8. Denoising loop
|
||||
num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
|
||||
with self.progress_bar(total=num_inference_steps) as progress_bar:
|
||||
|
||||
for i, t in enumerate(timesteps):
|
||||
|
||||
latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
|
||||
latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
|
||||
|
||||
# broadcast to batch dimension in a way that's compatible with ONNX/Core ML
|
||||
timestep = t.expand(latent_model_input.shape[0])
|
||||
|
||||
# predict noise model_output
|
||||
noise_pred = self.transformer(
|
||||
hidden_states=latent_model_input,
|
||||
encoder_hidden_states=prompt_embeds,
|
||||
timestep=timestep,
|
||||
image_rotary_emb=image_rotary_emb,
|
||||
attention_kwargs=attention_kwargs,
|
||||
return_dict=False,
|
||||
)[0]
|
||||
|
||||
if do_classifier_free_guidance:
|
||||
noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
|
||||
noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond)
|
||||
|
||||
k = 2
|
||||
lbd = 1 - (i / (num_inference_steps - 1)) ** k
|
||||
detail_lbd = 1 - (i / (num_inference_steps - 1))
|
||||
|
||||
if lbd > 0.2:
|
||||
|
||||
print(f"relight lbd = {lbd}")
|
||||
|
||||
## get pred x_0
|
||||
alpha_prod_t = self.scheduler.alphas_cumprod[t]
|
||||
beta_prod_t = 1 - alpha_prod_t
|
||||
pred_x0_latent = (alpha_prod_t**0.5) * latents - (beta_prod_t**0.5) * noise_pred
|
||||
|
||||
consist_target = self.decode_latents(pred_x0_latent)
|
||||
|
||||
## detail
|
||||
if i == 0:
|
||||
detail_diff = org_target - consist_target
|
||||
consist_target = consist_target + detail_lbd * detail_diff
|
||||
consist_target = rearrange(consist_target, "1 c f h w -> f c h w")
|
||||
|
||||
## fusion
|
||||
relight_target = relighter(consist_target)
|
||||
fusion_target = consist_target + lbd * (relight_target - consist_target)
|
||||
fusion_target = rearrange(fusion_target, "f c h w -> 1 c f h w")
|
||||
fusion_target = self.vae.encode(fusion_target).latent_dist.mode() * self.vae_scaling_factor_image
|
||||
fusion_target = fusion_target.to(fusion_target.dtype).permute(0, 2, 1, 3, 4)
|
||||
|
||||
## Progresive Light Fusion
|
||||
output = cog_ddim_step(self.scheduler,noise_pred, t, latents,fusion_target,**extra_step_kwargs, return_dict=False)
|
||||
|
||||
else:
|
||||
## org_scheduer step
|
||||
output = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)
|
||||
|
||||
latents = output[0]
|
||||
latents = latents.to(prompt_embeds.dtype)
|
||||
|
||||
# call the callback, if provided
|
||||
if callback_on_step_end is not None:
|
||||
callback_kwargs = {}
|
||||
for k in callback_on_step_end_tensor_inputs:
|
||||
callback_kwargs[k] = locals()[k]
|
||||
callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
|
||||
|
||||
latents = callback_outputs.pop("latents", latents)
|
||||
prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
|
||||
negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds)
|
||||
|
||||
if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
|
||||
progress_bar.update()
|
||||
|
||||
if not output_type == "latent":
|
||||
video = self.decode_latents(latents)
|
||||
video = self.video_processor.postprocess_video(video=video, output_type=output_type)
|
||||
else:
|
||||
video = latents
|
||||
|
||||
# Offload all models
|
||||
self.maybe_free_model_hooks()
|
||||
|
||||
if not return_dict:
|
||||
return (video,)
|
||||
|
||||
return CogVideoXPipelineOutput(frames=video)
|
||||
+10
-4
@@ -51,10 +51,16 @@ class Relighter:
|
||||
|
||||
conds, unconds = self.encode_prompt_pair(tokenizer, device, positive_prompt, negative_prompt)
|
||||
input_bg = self.create_background()
|
||||
bg = resize_and_center_crop(input_bg, self.image_width, self.image_height)
|
||||
bg_latent = numpy2pytorch([bg], device, vae.dtype)
|
||||
bg_latent = vae.encode(bg_latent).latent_dist.mode() * vae.config.scaling_factor
|
||||
|
||||
# bg = resize_and_center_crop(input_bg, self.image_width, self.image_height)
|
||||
# bg_latent = numpy2pytorch([bg], device, vae.dtype)
|
||||
# bg_latent = vae.encode(bg_latent).latent_dist.mode() * vae.config.scaling_factor
|
||||
if self.bg_source == BGSource.NONE:
|
||||
shape = (1, 4, self.image_width//8, self.image_height//8)
|
||||
bg_latent = torch.randn(shape, generator=generator, device=device, dtype=vae.dtype)
|
||||
else:
|
||||
bg = resize_and_center_crop(input_bg, self.image_width, self.image_height)
|
||||
bg_latent = numpy2pytorch([bg], device, vae.dtype)
|
||||
bg_latent = vae.encode(bg_latent).latent_dist.mode() * vae.config.scaling_factor
|
||||
self.bg_latent = bg_latent.repeat(self.num_frames, 1, 1, 1) ## 固定光源
|
||||
self.conds = conds.repeat(self.num_frames, 1, 1)
|
||||
self.unconds = unconds.repeat(self.num_frames, 1, 1)
|
||||
|
||||
@@ -0,0 +1,119 @@
|
||||
# Copyright 2024 Stability AI, Katherine Crowson and The HuggingFace Team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import math
|
||||
from dataclasses import dataclass
|
||||
from typing import List, Optional, Tuple, Union
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from diffusers.configuration_utils import ConfigMixin, register_to_config
|
||||
from diffusers.utils import BaseOutput, is_scipy_available, logging
|
||||
from diffusers.schedulers.scheduling_utils import SchedulerMixin
|
||||
|
||||
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
|
||||
|
||||
|
||||
@dataclass
|
||||
class FlowMatchEulerDiscreteSchedulerOutput(BaseOutput):
|
||||
"""
|
||||
Output class for the scheduler's `step` function output.
|
||||
|
||||
Args:
|
||||
prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images):
|
||||
Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the
|
||||
denoising loop.
|
||||
"""
|
||||
|
||||
prev_sample: torch.FloatTensor
|
||||
|
||||
def wan_flowmatch_step(
|
||||
self,
|
||||
model_output: torch.FloatTensor,
|
||||
timestep: Union[float, torch.FloatTensor],
|
||||
sample: torch.FloatTensor,
|
||||
fusion_sample: torch.FloatTensor,
|
||||
s_churn: float = 0.0,
|
||||
s_tmin: float = 0.0,
|
||||
s_tmax: float = float("inf"),
|
||||
s_noise: float = 1.0,
|
||||
generator: Optional[torch.Generator] = None,
|
||||
return_dict: bool = True,
|
||||
) -> Union[FlowMatchEulerDiscreteSchedulerOutput, Tuple]:
|
||||
"""
|
||||
Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
|
||||
process from the learned model outputs (most often the predicted noise).
|
||||
|
||||
Args:
|
||||
model_output (`torch.FloatTensor`):
|
||||
The direct output from learned diffusion model.
|
||||
timestep (`float`):
|
||||
The current discrete timestep in the diffusion chain.
|
||||
sample (`torch.FloatTensor`):
|
||||
A current instance of a sample created by the diffusion process.
|
||||
s_churn (`float`):
|
||||
s_tmin (`float`):
|
||||
s_tmax (`float`):
|
||||
s_noise (`float`, defaults to 1.0):
|
||||
Scaling factor for noise added to the sample.
|
||||
generator (`torch.Generator`, *optional*):
|
||||
A random number generator.
|
||||
return_dict (`bool`):
|
||||
Whether or not to return a
|
||||
[`~schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteSchedulerOutput`] or tuple.
|
||||
|
||||
Returns:
|
||||
[`~schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteSchedulerOutput`] or `tuple`:
|
||||
If return_dict is `True`,
|
||||
[`~schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteSchedulerOutput`] is returned,
|
||||
otherwise a tuple is returned where the first element is the sample tensor.
|
||||
"""
|
||||
|
||||
if (
|
||||
isinstance(timestep, int)
|
||||
or isinstance(timestep, torch.IntTensor)
|
||||
or isinstance(timestep, torch.LongTensor)
|
||||
):
|
||||
raise ValueError(
|
||||
(
|
||||
"Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to"
|
||||
" `FlowMatchEulerDiscreteScheduler.step()` is not supported. Make sure to pass"
|
||||
" one of the `scheduler.timesteps` as a timestep."
|
||||
),
|
||||
)
|
||||
|
||||
if self.step_index is None:
|
||||
self._init_step_index(timestep)
|
||||
|
||||
# Upcast to avoid precision issues when computing prev_sample
|
||||
sample = sample.to(torch.float32)
|
||||
|
||||
sigma = self.sigmas[self.step_index]
|
||||
sigma_next = self.sigmas[self.step_index + 1]
|
||||
|
||||
## new direction
|
||||
fusion_vector = (sample - fusion_sample) / sigma
|
||||
prev_sample = sample + (sigma_next - sigma) * fusion_vector
|
||||
|
||||
# Cast sample back to model compatible dtype
|
||||
prev_sample = prev_sample.to(model_output.dtype)
|
||||
|
||||
# upon completion increase step index by one
|
||||
self._step_index += 1
|
||||
|
||||
if not return_dict:
|
||||
return (prev_sample,)
|
||||
|
||||
return FlowMatchEulerDiscreteSchedulerOutput(prev_sample=prev_sample)
|
||||
+752
@@ -0,0 +1,752 @@
|
||||
# Copyright 2025 The Wan Team and The HuggingFace Team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import html
|
||||
from typing import Callable, Dict, List, Optional, Union
|
||||
|
||||
import ftfy
|
||||
import regex as re
|
||||
import torch
|
||||
from transformers import AutoTokenizer, UMT5EncoderModel
|
||||
|
||||
from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback
|
||||
from diffusers.models import AutoencoderKLWan, WanTransformer3DModel
|
||||
from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
|
||||
from diffusers.utils import is_torch_xla_available, logging, replace_example_docstring
|
||||
from diffusers.utils.torch_utils import randn_tensor
|
||||
from diffusers.video_processor import VideoProcessor
|
||||
from diffusers.pipelines.pipeline_utils import DiffusionPipeline
|
||||
|
||||
import PIL
|
||||
import numpy as np
|
||||
from diffusers.utils import BaseOutput
|
||||
from dataclasses import dataclass
|
||||
from einops import rearrange
|
||||
from .wan_flowmatch import wan_flowmatch_step
|
||||
from diffusers.utils import export_to_gif
|
||||
|
||||
PipelineImageInput = Union[
|
||||
PIL.Image.Image,
|
||||
np.ndarray,
|
||||
torch.Tensor,
|
||||
List[PIL.Image.Image],
|
||||
List[np.ndarray],
|
||||
List[torch.Tensor],
|
||||
]
|
||||
|
||||
@dataclass
|
||||
class WanPipelineOutput(BaseOutput):
|
||||
r"""
|
||||
Output class for Wan pipelines.
|
||||
|
||||
Args:
|
||||
frames (`torch.Tensor`, `np.ndarray`, or List[List[PIL.Image.Image]]):
|
||||
List of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing
|
||||
denoised PIL image sequences of length `num_frames.` It can also be a NumPy array or Torch tensor of shape
|
||||
`(batch_size, num_frames, channels, height, width)`.
|
||||
"""
|
||||
|
||||
frames: torch.Tensor
|
||||
|
||||
XLA_AVAILABLE = False
|
||||
|
||||
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
|
||||
|
||||
|
||||
EXAMPLE_DOC_STRING = """
|
||||
Examples:
|
||||
```python
|
||||
>>> import torch
|
||||
>>> from diffusers import AutoencoderKLWan, WanPipeline
|
||||
>>> from diffusers.utils import export_to_video
|
||||
|
||||
>>> # Available models: Wan-AI/Wan2.1-T2V-14B-Diffusers, Wan-AI/Wan2.1-T2V-1.3B-Diffusers
|
||||
>>> model_id = "Wan-AI/Wan2.1-T2V-14B-Diffusers"
|
||||
>>> vae = AutoencoderKLWan.from_pretrained(model_id, subfolder="vae", torch_dtype=torch.float32)
|
||||
>>> pipe = WanPipeline.from_pretrained(model_id, vae=vae, torch_dtype=torch.bfloat16)
|
||||
>>> pipe.to("cuda")
|
||||
|
||||
>>> prompt = "A cat walks on the grass, realistic"
|
||||
>>> negative_prompt = "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards"
|
||||
|
||||
>>> output = pipe(
|
||||
... prompt=prompt,
|
||||
... negative_prompt=negative_prompt,
|
||||
... height=480,
|
||||
... width=832,
|
||||
... num_frames=81,
|
||||
... guidance_scale=5.0,
|
||||
... ).frames[0]
|
||||
>>> export_to_video(output, "output.mp4", fps=15)
|
||||
```
|
||||
"""
|
||||
|
||||
import inspect
|
||||
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
|
||||
def retrieve_timesteps(
|
||||
scheduler,
|
||||
num_inference_steps: Optional[int] = None,
|
||||
device: Optional[Union[str, torch.device]] = None,
|
||||
timesteps: Optional[List[int]] = None,
|
||||
sigmas: Optional[List[float]] = None,
|
||||
**kwargs,
|
||||
):
|
||||
r"""
|
||||
Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
|
||||
custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
|
||||
|
||||
Args:
|
||||
scheduler (`SchedulerMixin`):
|
||||
The scheduler to get timesteps from.
|
||||
num_inference_steps (`int`):
|
||||
The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
|
||||
must be `None`.
|
||||
device (`str` or `torch.device`, *optional*):
|
||||
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
|
||||
timesteps (`List[int]`, *optional*):
|
||||
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
|
||||
`num_inference_steps` and `sigmas` must be `None`.
|
||||
sigmas (`List[float]`, *optional*):
|
||||
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
|
||||
`num_inference_steps` and `timesteps` must be `None`.
|
||||
|
||||
Returns:
|
||||
`Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
|
||||
second element is the number of inference steps.
|
||||
"""
|
||||
|
||||
if timesteps is not None and sigmas is not None:
|
||||
raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
|
||||
if timesteps is not None:
|
||||
accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
|
||||
if not accepts_timesteps:
|
||||
raise ValueError(
|
||||
f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
|
||||
f" timestep schedules. Please check whether you are using the correct scheduler."
|
||||
)
|
||||
scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
|
||||
timesteps = scheduler.timesteps
|
||||
num_inference_steps = len(timesteps)
|
||||
elif sigmas is not None:
|
||||
accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
|
||||
if not accept_sigmas:
|
||||
raise ValueError(
|
||||
f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
|
||||
f" sigmas schedules. Please check whether you are using the correct scheduler."
|
||||
)
|
||||
scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
|
||||
timesteps = scheduler.timesteps
|
||||
num_inference_steps = len(timesteps)
|
||||
else:
|
||||
scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
|
||||
timesteps = scheduler.timesteps
|
||||
return timesteps, num_inference_steps
|
||||
|
||||
def basic_clean(text):
|
||||
text = ftfy.fix_text(text)
|
||||
text = html.unescape(html.unescape(text))
|
||||
return text.strip()
|
||||
|
||||
|
||||
def whitespace_clean(text):
|
||||
text = re.sub(r"\s+", " ", text)
|
||||
text = text.strip()
|
||||
return text
|
||||
|
||||
def retrieve_latents(
|
||||
encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
|
||||
):
|
||||
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
|
||||
return encoder_output.latent_dist.sample(generator)
|
||||
elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax":
|
||||
return encoder_output.latent_dist.mode()
|
||||
elif hasattr(encoder_output, "latents"):
|
||||
return encoder_output.latents
|
||||
else:
|
||||
raise AttributeError("Could not access latents of provided encoder_output")
|
||||
|
||||
|
||||
def prompt_clean(text):
|
||||
text = whitespace_clean(basic_clean(text))
|
||||
return text
|
||||
|
||||
|
||||
class WanVideoToVideoPipeline(DiffusionPipeline):
|
||||
r"""
|
||||
Pipeline for text-to-video generation using Wan.
|
||||
|
||||
This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods
|
||||
implemented for all pipelines (downloading, saving, running on a particular device, etc.).
|
||||
|
||||
Args:
|
||||
tokenizer ([`T5Tokenizer`]):
|
||||
Tokenizer from [T5](https://huggingface.co/docs/transformers/en/model_doc/t5#transformers.T5Tokenizer),
|
||||
specifically the [google/umt5-xxl](https://huggingface.co/google/umt5-xxl) variant.
|
||||
text_encoder ([`T5EncoderModel`]):
|
||||
[T5](https://huggingface.co/docs/transformers/en/model_doc/t5#transformers.T5EncoderModel), specifically
|
||||
the [google/umt5-xxl](https://huggingface.co/google/umt5-xxl) variant.
|
||||
transformer ([`WanTransformer3DModel`]):
|
||||
Conditional Transformer to denoise the input latents.
|
||||
scheduler ([`UniPCMultistepScheduler`]):
|
||||
A scheduler to be used in combination with `transformer` to denoise the encoded image latents.
|
||||
vae ([`AutoencoderKLWan`]):
|
||||
Variational Auto-Encoder (VAE) Model to encode and decode videos to and from latent representations.
|
||||
"""
|
||||
|
||||
model_cpu_offload_seq = "text_encoder->transformer->vae"
|
||||
_callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
tokenizer: AutoTokenizer,
|
||||
text_encoder: UMT5EncoderModel,
|
||||
transformer: WanTransformer3DModel,
|
||||
vae: AutoencoderKLWan,
|
||||
scheduler: FlowMatchEulerDiscreteScheduler,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.register_modules(
|
||||
vae=vae,
|
||||
text_encoder=text_encoder,
|
||||
tokenizer=tokenizer,
|
||||
transformer=transformer,
|
||||
scheduler=scheduler,
|
||||
)
|
||||
|
||||
self.vae_scale_factor_temporal = 2 ** sum(self.vae.temperal_downsample) if getattr(self, "vae", None) else 4
|
||||
self.vae_scale_factor_spatial = 2 ** len(self.vae.temperal_downsample) if getattr(self, "vae", None) else 8
|
||||
self.video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial)
|
||||
|
||||
def _get_t5_prompt_embeds(
|
||||
self,
|
||||
prompt: Union[str, List[str]] = None,
|
||||
num_videos_per_prompt: int = 1,
|
||||
max_sequence_length: int = 226,
|
||||
device: Optional[torch.device] = None,
|
||||
dtype: Optional[torch.dtype] = None,
|
||||
):
|
||||
device = device or self._execution_device
|
||||
dtype = dtype or self.text_encoder.dtype
|
||||
|
||||
prompt = [prompt] if isinstance(prompt, str) else prompt
|
||||
prompt = [prompt_clean(u) for u in prompt]
|
||||
batch_size = len(prompt)
|
||||
|
||||
text_inputs = self.tokenizer(
|
||||
prompt,
|
||||
padding="max_length",
|
||||
max_length=max_sequence_length,
|
||||
truncation=True,
|
||||
add_special_tokens=True,
|
||||
return_attention_mask=True,
|
||||
return_tensors="pt",
|
||||
)
|
||||
text_input_ids, mask = text_inputs.input_ids, text_inputs.attention_mask
|
||||
seq_lens = mask.gt(0).sum(dim=1).long()
|
||||
|
||||
prompt_embeds = self.text_encoder(text_input_ids.to(device), mask.to(device)).last_hidden_state
|
||||
prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
|
||||
prompt_embeds = [u[:v] for u, v in zip(prompt_embeds, seq_lens)]
|
||||
prompt_embeds = torch.stack(
|
||||
[torch.cat([u, u.new_zeros(max_sequence_length - u.size(0), u.size(1))]) for u in prompt_embeds], dim=0
|
||||
)
|
||||
|
||||
# duplicate text embeddings for each generation per prompt, using mps friendly method
|
||||
_, seq_len, _ = prompt_embeds.shape
|
||||
prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1)
|
||||
prompt_embeds = prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1)
|
||||
|
||||
return prompt_embeds
|
||||
|
||||
# Copied from diffusers.pipelines.animatediff.pipeline_animatediff_video2video.AnimateDiffVideoToVideoPipeline.get_timesteps
|
||||
def get_timesteps(self, num_inference_steps, timesteps, strength, device):
|
||||
# get the original timestep using init_timestep
|
||||
init_timestep = min(int(num_inference_steps * strength), num_inference_steps)
|
||||
|
||||
t_start = max(num_inference_steps - init_timestep, 0)
|
||||
timesteps = timesteps[t_start * self.scheduler.order :]
|
||||
|
||||
return timesteps, num_inference_steps - t_start
|
||||
|
||||
def encode_prompt(
|
||||
self,
|
||||
prompt: Union[str, List[str]],
|
||||
negative_prompt: Optional[Union[str, List[str]]] = None,
|
||||
do_classifier_free_guidance: bool = True,
|
||||
num_videos_per_prompt: int = 1,
|
||||
prompt_embeds: Optional[torch.Tensor] = None,
|
||||
negative_prompt_embeds: Optional[torch.Tensor] = None,
|
||||
max_sequence_length: int = 226,
|
||||
device: Optional[torch.device] = None,
|
||||
dtype: Optional[torch.dtype] = None,
|
||||
):
|
||||
r"""
|
||||
Encodes the prompt into text encoder hidden states.
|
||||
|
||||
Args:
|
||||
prompt (`str` or `List[str]`, *optional*):
|
||||
prompt to be encoded
|
||||
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`).
|
||||
do_classifier_free_guidance (`bool`, *optional*, defaults to `True`):
|
||||
Whether to use classifier free guidance or not.
|
||||
num_videos_per_prompt (`int`, *optional*, defaults to 1):
|
||||
Number of videos that should be generated per prompt. torch device to place the resulting embeddings on
|
||||
prompt_embeds (`torch.Tensor`, *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.Tensor`, *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.
|
||||
device: (`torch.device`, *optional*):
|
||||
torch device
|
||||
dtype: (`torch.dtype`, *optional*):
|
||||
torch dtype
|
||||
"""
|
||||
device = device or self._execution_device
|
||||
|
||||
prompt = [prompt] if isinstance(prompt, str) else prompt
|
||||
if prompt is not None:
|
||||
batch_size = len(prompt)
|
||||
else:
|
||||
batch_size = prompt_embeds.shape[0]
|
||||
|
||||
if prompt_embeds is None:
|
||||
prompt_embeds = self._get_t5_prompt_embeds(
|
||||
prompt=prompt,
|
||||
num_videos_per_prompt=num_videos_per_prompt,
|
||||
max_sequence_length=max_sequence_length,
|
||||
device=device,
|
||||
dtype=dtype,
|
||||
)
|
||||
|
||||
if do_classifier_free_guidance and negative_prompt_embeds is None:
|
||||
negative_prompt = negative_prompt or ""
|
||||
negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt
|
||||
|
||||
if prompt is not None and type(prompt) is not type(negative_prompt):
|
||||
raise TypeError(
|
||||
f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
|
||||
f" {type(prompt)}."
|
||||
)
|
||||
elif batch_size != len(negative_prompt):
|
||||
raise ValueError(
|
||||
f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
|
||||
f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
|
||||
" the batch size of `prompt`."
|
||||
)
|
||||
|
||||
negative_prompt_embeds = self._get_t5_prompt_embeds(
|
||||
prompt=negative_prompt,
|
||||
num_videos_per_prompt=num_videos_per_prompt,
|
||||
max_sequence_length=max_sequence_length,
|
||||
device=device,
|
||||
dtype=dtype,
|
||||
)
|
||||
|
||||
return prompt_embeds, negative_prompt_embeds
|
||||
|
||||
def check_inputs(
|
||||
self,
|
||||
prompt,
|
||||
negative_prompt,
|
||||
height,
|
||||
width,
|
||||
prompt_embeds=None,
|
||||
negative_prompt_embeds=None,
|
||||
callback_on_step_end_tensor_inputs=None,
|
||||
):
|
||||
if height % 16 != 0 or width % 16 != 0:
|
||||
raise ValueError(f"`height` and `width` have to be divisible by 16 but are {height} and {width}.")
|
||||
|
||||
if callback_on_step_end_tensor_inputs is not None and not all(
|
||||
k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs
|
||||
):
|
||||
raise ValueError(
|
||||
f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}"
|
||||
)
|
||||
|
||||
if prompt is not None and prompt_embeds is not None:
|
||||
raise ValueError(
|
||||
f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
|
||||
" only forward one of the two."
|
||||
)
|
||||
elif negative_prompt is not None and negative_prompt_embeds is not None:
|
||||
raise ValueError(
|
||||
f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`: {negative_prompt_embeds}. Please make sure to"
|
||||
" only forward one of the two."
|
||||
)
|
||||
elif prompt is None and prompt_embeds is None:
|
||||
raise ValueError(
|
||||
"Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
|
||||
)
|
||||
elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
|
||||
raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
|
||||
elif negative_prompt is not None and (
|
||||
not isinstance(negative_prompt, str) and not isinstance(negative_prompt, list)
|
||||
):
|
||||
raise ValueError(f"`negative_prompt` has to be of type `str` or `list` but is {type(negative_prompt)}")
|
||||
|
||||
def prepare_latents(
|
||||
self,
|
||||
video: Optional[torch.Tensor] = None,
|
||||
batch_size: int = 1,
|
||||
num_channels_latents: int = 16,
|
||||
height: int = 720,
|
||||
width: int = 1280,
|
||||
num_frames: int = 81,
|
||||
dtype: Optional[torch.dtype] = None,
|
||||
device: Optional[torch.device] = None,
|
||||
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
|
||||
latents: Optional[torch.Tensor] = None,
|
||||
timestep: Optional[torch.Tensor] = None,
|
||||
) -> torch.Tensor:
|
||||
num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1
|
||||
latent_height = height // self.vae_scale_factor_spatial
|
||||
latent_width = width // self.vae_scale_factor_spatial
|
||||
|
||||
shape = (
|
||||
batch_size,
|
||||
num_channels_latents,
|
||||
num_latent_frames,
|
||||
latent_height,
|
||||
latent_width,
|
||||
)
|
||||
init_latents = [retrieve_latents(self.vae.encode(vid.unsqueeze(0)), generator) for vid in video]
|
||||
init_latents = torch.cat(init_latents, dim=0).to(dtype)
|
||||
|
||||
latents_mean = (
|
||||
torch.tensor(self.vae.config.latents_mean)
|
||||
.view(1, self.vae.config.z_dim, 1, 1, 1)
|
||||
.to(device, dtype)
|
||||
)
|
||||
latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, 1).to(
|
||||
device, dtype
|
||||
)
|
||||
init_latents = (init_latents - latents_mean) * latents_std
|
||||
|
||||
|
||||
noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
|
||||
latents = self.scheduler.scale_noise(init_latents, timestep, noise)
|
||||
|
||||
return latents
|
||||
|
||||
def encode_latents(self, video: torch.Tensor) -> torch.Tensor:
|
||||
|
||||
latents = self.vae.encode(video).latent_dist.mode()
|
||||
latents_mean = (
|
||||
torch.tensor(self.vae.config.latents_mean)
|
||||
.view(1, self.vae.config.z_dim, 1, 1, 1)
|
||||
.to(latents.device, latents.dtype)
|
||||
)
|
||||
latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, 1).to(
|
||||
latents.device, latents.dtype
|
||||
)
|
||||
|
||||
latents = (latents - latents_mean) * latents_std
|
||||
|
||||
return latents
|
||||
|
||||
def decode_latents(self, latents: torch.Tensor) -> torch.Tensor:
|
||||
|
||||
latents = latents.to(self.vae.dtype)
|
||||
latents_mean = (
|
||||
torch.tensor(self.vae.config.latents_mean)
|
||||
.view(1, self.vae.config.z_dim, 1, 1, 1)
|
||||
.to(latents.device, latents.dtype)
|
||||
)
|
||||
latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, 1).to(
|
||||
latents.device, latents.dtype
|
||||
)
|
||||
latents = latents / latents_std + latents_mean
|
||||
video = self.vae.decode(latents, return_dict=False)[0]
|
||||
|
||||
return video
|
||||
|
||||
@property
|
||||
def guidance_scale(self):
|
||||
return self._guidance_scale
|
||||
|
||||
@property
|
||||
def do_classifier_free_guidance(self):
|
||||
return self._guidance_scale > 1.0
|
||||
|
||||
@property
|
||||
def num_timesteps(self):
|
||||
return self._num_timesteps
|
||||
|
||||
@property
|
||||
def current_timestep(self):
|
||||
return self._current_timestep
|
||||
|
||||
@property
|
||||
def interrupt(self):
|
||||
return self._interrupt
|
||||
|
||||
@torch.no_grad()
|
||||
@replace_example_docstring(EXAMPLE_DOC_STRING)
|
||||
def __call__(
|
||||
self,
|
||||
ic_light_pipe=None,
|
||||
relight_prompt=None,
|
||||
bg_source=None,
|
||||
video: List[List[PipelineImageInput]] = None,
|
||||
strength = 0,
|
||||
prompt: Union[str, List[str]] = None,
|
||||
negative_prompt: Union[str, List[str]] = None,
|
||||
height: int = 720,
|
||||
width: int = 1280,
|
||||
num_frames: int = 81,
|
||||
timesteps: Optional[List[int]] = None,
|
||||
num_inference_steps: int = 50,
|
||||
guidance_scale: float = 5.0,
|
||||
num_videos_per_prompt: Optional[int] = 1,
|
||||
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
|
||||
latents: Optional[torch.Tensor] = None,
|
||||
prompt_embeds: Optional[torch.Tensor] = None,
|
||||
negative_prompt_embeds: Optional[torch.Tensor] = None,
|
||||
output_type: Optional[str] = "np",
|
||||
return_dict: bool = True,
|
||||
callback_on_step_end: Optional[
|
||||
Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
|
||||
] = None,
|
||||
callback_on_step_end_tensor_inputs: List[str] = ["latents"],
|
||||
max_sequence_length: int = 512,
|
||||
):
|
||||
r"""
|
||||
The call function to the pipeline for generation.
|
||||
|
||||
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.
|
||||
height (`int`, defaults to `720`):
|
||||
The height in pixels of the generated image.
|
||||
width (`int`, defaults to `1280`):
|
||||
The width in pixels of the generated image.
|
||||
num_frames (`int`, defaults to `129`):
|
||||
The number of frames in the generated video.
|
||||
num_inference_steps (`int`, 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`, defaults to `5.0`):
|
||||
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.
|
||||
num_videos_per_prompt (`int`, *optional*, defaults to 1):
|
||||
The number of images to generate per prompt.
|
||||
generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
|
||||
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
|
||||
generation deterministic.
|
||||
latents (`torch.Tensor`, *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 is generated by sampling using the supplied random `generator`.
|
||||
prompt_embeds (`torch.Tensor`, *optional*):
|
||||
Pre-generated text embeddings. Can be used to easily tweak text inputs (prompt weighting). If not
|
||||
provided, text embeddings are generated from the `prompt` input argument.
|
||||
output_type (`str`, *optional*, defaults to `"pil"`):
|
||||
The output format of the generated image. Choose between `PIL.Image` or `np.array`.
|
||||
return_dict (`bool`, *optional*, defaults to `True`):
|
||||
Whether or not to return a [`WanPipelineOutput`] instead of a plain tuple.
|
||||
callback_on_step_end (`Callable`, `PipelineCallback`, `MultiPipelineCallbacks`, *optional*):
|
||||
A function or a subclass of `PipelineCallback` or `MultiPipelineCallbacks` that is called at the end of
|
||||
each denoising step during the inference. with the following arguments: `callback_on_step_end(self:
|
||||
DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a
|
||||
list of all tensors as specified by `callback_on_step_end_tensor_inputs`.
|
||||
callback_on_step_end_tensor_inputs (`List`, *optional*):
|
||||
The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
|
||||
will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
|
||||
`._callback_tensor_inputs` attribute of your pipeline class.
|
||||
autocast_dtype (`torch.dtype`, *optional*, defaults to `torch.bfloat16`):
|
||||
The dtype to use for the torch.amp.autocast.
|
||||
|
||||
Examples:
|
||||
|
||||
Returns:
|
||||
[`~WanPipelineOutput`] or `tuple`:
|
||||
If `return_dict` is `True`, [`WanPipelineOutput`] is returned, otherwise a `tuple` is returned where
|
||||
the first element is a list with the generated images and the second element is a list of `bool`s
|
||||
indicating whether the corresponding generated image contains "not-safe-for-work" (nsfw) content.
|
||||
"""
|
||||
self._guidance_scale = guidance_scale
|
||||
self._current_timestep = None
|
||||
self._interrupt = False
|
||||
|
||||
device = self._execution_device
|
||||
|
||||
# 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]
|
||||
|
||||
# 3. Encode input prompt
|
||||
prompt_embeds, negative_prompt_embeds = self.encode_prompt(
|
||||
prompt=prompt,
|
||||
negative_prompt=negative_prompt,
|
||||
do_classifier_free_guidance=self.do_classifier_free_guidance,
|
||||
num_videos_per_prompt=num_videos_per_prompt,
|
||||
prompt_embeds=prompt_embeds,
|
||||
negative_prompt_embeds=negative_prompt_embeds,
|
||||
max_sequence_length=max_sequence_length,
|
||||
device=device,
|
||||
)
|
||||
|
||||
transformer_dtype = self.transformer.dtype
|
||||
prompt_embeds = prompt_embeds.to(transformer_dtype)
|
||||
if negative_prompt_embeds is not None:
|
||||
negative_prompt_embeds = negative_prompt_embeds.to(transformer_dtype)
|
||||
|
||||
# 4. Prepare timesteps
|
||||
timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps)
|
||||
timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, timesteps, strength, device)
|
||||
latent_timestep = timesteps[:1].repeat(batch_size * num_videos_per_prompt)
|
||||
self._num_timesteps = len(timesteps)
|
||||
|
||||
if latents is None:
|
||||
video = self.video_processor.preprocess_video(video, height=height, width=width)
|
||||
video = video.to(device=device, dtype=prompt_embeds.dtype)
|
||||
|
||||
# 5. Prepare latent variables
|
||||
num_channels_latents = self.transformer.config.in_channels
|
||||
|
||||
latents = self.prepare_latents(
|
||||
video,
|
||||
batch_size * num_videos_per_prompt,
|
||||
num_channels_latents,
|
||||
height,
|
||||
width,
|
||||
num_frames,
|
||||
prompt_embeds.dtype,
|
||||
device,
|
||||
generator,
|
||||
latents,
|
||||
latent_timestep,
|
||||
)
|
||||
|
||||
#################### Relighter ###################
|
||||
from .ic_light import Relighter
|
||||
num_frames = video.shape[2]
|
||||
relighter = Relighter(
|
||||
pipeline=ic_light_pipe,
|
||||
relight_prompt=relight_prompt,
|
||||
bg_source=bg_source,
|
||||
generator=generator,
|
||||
num_frames=num_frames,
|
||||
image_width=width,
|
||||
image_height=height,
|
||||
)
|
||||
|
||||
# 6. Denoising loop
|
||||
num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
|
||||
self._num_timesteps = len(timesteps)
|
||||
|
||||
with self.progress_bar(total=num_inference_steps) as progress_bar:
|
||||
for i, t in enumerate(timesteps):
|
||||
if self.interrupt:
|
||||
continue
|
||||
|
||||
self._current_timestep = t
|
||||
latent_model_input = latents.to(transformer_dtype)
|
||||
timestep = t.expand(latents.shape[0])
|
||||
|
||||
## init the step
|
||||
self.scheduler._init_step_index(timestep)
|
||||
|
||||
noise_pred = self.transformer(
|
||||
hidden_states=latent_model_input,
|
||||
timestep=timestep,
|
||||
encoder_hidden_states=prompt_embeds,
|
||||
return_dict=False,
|
||||
)[0]
|
||||
|
||||
if self.do_classifier_free_guidance:
|
||||
noise_uncond = self.transformer(
|
||||
hidden_states=latent_model_input,
|
||||
timestep=timestep,
|
||||
encoder_hidden_states=negative_prompt_embeds,
|
||||
return_dict=False,
|
||||
)[0]
|
||||
noise_pred = noise_uncond + guidance_scale * (noise_pred - noise_uncond)
|
||||
|
||||
## Progressive light fusion
|
||||
k = 1
|
||||
lbd = 1 - (i / (num_inference_steps - 1)) ** k
|
||||
|
||||
if lbd > 0.2:
|
||||
|
||||
## get pred_x0
|
||||
sigma = self.scheduler.sigmas[self.scheduler.step_index]
|
||||
pred_x0_latent = latents - sigma * noise_pred
|
||||
consist_target = self.decode_latents(pred_x0_latent)
|
||||
|
||||
consist_target = rearrange(consist_target, "1 c f h w -> f c h w")
|
||||
relight_target = relighter(consist_target)
|
||||
|
||||
fusion_target = consist_target + lbd * (relight_target - consist_target)
|
||||
fusion_target = rearrange(fusion_target, "f c h w -> 1 c f h w")
|
||||
|
||||
fusion_target = self.encode_latents(fusion_target)
|
||||
output = wan_flowmatch_step(self.scheduler, noise_pred, t, latents, fusion_target, return_dict=False)
|
||||
else:
|
||||
output = self.scheduler.step(noise_pred, t, latents, return_dict=False)
|
||||
|
||||
latents = output[0]
|
||||
|
||||
if callback_on_step_end is not None:
|
||||
callback_kwargs = {}
|
||||
for k in callback_on_step_end_tensor_inputs:
|
||||
callback_kwargs[k] = locals()[k]
|
||||
callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
|
||||
|
||||
latents = callback_outputs.pop("latents", latents)
|
||||
prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
|
||||
negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds)
|
||||
|
||||
# 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()
|
||||
|
||||
self._current_timestep = None
|
||||
|
||||
if not output_type == "latent":
|
||||
latents = latents.to(self.vae.dtype)
|
||||
latents_mean = (
|
||||
torch.tensor(self.vae.config.latents_mean)
|
||||
.view(1, self.vae.config.z_dim, 1, 1, 1)
|
||||
.to(latents.device, latents.dtype)
|
||||
)
|
||||
latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, 1).to(
|
||||
latents.device, latents.dtype
|
||||
)
|
||||
latents = latents / latents_std + latents_mean
|
||||
video = self.vae.decode(latents, return_dict=False)[0]
|
||||
video = self.video_processor.postprocess_video(video, output_type=output_type)
|
||||
|
||||
# Offload all models
|
||||
self.maybe_free_model_hooks()
|
||||
|
||||
if not return_dict:
|
||||
return (video,)
|
||||
|
||||
return WanPipelineOutput(frames=video)
|
||||
+148
@@ -0,0 +1,148 @@
|
||||
from PIL import Image,ImageSequence
|
||||
import numpy as np
|
||||
import torch
|
||||
from moviepy.editor import VideoFileClip
|
||||
import os
|
||||
import imageio
|
||||
import random
|
||||
from diffusers.utils import export_to_video
|
||||
|
||||
def resize_and_center_crop(image, target_width, target_height):
|
||||
pil_image = Image.fromarray(image)
|
||||
original_width, original_height = pil_image.size
|
||||
scale_factor = max(target_width / original_width, target_height / original_height)
|
||||
resized_width = int(round(original_width * scale_factor))
|
||||
resized_height = int(round(original_height * scale_factor))
|
||||
resized_image = pil_image.resize((resized_width, resized_height), Image.LANCZOS)
|
||||
|
||||
left = (resized_width - target_width) / 2
|
||||
top = (resized_height - target_height) / 2
|
||||
right = (resized_width + target_width) / 2
|
||||
bottom = (resized_height + target_height) / 2
|
||||
cropped_image = resized_image.crop((left, top, right, bottom))
|
||||
return np.array(cropped_image)
|
||||
|
||||
def numpy2pytorch(imgs, device, dtype):
|
||||
h = torch.from_numpy(np.stack(imgs, axis=0)).float() / 127.0 - 1.0
|
||||
h = h.movedim(-1, 1)
|
||||
return h.to(device=device, dtype=dtype)
|
||||
|
||||
def get_fg_video(video_list, mask_list, device, dtype):
|
||||
video_np = np.stack(video_list, axis=0)
|
||||
mask_np = np.stack(mask_list, axis=0)
|
||||
mask_bool = mask_np == 255
|
||||
video_fg = np.where(mask_bool, video_np, 127)
|
||||
|
||||
h = torch.from_numpy(video_fg).float() / 127.0 - 1.0
|
||||
h = h.movedim(-1, 1)
|
||||
return h.to(device=device, dtype=dtype)
|
||||
|
||||
|
||||
def pad(x, p, i):
|
||||
return x[:i] if len(x) >= i else x + [p] * (i - len(x))
|
||||
|
||||
def gif_to_mp4(gif_path, mp4_path):
|
||||
clip = VideoFileClip(gif_path)
|
||||
clip.write_videofile(mp4_path)
|
||||
|
||||
def generate_light_sequence(light_tensor, num_frames=16, direction="r"):
|
||||
|
||||
if direction in "l":
|
||||
target_tensor = torch.rot90(light_tensor, k=1, dims=(2, 3))
|
||||
elif direction in "r":
|
||||
target_tensor = torch.rot90(light_tensor, k=-1, dims=(2, 3))
|
||||
else:
|
||||
raise ValueError("direction must be either 'r' for right or 'l' for left")
|
||||
|
||||
# Generate the sequence
|
||||
out_list = []
|
||||
for frame_idx in range(num_frames):
|
||||
t = frame_idx / (num_frames - 1)
|
||||
interpolated_matrix = (1 - t) * light_tensor + t * target_tensor
|
||||
out_list.append(interpolated_matrix)
|
||||
|
||||
out_tensor = torch.stack(out_list, dim=0).squeeze(1)
|
||||
|
||||
return out_tensor
|
||||
|
||||
def tensor2vid(video: torch.Tensor, processor, output_type="np"):
|
||||
|
||||
batch_size, channels, num_frames, height, width = video.shape ## [1, 4, 16, 512, 512]
|
||||
outputs = []
|
||||
for batch_idx in range(batch_size):
|
||||
batch_vid = video[batch_idx].permute(1, 0, 2, 3)
|
||||
batch_output = processor.postprocess(batch_vid, output_type)
|
||||
|
||||
outputs.append(batch_output)
|
||||
|
||||
return outputs
|
||||
|
||||
def read_video(video_path:str, image_width, image_height):
|
||||
extension = video_path.split('.')[-1].lower()
|
||||
video_name = os.path.basename(video_path)
|
||||
video_list = []
|
||||
|
||||
if extension in "gif":
|
||||
## input from gif
|
||||
video = Image.open(video_path)
|
||||
for i, frame in enumerate(ImageSequence.Iterator(video)):
|
||||
frame = np.array(frame.convert("RGB"))
|
||||
frame = resize_and_center_crop(frame, image_width, image_height)
|
||||
video_list.append(frame)
|
||||
elif extension in "mp4":
|
||||
## input from mp4
|
||||
reader = imageio.get_reader(video_path)
|
||||
for frame in reader:
|
||||
frame = resize_and_center_crop(frame, image_width, image_height)
|
||||
video_list.append(frame)
|
||||
else:
|
||||
raise ValueError('Wrong input type')
|
||||
|
||||
video_list = [Image.fromarray(frame) for frame in video_list]
|
||||
|
||||
return video_list, video_name
|
||||
|
||||
def read_mask(mask_folder:str):
|
||||
mask_files = os.listdir(mask_folder)
|
||||
mask_files = sorted(mask_files)
|
||||
mask_list = []
|
||||
for mask_file in mask_files:
|
||||
mask_path = os.path.join(mask_folder, mask_file)
|
||||
mask = Image.open(mask_path).convert('RGB')
|
||||
mask_list.append(mask)
|
||||
|
||||
return mask_list
|
||||
|
||||
def decode_latents(vae, latents, decode_chunk_size: int = 16):
|
||||
|
||||
latents = 1 / vae.config.scaling_factor * latents
|
||||
video = []
|
||||
for i in range(0, latents.shape[0], decode_chunk_size):
|
||||
batch_latents = latents[i : i + decode_chunk_size]
|
||||
batch_latents = vae.decode(batch_latents).sample
|
||||
video.append(batch_latents)
|
||||
|
||||
video = torch.cat(video)
|
||||
|
||||
return video
|
||||
|
||||
def encode_video(vae, video, decode_chunk_size: int = 16) -> torch.Tensor:
|
||||
latents = []
|
||||
for i in range(0, len(video), decode_chunk_size):
|
||||
batch_video = video[i : i + decode_chunk_size]
|
||||
batch_video = vae.encode(batch_video).latent_dist.mode()
|
||||
latents.append(batch_video)
|
||||
return torch.cat(latents)
|
||||
|
||||
def vis_video(input_video, video_processor, save_path):
|
||||
## shape: 1, c, f, h, w
|
||||
relight_video = video_processor.postprocess_video(video=input_video, output_type="pil")
|
||||
export_to_video(relight_video[0], save_path)
|
||||
|
||||
def set_all_seed(seed):
|
||||
torch.manual_seed(seed)
|
||||
torch.cuda.manual_seed(seed)
|
||||
torch.cuda.manual_seed_all(seed)
|
||||
np.random.seed(seed)
|
||||
random.seed(seed)
|
||||
torch.backends.cudnn.deterministic = True
|
||||
Reference in New Issue
Block a user