Spaces:
Running on Zero
Running on Zero
File size: 5,114 Bytes
886038b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 | # Modified by the UniVEdit authors from Tencent-Hunyuan/OmniWeaving (HunyuanVideo-1.5).
import torch
import torchvision.transforms as transforms
import numpy as np
from PIL import Image
from decord import VideoReader
def prepare_custom_video(vr, total_subset=161, nframes=8):
if isinstance(vr, str):
vr = VideoReader(vr)
max_frames = min(len(vr), total_subset)
sample_indices = np.linspace(0, max_frames - 1, nframes).astype(int)
frames = vr.get_batch(sample_indices)
if hasattr(frames, 'asnumpy'):
frames = frames.asnumpy() # decord NDArray
else:
frames = frames.numpy() # PyTorch Tensor
frame_list = [Image.fromarray(frame) for frame in frames]
return frame_list
def get_semantic_images_np(video_path_list, nframes=8):
"""
Extract first frame from each video in video_path_list and return batchified numpy array.
Args:
video_path_list: List of video paths
Returns:
first_images_np: numpy array of shape (B, H, W, C) with dtype uint8 (0-255)
Returns None if video_path_list is None or empty
"""
if video_path_list is None or len(video_path_list) == 0:
return None
first_frames = []
sampled_frames = []
for video_path in video_path_list:
# Read first frame from video
if isinstance(video_path, str):
video_reader = VideoReader(video_path)
else:
assert isinstance(video_path, VideoReader), f"video_path must be either str or VideoReader, got {type(video_path)}"
video_reader = video_path
if len(video_reader) == 0:
raise ValueError(f"Video {video_path} has no frames")
sample_frame_list = prepare_custom_video(video_reader, nframes=nframes)
sampled_frames.append(sample_frame_list)
# Get first frame (index 0)
first_frame = video_reader.get_batch([0])
# Convert to numpy: shape is (1, H, W, C)
if hasattr(first_frame, 'asnumpy'):
first_frame_np = first_frame.asnumpy() # decord NDArray
else:
first_frame_np = first_frame.numpy() # PyTorch Tensor
# Remove batch dimension: shape becomes (H, W, C)
first_frame_np = first_frame_np[0]
first_frames.append(first_frame_np)
del video_reader
# Stack all first frames into batch: shape becomes (B, H, W, C)
first_images_np = np.stack(first_frames, axis=0)
# Ensure uint8 format (0-255)
first_images_np = first_images_np.astype(np.uint8)
return first_images_np, sampled_frames
def get_semantic_images_np2(video_path_list, first_image_pil, vae, F, nframes=8, is_mode=True):
"""
Extract first frame from each video in video_path_list and return batchified numpy array.
Also encode first_image_pil list into cond_latents with values filled at F index=1.
Args:
video_path_list: List of video paths or VideoReader objects
first_image_pil: List of PIL.Image.Image
args: Arguments object containing vae_precision
vae: VAE model
F: Time dimension for cond_latents (must be >= 2)
nframes: Number of frames to sample for each video (default: 8)
is_mode: If True, use mode() for latent distribution, else use sample()
Returns:
first_images_np: numpy array of shape (B, H, W, C) with dtype uint8 (0-255)
sampled_frames: list of sampled frame lists per video
cond_latents: tensor of shape (B, C, F, H, W), values at index=1, others 0
first_image_pil_np: numpy array of shape (B, H, W, C) with dtype uint8 (0-255)
"""
first_images_np, sampled_frames = get_semantic_images_np(video_path_list, nframes=nframes)
if first_image_pil is None or len(first_image_pil) == 0:
raise ValueError("first_image_pil must be a non-empty list of PIL images")
if F < 2:
raise ValueError(f"F must be >= 2 to fill index=1, got {F}")
first_image_pil_np = np.array([np.array(img, dtype=np.uint8) for img in first_image_pil])
image_transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize([0.5], [0.5])
])
first_images_pixel_values = [image_transform(image).unsqueeze(0) for image in first_image_pil]
first_images_pixel_values = torch.cat(first_images_pixel_values).unsqueeze(2).to(vae.device)
with torch.autocast(device_type="cuda", dtype=torch.float16, enabled=True):
if is_mode:
encoded_latents = vae.encode(first_images_pixel_values).latent_dist.mode()
else:
encoded_latents = vae.encode(first_images_pixel_values).latent_dist.sample()
encoded_latents.mul_(vae.config.scaling_factor)
B = encoded_latents.shape[0]
C, H, W = encoded_latents.shape[1], encoded_latents.shape[3], encoded_latents.shape[4]
cond_latents = torch.zeros(B, C, F, H, W, dtype=encoded_latents.dtype, device=encoded_latents.device)
cond_latents[:, :, 1, :, :] = encoded_latents.squeeze(2)
return first_images_np, sampled_frames, cond_latents, first_image_pil_np
|