diffusion: remove unused workflows folder (#13114)
This commit is contained in:
Mick
@@ -1 +0,0 @@
|
||||
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
|
||||
@@ -1 +0,0 @@
|
||||
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
|
||||
@@ -1,341 +0,0 @@
|
||||
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
|
||||
|
||||
import dataclasses
|
||||
import gc
|
||||
import os
|
||||
import random
|
||||
from collections.abc import Callable
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import pyarrow as pa
|
||||
import torch
|
||||
from datasets import Dataset, Video, load_dataset
|
||||
|
||||
from sglang.multimodal_gen.configs.configs import (
|
||||
DatasetType,
|
||||
PreprocessConfig,
|
||||
VideoLoaderType,
|
||||
)
|
||||
from sglang.multimodal_gen.configs.sample.base import DataType
|
||||
from sglang.multimodal_gen.dataset.dataloader.parquet_io import (
|
||||
ParquetDatasetWriter,
|
||||
records_to_table,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.distributed.parallel_state import (
|
||||
get_world_rank,
|
||||
get_world_size,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.pipelines.pipeline_batch_info import PreprocessBatch
|
||||
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class PreprocessingDataValidator:
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
max_height: int = 1024,
|
||||
max_width: int = 1024,
|
||||
max_h_div_w_ratio: float = 17 / 16,
|
||||
min_h_div_w_ratio: float = 8 / 16,
|
||||
num_frames: int = 16,
|
||||
train_fps: int = 24,
|
||||
speed_factor: float = 1.0,
|
||||
video_length_tolerance_range: float = 5.0,
|
||||
drop_short_ratio: float = 0.0,
|
||||
hw_aspect_threshold: float = 1.5,
|
||||
):
|
||||
self.max_height = max_height
|
||||
self.max_width = max_width
|
||||
self.max_h_div_w_ratio = max_h_div_w_ratio
|
||||
self.min_h_div_w_ratio = min_h_div_w_ratio
|
||||
self.num_frames = num_frames
|
||||
self.train_fps = train_fps
|
||||
self.speed_factor = speed_factor
|
||||
self.video_length_tolerance_range = video_length_tolerance_range
|
||||
self.drop_short_ratio = drop_short_ratio
|
||||
self.hw_aspect_threshold = hw_aspect_threshold
|
||||
self.validators: dict[str, Callable[[dict[str, Any]], bool]] = {}
|
||||
self.filter_counts: dict[str, int] = {}
|
||||
|
||||
self.num_items_before_filtering = 0
|
||||
self.num_items_after_filtering = 0
|
||||
|
||||
self.register_validators()
|
||||
|
||||
def register_validators(self) -> None:
|
||||
self.add_validator("data_type_validator", self._validate_data_type)
|
||||
self.add_validator("resolution_validator", self._validate_resolution)
|
||||
self.add_validator("frame_sampling_validator", self._validate_frame_sampling)
|
||||
|
||||
def add_validator(
|
||||
self, name: str, validator: Callable[[dict[str, Any]], bool]
|
||||
) -> None:
|
||||
self.validators[name] = validator
|
||||
self.filter_counts[name] = 0
|
||||
|
||||
def __call__(self, batch: dict[str, Any]) -> bool:
|
||||
"""
|
||||
Validate whether the preprocessing data batch is valid.
|
||||
"""
|
||||
self.num_items_before_filtering += 1
|
||||
|
||||
for name, validator in self.validators.items():
|
||||
if not validator(batch):
|
||||
self.filter_counts[name] += 1
|
||||
return False
|
||||
|
||||
self.num_items_after_filtering += 1
|
||||
return True
|
||||
|
||||
def _validate_data_type(self, batch: dict[str, Any]) -> bool:
|
||||
"""Validate basic validity of data items"""
|
||||
return not (
|
||||
batch["caption"] is None
|
||||
or batch["caption"] == ""
|
||||
or batch["fps"] is None
|
||||
or batch["fps"] <= 0
|
||||
or batch["num_frames"] is None
|
||||
or batch["num_frames"] <= 0
|
||||
)
|
||||
|
||||
def _validate_resolution(self, batch: dict[str, Any]) -> bool:
|
||||
"""Validate resolution constraints"""
|
||||
|
||||
aspect = self.max_height / self.max_width
|
||||
if batch["resolution"] is not None:
|
||||
height = batch["resolution"].get("height", None)
|
||||
width = batch["resolution"].get("width", None)
|
||||
|
||||
if height is None or width is None:
|
||||
return False
|
||||
|
||||
return self._filter_resolution(
|
||||
height,
|
||||
width,
|
||||
max_h_div_w_ratio=self.hw_aspect_threshold * aspect,
|
||||
min_h_div_w_ratio=1 / self.hw_aspect_threshold * aspect,
|
||||
)
|
||||
|
||||
def _filter_resolution(
|
||||
self, h: int, w: int, max_h_div_w_ratio: float, min_h_div_w_ratio: float
|
||||
) -> bool:
|
||||
"""Filter based on aspect ratio"""
|
||||
return (min_h_div_w_ratio <= h / w <= max_h_div_w_ratio) and (
|
||||
self.min_h_div_w_ratio <= h / w <= self.max_h_div_w_ratio
|
||||
)
|
||||
|
||||
def _validate_frame_sampling(self, batch: dict[str, Any]) -> bool:
|
||||
"""Validate frame sampling constraints"""
|
||||
|
||||
if batch["num_frames"] / batch["fps"] > self.video_length_tolerance_range * (
|
||||
self.num_frames / self.train_fps * self.speed_factor
|
||||
):
|
||||
return False
|
||||
|
||||
frame_interval = batch["fps"] / self.train_fps
|
||||
start_frame_idx = 0
|
||||
frame_indices = np.arange(
|
||||
start_frame_idx, batch["num_frames"], frame_interval
|
||||
).astype(int)
|
||||
return not (
|
||||
len(frame_indices) < self.num_frames
|
||||
and random.random() < self.drop_short_ratio
|
||||
)
|
||||
|
||||
def log_validation_stats(self):
|
||||
info = ""
|
||||
for name, count in self.filter_counts.items():
|
||||
info += f"failed in {name}: {count}, "
|
||||
info += f"number of items before filtering: {self.num_items_before_filtering}, "
|
||||
info += f"number of items after filtering: {self.num_items_after_filtering}"
|
||||
|
||||
logger.info(info)
|
||||
|
||||
|
||||
class VideoForwardBatchBuilder:
|
||||
|
||||
def __init__(self, seed: int):
|
||||
self.seed = seed
|
||||
|
||||
def __call__(self, batch: list) -> PreprocessBatch:
|
||||
forward_batch = PreprocessBatch(
|
||||
video_loader=[item["video"] for item in batch],
|
||||
video_file_name=[item["name"] for item in batch],
|
||||
height=[item["resolution"]["height"] for item in batch],
|
||||
width=[item["resolution"]["width"] for item in batch],
|
||||
fps=[item["fps"] for item in batch],
|
||||
num_frames=[item["num_frames"] for item in batch],
|
||||
prompt=[item["caption"] for item in batch],
|
||||
prompt_attention_mask=[],
|
||||
data_type=DataType.VIDEO,
|
||||
generator=torch.Generator("cpu").manual_seed(self.seed),
|
||||
)
|
||||
return forward_batch
|
||||
|
||||
|
||||
class ParquetDatasetSaver:
|
||||
"""Component for saving and writing Parquet datasets using shared parquet_io."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
flush_frequency: int,
|
||||
samples_per_file: int,
|
||||
schema: pa.Schema,
|
||||
record_creator: Callable[..., list[dict[str, Any]]],
|
||||
):
|
||||
self.flush_frequency = flush_frequency
|
||||
self.samples_per_file = samples_per_file
|
||||
self.schema = schema
|
||||
self.create_records_from_batch = record_creator
|
||||
self.num_processed_samples: int = 0
|
||||
self._writer: ParquetDatasetWriter | None = None
|
||||
|
||||
def save_and_write_parquet_batch(
|
||||
self,
|
||||
batch: PreprocessBatch,
|
||||
output_dir: str,
|
||||
extra_features: dict[str, Any] | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
Save and write Parquet dataset batch
|
||||
|
||||
Args:
|
||||
batch: PreprocessBatch containing video and metadata information
|
||||
output_dir: Output directory
|
||||
extra_features: Extra features
|
||||
|
||||
Returns:
|
||||
Number of processed samples
|
||||
"""
|
||||
assert isinstance(batch.latents, torch.Tensor)
|
||||
assert isinstance(batch.prompt_embeds, list)
|
||||
assert isinstance(batch.prompt_attention_mask, list)
|
||||
|
||||
# Process non-padded embeddings (if needed)
|
||||
if batch.prompt_attention_mask is not None:
|
||||
batch.prompt_embeds = self._process_non_padded_embeddings(
|
||||
batch.prompt_embeds[0], batch.prompt_attention_mask[0]
|
||||
)
|
||||
else:
|
||||
raise ValueError("prompt_attention_mask is None")
|
||||
|
||||
# Prepare batch data for Parquet dataset
|
||||
batch_data: list[dict[str, Any]] = []
|
||||
|
||||
for key in dataclasses.fields(batch):
|
||||
value = getattr(batch, key.name)
|
||||
if isinstance(value, list):
|
||||
for idx in range(len(value)):
|
||||
if isinstance(value[idx], torch.Tensor):
|
||||
value[idx] = value[idx].cpu().numpy()
|
||||
elif isinstance(value, torch.Tensor):
|
||||
value = value.cpu().numpy()
|
||||
setattr(batch, key.name, value)
|
||||
|
||||
# Create record for Parquet dataset
|
||||
records = self.create_records_from_batch(batch)
|
||||
batch_data.extend(records)
|
||||
|
||||
if batch_data:
|
||||
self.num_processed_samples += len(batch_data)
|
||||
table = records_to_table(batch_data, self.schema)
|
||||
if self._writer is None:
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
self._writer = ParquetDatasetWriter(
|
||||
out_dir=output_dir, samples_per_file=self.samples_per_file
|
||||
)
|
||||
self._writer.append_table(table)
|
||||
logger.debug("Collected batch with %s samples", len(table))
|
||||
|
||||
# If flush is needed
|
||||
if self.num_processed_samples >= self.flush_frequency:
|
||||
self.flush_tables()
|
||||
|
||||
def _process_non_padded_embeddings(
|
||||
self, prompt_embeds: torch.Tensor, prompt_attention_mask: torch.Tensor
|
||||
) -> list[torch.Tensor]:
|
||||
"""Process non-padded embeddings"""
|
||||
assert isinstance(prompt_embeds, torch.Tensor)
|
||||
assert isinstance(prompt_attention_mask, torch.Tensor)
|
||||
assert prompt_embeds.shape[0] == prompt_attention_mask.shape[0]
|
||||
|
||||
# Get sequence lengths from attention masks (number of 1s)
|
||||
seq_lens = prompt_attention_mask.sum(dim=1)
|
||||
|
||||
non_padded_embeds = []
|
||||
|
||||
# Process each item in the batch
|
||||
for i in range(prompt_embeds.size(0)):
|
||||
seq_len = seq_lens[i].item()
|
||||
# Slice the embeddings and masks to keep only non-padding parts
|
||||
non_padded_embeds.append(prompt_embeds[i, :seq_len])
|
||||
|
||||
return non_padded_embeds
|
||||
|
||||
def flush_tables(self, write_remainder: bool = False):
|
||||
"""Flush buffered records to disk.
|
||||
|
||||
Args:
|
||||
output_dir: Directory where parquet files are written. Kept for API
|
||||
symmetry (writer already configured with this path).
|
||||
write_remainder: If True, also write any leftover rows smaller than
|
||||
``samples_per_file`` as a final small file. Useful for the last flush.
|
||||
"""
|
||||
if self._writer is None:
|
||||
return
|
||||
_ = self._writer.flush(write_remainder=write_remainder)
|
||||
# Reset processed sample count modulo samples_per_file
|
||||
remainder = self.num_processed_samples % self.samples_per_file
|
||||
self.num_processed_samples = 0 if write_remainder else remainder
|
||||
|
||||
def clean_up(self) -> None:
|
||||
"""Clean up all tables"""
|
||||
self.flush_tables(write_remainder=True)
|
||||
self._writer = None
|
||||
self.num_processed_samples = 0
|
||||
gc.collect()
|
||||
|
||||
def __del__(self):
|
||||
self.clean_up()
|
||||
|
||||
|
||||
def build_dataset(
|
||||
preprocess_config: PreprocessConfig,
|
||||
split: str,
|
||||
validator: Callable[[dict[str, Any]], bool],
|
||||
) -> Dataset:
|
||||
if preprocess_config.dataset_type == DatasetType.HF:
|
||||
dataset = load_dataset(preprocess_config.dataset_path, split=split)
|
||||
dataset = dataset.filter(validator)
|
||||
dataset = dataset.shard(num_shards=get_world_size(), index=get_world_rank())
|
||||
elif preprocess_config.dataset_type == DatasetType.MERGED:
|
||||
metadata_json_path = os.path.join(
|
||||
preprocess_config.dataset_path, "videos2caption.json"
|
||||
)
|
||||
video_folder = os.path.join(preprocess_config.dataset_path, "videos")
|
||||
dataset = load_dataset("json", data_files=metadata_json_path, split=split)
|
||||
column_names = dataset.column_names
|
||||
# rename columns to match the schema
|
||||
if "cap" in column_names:
|
||||
dataset = dataset.rename_column("cap", "caption")
|
||||
if "path" in column_names:
|
||||
dataset = dataset.rename_column("path", "name")
|
||||
|
||||
dataset = dataset.filter(validator)
|
||||
dataset = dataset.shard(num_shards=get_world_size(), index=get_world_rank())
|
||||
|
||||
# add video column
|
||||
def add_video_column(item: dict[str, Any]) -> dict[str, Any]:
|
||||
item["video"] = os.path.join(video_folder, item["name"])
|
||||
return item
|
||||
|
||||
dataset = dataset.map(add_video_column)
|
||||
if preprocess_config.video_loader_type == VideoLoaderType.TORCHCODEC:
|
||||
dataset = dataset.cast_column("video", Video())
|
||||
else:
|
||||
raise ValueError(f"Invalid dataset type: {preprocess_config.dataset_type}")
|
||||
|
||||
return dataset
|
||||
@@ -1,140 +0,0 @@
|
||||
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
|
||||
|
||||
import os
|
||||
from typing import cast
|
||||
|
||||
from torch.utils.data import DataLoader
|
||||
|
||||
from sglang.multimodal_gen.configs.configs import PreprocessConfig
|
||||
from sglang.multimodal_gen.dataset.dataloader.record_schema import (
|
||||
basic_t2v_record_creator,
|
||||
i2v_record_creator,
|
||||
)
|
||||
from sglang.multimodal_gen.dataset.dataloader.schema import (
|
||||
pyarrow_schema_i2v,
|
||||
pyarrow_schema_t2v,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.distributed.parallel_state import get_world_rank
|
||||
from sglang.multimodal_gen.runtime.server_args import ServerArgs, WorkloadType
|
||||
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
|
||||
from sglang.multimodal_gen.runtime.workflow.preprocess.components import (
|
||||
ParquetDatasetSaver,
|
||||
PreprocessingDataValidator,
|
||||
VideoForwardBatchBuilder,
|
||||
build_dataset,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.workflow.workflow_base import WorkflowBase
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class PreprocessWorkflow(WorkflowBase):
|
||||
|
||||
def register_pipelines(self) -> None:
|
||||
self.add_pipeline_config("preprocess_pipeline", self.server_args)
|
||||
|
||||
def register_components(self) -> None:
|
||||
assert self.server_args.preprocess_config is not None
|
||||
preprocess_config: PreprocessConfig = self.server_args.preprocess_config
|
||||
|
||||
# raw data validator
|
||||
raw_data_validator = PreprocessingDataValidator(
|
||||
max_height=preprocess_config.max_height,
|
||||
max_width=preprocess_config.max_width,
|
||||
num_frames=preprocess_config.num_frames,
|
||||
train_fps=preprocess_config.train_fps,
|
||||
speed_factor=preprocess_config.speed_factor,
|
||||
video_length_tolerance_range=preprocess_config.video_length_tolerance_range,
|
||||
drop_short_ratio=preprocess_config.drop_short_ratio,
|
||||
)
|
||||
self.add_component("raw_data_validator", raw_data_validator)
|
||||
|
||||
# training dataset
|
||||
training_dataset = build_dataset(
|
||||
preprocess_config, split="train", validator=raw_data_validator
|
||||
)
|
||||
# we do not use collate_fn here because we use iterable-style Dataset
|
||||
# and want to keep the original type of the dataset
|
||||
training_dataloader = DataLoader(
|
||||
training_dataset,
|
||||
batch_size=preprocess_config.preprocess_video_batch_size,
|
||||
num_workers=preprocess_config.dataloader_num_workers,
|
||||
collate_fn=lambda x: x,
|
||||
)
|
||||
self.add_component("training_dataloader", training_dataloader)
|
||||
|
||||
# try to load validation dataset if it exists
|
||||
try:
|
||||
validation_dataset = build_dataset(
|
||||
preprocess_config, split="validation", validator=raw_data_validator
|
||||
)
|
||||
validation_dataloader = DataLoader(
|
||||
validation_dataset,
|
||||
batch_size=preprocess_config.preprocess_video_batch_size,
|
||||
num_workers=preprocess_config.dataloader_num_workers,
|
||||
collate_fn=lambda x: x,
|
||||
)
|
||||
except ValueError:
|
||||
logger.warning(
|
||||
"Validation dataset not found, skipping validation dataset preprocessing."
|
||||
)
|
||||
validation_dataloader = None
|
||||
|
||||
self.add_component("validation_dataloader", validation_dataloader)
|
||||
|
||||
# forward batch builder
|
||||
video_forward_batch_builder = VideoForwardBatchBuilder(
|
||||
seed=self.server_args.preprocess_config.seed
|
||||
)
|
||||
self.add_component("video_forward_batch_builder", video_forward_batch_builder)
|
||||
|
||||
# record creator
|
||||
if self.server_args.workload_type == WorkloadType.I2V:
|
||||
record_creator = i2v_record_creator
|
||||
schema = pyarrow_schema_i2v
|
||||
else:
|
||||
record_creator = basic_t2v_record_creator
|
||||
schema = pyarrow_schema_t2v
|
||||
processed_dataset_saver = ParquetDatasetSaver(
|
||||
flush_frequency=self.server_args.preprocess_config.flush_frequency,
|
||||
samples_per_file=self.server_args.preprocess_config.samples_per_file,
|
||||
schema=schema,
|
||||
record_creator=record_creator,
|
||||
)
|
||||
self.add_component("processed_dataset_saver", processed_dataset_saver)
|
||||
|
||||
def prepare_system_environment(self) -> None:
|
||||
assert self.server_args.preprocess_config is not None
|
||||
dataset_output_dir = self.server_args.preprocess_config.dataset_output_dir
|
||||
os.makedirs(dataset_output_dir, exist_ok=True)
|
||||
|
||||
validation_dataset_output_dir = os.path.join(
|
||||
dataset_output_dir, "validation_dataset", f"worker_{get_world_rank()}"
|
||||
)
|
||||
os.makedirs(validation_dataset_output_dir, exist_ok=True)
|
||||
self.validation_dataset_output_dir = validation_dataset_output_dir
|
||||
|
||||
training_dataset_output_dir = os.path.join(
|
||||
dataset_output_dir, "training_dataset", f"worker_{get_world_rank()}"
|
||||
)
|
||||
os.makedirs(training_dataset_output_dir, exist_ok=True)
|
||||
self.training_dataset_output_dir = training_dataset_output_dir
|
||||
|
||||
@classmethod
|
||||
def get_workflow_cls(cls, server_args: ServerArgs) -> "PreprocessWorkflow":
|
||||
if server_args.workload_type == WorkloadType.T2V:
|
||||
from sglang.multimodal_gen.runtime.workflow.preprocess.preprocess_workflow_t2v import (
|
||||
PreprocessWorkflowT2V,
|
||||
)
|
||||
|
||||
return cast(PreprocessWorkflow, PreprocessWorkflowT2V)
|
||||
elif server_args.workload_type == WorkloadType.I2V:
|
||||
from sglang.multimodal_gen.runtime.workflow.preprocess.preprocess_workflow_i2v import (
|
||||
PreprocessWorkflowI2V,
|
||||
)
|
||||
|
||||
return cast(PreprocessWorkflow, PreprocessWorkflowI2V)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Workload type: {server_args.workload_type} is not supported in preprocessing workflow."
|
||||
)
|
||||
@@ -1,70 +0,0 @@
|
||||
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from tqdm import tqdm
|
||||
|
||||
from sglang.multimodal_gen.dataset.preprocessing_datasets import PreprocessBatch
|
||||
from sglang.multimodal_gen.runtime.workflow.preprocess.components import (
|
||||
ParquetDatasetSaver,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.workflow.preprocess.preprocess_workflow import (
|
||||
PreprocessWorkflow,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from torch.utils.data import DataLoader
|
||||
|
||||
from sglang.multimodal_gen.runtime.pipelines.composed_pipeline_base import (
|
||||
ComposedPipelineBase,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.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.server_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.server_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()
|
||||
@@ -1,70 +0,0 @@
|
||||
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
|
||||
|
||||
from typing import TYPE_CHECKING, Optional
|
||||
|
||||
from tqdm import tqdm
|
||||
|
||||
from sglang.multimodal_gen.runtime.pipelines.pipeline_batch_info import PreprocessBatch
|
||||
from sglang.multimodal_gen.runtime.workflow.preprocess.components import (
|
||||
ParquetDatasetSaver,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.workflow.preprocess.preprocess_workflow import (
|
||||
PreprocessWorkflow,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from torch.utils.data import DataLoader
|
||||
|
||||
from sglang.multimodal_gen.runtime.pipelines.composed_pipeline_base import (
|
||||
ComposedPipelineBase,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.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.server_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.server_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()
|
||||
@@ -1,188 +0,0 @@
|
||||
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Any, Optional
|
||||
|
||||
from sglang.multimodal_gen.runtime.pipelines import ComposedPipelineBase, build_pipeline
|
||||
from sglang.multimodal_gen.runtime.pipelines.pipeline_registry import PipelineType
|
||||
from sglang.multimodal_gen.runtime.server_args import ExecutionMode, ServerArgs
|
||||
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class WorkflowBase(ABC):
|
||||
"""
|
||||
Abstract base class for defining video processing workflows.
|
||||
|
||||
A workflow serves as the top-level orchestrator that coordinates multiple pipelines
|
||||
and components to accomplish a specific video processing task. The workflow pattern
|
||||
provides several key benefits:
|
||||
|
||||
1. **Separation of Concerns**: Workflows separate high-level orchestration logic
|
||||
from low-level processing implementations in pipelines.
|
||||
|
||||
2. **Modularity**: Different workflows can be created for different execution modes
|
||||
(preprocess, inference, etc.) while sharing common pipeline components.
|
||||
|
||||
3. **Configuration Management**: Workflows manage the configuration and initialization
|
||||
of multiple related pipelines and components in a centralized manner.
|
||||
|
||||
4. **Environment Setup**: Workflows handle system-level setup and resource
|
||||
allocation before pipeline execution begins.
|
||||
|
||||
5. **Lifecycle Management**: Workflows control the complete lifecycle from
|
||||
initialization through execution to cleanup.
|
||||
|
||||
The workflow acts as a factory and coordinator, creating the appropriate pipelines
|
||||
based on configuration, setting up the execution environment, and orchestrating
|
||||
the overall processing flow.
|
||||
"""
|
||||
|
||||
def __init__(self, server_args: ServerArgs):
|
||||
"""
|
||||
Initialize the workflow with configuration arguments.
|
||||
|
||||
Args:
|
||||
server_args: Configuration object containing all parameters
|
||||
needed for workflow and pipeline setup.
|
||||
"""
|
||||
self.server_args = server_args
|
||||
|
||||
# TODO: pipeline_config should be: dict[str, PipelineConfig]
|
||||
# pipeline_type should be included in the PipelineConfig
|
||||
# pipeline_config[pipeline_name] = (pipeline_type, server_args)
|
||||
self._pipeline_configs: dict[str, tuple[PipelineType, ServerArgs]] = {}
|
||||
self._pipelines: dict[str, ComposedPipelineBase] = {}
|
||||
self._components: dict[str, Any] = {}
|
||||
self.register_pipelines()
|
||||
self.register_components()
|
||||
|
||||
self.prepare_system_environment()
|
||||
self.load_pipelines()
|
||||
|
||||
def load_pipelines(self) -> None:
|
||||
"""
|
||||
Create and initialize all registered pipelines.
|
||||
|
||||
This method instantiates pipeline objects from their configurations
|
||||
and makes them available as both dictionary entries and instance
|
||||
attributes for convenient access.
|
||||
"""
|
||||
for pipeline_name, pipeline_config in self._pipeline_configs.items():
|
||||
pipeline_type, server_args = pipeline_config
|
||||
pipeline = build_pipeline(server_args, pipeline_type)
|
||||
self._pipelines[pipeline_name] = pipeline
|
||||
setattr(self, pipeline_name, pipeline)
|
||||
|
||||
def add_pipeline_config(
|
||||
self, pipeline_name: str, pipeline_config: tuple[PipelineType, ServerArgs]
|
||||
) -> None:
|
||||
"""
|
||||
Register a pipeline configuration for later instantiation.
|
||||
|
||||
Args:
|
||||
pipeline_name: Unique identifier for the pipeline.
|
||||
pipeline_config: Tuple containing the pipeline type and
|
||||
configuration arguments.
|
||||
"""
|
||||
self._pipeline_configs[pipeline_name] = pipeline_config
|
||||
|
||||
def add_component(self, component_name: str, component: Any) -> None:
|
||||
"""
|
||||
Register a component instance with the workflow.
|
||||
|
||||
Components are auxiliary objects that may be shared across pipelines
|
||||
or used for workflow-level functionality (e.g., databases, caches,
|
||||
external services).
|
||||
|
||||
Args:
|
||||
component_name: Unique identifier for the component.
|
||||
component: The component instance to register.
|
||||
"""
|
||||
self._components[component_name] = component
|
||||
setattr(self, component_name, component)
|
||||
|
||||
def get_component(self, component_name: str) -> Any:
|
||||
"""
|
||||
Retrieve a registered component by name.
|
||||
|
||||
Args:
|
||||
component_name: The name of the component to retrieve.
|
||||
|
||||
Returns:
|
||||
The component instance.
|
||||
"""
|
||||
return self._components[component_name]
|
||||
|
||||
@abstractmethod
|
||||
def register_components(self) -> None:
|
||||
"""
|
||||
Register workflow-specific components.
|
||||
|
||||
Subclasses must implement this method to register any components
|
||||
needed for their specific workflow (e.g., databases, external APIs,
|
||||
shared resources).
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def register_pipelines(self) -> None:
|
||||
"""
|
||||
Register workflow-specific pipelines.
|
||||
|
||||
Subclasses must implement this method to define which pipelines
|
||||
are needed for their specific workflow and how they should be
|
||||
configured.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def prepare_system_environment(self) -> None:
|
||||
"""
|
||||
Prepare the system environment for workflow execution.
|
||||
|
||||
Subclasses must implement this method to handle any system-level
|
||||
setup required before pipeline execution (e.g., GPU initialization,
|
||||
temporary directories, resource allocation).
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def run(self):
|
||||
"""
|
||||
Execute the main workflow logic.
|
||||
|
||||
Subclasses must implement this method to define the specific
|
||||
execution flow for their workflow, coordinating the registered
|
||||
pipelines and components to accomplish the desired task.
|
||||
"""
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
def get_workflow_cls(cls, server_args: ServerArgs) -> Optional["WorkflowBase"]:
|
||||
"""
|
||||
Factory method to get the appropriate workflow class based on execution mode.
|
||||
|
||||
This method acts as a workflow factory, returning the appropriate
|
||||
workflow class implementation based on the specified execution mode
|
||||
in the configuration arguments.
|
||||
|
||||
Args:
|
||||
server_args: Configuration object containing the execution mode
|
||||
and other parameters.
|
||||
|
||||
Returns:
|
||||
The appropriate workflow class for the specified execution mode,
|
||||
or None if no workflow is available for the given mode.
|
||||
"""
|
||||
if server_args.mode == ExecutionMode.PREPROCESS:
|
||||
from sglang.multimodal_gen.runtime.workflow.preprocess.preprocess_workflow import (
|
||||
PreprocessWorkflow,
|
||||
)
|
||||
|
||||
return PreprocessWorkflow.get_workflow_cls(server_args)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Execution mode: {server_args.mode} is not supported in workflow."
|
||||
)
|
||||
Reference in New Issue
Block a user