46 lines
2.0 KiB
Python
46 lines
2.0 KiB
Python
from typing import TYPE_CHECKING
|
|
|
|
from tqdm import tqdm
|
|
|
|
from fastvideo.dataset.preprocessing_datasets import PreprocessBatch
|
|
from fastvideo.workflow.preprocess.components import ParquetDatasetSaver
|
|
from fastvideo.workflow.preprocess.preprocess_workflow import PreprocessWorkflow
|
|
|
|
if TYPE_CHECKING:
|
|
from torch.utils.data import DataLoader
|
|
|
|
from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase
|
|
from fastvideo.workflow.preprocess.components import (VideoForwardBatchBuilder)
|
|
|
|
|
|
class PreprocessWorkflowI2V(PreprocessWorkflow):
|
|
training_dataloader: "DataLoader"
|
|
validation_dataloader: "DataLoader"
|
|
preprocess_pipeline: "ComposedPipelineBase"
|
|
processed_dataset_saver: "ParquetDatasetSaver"
|
|
video_forward_batch_builder: "VideoForwardBatchBuilder"
|
|
|
|
def run(self) -> None:
|
|
# Training dataset preprocessing
|
|
for batch in tqdm(self.training_dataloader, desc="Preprocessing training dataset", unit="batch"):
|
|
forward_batch: PreprocessBatch = self.video_forward_batch_builder(batch)
|
|
|
|
forward_batch = self.preprocess_pipeline.forward(forward_batch, self.fastvideo_args)
|
|
|
|
self.processed_dataset_saver.save_and_write_parquet_batch(forward_batch, self.training_dataset_output_dir)
|
|
|
|
self.processed_dataset_saver.flush_tables()
|
|
self.processed_dataset_saver.clean_up()
|
|
|
|
# Validation dataset preprocessing
|
|
if self.validation_dataloader is not None:
|
|
for batch in tqdm(self.validation_dataloader, desc="Preprocessing validation dataset", unit="batch"):
|
|
forward_batch = self.video_forward_batch_builder(batch)
|
|
|
|
forward_batch = self.preprocess_pipeline.forward(forward_batch, self.fastvideo_args)
|
|
|
|
self.processed_dataset_saver.save_and_write_parquet_batch(forward_batch,
|
|
self.validation_dataset_output_dir)
|
|
self.processed_dataset_saver.flush_tables()
|
|
self.processed_dataset_saver.clean_up()
|