support for parallel inference using xfuser
This commit is contained in:
mengli.cml
@@ -1145,4 +1145,4 @@ class EasyAnimateDiTBlock(nn.Module):
|
||||
norm_encoder_hidden_states = self.ff(norm_encoder_hidden_states)
|
||||
hidden_states = hidden_states + gate_ff * norm_hidden_states
|
||||
encoder_hidden_states = encoder_hidden_states + enc_gate_ff * norm_encoder_hidden_states
|
||||
return hidden_states, encoder_hidden_states
|
||||
return hidden_states, encoder_hidden_states
|
||||
|
||||
@@ -6,6 +6,21 @@ from diffusers.models.attention import Attention
|
||||
from diffusers.models.embeddings import apply_rotary_emb
|
||||
from einops import rearrange, repeat
|
||||
|
||||
try:
|
||||
import xfuser
|
||||
from xfuser.core.distributed import (
|
||||
get_sequence_parallel_world_size,
|
||||
get_sequence_parallel_rank,
|
||||
get_sp_group,
|
||||
initialize_model_parallel,
|
||||
init_distributed_environment
|
||||
)
|
||||
from xfuser.core.long_ctx_attention import xFuserLongContextAttention
|
||||
except Exception as ex:
|
||||
get_sequence_parallel_world_size = None
|
||||
get_sequence_parallel_rank = None
|
||||
xFuserLongContextAttention = None
|
||||
|
||||
|
||||
class HunyuanAttnProcessor2_0:
|
||||
r"""
|
||||
@@ -217,7 +232,14 @@ class LazyKVCompressionProcessor2_0:
|
||||
|
||||
class EasyAnimateAttnProcessor2_0:
|
||||
def __init__(self):
|
||||
pass
|
||||
if xFuserLongContextAttention is not None:
|
||||
try:
|
||||
get_sequence_parallel_world_size()
|
||||
self.hybrid_seq_parallel_attn = xFuserLongContextAttention()
|
||||
except Exception:
|
||||
self.hybrid_seq_parallel_attn = None
|
||||
else:
|
||||
self.hybrid_seq_parallel_attn = None
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
@@ -284,11 +306,30 @@ class EasyAnimateAttnProcessor2_0:
|
||||
if not attn.is_cross_attention:
|
||||
key[:, :, text_seq_length:] = apply_rotary_emb(key[:, :, text_seq_length:], image_rotary_emb)
|
||||
|
||||
hidden_states = F.scaled_dot_product_attention(
|
||||
query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
|
||||
)
|
||||
if self.hybrid_seq_parallel_attn is None:
|
||||
hidden_states = F.scaled_dot_product_attention(
|
||||
query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
|
||||
)
|
||||
hidden_states = hidden_states.transpose(1, 2)
|
||||
else:
|
||||
sp_world_rank = get_sequence_parallel_rank()
|
||||
sp_world_size = get_sequence_parallel_world_size()
|
||||
|
||||
hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
|
||||
img_q = query[:, :, text_seq_length:].transpose(1,2)
|
||||
txt_q = query[:, :, :text_seq_length].transpose(1,2)
|
||||
img_k = key[:, :, text_seq_length:].transpose(1,2)
|
||||
txt_k = key[:, :, :text_seq_length].transpose(1,2)
|
||||
img_v = value[:, :, text_seq_length:].transpose(1,2)
|
||||
txt_v = value[:, :, :text_seq_length].transpose(1,2)
|
||||
|
||||
hidden_states = self.hybrid_seq_parallel_attn(None,
|
||||
img_q, img_k, img_v, dropout_p=0.0, causal=False,
|
||||
joint_tensor_query=txt_q,
|
||||
joint_tensor_key=txt_k,
|
||||
joint_tensor_value=txt_v,
|
||||
joint_strategy='front',)
|
||||
|
||||
hidden_states = hidden_states.reshape(batch_size, -1, attn.heads * head_dim)
|
||||
|
||||
if attn2 is None:
|
||||
# linear proj
|
||||
|
||||
@@ -51,6 +51,23 @@ except:
|
||||
from diffusers.models.embeddings import \
|
||||
CaptionProjection as PixArtAlphaTextProjection
|
||||
|
||||
try:
|
||||
import xfuser
|
||||
from xfuser.core.distributed import (
|
||||
get_sequence_parallel_world_size,
|
||||
get_sequence_parallel_rank,
|
||||
get_sp_group,
|
||||
initialize_model_parallel,
|
||||
init_distributed_environment
|
||||
)
|
||||
except Exception as ex:
|
||||
xfuser = None
|
||||
get_sequence_parallel_world_size = None
|
||||
get_sequence_parallel_rank = None
|
||||
get_sp_group = None
|
||||
initialize_model_parallel = None
|
||||
init_distributed_environment = None
|
||||
|
||||
|
||||
class CLIPProjection(nn.Module):
|
||||
"""
|
||||
@@ -1375,6 +1392,14 @@ class EasyAnimateTransformer3DModel(ModelMixin, ConfigMixin):
|
||||
|
||||
self.gradient_checkpointing = False
|
||||
|
||||
try:
|
||||
self.sp_world_size = get_sequence_parallel_world_size()
|
||||
self.sp_world_rank = get_sequence_parallel_rank()
|
||||
except Exception:
|
||||
self.sp_world_size = 1
|
||||
self.sp_world_rank = 0
|
||||
xfuser = None
|
||||
|
||||
def _set_gradient_checkpointing(self, module, value=False):
|
||||
self.gradient_checkpointing = value
|
||||
|
||||
@@ -1399,6 +1424,40 @@ class EasyAnimateTransformer3DModel(ModelMixin, ConfigMixin):
|
||||
):
|
||||
batch_size, channels, video_length, height, width = hidden_states.size()
|
||||
|
||||
if xfuser is not None and self.sp_world_size > 1:
|
||||
if hidden_states.shape[-2] // self.patch_size % self.sp_world_size == 0:
|
||||
split_height = height // self.sp_world_size
|
||||
split_dim = -2
|
||||
elif hidden_states.shape[-2] // self.patch_size % self.sp_world_size == 0:
|
||||
split_width = width // self.sp_world_size
|
||||
split_dim = -1
|
||||
else:
|
||||
raise ValueError("Cannot split video sequence into ulysses_degree x ring_degree=%d parts evenly, hidden_states.shape=%s" % (self.sp_world_size, str(hidden_states.shape)))
|
||||
|
||||
|
||||
hidden_states = torch.chunk(hidden_states, self.sp_world_size, dim=split_dim)[self.sp_world_rank]
|
||||
if inpaint_latents is not None:
|
||||
inpaint_latents = torch.chunk(inpaint_latents, self.sp_world_size, dim=split_dim)[self.sp_world_rank]
|
||||
|
||||
if image_rotary_emb is not None:
|
||||
embed_dim = image_rotary_emb[0].shape[-1]
|
||||
freq_cos = image_rotary_emb[0].reshape(video_length, height // self.patch_size, width // self.patch_size, embed_dim)
|
||||
freq_sin = image_rotary_emb[1].reshape(video_length, height // self.patch_size, width // self.patch_size, embed_dim)
|
||||
|
||||
freq_cos = torch.chunk(freq_cos, self.sp_world_size, dim=split_dim-1)[self.sp_world_rank]
|
||||
freq_sin = torch.chunk(freq_sin, self.sp_world_size, dim=split_dim-1)[self.sp_world_rank]
|
||||
|
||||
freq_cos = freq_cos.reshape(-1, embed_dim)
|
||||
freq_sin = freq_sin.reshape(-1, embed_dim)
|
||||
|
||||
image_rotary_emb = (freq_cos, freq_sin)
|
||||
|
||||
if split_dim == -2:
|
||||
height = split_height
|
||||
elif split_dim == -1:
|
||||
width = split_width
|
||||
|
||||
|
||||
# 1. Time embedding
|
||||
temb = self.time_proj(timestep).to(dtype=hidden_states.dtype)
|
||||
temb = self.time_embedding(temb, timestep_cond)
|
||||
@@ -1486,6 +1545,9 @@ class EasyAnimateTransformer3DModel(ModelMixin, ConfigMixin):
|
||||
output = hidden_states.reshape(batch_size, video_length, height // p, width // p, channels, p, p)
|
||||
output = output.permute(0, 4, 1, 2, 5, 3, 6).flatten(5, 6).flatten(3, 4)
|
||||
|
||||
if xfuser is not None and self.sp_world_size > 1:
|
||||
output = get_sp_group().all_gather(output, dim=split_dim)
|
||||
|
||||
if not return_dict:
|
||||
return (output,)
|
||||
return Transformer2DModelOutput(sample=output)
|
||||
@@ -1606,4 +1668,4 @@ class EasyAnimateTransformer3DModel(ModelMixin, ConfigMixin):
|
||||
print(f"### attn1 Parameters: {sum(params) / 1e6} M")
|
||||
|
||||
model = model.to(torch_dtype)
|
||||
return model
|
||||
return model
|
||||
|
||||
@@ -0,0 +1,401 @@
|
||||
import os
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
from diffusers import (DDIMScheduler,
|
||||
DPMSolverMultistepScheduler,
|
||||
EulerAncestralDiscreteScheduler, EulerDiscreteScheduler,
|
||||
PNDMScheduler)
|
||||
from omegaconf import OmegaConf
|
||||
from PIL import Image
|
||||
from transformers import (BertModel, BertTokenizer, CLIPImageProcessor,
|
||||
CLIPVisionModelWithProjection,
|
||||
T5EncoderModel, T5Tokenizer)
|
||||
|
||||
from easyanimate.models import (name_to_autoencoder_magvit,
|
||||
name_to_transformer3d)
|
||||
from easyanimate.pipeline.pipeline_easyanimate_inpaint import \
|
||||
EasyAnimateInpaintPipeline
|
||||
from easyanimate.pipeline.pipeline_easyanimate_multi_text_encoder_inpaint import \
|
||||
EasyAnimatePipeline_Multi_Text_Encoder_Inpaint
|
||||
from easyanimate.utils.lora_utils import merge_lora, unmerge_lora
|
||||
from easyanimate.utils.utils import get_image_to_video_latent, save_videos_grid
|
||||
from easyanimate.utils.fp8_optimization import convert_weight_dtype_wrapper
|
||||
|
||||
|
||||
try:
|
||||
import xfuser
|
||||
from xfuser.core.distributed import (
|
||||
get_sequence_parallel_world_size,
|
||||
get_sequence_parallel_rank,
|
||||
get_sp_group,
|
||||
initialize_model_parallel,
|
||||
init_distributed_environment
|
||||
)
|
||||
except:
|
||||
xfuser = None
|
||||
get_sequence_parallel_world_size = None
|
||||
get_sequence_parallel_rank = None
|
||||
get_sp_group = None
|
||||
initialize_model_parallel = None
|
||||
init_distributed_environment = None
|
||||
|
||||
ulysses_degree = 2
|
||||
ring_degree = 2
|
||||
|
||||
if ulysses_degree > 1 or ring_degree > 1:
|
||||
dist.init_process_group("nccl")
|
||||
print('parallel inference enabled: ulysses_degree=%d ring_degree=%d rank=%d world_size=%d' % (
|
||||
ulysses_degree, ring_degree, dist.get_rank(),
|
||||
dist.get_world_size()))
|
||||
assert dist.get_world_size() == ring_degree * ulysses_degree, \
|
||||
"number of GPUs should be equal to ring_degree * ulysses_degree."
|
||||
init_distributed_environment(rank=dist.get_rank(), world_size=dist.get_world_size())
|
||||
initialize_model_parallel(sequence_parallel_degree=dist.get_world_size(),
|
||||
ring_degree=ring_degree,
|
||||
ulysses_degree=ulysses_degree)
|
||||
device = torch.device("cuda:%d" % dist.get_rank())
|
||||
print('rank=%d device=%s' % (dist.get_rank(), str(device)))
|
||||
else:
|
||||
device = "cuda"
|
||||
|
||||
# GPU memory mode, which can be choosen in [model_cpu_offload, model_cpu_offload_and_qfloat8, sequential_cpu_offload].
|
||||
# model_cpu_offload means that the entire model will be moved to the CPU after use, which can save some GPU memory.
|
||||
#
|
||||
# model_cpu_offload_and_qfloat8 indicates that the entire model will be moved to the CPU after use,
|
||||
# and the transformer model has been quantized to float8, which can save more GPU memory.
|
||||
#
|
||||
# sequential_cpu_offload means that each layer of the model will be moved to the CPU after use,
|
||||
# resulting in slower speeds but saving a large amount of GPU memory.
|
||||
GPU_memory_mode = "model_cpu_offload"
|
||||
|
||||
# Config and model path
|
||||
config_path = "config/easyanimate_video_v5_magvit_multi_text_encoder.yaml"
|
||||
model_name = "models/Diffusion_Transformer/EasyAnimateV5-12b-zh-InP"
|
||||
|
||||
# Choose the sampler in "Euler" "Euler A" "DPM++" "PNDM" and "DDIM"
|
||||
# EasyAnimateV1, V2 and V3 cannot use DDIM.
|
||||
# EasyAnimateV4 and V5 support DDIM.
|
||||
sampler_name = "DDIM"
|
||||
|
||||
# Load pretrained model if need
|
||||
transformer_path = None
|
||||
# Only V1 does need a motion module
|
||||
motion_module_path = None
|
||||
vae_path = None
|
||||
lora_path = None
|
||||
|
||||
# Other params
|
||||
sample_size = [384, 672]
|
||||
# In EasyAnimateV1, the video_length of video is 40 ~ 80.
|
||||
# In EasyAnimateV2, V3, V4, the video_length of video is 1 ~ 144.
|
||||
# In EasyAnimateV5, the video_length of video is 1 ~ 49.
|
||||
# If u want to generate a image, please set the video_length = 1.
|
||||
video_length = 49
|
||||
fps = 8
|
||||
|
||||
# If you want to generate ultra long videos, please set partial_video_length as the length of each sub video segment
|
||||
partial_video_length = None
|
||||
overlap_video_length = 4
|
||||
|
||||
# Use torch.float16 if GPU does not support torch.bfloat16
|
||||
# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
|
||||
weight_dtype = torch.bfloat16
|
||||
# If you want to generate from text, please set the validation_image_start = None and validation_image_end = None
|
||||
validation_image_start = "asset/1.png"
|
||||
validation_image_end = None
|
||||
|
||||
# EasyAnimateV1, V2 and V3 support English.
|
||||
# EasyAnimateV4 and V5 support English and Chinese.
|
||||
# 使用更长的neg prompt如"模糊，突变，变形，失真，画面暗，文本字幕，画面固定，连环画，漫画，线稿，没有主体。"，可以增加稳定性
|
||||
# 在neg prompt中添加"安静，固定"等词语可以增加动态性。
|
||||
prompt = "一条狗正在摇头。质量高、杰作、最佳品质、高分辨率、超精细、梦幻般。"
|
||||
negative_prompt = "扭曲的身体，肢体残缺，文本字幕，漫画，静止，丑陋，错误，乱码。"
|
||||
#
|
||||
# Using longer neg prompt such as "Blurring, mutation, deformation, distortion, dark and solid, comics, text subtitles, line art." can increase stability
|
||||
# Adding words such as "quiet, solid" to the neg prompt can increase dynamism.
|
||||
# prompt = "The dog is shaking head. The video is of high quality, and the view is very clear. High quality, masterpiece, best quality, highres, ultra-detailed, fantastic."
|
||||
# negative_prompt = "Twisted body, limb deformities, text captions, comic, static, ugly, error, messy code."
|
||||
guidance_scale = 6.0
|
||||
seed = 43
|
||||
num_inference_steps = 50
|
||||
lora_weight = 0.60
|
||||
save_path = "samples/easyanimate-videos_i2v"
|
||||
|
||||
config = OmegaConf.load(config_path)
|
||||
|
||||
# Get Transformer
|
||||
Choosen_Transformer3DModel = name_to_transformer3d[
|
||||
config['transformer_additional_kwargs'].get('transformer_type', 'Transformer3DModel')
|
||||
]
|
||||
|
||||
transformer_additional_kwargs = OmegaConf.to_container(config['transformer_additional_kwargs'])
|
||||
if weight_dtype == torch.float16:
|
||||
transformer_additional_kwargs["upcast_attention"] = True
|
||||
|
||||
transformer = Choosen_Transformer3DModel.from_pretrained_2d(
|
||||
model_name,
|
||||
subfolder="transformer",
|
||||
transformer_additional_kwargs=transformer_additional_kwargs,
|
||||
torch_dtype=torch.float8_e4m3fn if GPU_memory_mode == "model_cpu_offload_and_qfloat8" else weight_dtype,
|
||||
low_cpu_mem_usage=True,
|
||||
)
|
||||
|
||||
transformer = transformer.to(device)
|
||||
|
||||
if transformer_path is not None:
|
||||
print(f"From checkpoint: {transformer_path}")
|
||||
if transformer_path.endswith("safetensors"):
|
||||
from safetensors.torch import load_file, safe_open
|
||||
state_dict = load_file(transformer_path)
|
||||
else:
|
||||
state_dict = torch.load(transformer_path, map_location="cpu")
|
||||
state_dict = state_dict["state_dict"] if "state_dict" in state_dict else state_dict
|
||||
|
||||
m, u = transformer.load_state_dict(state_dict, strict=False)
|
||||
print(f"missing keys: {len(m)}, unexpected keys: {len(u)}")
|
||||
|
||||
if motion_module_path is not None:
|
||||
print(f"From Motion Module: {motion_module_path}")
|
||||
if motion_module_path.endswith("safetensors"):
|
||||
from safetensors.torch import load_file, safe_open
|
||||
state_dict = load_file(motion_module_path)
|
||||
else:
|
||||
state_dict = torch.load(motion_module_path, map_location="cpu")
|
||||
state_dict = state_dict["state_dict"] if "state_dict" in state_dict else state_dict
|
||||
|
||||
m, u = transformer.load_state_dict(state_dict, strict=False)
|
||||
print(f"missing keys: {len(m)}, unexpected keys: {len(u)}, {u}")
|
||||
|
||||
# Get Vae
|
||||
Choosen_AutoencoderKL = name_to_autoencoder_magvit[
|
||||
config['vae_kwargs'].get('vae_type', 'AutoencoderKL')
|
||||
]
|
||||
vae = Choosen_AutoencoderKL.from_pretrained(
|
||||
model_name,
|
||||
subfolder="vae",
|
||||
vae_additional_kwargs=OmegaConf.to_container(config['vae_kwargs'])
|
||||
).to(weight_dtype).to(device)
|
||||
if config['vae_kwargs'].get('vae_type', 'AutoencoderKL') == 'AutoencoderKLMagvit' and weight_dtype == torch.float16:
|
||||
vae.upcast_vae = True
|
||||
|
||||
if vae_path is not None:
|
||||
print(f"From checkpoint: {vae_path}")
|
||||
if vae_path.endswith("safetensors"):
|
||||
from safetensors.torch import load_file, safe_open
|
||||
state_dict = load_file(vae_path)
|
||||
else:
|
||||
state_dict = torch.load(vae_path, map_location="cpu")
|
||||
state_dict = state_dict["state_dict"] if "state_dict" in state_dict else state_dict
|
||||
|
||||
m, u = vae.load_state_dict(state_dict, strict=False)
|
||||
print(f"missing keys: {len(m)}, unexpected keys: {len(u)}")
|
||||
|
||||
if config['text_encoder_kwargs'].get('enable_multi_text_encoder', False):
|
||||
tokenizer = BertTokenizer.from_pretrained(
|
||||
model_name, subfolder="tokenizer"
|
||||
)
|
||||
tokenizer_2 = T5Tokenizer.from_pretrained(
|
||||
model_name, subfolder="tokenizer_2"
|
||||
)
|
||||
else:
|
||||
tokenizer = T5Tokenizer.from_pretrained(
|
||||
model_name, subfolder="tokenizer"
|
||||
)
|
||||
tokenizer_2 = None
|
||||
|
||||
if config['text_encoder_kwargs'].get('enable_multi_text_encoder', False):
|
||||
text_encoder = BertModel.from_pretrained(
|
||||
model_name, subfolder="text_encoder", torch_dtype=weight_dtype
|
||||
).to(device)
|
||||
text_encoder_2 = T5EncoderModel.from_pretrained(
|
||||
model_name, subfolder="text_encoder_2", torch_dtype=weight_dtype
|
||||
).to(device)
|
||||
else:
|
||||
text_encoder = T5EncoderModel.from_pretrained(
|
||||
model_name, subfolder="text_encoder", torch_dtype=weight_dtype
|
||||
).to(device)
|
||||
text_encoder_2 = None
|
||||
|
||||
if transformer.config.in_channels != vae.config.latent_channels and config['transformer_additional_kwargs'].get('enable_clip_in_inpaint', True):
|
||||
clip_image_encoder = CLIPVisionModelWithProjection.from_pretrained(
|
||||
model_name, subfolder="image_encoder"
|
||||
).to(device, weight_dtype)
|
||||
clip_image_processor = CLIPImageProcessor.from_pretrained(
|
||||
model_name, subfolder="image_encoder"
|
||||
)
|
||||
else:
|
||||
clip_image_encoder = None
|
||||
clip_image_processor = None
|
||||
|
||||
# Get Scheduler
|
||||
Choosen_Scheduler = scheduler_dict = {
|
||||
"Euler": EulerDiscreteScheduler,
|
||||
"Euler A": EulerAncestralDiscreteScheduler,
|
||||
"DPM++": DPMSolverMultistepScheduler,
|
||||
"PNDM": PNDMScheduler,
|
||||
"DDIM": DDIMScheduler,
|
||||
}[sampler_name]
|
||||
|
||||
scheduler = Choosen_Scheduler.from_pretrained(
|
||||
model_name,
|
||||
subfolder="scheduler"
|
||||
)
|
||||
if config['text_encoder_kwargs'].get('enable_multi_text_encoder', False):
|
||||
pipeline = EasyAnimatePipeline_Multi_Text_Encoder_Inpaint.from_pretrained(
|
||||
model_name,
|
||||
text_encoder=text_encoder,
|
||||
text_encoder_2=text_encoder_2,
|
||||
tokenizer=tokenizer,
|
||||
tokenizer_2=tokenizer_2,
|
||||
vae=vae,
|
||||
transformer=transformer,
|
||||
scheduler=scheduler,
|
||||
torch_dtype=weight_dtype,
|
||||
clip_image_encoder=clip_image_encoder,
|
||||
clip_image_processor=clip_image_processor,
|
||||
)
|
||||
else:
|
||||
pipeline = EasyAnimateInpaintPipeline.from_pretrained(
|
||||
model_name,
|
||||
text_encoder=text_encoder,
|
||||
tokenizer=tokenizer,
|
||||
vae=vae,
|
||||
transformer=transformer,
|
||||
scheduler=scheduler,
|
||||
torch_dtype=weight_dtype,
|
||||
clip_image_encoder=clip_image_encoder,
|
||||
clip_image_processor=clip_image_processor,
|
||||
)
|
||||
|
||||
if GPU_memory_mode == "sequential_cpu_offload":
|
||||
pipeline.enable_sequential_cpu_offload(device=device)
|
||||
elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
|
||||
pipeline.enable_model_cpu_offload(device=device)
|
||||
convert_weight_dtype_wrapper(transformer, weight_dtype)
|
||||
else:
|
||||
pipeline.enable_model_cpu_offload(device=device)
|
||||
|
||||
# print('pipeline to device=%s' % str(device))
|
||||
# pipeline.to(device)
|
||||
# print('pipeline.device=%s' % str(pipeline.device))
|
||||
|
||||
|
||||
generator = torch.Generator(device=device).manual_seed(seed)
|
||||
|
||||
if lora_path is not None:
|
||||
pipeline = merge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype)
|
||||
|
||||
if partial_video_length is not None:
|
||||
init_frames = 0
|
||||
last_frames = init_frames + partial_video_length
|
||||
while init_frames < video_length:
|
||||
if last_frames >= video_length:
|
||||
if pipeline.vae.quant_conv.weight.ndim==5:
|
||||
mini_batch_encoder = pipeline.vae.mini_batch_encoder
|
||||
_partial_video_length = video_length - init_frames
|
||||
if vae.cache_mag_vae:
|
||||
_partial_video_length = int((_partial_video_length - 1) // vae.mini_batch_encoder * vae.mini_batch_encoder) + 1
|
||||
else:
|
||||
_partial_video_length = int(_partial_video_length // vae.mini_batch_encoder * vae.mini_batch_encoder)
|
||||
else:
|
||||
_partial_video_length = video_length - init_frames
|
||||
|
||||
if _partial_video_length <= 0:
|
||||
break
|
||||
else:
|
||||
_partial_video_length = partial_video_length
|
||||
|
||||
input_video, input_video_mask, clip_image = get_image_to_video_latent(validation_image, None, video_length=_partial_video_length, sample_size=sample_size)
|
||||
|
||||
with torch.no_grad():
|
||||
sample = pipeline(
|
||||
prompt,
|
||||
video_length = _partial_video_length,
|
||||
negative_prompt = negative_prompt,
|
||||
height = sample_size[0],
|
||||
width = sample_size[1],
|
||||
generator = generator,
|
||||
guidance_scale = guidance_scale,
|
||||
num_inference_steps = num_inference_steps,
|
||||
|
||||
video = input_video,
|
||||
mask_video = input_video_mask,
|
||||
clip_image = clip_image,
|
||||
).videos
|
||||
|
||||
if init_frames != 0:
|
||||
mix_ratio = torch.from_numpy(
|
||||
np.array([float(_index) / float(overlap_video_length) for _index in range(overlap_video_length)], np.float32)
|
||||
).unsqueeze(0).unsqueeze(0).unsqueeze(-1).unsqueeze(-1)
|
||||
|
||||
new_sample[:, :, -overlap_video_length:] = new_sample[:, :, -overlap_video_length:] * (1 - mix_ratio) + \
|
||||
sample[:, :, :overlap_video_length] * mix_ratio
|
||||
new_sample = torch.cat([new_sample, sample[:, :, overlap_video_length:]], dim = 2)
|
||||
|
||||
sample = new_sample
|
||||
else:
|
||||
new_sample = sample
|
||||
|
||||
if last_frames >= video_length:
|
||||
break
|
||||
|
||||
validation_image = [
|
||||
Image.fromarray(
|
||||
(sample[0, :, _index].transpose(0, 1).transpose(1, 2) * 255).numpy().astype(np.uint8)
|
||||
) for _index in range(-overlap_video_length, 0)
|
||||
]
|
||||
|
||||
init_frames = init_frames + _partial_video_length - overlap_video_length
|
||||
last_frames = init_frames + _partial_video_length
|
||||
else:
|
||||
if vae.cache_mag_vae:
|
||||
video_length = int((video_length - 1) // vae.mini_batch_encoder * vae.mini_batch_encoder) + 1 if video_length != 1 else 1
|
||||
else:
|
||||
video_length = int(video_length // vae.mini_batch_encoder * vae.mini_batch_encoder) if video_length != 1 else 1
|
||||
input_video, input_video_mask, clip_image = get_image_to_video_latent(validation_image_start, validation_image_end, video_length=video_length, sample_size=sample_size)
|
||||
|
||||
with torch.no_grad():
|
||||
sample = pipeline(
|
||||
prompt,
|
||||
video_length = video_length,
|
||||
negative_prompt = negative_prompt,
|
||||
height = sample_size[0],
|
||||
width = sample_size[1],
|
||||
generator = generator,
|
||||
guidance_scale = guidance_scale,
|
||||
num_inference_steps = num_inference_steps,
|
||||
|
||||
video = input_video,
|
||||
mask_video = input_video_mask,
|
||||
clip_image = clip_image,
|
||||
).videos
|
||||
|
||||
if lora_path is not None:
|
||||
pipeline = unmerge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype)
|
||||
|
||||
if not os.path.exists(save_path):
|
||||
os.makedirs(save_path, exist_ok=True)
|
||||
|
||||
index = len([path for path in os.listdir(save_path)]) + 1
|
||||
prefix = str(index).zfill(8)
|
||||
|
||||
if video_length == 1:
|
||||
save_sample_path = os.path.join(save_path, prefix + f".png")
|
||||
|
||||
image = sample[0, :, 0]
|
||||
image = image.transpose(0, 1).transpose(1, 2)
|
||||
image = (image * 255).numpy().astype(np.uint8)
|
||||
image = Image.fromarray(image)
|
||||
image.save(save_sample_path)
|
||||
else:
|
||||
if ulysses_degree * ring_degree > 1:
|
||||
if dist.get_rank() == 0:
|
||||
video_path = os.path.join(save_path, prefix + ".mp4")
|
||||
save_videos_grid(sample, video_path, fps=fps)
|
||||
print('save video to %s' % video_path)
|
||||
else:
|
||||
video_path = os.path.join(save_path, prefix + ".mp4")
|
||||
save_videos_grid(sample, video_path, fps=fps)
|
||||
print('save video to %s' % video_path)
|
||||
Reference in New Issue
Block a user