| 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: |
| |
| 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() |
|
|
| |
| 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() |
|
|
