570 lines
23 KiB
Python
570 lines
23 KiB
Python
|
|
import gc
|
|
import os
|
|
import subprocess
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from accelerate.logging import get_logger
|
|
from diffusers import FlowMatchEulerDiscreteScheduler
|
|
from einops import rearrange
|
|
from omegaconf import OmegaConf
|
|
from PIL import Image
|
|
from tqdm.auto import tqdm
|
|
from transformers import AutoTokenizer
|
|
import librosa
|
|
from pathlib import Path
|
|
import imageio
|
|
import torchvision
|
|
from transformers import Wav2Vec2Model, Wav2Vec2Processor
|
|
import math
|
|
|
|
from .StableAvatar.fm_solvers_unipc import FlowUniPCMultistepScheduler
|
|
from .StableAvatar.wan.models.cache_utils import get_teacache_coefficients
|
|
from .StableAvatar.wan.models.wan_fantasy_transformer3d_1B import WanTransformer3DFantasyModel
|
|
from .StableAvatar.wan.models.wan_vae import AutoencoderKLWan
|
|
from .StableAvatar.wan.pipeline.wan_inference_long_pipeline import WanI2VTalkingInferenceLongPipeline
|
|
from .StableAvatar.wan.utils.fp8_optimization import replace_parameters_by_name, convert_weight_dtype_wrapper, convert_model_weight_to_float8
|
|
from .StableAvatar.wan.utils.utils import get_image_to_video_latent
|
|
|
|
logger = get_logger(__name__, log_level="INFO")
|
|
|
|
|
|
|
|
def save_video_ffmpeg(gen_video_samples, save_path, vocal_audio_path, fps=25, quality=10):
|
|
def save_video(frames, save_path, fps, quality=9, ffmpeg_params=None, saved_frames_dir=None):
|
|
writer = imageio.get_writer(
|
|
save_path, fps=fps, quality=quality, ffmpeg_params=ffmpeg_params
|
|
)
|
|
idx = 0
|
|
for frame in tqdm(frames, desc="Saving video"):
|
|
frame = np.array(frame)
|
|
frame_path = os.path.join(saved_frames_dir, f"frame_{idx}.png")
|
|
idx = idx + 1
|
|
imageio.imwrite(frame_path, frame)
|
|
writer.append_data(frame)
|
|
writer.close()
|
|
|
|
save_path_tmp = os.path.join(save_path, "video_without_audio.mp4")
|
|
saved_frames_dir = os.path.join(save_path, "animated_images")
|
|
os.makedirs(saved_frames_dir, exist_ok=True)
|
|
|
|
# video_audio = (gen_video_samples + 1) / 2 # C T H W
|
|
video_audio = (gen_video_samples / 2 + 0.5).clamp(0, 1)
|
|
video_audio = video_audio.permute(1, 2, 3, 0).cpu().numpy()
|
|
video_audio = np.clip(video_audio * 255, 0, 255).astype(np.uint8) # to [0, 255]
|
|
save_video(video_audio, save_path_tmp, fps=fps, quality=quality, saved_frames_dir=saved_frames_dir)
|
|
|
|
# crop audio according to video length
|
|
_, T, _, _ = gen_video_samples.shape
|
|
duration = T / fps
|
|
save_path_crop_audio = os.path.join(save_path, "cropped_audio.wav")
|
|
final_command = [
|
|
"ffmpeg",
|
|
"-i",
|
|
vocal_audio_path,
|
|
"-t",
|
|
f'{duration}',
|
|
save_path_crop_audio,
|
|
]
|
|
subprocess.run(final_command, check=True)
|
|
|
|
|
|
def filter_kwargs(cls, kwargs):
|
|
import inspect
|
|
sig = inspect.signature(cls.__init__)
|
|
valid_params = set(sig.parameters.keys()) - {'self', 'cls'}
|
|
filtered_kwargs = {k: v for k, v in kwargs.items() if k in valid_params}
|
|
return filtered_kwargs
|
|
|
|
|
|
def get_random_downsample_ratio(sample_size, image_ratio=[],
|
|
all_choices=False, rng=None):
|
|
def _create_special_list(length):
|
|
if length == 1:
|
|
return [1.0]
|
|
if length >= 2:
|
|
first_element = 0.75
|
|
remaining_sum = 1.0 - first_element
|
|
other_elements_value = remaining_sum / (length - 1)
|
|
special_list = [first_element] + [other_elements_value] * (length - 1)
|
|
return special_list
|
|
|
|
if sample_size >= 1536:
|
|
number_list = [1, 1.25, 1.5, 2, 2.5, 3] + image_ratio
|
|
elif sample_size >= 1024:
|
|
number_list = [1, 1.25, 1.5, 2] + image_ratio
|
|
elif sample_size >= 768:
|
|
number_list = [1, 1.25, 1.5] + image_ratio
|
|
elif sample_size >= 512:
|
|
number_list = [1] + image_ratio
|
|
else:
|
|
number_list = [1]
|
|
|
|
if all_choices:
|
|
return number_list
|
|
|
|
number_list_prob = np.array(_create_special_list(len(number_list)))
|
|
if rng is None:
|
|
return np.random.choice(number_list, p=number_list_prob)
|
|
else:
|
|
return rng.choice(number_list, p=number_list_prob)
|
|
|
|
|
|
def resize_mask(mask, latent, process_first_frame_only=True):
|
|
latent_size = latent.size()
|
|
batch_size, channels, num_frames, height, width = mask.shape
|
|
|
|
if process_first_frame_only:
|
|
target_size = list(latent_size[2:])
|
|
target_size[0] = 1
|
|
first_frame_resized = F.interpolate(
|
|
mask[:, :, 0:1, :, :],
|
|
size=target_size,
|
|
mode='trilinear',
|
|
align_corners=False
|
|
)
|
|
|
|
target_size = list(latent_size[2:])
|
|
target_size[0] = target_size[0] - 1
|
|
if target_size[0] != 0:
|
|
remaining_frames_resized = F.interpolate(
|
|
mask[:, :, 1:, :, :],
|
|
size=target_size,
|
|
mode='trilinear',
|
|
align_corners=False
|
|
)
|
|
resized_mask = torch.cat([first_frame_resized, remaining_frames_resized], dim=2)
|
|
else:
|
|
resized_mask = first_frame_resized
|
|
else:
|
|
target_size = list(latent_size[2:])
|
|
resized_mask = F.interpolate(
|
|
mask,
|
|
size=target_size,
|
|
mode='trilinear',
|
|
align_corners=False
|
|
)
|
|
return resized_mask
|
|
|
|
|
|
# # Will error if the minimal version of diffusers is not installed. Remove at your own risks.
|
|
# check_min_version("0.18.0.dev0")
|
|
|
|
|
|
|
|
|
|
def save_videos_from_pil(pil_images, path, fps=8):
|
|
save_fmt = Path(path).suffix
|
|
os.makedirs(os.path.dirname(path), exist_ok=True)
|
|
|
|
if save_fmt == ".mp4":
|
|
with imageio.get_writer(path, fps=fps) as writer:
|
|
for img in pil_images:
|
|
img_array = np.array(img) # Convert PIL Image to numpy array
|
|
writer.append_data(img_array)
|
|
|
|
elif save_fmt == ".gif":
|
|
pil_images[0].save(
|
|
fp=path,
|
|
format="GIF",
|
|
append_images=pil_images[1:],
|
|
save_all=True,
|
|
duration=(1 / fps * 1000),
|
|
loop=0,
|
|
optimize=False,
|
|
lossless=True
|
|
)
|
|
else:
|
|
raise ValueError("Unsupported file type. Use .mp4 or .gif.")
|
|
|
|
|
|
def save_videos_grid_png_and_mp4(videos: torch.Tensor, rescale=False, n_rows=6, save_frames_path=None):
|
|
videos = rearrange(videos, "b c t h w -> t b c h w")
|
|
height, width = videos.shape[-2:]
|
|
outputs = []
|
|
|
|
for i, x in enumerate(videos):
|
|
x = torchvision.utils.make_grid(x, nrow=n_rows) # (c h w)
|
|
x = x.transpose(0, 1).transpose(1, 2).squeeze(-1) # (h w c)
|
|
if rescale:
|
|
x = (x + 1.0) / 2.0 # -1,1 -> 0,1
|
|
x = (x * 255).numpy().astype(np.uint8)
|
|
x = Image.fromarray(x)
|
|
outputs.append(x)
|
|
|
|
pil_frames = [Image.fromarray(frame) if isinstance(frame, np.ndarray) else frame for frame in outputs]
|
|
num_frames = len(pil_frames)
|
|
for i in range(num_frames):
|
|
pil_frame = pil_frames[i]
|
|
save_path = os.path.join(save_frames_path, f'frame_{i}.png')
|
|
pil_frame.save(save_path)
|
|
|
|
|
|
def linear_decay(initial_value, final_value, total_steps, current_step):
|
|
if current_step >= total_steps:
|
|
return final_value
|
|
current_step = max(0, current_step)
|
|
step_size = (final_value - initial_value) / total_steps
|
|
current_value = initial_value + step_size * current_step
|
|
return current_value
|
|
|
|
|
|
def generate_timestep_with_lognorm(low, high, shape, device="cpu", generator=None):
|
|
u = torch.normal(mean=0.0, std=1.0, size=shape, device=device, generator=generator)
|
|
t = 1 / (1 + torch.exp(-u)) * (high - low) + low
|
|
return torch.clip(t.to(torch.int32), low, high - 1)
|
|
|
|
|
|
|
|
def load_StableAvatar_model(args,vae_path,config,device,weight_dtype,use_mmgp,lora_path):
|
|
|
|
sampler_name = "Flow" if not lora_path else "Flow_Unipc"
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained(os.path.join(args.pretrained_model_name_or_path, config['text_encoder_kwargs'].get('tokenizer_subpath', 'tokenizer')), )
|
|
|
|
vae = AutoencoderKLWan.from_pretrained(
|
|
vae_path,
|
|
additional_kwargs=OmegaConf.to_container(config['vae_kwargs']),
|
|
)
|
|
wav2vec_processor = Wav2Vec2Processor.from_pretrained(args.pretrained_wav2vec_path)
|
|
wav2vec = Wav2Vec2Model.from_pretrained(args.pretrained_wav2vec_path).to("cpu")
|
|
|
|
|
|
transformer3d = WanTransformer3DFantasyModel.from_pretrained(
|
|
args.pretrained_dit_path,
|
|
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
|
|
low_cpu_mem_usage=False,
|
|
torch_dtype=weight_dtype,
|
|
)
|
|
if args.transformer_path is not None:
|
|
print(f"From checkpoint: {args.transformer_path}")
|
|
state_dict = torch.load(args.transformer_path, map_location="cpu", weights_only=False)
|
|
state_dict = state_dict["state_dict"] if "state_dict" in state_dict else state_dict
|
|
m, u = transformer3d.load_state_dict(state_dict, strict=False)
|
|
print(f"missing keys: {len(m)}, unexpected keys: {len(u)}")
|
|
|
|
if lora_path:
|
|
from .lora_adapter import WanLoraWrapper
|
|
lora_wrapper = WanLoraWrapper(transformer3d)
|
|
lora_name = lora_wrapper.load_lora(lora_path)
|
|
lora_wrapper.apply_lora(lora_name, 1.0)
|
|
transformer3d=lora_wrapper.model
|
|
|
|
|
|
Choosen_Scheduler = scheduler_dict = {
|
|
"Flow": FlowMatchEulerDiscreteScheduler,
|
|
"Flow_Unipc": FlowUniPCMultistepScheduler,
|
|
}[sampler_name]
|
|
|
|
if sampler_name=="Flow_Unipc":
|
|
config['scheduler_kwargs']["solver_order"]=2
|
|
config['scheduler_kwargs']["lower_order_final"]=True
|
|
config['scheduler_kwargs']["solver_type"]="bh2"
|
|
config['scheduler_kwargs']["predict_x0"]=True
|
|
config['scheduler_kwargs']["sample_max_value"]=1.0
|
|
config['scheduler_kwargs']["dynamic_thresholding_ratio"]=0.995
|
|
config['scheduler_kwargs']["prediction_type"]="flow_prediction"
|
|
config['scheduler_kwargs']["thresholding"]=False
|
|
scheduler = Choosen_Scheduler(
|
|
**filter_kwargs(Choosen_Scheduler, OmegaConf.to_container(config['scheduler_kwargs']))
|
|
)
|
|
pipeline = WanI2VTalkingInferenceLongPipeline(
|
|
tokenizer=tokenizer,
|
|
#text_encoder=text_encoder,
|
|
vae=vae,
|
|
transformer=transformer3d,
|
|
#clip_image_encoder=clip_image_encoder,
|
|
scheduler=scheduler,
|
|
wav2vec_processor=wav2vec_processor,
|
|
wav2vec=wav2vec,
|
|
)
|
|
# if args.ulysses_degree > 1 or args.ring_degree > 1:
|
|
# transformer3d.enable_multi_gpus_inference()
|
|
# if args.fsdp_dit:
|
|
# shard_fn = partial(shard_model, device_id=device, param_dtype=weight_dtype)
|
|
# pipeline.transformer = shard_fn(pipeline.transformer)
|
|
quantizeTransformer=False if args.GPU_memory_mode=="None" else True
|
|
if use_mmgp!="None":
|
|
from mmgp import offload, profile_type
|
|
pipeline.to("cpu")
|
|
if use_mmgp=="VerylowRAM_LowVRAM":
|
|
offload.profile(pipeline, profile_type.VerylowRAM_LowVRAM,quantizeTransformer=quantizeTransformer)
|
|
elif use_mmgp=="LowRAM_LowVRAM":
|
|
offload.profile(pipeline, profile_type.LowRAM_LowVRAM,quantizeTransformer=quantizeTransformer)
|
|
elif use_mmgp=="LowRAM_HighVRAM":
|
|
offload.profile(pipeline, profile_type.LowRAM_HighVRAM,quantizeTransformer=quantizeTransformer)
|
|
elif use_mmgp=="HighRAM_LowVRAM":
|
|
offload.profile(pipeline, profile_type.HighRAM_LowVRAM,quantizeTransformer=quantizeTransformer)
|
|
elif use_mmgp=="HighRAM_HighVRAM":
|
|
offload.profile(pipeline, profile_type.HighRAM_HighVRAM,quantizeTransformer=quantizeTransformer)
|
|
elif args.GPU_memory_mode == "sequential_cpu_offload":
|
|
replace_parameters_by_name(transformer3d, ["modulation", ], device=device)
|
|
transformer3d.freqs = transformer3d.freqs.to(device=device)
|
|
pipeline.enable_sequential_cpu_offload(device=device)
|
|
elif args.GPU_memory_mode == "model_cpu_offload_and_qfloat8":
|
|
convert_model_weight_to_float8(transformer3d, exclude_module_name=["modulation", ])
|
|
convert_weight_dtype_wrapper(transformer3d, weight_dtype)
|
|
pipeline.enable_model_cpu_offload(device=device)
|
|
elif args.GPU_memory_mode == "model_cpu_offload":
|
|
pipeline.enable_model_cpu_offload(device=device)
|
|
else:
|
|
pipeline.to(device=device)
|
|
|
|
|
|
temporal_compression_ratio=vae.config.temporal_compression_ratio
|
|
return pipeline,tokenizer,temporal_compression_ratio,sampler_name
|
|
|
|
def pre_data_process(text_encoder,clip_image_encoder,tokenizer,prompt,negative_prompt,infer_img,device,width,height,args,weight_dtype):
|
|
|
|
|
|
clip_sample_n_frames = args.clip_sample_n_frames
|
|
temporal_compression_ratio=args.temporal_compression_ratio
|
|
do_classifier_free_guidance=True
|
|
with torch.no_grad():
|
|
prompt_embeds, negative_prompt_embeds=encode_prompt(text_encoder,tokenizer,prompt,negative_prompt,do_classifier_free_guidance,1,device=device,dtype=weight_dtype)
|
|
|
|
video_length = int((clip_sample_n_frames - 1) // temporal_compression_ratio * temporal_compression_ratio) + 1 if clip_sample_n_frames != 1 else 1
|
|
input_video, input_video_mask, _ = get_image_to_video_latent([infer_img], None, video_length=video_length, sample_size=[width, height]) # 首尾帧的处理流程
|
|
sr = 16000
|
|
vocal_input, sample_rate = librosa.load(args.validation_driven_audio_path, sr=sr)
|
|
|
|
|
|
#clip_image = cond_image = Image.open(cond_file_path).convert('RGB')
|
|
clip_image = cond_image = infer_img
|
|
cond_image = cond_image.resize([width, height])
|
|
clip_image = clip_image.resize([width, height])
|
|
clip_image = torch.from_numpy(np.array(clip_image)).permute(2, 0, 1)
|
|
clip_image = clip_image / 255
|
|
clip_image = (clip_image - 0.5) * 2 # C H W
|
|
cond_image = torch.from_numpy(np.array(cond_image)).permute(2, 0, 1).unsqueeze(1).unsqueeze(0)
|
|
cond_image = cond_image / 255
|
|
cond_image = (cond_image - 0.5) * 2 # normalization
|
|
cond_image = cond_image.to(device) # 1 C 1 H W
|
|
|
|
clip_image = clip_image.to(device, weight_dtype)
|
|
# clip_context = clip_image_encoder([clip_image[:, None, :, :]])
|
|
#print("clip_image:",clip_image.shape) #clip_image: torch.Size([3, 720, 480])
|
|
clip_image=clip_image.permute(1, 2, 0).unsqueeze(0) #comfy need [B,C,H,W] -->[B,H,W,C]
|
|
clip_dict=clip_image_encoder.encode_image(clip_image)
|
|
clip_context =clip_dict["penultimate_hidden_states"].to(device, weight_dtype)
|
|
clip_context = (torch.cat([clip_context, clip_context, clip_context], dim=0) if do_classifier_free_guidance else clip_context) #LCM need [B,C,H,W]
|
|
|
|
clip_image_encoder.patcher.cleanup()
|
|
|
|
gc.collect()
|
|
emb={"prompt_embeds":prompt_embeds,"negative_prompt_embeds":negative_prompt_embeds,"vocal_input":vocal_input,"video_length":video_length,
|
|
"clip_image_tensor":cond_image,"clip_context":clip_context,
|
|
"sample_rate":sample_rate,"input_video":input_video,"input_video_mask":input_video_mask,"sr":sr,
|
|
}
|
|
return emb
|
|
|
|
|
|
def infer_StableAvatar(pipeline,args,seed,cfg,device,steps,frame_rate,sample_text_guide_scale,sample_audio_guide_scale,overlap_window_length,weight_dtype,overlapping_weight_scheme):
|
|
|
|
coefficients = get_teacache_coefficients(args.get("pretrained_model_name_or_path")) if args.get("enable_teacache") else None
|
|
if coefficients is not None:
|
|
#print(f"Enable TeaCache with threshold {args.teacache_threshold} and skip the first {args.get(num_skip_start_steps} steps.")
|
|
pipeline.transformer.enable_teacache(
|
|
coefficients,
|
|
args.get("sample_steps"),
|
|
args.get("teacache_threshold"),
|
|
num_skip_start_steps=args.get("num_skip_start_steps"),
|
|
offload=args.get("teacache_offload")
|
|
)
|
|
|
|
if 4 == steps and args.get("sampler_name")=="Flow_Unipc":
|
|
|
|
print("###### Using LCM schedulers ######")
|
|
target_height = args.get("height", 480)
|
|
target_width = args.get("width", 832)
|
|
target_video_length = args.get("video_length", 81)
|
|
vae = pipeline.vae
|
|
compressed_height = target_height // vae.spacial_compression_ratio
|
|
compressed_width = target_width // vae.spacial_compression_ratio
|
|
compressed_length = (target_video_length - 1) // vae.temporal_compression_ratio + 1
|
|
|
|
lcm_config = {
|
|
"infer_steps": 4,
|
|
"target_video_length": 81,
|
|
"target_height": 480,
|
|
"target_width": 832,
|
|
"self_attn_1_type": "flash_attn3",
|
|
"cross_attn_1_type": "flash_attn3",
|
|
"cross_attn_2_type": "flash_attn3",
|
|
"seed": seed,
|
|
"sample_guide_scale": 5,
|
|
"denoising_step_list": [1000, 750, 500, 250],
|
|
"sample_shift": 5,
|
|
"enable_cfg": False,
|
|
"cpu_offload": False,
|
|
"target_shape":[1, vae.config.latent_channels, compressed_length, compressed_height, compressed_width],
|
|
"task":"t2v",
|
|
"patch_size": (1, 2, 2),
|
|
"vae_stride": (4, 8, 8),
|
|
}
|
|
|
|
config_ =OmegaConf.create(lcm_config)
|
|
from .StableAvatar.flow_match_lcm import WanStepDistillScheduler
|
|
pipeline.scheduler = WanStepDistillScheduler(config_)
|
|
|
|
|
|
|
|
generator = torch.Generator(device=device).manual_seed(seed)
|
|
|
|
with torch.no_grad():
|
|
sample = pipeline(
|
|
None,
|
|
num_frames=args.get("video_length"),
|
|
negative_prompt=None,
|
|
height=args.get("height"),
|
|
width=args.get("width"),
|
|
guidance_scale=cfg,
|
|
generator=generator,
|
|
num_inference_steps=steps,
|
|
video=args.get("input_video"),
|
|
mask_video=args.get("input_video_mask"),
|
|
prompt_embeds=args.get("prompt_embeds"),
|
|
negative_prompt_embeds=args.get("negative_prompt_embeds"),
|
|
clip_image=None,
|
|
text_guide_scale=sample_text_guide_scale,
|
|
audio_guide_scale=sample_audio_guide_scale,
|
|
vocal_input_values=args.get("vocal_input"),
|
|
motion_frame=args.get("motion_frame"),
|
|
fps=frame_rate,
|
|
sr=args.get("sr"),
|
|
cond_file_path=None,
|
|
seed=seed,
|
|
overlap_window_length=overlap_window_length,
|
|
clip_image_tensor=args.get("clip_image_tensor"),
|
|
clip_context=args.get("clip_context") ,
|
|
weight_dtype=weight_dtype,
|
|
overlapping_weight_scheme=overlapping_weight_scheme,
|
|
).videos
|
|
|
|
|
|
del pipeline
|
|
|
|
return sample
|
|
|
|
|
|
|
|
def encode_prompt(
|
|
text_encoder,
|
|
tokenizer,
|
|
prompt,
|
|
negative_prompt = None,
|
|
do_classifier_free_guidance: bool = True,
|
|
num_videos_per_prompt: int = 1,
|
|
prompt_embeds = None,
|
|
negative_prompt_embeds = None,
|
|
max_sequence_length: int = 512,
|
|
device= None,
|
|
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
|
|
|
|
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,
|
|
# )
|
|
text_inputs = tokenizer(
|
|
prompt,
|
|
padding="max_length",
|
|
max_length=max_sequence_length,
|
|
truncation=True,
|
|
add_special_tokens=True,
|
|
return_tensors="pt",
|
|
)
|
|
|
|
prompt_attention_mask = text_inputs.attention_mask
|
|
|
|
prompt_embeds=cf_clip(prompt, text_encoder,prompt_attention_mask,device, dtype)[0]
|
|
|
|
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`."
|
|
)
|
|
text_inputs = tokenizer(
|
|
negative_prompt,
|
|
padding="max_length",
|
|
max_length=max_sequence_length,
|
|
truncation=True,
|
|
add_special_tokens=True,
|
|
return_tensors="pt",
|
|
)
|
|
|
|
prompt_attention_mask = text_inputs.attention_mask
|
|
# 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,
|
|
# )
|
|
negative_prompt_embeds=cf_clip(negative_prompt, text_encoder,prompt_attention_mask,device, dtype)[0]
|
|
|
|
return prompt_embeds, negative_prompt_embeds
|
|
|
|
def cf_clip(txt_list, clip,prompt_attention_mask,device, dtype):
|
|
seq_lens = prompt_attention_mask.gt(0).sum(dim=1).long()
|
|
pos_cond_list = []
|
|
for i in txt_list:
|
|
tokens_p = clip.tokenize(i)
|
|
output_p = clip.encode_from_tokens(tokens_p, return_dict=True) # {"pooled_output":tensor}
|
|
cond_p = output_p.pop("cond").to(device, dtype)
|
|
#print(cond_p.shape) #torch.Size([1, 231, 768])
|
|
positive=[u[:v] for u, v in zip(cond_p, seq_lens)]
|
|
pos_cond_list.append(positive)
|
|
|
|
return pos_cond_list
|
|
|
|
|