44 lines
1.5 KiB
Python
44 lines
1.5 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
from fastvideo.dataset.lmdb_utils import get_array_shape_from_lmdb, retrieve_row_from_lmdb
|
|
from torch.utils.data import Dataset
|
|
import numpy as np
|
|
import torch
|
|
import lmdb
|
|
|
|
# from Self-Forcing: https://github.com/guandeh17/Self-Forcing/blob/main/utils/dataset.py
|
|
class ODERegressionLMDBDataset(Dataset):
|
|
def __init__(self, data_path: str, max_pair: int = int(1e8)):
|
|
print(f"data_path: {data_path}")
|
|
self.env = lmdb.open(data_path, readonly=True,
|
|
lock=False, readahead=False, meminit=False)
|
|
|
|
self.latents_shape = get_array_shape_from_lmdb(self.env, 'latents')
|
|
self.max_pair = max_pair
|
|
|
|
def __len__(self):
|
|
return min(self.latents_shape[0], self.max_pair)
|
|
|
|
def __getitem__(self, idx):
|
|
"""
|
|
Outputs:
|
|
- prompts: List of Strings
|
|
- latents: Tensor of shape (num_denoising_steps, num_frames, num_channels, height, width). It is ordered from pure noise to clean image.
|
|
"""
|
|
latents = retrieve_row_from_lmdb(
|
|
self.env,
|
|
"latents", np.float16, idx, shape=self.latents_shape[1:]
|
|
)
|
|
|
|
if len(latents.shape) == 4:
|
|
latents = latents[None, ...]
|
|
|
|
prompts = retrieve_row_from_lmdb(
|
|
self.env,
|
|
"prompts", str, idx
|
|
)
|
|
return {
|
|
"prompts": prompts,
|
|
"ode_latent": torch.tensor(latents, dtype=torch.float32)
|
|
}
|
|
|