import os import cv2 import math import json import torch import random import librosa import traceback import torchvision import numpy as np import pandas as pd from PIL import Image from einops import rearrange from torch.utils.data import Dataset from decord import VideoReader, cpu from transformers import CLIPImageProcessor import torchvision.transforms as transforms from torchvision.transforms import ToPILImage def get_audio_feature(feature_extractor, audio_path): audio_input, sampling_rate = librosa.load(audio_path, sr=16000) assert sampling_rate == 16000 audio_features = [] window = 750*640 for i in range(0, len(audio_input), window): audio_feature = feature_extractor(audio_input[i:i+window], sampling_rate=sampling_rate, return_tensors="pt", ).input_features audio_features.append(audio_feature) audio_features = torch.cat(audio_features, dim=-1) return audio_features, len(audio_input) // 640 class VideoAudioTextLoaderVal(Dataset): def __init__( self, image_size: int, meta_file: str, **kwargs, ): super().__init__() self.meta_file = meta_file self.image_size = image_size self.text_encoder = kwargs.get("text_encoder", None) # llava_text_encoder self.text_encoder_2 = kwargs.get("text_encoder_2", None) # clipL_text_encoder self.feature_extractor = kwargs.get("feature_extractor", None) self.meta_files = [] csv_data = pd.read_csv(meta_file) for idx in range(len(csv_data)): self.meta_files.append( { "videoid": str(csv_data["videoid"][idx]), "image_path": str(csv_data["image"][idx]), "audio_path": str(csv_data["audio"][idx]), "prompt": str(csv_data["prompt"][idx]), "fps": float(csv_data["fps"][idx]) } ) self.llava_transform = transforms.Compose( [ transforms.Resize((336, 336), interpolation=transforms.InterpolationMode.BILINEAR), transforms.ToTensor(), transforms.Normalize((0.48145466, 0.4578275, 0.4082107), (0.26862954, 0.26130258, 0.27577711)), ] ) self.clip_image_processor = CLIPImageProcessor() self.device = torch.device("cuda") self.weight_dtype = torch.float16 def __len__(self): return len(self.meta_files) @staticmethod def get_text_tokens(text_encoder, description, dtype_encode="video"): text_inputs = text_encoder.text2tokens(description, data_type=dtype_encode) text_ids = text_inputs["input_ids"].squeeze(0) text_mask = text_inputs["attention_mask"].squeeze(0) return text_ids, text_mask def get_batch_data(self, idx): meta_file = self.meta_files[idx] videoid = meta_file["videoid"] image_path = meta_file["image_path"] audio_path = meta_file["audio_path"] prompt = "Authentic, Realistic, Natural, High-quality, Lens-Fixed, " + meta_file["prompt"] fps = meta_file["fps"] img_size = self.image_size ref_image = Image.open(image_path).convert('RGB') # Resize reference image w, h = ref_image.size scale = img_size / min(w, h) new_w = round(w * scale / 64) * 64 new_h = round(h * scale / 64) * 64 if img_size == 704: img_size_long = 1216 if new_w * new_h > img_size * img_size_long: import math scale = math.sqrt(img_size * img_size_long / w / h) new_w = round(w * scale / 64) * 64 new_h = round(h * scale / 64) * 64 ref_image = ref_image.resize((new_w, new_h), Image.LANCZOS) ref_image = np.array(ref_image) ref_image = torch.from_numpy(ref_image) audio_input, audio_len = get_audio_feature(self.feature_extractor, audio_path) audio_prompts = audio_input[0] motion_bucket_id_heads = np.array([25] * 4) motion_bucket_id_exps = np.array([30] * 4) motion_bucket_id_heads = torch.from_numpy(motion_bucket_id_heads) motion_bucket_id_exps = torch.from_numpy(motion_bucket_id_exps) fps = torch.from_numpy(np.array(fps)) to_pil = ToPILImage() pixel_value_ref = rearrange(ref_image.clone().unsqueeze(0), "b h w c -> b c h w") # (b c h w) pixel_value_ref_llava = [self.llava_transform(to_pil(image)) for image in pixel_value_ref] pixel_value_ref_llava = torch.stack(pixel_value_ref_llava, dim=0) pixel_value_ref_clip = self.clip_image_processor( images=Image.fromarray((pixel_value_ref[0].permute(1,2,0)).data.cpu().numpy().astype(np.uint8)), return_tensors="pt" ).pixel_values[0] pixel_value_ref_clip = pixel_value_ref_clip.unsqueeze(0) # Encode text prompts text_ids, text_mask = self.get_text_tokens(self.text_encoder, prompt) text_ids_2, text_mask_2 = self.get_text_tokens(self.text_encoder_2, prompt) # Output batch batch = { "text_prompt": prompt, # "videoid": videoid, "pixel_value_ref": pixel_value_ref.to(dtype=torch.float16), # 参考图，用于vae提特征 (1, 3, h, w), 取值范围(0, 255) "pixel_value_ref_llava": pixel_value_ref_llava.to(dtype=torch.float16), # 参考图，用于llava提特征 (1, 3, 336, 336), 取值范围 = CLIP取值范围 "pixel_value_ref_clip": pixel_value_ref_clip.to(dtype=torch.float16), # 参考图，用于clip_image_encoder提特征 (1, 3, 244, 244), 取值范围 = CLIP取值范围 "audio_prompts": audio_prompts.to(dtype=torch.float16), "motion_bucket_id_heads": motion_bucket_id_heads.to(dtype=text_ids.dtype), "motion_bucket_id_exps": motion_bucket_id_exps.to(dtype=text_ids.dtype), "fps": fps.to(dtype=torch.float16), "text_ids": text_ids.clone(), # 对应llava_text_encoder "text_mask": text_mask.clone(), # 对应llava_text_encoder "text_ids_2": text_ids_2.clone(), # 对应clip_text_encoder "text_mask_2": text_mask_2.clone(), # 对应clip_text_encoder "audio_len": audio_len, "image_path": image_path, "audio_path": audio_path, } return batch def __getitem__(self, idx): return self.get_batch_data(idx)