Skip to content

full video can be generated when the args.L parameter is small #36

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wailovet wants to merge 2 commits into
base: main
Choose a base branch
from
Open

full video can be generated when the args.L parameter is small #36

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
01d8ff8
can generate full video
wailovet Mar 28, 2024
34f837f
full video can be generated when the args.L parameter is small
wailovet Mar 28, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
159 changes: 96 additions & 63 deletions scripts/vid2vid.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -62,6 +62,7 @@ def main():
config.pretrained_base_model_path,
subfolder="unet",
).to(dtype=weight_dtype, device="cuda")


inference_config_path = config.inference_config
infer_config = OmegaConf.load(inference_config_path)
Expand Down Expand Up @@ -117,7 +118,6 @@ def main():

lmk_extractor = LMKExtractor()
vis = FaceMeshVisualizer(forehead_edge=False)


for ref_image_path in config["test_cases"].keys():
# Each ref_image may correspond to multiple actions
Expand All @@ -143,73 +143,106 @@ def main():
[transforms.Resize((height, width)), transforms.ToTensor()]
)

step = 1
if src_fps == 60:
src_fps = 30
step = 2
# step = 1
# if src_fps == 60:
# src_fps = 30
# step = 2

pose_trans_list = []
verts_list = []
bs_list = []
src_tensor_list = []
args_L = len(source_images) if args.L is None else args.L*step
for src_image_pil in source_images[: args_L: step]:
src_tensor_list.append(pose_transform(src_image_pil))
src_img_np = cv2.cvtColor(np.array(src_image_pil), cv2.COLOR_RGB2BGR)
frame_height, frame_width, _ = src_img_np.shape
src_img_result = lmk_extractor(src_img_np)
if src_img_result is None:
break
pose_trans_list.append(src_img_result['trans_mat'])
verts_list.append(src_img_result['lmks3d'])
bs_list.append(src_img_result['bs'])
args_L = len(source_images) if args.L is None else args.L

total_frames = len(source_images)
batch_size = total_frames // args_L
if total_frames % args_L != 0:
batch_size += 1

pose_arr = np.array(pose_trans_list)
verts_arr = np.array(verts_list)
bs_arr = np.array(bs_list)
min_bs_idx = np.argmin(bs_arr.sum(1))

# face retarget
verts_arr = verts_arr - verts_arr[min_bs_idx] + face_result['lmks3d']
# project 3D mesh to 2D landmark
projected_vertices = project_points_with_trans(verts_arr, pose_arr, [frame_height, frame_width])

pose_list = []
for i, verts in enumerate(projected_vertices):
lmk_img = vis.draw_landmarks((frame_width, frame_height), verts, normed=False)
pose_image_np = cv2.resize(lmk_img, (width, height))
pose_list.append(pose_image_np)

pose_list = np.array(pose_list)

video_length = len(src_tensor_list)

ref_image_tensor = pose_transform(ref_image_pil) # (c, h, w)
ref_image_tensor = ref_image_tensor.unsqueeze(1).unsqueeze(
0
) # (1, c, 1, h, w)
ref_image_tensor = repeat(
ref_image_tensor, "b c f h w -> b c (repeat f) h w", repeat=video_length
)
video_list = None
ref_image_list = None
ref_src_tensor_list = None

for batch_num in range(batch_size):
print(f"Processing batch {batch_num+1}/{batch_size}")
pose_trans_list = []
verts_list = []
bs_list = []
src_tensor_list = []

# for src_image_pil in source_images[: args_L: step]:
for j in range(args_L):
if batch_num*args_L+j >= total_frames:
break
src_image_pil = source_images[batch_num*args_L+j]
src_tensor_list.append(pose_transform(src_image_pil))
src_img_np = cv2.cvtColor(np.array(src_image_pil), cv2.COLOR_RGB2BGR)
frame_height, frame_width, _ = src_img_np.shape
src_img_result = lmk_extractor(src_img_np)
if src_img_result is None:
break
pose_trans_list.append(src_img_result['trans_mat'])
verts_list.append(src_img_result['lmks3d'])
bs_list.append(src_img_result['bs'])


pose_arr = np.array(pose_trans_list)
verts_arr = np.array(verts_list)
bs_arr = np.array(bs_list)
min_bs_idx = np.argmin(bs_arr.sum(1))

# face retarget
verts_arr = verts_arr - verts_arr[min_bs_idx] + face_result['lmks3d']
# project 3D mesh to 2D landmark
projected_vertices = project_points_with_trans(verts_arr, pose_arr, [frame_height, frame_width])

pose_list = []
for i, verts in enumerate(projected_vertices):
lmk_img = vis.draw_landmarks((frame_width, frame_height), verts, normed=False)
pose_image_np = cv2.resize(lmk_img, (width, height))
pose_list.append(pose_image_np)

pose_list = np.array(pose_list)

video_length = len(src_tensor_list)

ref_image_tensor = pose_transform(ref_image_pil) # (c, h, w)
ref_image_tensor = ref_image_tensor.unsqueeze(1).unsqueeze(
0
) # (1, c, 1, h, w)
ref_image_tensor = repeat(
ref_image_tensor, "b c f h w -> b c (repeat f) h w", repeat=video_length
)

src_tensor = torch.stack(src_tensor_list, dim=0) # (f, c, h, w)
src_tensor = src_tensor.transpose(0, 1)
src_tensor = src_tensor.unsqueeze(0)
# print(f"pose_list: {pose_list.shape}, ref_pose: {ref_pose.shape}, ref_image_tensor: {ref_image_tensor.shape}, src_tensor: {src_tensor.shape}")
# input("Press Enter to continue...")
video = pipe(
ref_image_pil,
pose_list,
ref_pose,
width,
height,
video_length,
args.steps,
args.cfg,
generator=generator,
).videos
if video_list is None:
video_list = video
ref_image_list = ref_image_tensor
ref_src_tensor_list = src_tensor

else:
video_list = torch.cat([video_list, video], dim=2)
ref_image_list = torch.cat([ref_image_list, ref_image_tensor], dim=2)
ref_src_tensor_list = torch.cat([ref_src_tensor_list, src_tensor], dim=2)





video = torch.cat([ref_image_list, video_list, ref_src_tensor_list], dim=0)

src_tensor = torch.stack(src_tensor_list, dim=0) # (f, c, h, w)
src_tensor = src_tensor.transpose(0, 1)
src_tensor = src_tensor.unsqueeze(0)

video = pipe(
ref_image_pil,
pose_list,
ref_pose,
width,
height,
video_length,
args.steps,
args.cfg,
generator=generator,
).videos

video = torch.cat([ref_image_tensor, video, src_tensor], dim=0)
save_path = f"{save_dir}/{ref_name}_{pose_name}_{args.H}x{args.W}_{int(args.cfg)}_{time_str}_noaudio.mp4"
save_videos_grid(
video,
Expand Down