SolitaryThinker
fa40553afb
checkpoint ode checkpoint fix t2v to i2v lint chekpt checkpoint hacked but working ode_init scripts WIP fixing time embedding WIP fixing time embedding checkpoint update fix revert revert revert update update visualize
229 lines
8.7 KiB
Python
229 lines
8.7 KiB
Python
import random
|
|
from typing import Any, cast
|
|
|
|
import numpy as np
|
|
import torch
|
|
from fastvideo.logger import init_logger
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
def pad(t: torch.Tensor, padding_length: int) -> tuple[torch.Tensor, torch.Tensor]:
|
|
"""
|
|
Pad or crop an embedding [L, D] to exactly padding_length tokens.
|
|
Return:
|
|
- [L, D] tensor in pinned CPU memory
|
|
- [L] attention mask in pinned CPU memory
|
|
"""
|
|
L, D = t.shape
|
|
if padding_length > L: # pad
|
|
pad = torch.zeros(padding_length - L, D, dtype=t.dtype, device=t.device)
|
|
return torch.cat([t, pad], 0), torch.cat(
|
|
[torch.ones(L), torch.zeros(padding_length - L)], 0)
|
|
else: # crop
|
|
return t[:padding_length], torch.ones(padding_length)
|
|
|
|
|
|
def get_torch_tensors_from_row_dict(row_dict,
|
|
keys,
|
|
cfg_rate,
|
|
rng=None) -> dict[str, Any]:
|
|
"""
|
|
Get the latents and prompts from a row dictionary.
|
|
"""
|
|
return_dict = {}
|
|
for key in keys:
|
|
shape, bytes = None, None
|
|
if isinstance(key, tuple):
|
|
for k in key:
|
|
try:
|
|
shape = row_dict[f"{k}_shape"]
|
|
bytes = row_dict[f"{k}_bytes"]
|
|
except KeyError:
|
|
continue
|
|
key = key[0]
|
|
if shape is None or bytes is None:
|
|
raise ValueError(f"Key {key} not found in row_dict")
|
|
else:
|
|
shape = row_dict[f"{key}_shape"]
|
|
bytes = row_dict[f"{key}_bytes"]
|
|
|
|
# TODO (peiyuan): read precision
|
|
if key == 'text_embedding' and (rng.random()
|
|
if rng else random.random()) < cfg_rate:
|
|
data = np.zeros((512, 4096), dtype=np.float32)
|
|
else:
|
|
data = np.frombuffer(bytes, dtype=np.float32).reshape(shape).copy()
|
|
data = torch.from_numpy(data)
|
|
if len(data.shape) == 3:
|
|
B, L, D = data.shape
|
|
assert B == 1, "Batch size must be 1"
|
|
data = data.squeeze(0)
|
|
return_dict[key] = data
|
|
return return_dict
|
|
|
|
|
|
def collate_latents_embs_masks(
|
|
batch_to_process,
|
|
text_padding_length,
|
|
keys,
|
|
cfg_rate=0.0,
|
|
rng=None) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[str]]:
|
|
# Initialize tensors to hold padded embeddings and masks
|
|
all_latents = []
|
|
all_embs = []
|
|
all_masks = []
|
|
caption_text = []
|
|
# Process each row individually
|
|
for i, row in enumerate(batch_to_process):
|
|
# Get tensors from row
|
|
data = get_torch_tensors_from_row_dict(row, keys, cfg_rate, rng)
|
|
latents, emb = data["vae_latent"], data["text_embedding"]
|
|
|
|
padded_emb, mask = pad(emb, text_padding_length)
|
|
# Store in batch tensors
|
|
all_latents.append(latents)
|
|
all_embs.append(padded_emb)
|
|
all_masks.append(mask)
|
|
# TODO(py): remove this once we fix preprocess
|
|
try:
|
|
caption_text.append(row["prompt"])
|
|
except KeyError:
|
|
caption_text.append(row["caption"])
|
|
|
|
# Pin memory for faster transfer to GPU
|
|
all_latents = torch.stack(all_latents)
|
|
all_embs = torch.stack(all_embs)
|
|
all_masks = torch.stack(all_masks)
|
|
|
|
return all_latents, all_embs, all_masks, caption_text
|
|
|
|
|
|
def collate_rows_from_parquet_schema(rows,
|
|
parquet_schema,
|
|
text_padding_length,
|
|
cfg_rate=0.0,
|
|
rng=None) -> dict[str, Any]:
|
|
"""
|
|
Collate rows from parquet files based on the provided schema.
|
|
Dynamically processes tensor fields based on schema and returns batched data.
|
|
|
|
Args:
|
|
rows: List of row dictionaries from parquet files
|
|
parquet_schema: PyArrow schema defining the structure of the data
|
|
|
|
Returns:
|
|
Dict containing batched tensors and metadata
|
|
"""
|
|
if not rows:
|
|
return cast(dict[str, Any], {})
|
|
|
|
# Initialize containers for different data types
|
|
batch_data: dict[str, Any] = {}
|
|
|
|
# Get tensor and metadata field names from schema (fields ending with '_bytes')
|
|
tensor_fields = []
|
|
metadata_fields = []
|
|
for field in parquet_schema.names:
|
|
if field.endswith('_bytes'):
|
|
shape_field = field.replace('_bytes', '_shape')
|
|
dtype_field = field.replace('_bytes', '_dtype')
|
|
tensor_name = field.replace('_bytes', '')
|
|
tensor_fields.append(tensor_name)
|
|
assert shape_field in parquet_schema.names, f"Shape field {shape_field} not found in schema for field {field}. Currently we only support *_bytes fields for tensors."
|
|
assert dtype_field in parquet_schema.names, f"Dtype field {dtype_field} not found in schema for field {field}. Currently we only support *_bytes fields for tensors."
|
|
elif not field.endswith('_shape') and not field.endswith('_dtype'):
|
|
# Only add actual metadata fields, not the shape/dtype helper fields
|
|
metadata_fields.append(field)
|
|
|
|
# Process each tensor field
|
|
for tensor_name in tensor_fields:
|
|
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"
|
|
|
|
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 (rng.random(
|
|
) if rng else 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))
|
|
|
|
# Stack tensors with special handling for text embeddings
|
|
if tensor_name == 'text_embedding':
|
|
# Handle text embeddings with padding
|
|
padded_tensors = []
|
|
attention_masks = []
|
|
|
|
for tensor in tensor_list:
|
|
if tensor.numel() > 0:
|
|
padded_tensor, mask = pad(tensor, text_padding_length)
|
|
padded_tensors.append(padded_tensor)
|
|
attention_masks.append(mask)
|
|
else:
|
|
# Handle empty embeddings - assume default embedding dimension
|
|
padded_tensors.append(
|
|
torch.zeros(text_padding_length,
|
|
768,
|
|
dtype=torch.bfloat16))
|
|
attention_masks.append(torch.zeros(text_padding_length))
|
|
|
|
batch_data[tensor_name] = torch.stack(padded_tensors)
|
|
batch_data['text_attention_mask'] = torch.stack(attention_masks)
|
|
else:
|
|
# Stack all tensors to preserve batch consistency
|
|
# Don't filter out None or empty tensors as this breaks batch sizing
|
|
try:
|
|
batch_data[tensor_name] = torch.stack(tensor_list)
|
|
except ValueError as e:
|
|
shapes = [
|
|
t.shape
|
|
if t is not None and hasattr(t, 'shape') else 'None/Invalid'
|
|
for t in tensor_list
|
|
]
|
|
raise ValueError(
|
|
f"Failed to stack tensors for field '{tensor_name}'. "
|
|
f"Tensor shapes: {shapes}. "
|
|
f"All tensors in a batch must have compatible shapes. "
|
|
f"Original error: {e}") from e
|
|
|
|
# Process metadata fields into info_list
|
|
info_list = []
|
|
for row in rows:
|
|
info = {}
|
|
for field in metadata_fields:
|
|
info[field] = row.get(field, "")
|
|
|
|
# Add prompt field for backward compatibility
|
|
info["prompt"] = info.get("caption", "")
|
|
info_list.append(info)
|
|
|
|
batch_data['info_list'] = info_list
|
|
|
|
# Add caption_text for backward compatibility
|
|
if info_list and 'caption' in info_list[0]:
|
|
batch_data['caption_text'] = [info['caption'] for info in info_list]
|
|
|
|
return batch_data
|