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