ALSv/self-forcing

1import numpy as np2 3 4def get_array_shape_from_lmdb(env, array_name):5    with env.begin() as txn:6        image_shape = txn.get(f"{array_name}_shape".encode()).decode()7        image_shape = tuple(map(int, image_shape.split()))8    return image_shape9 10 11def store_arrays_to_lmdb(env, arrays_dict, start_index=0):12    """13    Store rows of multiple numpy arrays in a single LMDB.14    Each row is stored separately with a naming convention.15    """16    with env.begin(write=True) as txn:17        for array_name, array in arrays_dict.items():18            for i, row in enumerate(array):19                # Convert row to bytes20                if isinstance(row, str):21                    row_bytes = row.encode()22                else:23                    row_bytes = row.tobytes()24 25                data_key = f'{array_name}_{start_index + i}_data'.encode()26 27                txn.put(data_key, row_bytes)28 29 30def process_data_dict(data_dict, seen_prompts):31    output_dict = {}32 33    all_videos = []34    all_prompts = []35    for prompt, video in data_dict.items():36        if prompt in seen_prompts:37            continue38        else:39            seen_prompts.add(prompt)40 41        video = video.half().numpy()42        all_videos.append(video)43        all_prompts.append(prompt)44 45    if len(all_videos) == 0:46        return {"latents": np.array([]), "prompts": np.array([])}47 48    all_videos = np.concatenate(all_videos, axis=0)49 50    output_dict['latents'] = all_videos51    output_dict['prompts'] = np.array(all_prompts)52 53    return output_dict54 55 56def retrieve_row_from_lmdb(lmdb_env, array_name, dtype, row_index, shape=None):57    """58    Retrieve a specific row from a specific array in the LMDB.59    """60    data_key = f'{array_name}_{row_index}_data'.encode()61 62    with lmdb_env.begin() as txn:63        row_bytes = txn.get(data_key)64 65    if dtype == str:66        array = row_bytes.decode()67    else:68        array = np.frombuffer(row_bytes, dtype=dtype)69 70    if shape is not None and len(shape) > 0:71        array = array.reshape(shape)72    return array73