mirror of
https://github.com/storytold/FineTrainers-Conditioning.git
synced 2026-10-09 00:09:45 +00:00
sayakpaul
51 lines
1.5 KiB
Python
51 lines
1.5 KiB
Python
"""
|
|
Taken from
|
|
https://github.com/genmoai/mochi/blob/main/demos/fine_tuner/dataset.py
|
|
"""
|
|
|
|
from pathlib import Path
|
|
|
|
import click
|
|
import torch
|
|
from torch.utils.data import DataLoader, Dataset
|
|
|
|
|
|
def load_to_cpu(x):
|
|
return torch.load(x, map_location=torch.device("cpu"), weights_only=True)
|
|
|
|
|
|
class LatentEmbedDataset(Dataset):
|
|
def __init__(self, file_paths, repeat=1):
|
|
self.items = [
|
|
(Path(p).with_suffix(".latent.pt"), Path(p).with_suffix(".embed.pt"))
|
|
for p in file_paths
|
|
if Path(p).with_suffix(".latent.pt").is_file() and Path(p).with_suffix(".embed.pt").is_file()
|
|
]
|
|
self.items = self.items * repeat
|
|
print(f"Loaded {len(self.items)}/{len(file_paths)} valid file pairs.")
|
|
|
|
def __len__(self):
|
|
return len(self.items)
|
|
|
|
def __getitem__(self, idx):
|
|
latent_path, embed_path = self.items[idx]
|
|
return load_to_cpu(latent_path), load_to_cpu(embed_path)
|
|
|
|
|
|
@click.command()
|
|
@click.argument("directory", type=click.Path(exists=True, file_okay=False))
|
|
def process_videos(directory):
|
|
dir_path = Path(directory)
|
|
mp4_files = [str(f) for f in dir_path.glob("**/*.mp4") if not f.name.endswith(".recon.mp4")]
|
|
assert mp4_files, f"No mp4 files found"
|
|
|
|
dataset = LatentEmbedDataset(mp4_files)
|
|
dataloader = DataLoader(dataset, batch_size=4, shuffle=True)
|
|
|
|
for latents, embeds in dataloader:
|
|
print([(k, v.shape) for k, v in latents.items()])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
process_videos()
|