File size: 2,307 Bytes
74da989 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 | from typing import TYPE_CHECKING, Optional
from tqdm import tqdm
from trainer.pipelines.pipeline_batch_info import PreprocessBatch
from trainer.workflow.preprocess.components import ParquetDatasetSaver
from trainer.workflow.preprocess.preprocess_workflow import PreprocessWorkflow
if TYPE_CHECKING:
from torch.utils.data import DataLoader
from trainer.pipelines.composed_pipeline_base import ComposedPipelineBase
from trainer.workflow.preprocess.components import (
VideoForwardBatchBuilder)
class PreprocessWorkflowT2V(PreprocessWorkflow):
training_dataloader: "DataLoader"
validation_dataloader: Optional["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.trainer_args)
self.processed_dataset_saver.save_and_write_parquet_batch(
forward_batch, self.training_dataset_output_dir)
self.processed_dataset_saver.flush_tables(
self.training_dataset_output_dir)
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.trainer_args)
self.processed_dataset_saver.save_and_write_parquet_batch(
forward_batch, self.validation_dataset_output_dir)
self.processed_dataset_saver.flush_tables(
self.validation_dataset_output_dir)
self.processed_dataset_saver.clean_up()
|