From b0a104fef6ed145c5b2b17a00a04406d1d03b16f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9CBrianChen1129=E2=80=9D?= Date: Wed, 2 Jul 2025 04:38:01 +0000 Subject: [PATCH] update --- fastvideo/v1/dataset/utils.py | 52 ++++++++++++++++++----------------- 1 file changed, 27 insertions(+), 25 deletions(-) diff --git a/fastvideo/v1/dataset/utils.py b/fastvideo/v1/dataset/utils.py index 49f1152da..57205a959 100644 --- a/fastvideo/v1/dataset/utils.py +++ b/fastvideo/v1/dataset/utils.py @@ -133,35 +133,37 @@ def collate_rows_from_parquet_schema(rows, tensor_list = [] for row in rows: - # Get tensor data from row using the existing helper function pattern - shape_key = f"{tensor_name}_shape" - bytes_key = f"{tensor_name}_bytes" + try: + # Get tensor data from row using the existing helper function pattern + shape_key = f"{tensor_name}_shape" + bytes_key = f"{tensor_name}_bytes" - if shape_key in row and bytes_key in row: - shape = row[shape_key] - bytes_data = row[bytes_key] + if shape_key in row and bytes_key in row: + shape = row[shape_key] + bytes_data = row[bytes_key] - if len(bytes_data) == 0: - tensor = torch.zeros(0, dtype=torch.bfloat16) - else: - # Convert bytes to tensor using float32 as default - if tensor_name == 'text_embedding' and random.random( - ) < cfg_rate: - data = np.zeros((512, 4096), dtype=np.float32) + if len(bytes_data) == 0: + tensor = torch.zeros(0, dtype=torch.bfloat16) else: - data = np.frombuffer( - bytes_data, dtype=np.float32).reshape(shape).copy() - tensor = torch.from_numpy(data) - # if len(data.shape) == 3: - # B, L, D = tensor.shape - # assert B == 1, "Batch size must be 1" - # tensor = tensor.squeeze(0) - - tensor_list.append(tensor) - else: - # Handle missing tensor data - tensor_list.append(torch.zeros(0, dtype=torch.bfloat16)) + # Convert bytes to tensor using float32 as default + if tensor_name == 'text_embedding' and random.random( + ) < cfg_rate: + data = np.zeros((512, 4096), dtype=np.float32) + else: + data = np.frombuffer( + bytes_data, dtype=np.float32).reshape(shape).copy() + tensor = torch.from_numpy(data) + # if len(data.shape) == 3: + # B, L, D = tensor.shape + # assert B == 1, "Batch size must be 1" + # tensor = tensor.squeeze(0) + tensor_list.append(tensor) + else: + # Handle missing tensor data + tensor_list.append(torch.zeros(0, dtype=torch.bfloat16)) + except: + continue # Stack tensors with special handling for text embeddings if tensor_name == 'text_embedding': # Handle text embeddings with padding