From 47a4390ce0db9212d5ffd173a996cc4dad867991 Mon Sep 17 00:00:00 2001 From: hengtaoguo Date: Fri, 24 Jul 2026 23:13:10 +0000 Subject: [PATCH] End2End Video SFT with Varied Size/Duration --- .../sft-vision-llava-video-178k.yml | 23 +++ .../input_pipeline/hf_data_processing.py | 21 +- .../input_pipeline/input_pipeline_utils.py | 126 ++++++++---- src/maxtext/multimodal/processor.py | 13 +- .../multimodal/processor_qwen3_omni.py | 187 ++++++++++++++---- src/maxtext/trainers/pre_train/train.py | 24 ++- src/maxtext/utils/maxtext_utils.py | 25 ++- .../prepare_llava_video_178k.py | 127 ++++++++++++ 8 files changed, 450 insertions(+), 96 deletions(-) create mode 100644 src/maxtext/configs/post_train/sft-vision-llava-video-178k.yml create mode 100644 tools/data_generation/prepare_llava_video_178k.py diff --git a/src/maxtext/configs/post_train/sft-vision-llava-video-178k.yml b/src/maxtext/configs/post_train/sft-vision-llava-video-178k.yml new file mode 100644 index 0000000000..c6694a5554 --- /dev/null +++ b/src/maxtext/configs/post_train/sft-vision-llava-video-178k.yml @@ -0,0 +1,23 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +base_config: "sft-vision-chartqa.yml" + +# Dataset +hf_path: "parquet" +hf_train_files: "/mounted/LLaVA-Video-178K/0_30_s_academic_v0_1/*.parquet" +train_data_columns: ["query", "response"] +train_image_column: "video" +eval_data_columns: ["query", "response"] +eval_image_column: "video" diff --git a/src/maxtext/input_pipeline/hf_data_processing.py b/src/maxtext/input_pipeline/hf_data_processing.py index 370f1895bd..2a337e35cf 100644 --- a/src/maxtext/input_pipeline/hf_data_processing.py +++ b/src/maxtext/input_pipeline/hf_data_processing.py @@ -42,6 +42,21 @@ def _get_pad_id(tokenizer): return pad_id +def format_sharegpt(example): + """Convert a ShareGPT conversation into query and response fields.""" + query = "" + response = "" + for turn in example["conversations"]: + if turn["from"] == "human": + query = turn["value"] + elif turn["from"] == "gpt": + response = turn["value"] + break + example["query"] = query + example["response"] = response + return example + + def vision_sft_preprocessing_pipeline( dataset, config, @@ -54,6 +69,9 @@ def vision_sft_preprocessing_pipeline( ): """pipeline for multimodal SFT with HF dataset""" + if "conversations" in dataset.features and len(text_columns) == 2 and text_columns[0] not in dataset.features: + dataset = dataset.map(format_sharegpt) + assert len(text_columns) == 2, f"Need two text_columns for query and response, received {text_columns=}" # Tunix GA requires per-micro-batch slicing at the data level, # whereas Native GA processes the full batch and splits it internally. @@ -99,6 +117,7 @@ def vision_sft_preprocessing_pipeline( "column": text_columns[0], "image_placeholder": config.image_placeholder, "model_name": config.model_name, + "video_placeholder": getattr(config, "video_placeholder", "<|video|>"), }, ) dataset = dataset.map( @@ -181,7 +200,7 @@ def vision_sft_preprocessing_pipeline( data_source=dataset, operations=operations, sampler=dummy_index_sampler, - worker_count=1, # only supports <=1 for now, more workers results in duplicated data + worker_count=1, worker_buffer_size=1, read_options=grain.ReadOptions(num_threads=1, prefetch_buffer_size=batch_size * 4), ) diff --git a/src/maxtext/input_pipeline/input_pipeline_utils.py b/src/maxtext/input_pipeline/input_pipeline_utils.py index d996a88908..477f47c59c 100644 --- a/src/maxtext/input_pipeline/input_pipeline_utils.py +++ b/src/maxtext/input_pipeline/input_pipeline_utils.py @@ -90,19 +90,56 @@ def _process_string(string_tensor): ########## Functions used by HF pipeline -def reformat_prompt(example, column, image_placeholder, model_name): +def is_video_file(path): + if not isinstance(path, str): + return False + video_extensions = (".mp4", ".avi", ".mkv", ".webm", ".mov", ".gif") + return path.lower().endswith(video_extensions) + + +def reformat_prompt( + example, + column, + image_placeholder, + model_name, + video_placeholder="<|video|>", + num_videos=0, + num_image_tokens=None, + num_video_tokens=None, +): """reformat prompt for multimodal SFT""" if isinstance(example["images"], list): - num_images = len(example["images"]) + num_images = sum(1 for img in example["images"] if not is_video_file(img)) + num_videos_in_example = sum(1 for img in example["images"] if is_video_file(img)) + if num_videos_in_example > 0: + num_videos = num_videos_in_example else: - num_images = 1 - example[column] = mm_processor.reformat_prompt(example[column], image_placeholder, model_name, num_images) + if is_video_file(example["images"]): + num_images = 0 + num_videos = 1 + else: + num_images = 1 + num_videos = 0 + example[column] = mm_processor.reformat_prompt( + example[column], + image_placeholder=image_placeholder, + model_name=model_name, + num_images=num_images, + video_placeholder=video_placeholder, + num_videos=num_videos, + num_image_tokens=num_image_tokens, + num_video_tokens=num_video_tokens, + ) return example def reformat_response(example, column, model_name): """reformat response for multimodal SFT""" - example[column] = mm_processor.reformat_response(example[column][0], model_name) + if isinstance(example[column], (list, np.ndarray)): + response = example[column][0] + else: + response = example[column] + example[column] = mm_processor.reformat_response(response, model_name) return example @@ -123,12 +160,16 @@ def pre_process_image_sft(example, image_column, config): """pre-process image for multimodal SFT""" def _process_image_fn(image): - if isinstance(image, list): + if isinstance(image, str) and is_video_file(image): + image = mm_processor.preprocess_image_for_training(image, config) + elif isinstance(image, list) and len(image) > 0 and isinstance(image[0], str) and is_video_file(image[0]): + image = mm_processor.preprocess_image_for_training(image, config) + elif isinstance(image, list): image = [np.array(mm_utils.convert_to_RGB(img)) for img in image] + image = mm_processor.preprocess_image_for_training(image, config) else: image = np.array(mm_utils.convert_to_RGB(image)) - - image = mm_processor.preprocess_image_for_training(image, config) + image = mm_processor.preprocess_image_for_training(image, config) return image example[image_column] = _process_image_fn(example[image_column]) @@ -453,7 +494,10 @@ def __getitem__(self, index): The next item in the iterator is returned.""" if not self.data_iters: self.data_iters = [iter(x) for x in self.datasets] - idx = int(current_thread().name.split("_")[1]) + try: + idx = int(current_thread().name.split("_")[1]) + except (IndexError, ValueError): + idx = 0 while True: try: @@ -777,6 +821,9 @@ def _pad_image_and_mask(self, preprocessed_image: mm_utils.PreprocessorOutput) - if not isinstance(preprocessed_image, mm_utils.PreprocessorOutput): raise TypeError(f"Input must be multimodal_utils.PreprocessorOutput, but got {type(preprocessed_image)}") + if getattr(preprocessed_image, "video_values", None) is not None: + return preprocessed_image + if preprocessed_image.pixel_values is None: raise ValueError("Input preprocessed_image must have pixel_values to pad images.") @@ -870,16 +917,7 @@ def map( @dataclasses.dataclass class ExtractImagesAndMasks(grain.MapTransform): - """Extracts images and masks from a PreprocessorOutput object. - - This transform is used in multi-modal data pipelines to extract the image - tensors and their corresponding masks from a PreprocessorOutput object. - The extracted images and masks are then added to the data element under - the keys 'images' and 'image_masks', respectively. - - If the 'images' key is not present in the input element, the transform - returns the element unchanged. - """ + """Extracts images and masks from a PreprocessorOutput object.""" def map(self, element: dict[str, np.ndarray]) -> dict[str, np.ndarray]: """Applies the extraction transformation to the 'images' field if present.""" @@ -891,28 +929,25 @@ def map(self, element: dict[str, np.ndarray]) -> dict[str, np.ndarray]: raise TypeError(f"'images' must be of type PreprocessorOutput, but got {type(preprocessed_image)}") output = element.copy() - output["images"] = preprocessed_image.pixel_values # pyrefly: ignore[unsupported-operation] - if preprocessed_image.pixel_mask is not None: - output["image_masks"] = preprocessed_image.pixel_mask + if getattr(preprocessed_image, "video_values", None) is not None: + output["images"] = preprocessed_image.video_values + if getattr(preprocessed_image, "video_mask", None) is not None: + output["image_masks"] = preprocessed_image.video_mask + if getattr(preprocessed_image, "video_grid_thw", None) is not None: + output["video_grid_thw"] = preprocessed_image.video_grid_thw + if getattr(preprocessed_image, "video_second_per_grid", None) is not None: + output["second_per_grids"] = preprocessed_image.video_second_per_grid + else: + output["images"] = preprocessed_image.pixel_values + if preprocessed_image.pixel_mask is not None: + output["image_masks"] = preprocessed_image.pixel_mask return output @dataclasses.dataclass class FoldImagesIntoBatch(grain.MapTransform): - """Folds the 'image' dimension into the batch dimension. - - This transform is used in multi-modal data pipelines where each data example - might have multiple associated images. For model processing, it's often - efficient to treat each image as a separate item in a larger batch. - - This operation reshapes the 'images' tensor from a shape like - (B, N, T, H, W, C) to (B * N, T, H, W, C), where B is the batch size, N is - the number of images per example, and T is the number of image tiles. - - The transformation is triggered only if the input 'images' tensor has more - dimensions than the expected batched image tensor. - """ + """Folds the 'image' dimension into the batch dimension.""" model_name: str | None = None @@ -926,16 +961,21 @@ def map(self, element: dict[str, np.ndarray]) -> dict[str, np.ndarray]: if images is None: return element - # If ndim is greater than the expected ndim for a batched image tensor, - # it implies an extra dimension (e.g., number of images per example) - # that needs to be folded into the batch dimension. + image_masks = element.get("image_masks") + video_grid_thw = element.get("video_grid_thw") + second_per_grids = element.get("second_per_grids") + if images.ndim > len(self.target_shape): - # Compute the new shape by merging the batch and image count dimensions. - trailing_dims = self.target_shape[1:] + element["images"] = images.reshape(-1, *images.shape[2:]) + + if image_masks is not None: + element["image_masks"] = image_masks.reshape(-1, *image_masks.shape[2:]) + + if video_grid_thw is not None: + element["video_grid_thw"] = video_grid_thw.reshape(-1, video_grid_thw.shape[-1]) - # Reshape merges the leading dimensions (B, N) into one (-1) and - # appends the correct trailing dimensions. - element["images"] = images.reshape(-1, *trailing_dims) + if second_per_grids is not None: + element["second_per_grids"] = second_per_grids.reshape(-1) return element diff --git a/src/maxtext/multimodal/processor.py b/src/maxtext/multimodal/processor.py index 453ef026a0..cd1c3e781b 100644 --- a/src/maxtext/multimodal/processor.py +++ b/src/maxtext/multimodal/processor.py @@ -98,7 +98,16 @@ def get_image_offsets(config, processor_output: mm_utils.PreprocessorOutput | No return 0 -def reformat_prompt(prompt, image_placeholder, model_name, num_images, video_placeholder="<|video|>", num_videos=0): +def reformat_prompt( + prompt, + image_placeholder, + model_name, + num_images, + video_placeholder="<|video|>", + num_videos=0, + num_image_tokens=None, + num_video_tokens=None, +): """Reformat prompt for different models.""" if model_name in ["gemma3-4b", "gemma3-12b", "gemma3-27b"]: from maxtext.multimodal.processor_gemma3 import reformat_prompt_gemma3 # pylint: disable=import-outside-toplevel @@ -121,6 +130,8 @@ def reformat_prompt(prompt, image_placeholder, model_name, num_images, video_pla num_images=num_images, video_placeholder=video_placeholder, num_videos=num_videos, + num_image_tokens=num_image_tokens, + num_video_tokens=num_video_tokens, ) else: return prompt diff --git a/src/maxtext/multimodal/processor_qwen3_omni.py b/src/maxtext/multimodal/processor_qwen3_omni.py index c2ad512f45..a6e60b2b08 100644 --- a/src/maxtext/multimodal/processor_qwen3_omni.py +++ b/src/maxtext/multimodal/processor_qwen3_omni.py @@ -17,8 +17,10 @@ Original implementation from HuggingFace: Qwen/Qwen3-Omni-30B-A3B-Instruct. """ +import logging import math import os +import tempfile from dataclasses import dataclass import numpy as np @@ -32,6 +34,7 @@ decord = None from maxtext.multimodal import utils as mm_utils +from maxtext.utils import gcs_utils from maxtext.utils import max_logging # Image constants. @@ -471,25 +474,29 @@ def _read_video_decord(video_path, video_start=0.0, video_end=None) -> tuple[np. } try: vr = decord.VideoReader(video_path) - except Exception as e: - raise RuntimeError(f"Failed to read video from {video_path}: {e}") from e - total_frames, video_fps = len(vr), vr.get_avg_fps() - start_frame, end_frame, total_frames = calculate_video_frame_range( - video_config, - total_frames, - video_fps, - ) - nframes = smart_nframes(video_config, total_frames=total_frames, video_fps=video_fps) + total_frames, video_fps = len(vr), vr.get_avg_fps() + start_frame, end_frame, total_frames = calculate_video_frame_range( + video_config, + total_frames, + video_fps, + ) + nframes = smart_nframes(video_config, total_frames=total_frames, video_fps=video_fps) - # Use numpy linspace instead of torch.linspace - idx = np.linspace(start_frame, end_frame, nframes).round().astype(int).tolist() + # Use numpy linspace instead of torch.linspace + idx = np.linspace(start_frame, end_frame, nframes).round().astype(int).tolist() - video = vr.get_batch(idx).asnumpy() - # Convert from THWC to TCHW format using numpy - video = np.transpose(video, (0, 3, 1, 2)) + video = vr.get_batch(idx).asnumpy() + # Convert from THWC to TCHW format using numpy + video = np.transpose(video, (0, 3, 1, 2)) - sample_fps = nframes / max(total_frames, 1e-6) * video_fps - return video, sample_fps + sample_fps = nframes / max(total_frames, 1e-6) * video_fps + return video, sample_fps + except Exception as e: # pylint: disable=broad-exception-caught + logging.warning("Failed to read/decode video %s: %s. Using dummy video.", video_path, e) + # Return dummy video: 4 frames of 224x224 black pixels + # 224 is a multiple of 28 (patch_size * merge_size = 14 * 2 = 28) + dummy_video = np.zeros((4, 3, 224, 224), dtype=np.uint8) + return dummy_video, 2.0 def preprocess_video(video, config): @@ -634,31 +641,114 @@ def pre_process_audio_qwen3_omni(audio_array): return audio_features, audio_features_mask +def is_video_file(path): + if not isinstance(path, str): + return False + video_extensions = (".mp4", ".avi", ".mkv", ".webm", ".mov", ".gif") + return path.lower().endswith(video_extensions) + + def preprocess_mm_data_qwen3_omni_for_training(images, config): - """Preprocesses image(s) for Qwen3-Omni SFT training using model config constants.""" - images_in = [images] if isinstance(images, np.ndarray) else images - if config.image_size_for_vit is None: - force_resize = None - elif isinstance(config.image_size_for_vit, (list, tuple)): - force_resize = tuple(config.image_size_for_vit) + """Preprocesses image(s) or video for Qwen3-Omni SFT training using model config constants.""" + is_video = False + if isinstance(images, str) and is_video_file(images): + is_video = True + elif isinstance(images, list) and len(images) > 0 and isinstance(images[0], str) and is_video_file(images[0]): + is_video = True + + if is_video: + video_path = images + if isinstance(video_path, list): + if len(video_path) > 1: + raise ValueError("Only 1 video per example is supported for now.") + video_path = video_path[0] + + # Auto-assemble video directory from train or eval files + video_dir = None + if config.hf_train_files: + video_dir = os.path.dirname(config.hf_train_files) + elif config.hf_eval_files: + video_dir = os.path.dirname(config.hf_eval_files) + + if video_dir: + video_path = os.path.join(video_dir, video_path) + + is_gcs_video = video_path.startswith("gs://") + temp_local_path = None + + if is_gcs_video: + ext = os.path.splitext(video_path)[1] + with tempfile.NamedTemporaryFile(suffix=ext, prefix="maxtext_video_sft_", delete=False) as temp_file: + temp_local_path = temp_file.name + + max_logging.log(f"Downloading remote video {video_path} to temp local path {temp_local_path}...") + try: + bucket_name, prefix_name = gcs_utils.parse_gcs_bucket_and_prefix(video_path) + storage_client = gcs_utils.storage.Client() + bucket = storage_client.get_bucket(bucket_name) + blob = bucket.blob(prefix_name) + blob.download_to_filename(temp_local_path) + video_path = temp_local_path + except Exception as e: + if os.path.exists(temp_local_path): + os.remove(temp_local_path) + raise RuntimeError(f"Failed to download video from GCS: {e}") from e + else: + if not os.path.exists(video_path): + raw_path = images if isinstance(images, str) else images[0] + if os.path.exists(raw_path): + video_path = raw_path + else: + raise FileNotFoundError(f"Video file not found: {video_path} (original: {raw_path})") + + try: + video_array, _ = _read_video_decord(video_path) + finally: + if temp_local_path and os.path.exists(temp_local_path): + max_logging.log(f"Cleaning up temp local video file {temp_local_path}...") + os.remove(temp_local_path) + video_processed, video_grid_thw = preprocess_video(video_array, config) + video_values = np.reshape( + video_processed, + ( + 1, + config.num_channels_for_vit, + config.temporal_patch_size_for_vit * video_grid_thw[0, 0], + config.patch_size_for_vit * video_grid_thw[0, 1], + config.patch_size_for_vit * video_grid_thw[0, 2], + ), + ) + video_values, video_grid_thw, video_mask = maybe_pad_video_values_to_max_grid(video_values, video_grid_thw, config) + return Qwen3OmniPreprocessorOutput( + num_videos=1, + video_values=video_values, + video_grid_thw=video_grid_thw, + video_mask=video_mask, + ) else: - force_resize = (config.image_size_for_vit, config.image_size_for_vit) - pixel_values, pixel_grid_thw = pre_process_qwen3_image(images_in, config, force_resize=force_resize) - pixel_values = np.reshape( - pixel_values, - ( - len(images_in), - config.num_channels_for_vit, - config.temporal_patch_size_for_vit * pixel_grid_thw[0, 0], - config.patch_size_for_vit * pixel_grid_thw[0, 1], - config.patch_size_for_vit * pixel_grid_thw[0, 2], - ), - ) - return Qwen3OmniPreprocessorOutput( - num_images=len(images_in), - pixel_values=pixel_values, - pixel_grid_thw=pixel_grid_thw, - ) + images_in = [images] if isinstance(images, np.ndarray) else images + if config.image_size_for_vit is None: + force_resize = None + elif isinstance(config.image_size_for_vit, (list, tuple)): + force_resize = tuple(config.image_size_for_vit) + else: + force_resize = (config.image_size_for_vit, config.image_size_for_vit) + pixel_values, pixel_grid_thw = pre_process_qwen3_image(images_in, config, force_resize=force_resize) + pixel_values = np.reshape( + pixel_values, + ( + len(images_in), + config.num_channels_for_vit, + config.temporal_patch_size_for_vit * pixel_grid_thw[0, 0], + config.patch_size_for_vit * pixel_grid_thw[0, 1], + config.patch_size_for_vit * pixel_grid_thw[0, 2], + ), + ) + return Qwen3OmniPreprocessorOutput( + num_images=len(images_in), + pixel_values=pixel_values, + pixel_grid_thw=pixel_grid_thw, + ) def preprocess_mm_data_qwen3_omni(config): @@ -1280,14 +1370,27 @@ def get_rope_index( def reformat_prompt_qwen3_omni( - prompt, image_placeholder="<|image|>", num_images=0, video_placeholder="<|video|>", num_videos=0 + prompt, + image_placeholder="<|image|>", + num_images=0, + video_placeholder="<|video|>", + num_videos=0, + num_image_tokens=None, + num_video_tokens=None, ): """Reformat the prompt for Qwen3-Omni model.""" # Qwen3-Omni vision format: <|vision_start|><|image_pad|><|vision_end|> # Qwen3-Omni mm token order: image_pad, video_pad, audio_pad (standalone audios), then text tokens. # use_audio_in_video mode: such audio tokens are interleaved within video tokens. - qwen3_image_placeholder = "<|vision_start|><|image_pad|><|vision_end|>" - qwen3_video_placeholder = "<|vision_start|><|video_pad|><|vision_end|>" + if num_image_tokens is not None: + qwen3_image_placeholder = f"<|vision_start|>{'<|image_pad|>' * num_image_tokens}<|vision_end|>" + else: + qwen3_image_placeholder = "<|vision_start|><|image_pad|><|vision_end|>" + + if num_video_tokens is not None: + qwen3_video_placeholder = f"<|vision_start|>{'<|video_pad|>' * num_video_tokens}<|vision_end|>" + else: + qwen3_video_placeholder = "<|vision_start|><|video_pad|><|vision_end|>" if video_placeholder in prompt: prompt = prompt.replace(video_placeholder, qwen3_video_placeholder) diff --git a/src/maxtext/trainers/pre_train/train.py b/src/maxtext/trainers/pre_train/train.py index ee92b93435..3c62826ecc 100644 --- a/src/maxtext/trainers/pre_train/train.py +++ b/src/maxtext/trainers/pre_train/train.py @@ -137,13 +137,21 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr model_vars["batch_stats"] = sparsity_state else: model_vars = params + is_video = "video_grid_thw" in data logits, intermediate_outputs = model.apply( model_vars, data["inputs"], data["inputs_position"], decoder_segment_ids=data["inputs_segmentation"], - encoder_images=data["images"] if config.use_multimodal else None, - encoder_image_masks=data["image_masks"] if config.use_multimodal and "image_masks" in data else None, + encoder_images=data["images"] if config.use_multimodal and not is_video else None, + encoder_image_masks=data["image_masks"] + if config.use_multimodal and not is_video and "image_masks" in data + else None, + encoder_videos=data["images"] if config.use_multimodal and is_video else None, + encoder_video_masks=data["image_masks"] if config.use_multimodal and is_video and "image_masks" in data else None, + encoder_video_grid_thw=data["video_grid_thw"] + if config.use_multimodal and is_video and "video_grid_thw" in data + else None, enable_dropout=config.enable_dropout if is_train else False, rngs={"dropout": rng1, "params": aqt_rng}, # pyrefly: ignore[bad-argument-type] mutable=mutable_collections, @@ -187,12 +195,20 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr total_z_loss = jnp.sum(z_loss) else: # Flax NNX model: forward pass, then pop Intermediates sown during it. + is_video = "video_grid_thw" in data logits = model( decoder_input_tokens=data["inputs"], decoder_positions=data["inputs_position"], decoder_segment_ids=data["inputs_segmentation"], - encoder_images=data["images"] if config.use_multimodal else None, - encoder_image_masks=data["image_masks"] if config.use_multimodal and "image_masks" in data else None, + encoder_images=data["images"] if config.use_multimodal and not is_video else None, + encoder_image_masks=data["image_masks"] + if config.use_multimodal and not is_video and "image_masks" in data + else None, + encoder_videos=data["images"] if config.use_multimodal and is_video else None, + encoder_video_masks=data["image_masks"] if config.use_multimodal and is_video and "image_masks" in data else None, + encoder_video_grid_thw=data["video_grid_thw"] + if config.use_multimodal and is_video and "video_grid_thw" in data + else None, enable_dropout=config.enable_dropout if is_train else False, decoder_target_tokens=data["targets"], decoder_target_mask=data["targets_segmentation"], diff --git a/src/maxtext/utils/maxtext_utils.py b/src/maxtext/utils/maxtext_utils.py index c7530910eb..c24107b64d 100644 --- a/src/maxtext/utils/maxtext_utils.py +++ b/src/maxtext/utils/maxtext_utils.py @@ -162,11 +162,26 @@ def get_shaped_batch(config, batch_sharding=None): shaped_batch["targets_position"] = jax.ShapeDtypeStruct(batch_shape, jnp.int32, sharding=batch_sharding) shaped_batch["targets_segmentation"] = jax.ShapeDtypeStruct(batch_shape, jnp.int32, sharding=batch_sharding) if config.use_multimodal: - image_shape = mm_processor.get_dummy_image_shape_for_init( - config.model_name, batch_size=config.micro_batch_size_to_train_on - ) - shaped_batch["images"] = jax.ShapeDtypeStruct(image_shape, jnp.int32, sharding=batch_sharding) - shaped_batch["image_masks"] = jax.ShapeDtypeStruct(image_shape[:2], jnp.int32, sharding=batch_sharding) + is_video = getattr(config, "video_max_grid_t", None) is not None + if is_video: + max_t = config.video_max_grid_t + max_h = config.video_max_grid_h + max_w = config.video_max_grid_w + tps = config.temporal_patch_size_for_vit + patch = config.patch_size_for_vit + channels = config.num_channels_for_vit + batch_size = config.micro_batch_size_to_train_on + video_shape = (batch_size, channels, max_t * tps, max_h * patch, max_w * patch) + video_mask_shape = (batch_size, 1, max_t * tps, max_h * patch, max_w * patch) + shaped_batch["images"] = jax.ShapeDtypeStruct(video_shape, jnp.float32, sharding=batch_sharding) + shaped_batch["image_masks"] = jax.ShapeDtypeStruct(video_mask_shape, jnp.int32, sharding=batch_sharding) + shaped_batch["video_grid_thw"] = jax.ShapeDtypeStruct((batch_size, 3), jnp.int32, sharding=batch_sharding) + else: + image_shape = mm_processor.get_dummy_image_shape_for_init( + config.model_name, batch_size=config.micro_batch_size_to_train_on + ) + shaped_batch["images"] = jax.ShapeDtypeStruct(image_shape, jnp.int32, sharding=batch_sharding) + shaped_batch["image_masks"] = jax.ShapeDtypeStruct(image_shape[:2], jnp.int32, sharding=batch_sharding) if config.use_audio: audio_shape = mm_processor.get_dummy_audio_shape_for_init(config) shaped_batch["audios"] = jax.ShapeDtypeStruct(audio_shape, jnp.float32, sharding=batch_sharding) diff --git a/tools/data_generation/prepare_llava_video_178k.py b/tools/data_generation/prepare_llava_video_178k.py new file mode 100644 index 0000000000..1b959f81a1 --- /dev/null +++ b/tools/data_generation/prepare_llava_video_178k.py @@ -0,0 +1,127 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Script to download and prepare LLaVA-Video-178K dataset for MaxText.""" + +import argparse +import os +import shutil +import subprocess +import tarfile + +import datasets +from huggingface_hub import hf_hub_download + + +def main(): + parser = argparse.ArgumentParser(description="Prepare LLaVA-Video-178K dataset for MaxText.") + parser.add_argument("--output_dir", type=str, required=True, help="Directory to save the prepared dataset.") + parser.add_argument("--fold", type=str, default="0_30_s_academic_v0_1", help="Fold to download.") + parser.add_argument("--download_videos", action="store_true", help="Whether to download videos (large).") + parser.add_argument( + "--video_subset", + type=int, + default=None, + help="Number of video tars to download (1-8). If None, downloads all.", + ) + parser.add_argument( + "--local_temp_dir", + type=str, + default="/tmp/llava_video_178k_temp", + help="Local directory for temporary storage when uploading to GCS.", + ) + args = parser.parse_args() + + is_gcs = args.output_dir.startswith("gs://") + if is_gcs: + local_output_dir = args.local_temp_dir + print(f"GCS output directory detected. Using local temp directory: {local_output_dir}") + else: + local_output_dir = args.output_dir + + fold_dir = os.path.join(local_output_dir, args.fold) + os.makedirs(fold_dir, exist_ok=True) + + repo_id = "lmms-lab/LLaVA-Video-178K" + + # 1. Download and convert metadata + print("Downloading metadata...") + metadata_file = f"{args.fold}_cap_processed.json" + local_metadata_path = hf_hub_download( + repo_id=repo_id, + filename=f"{args.fold}/{metadata_file}", + repo_type="dataset", + local_dir=local_output_dir, + local_dir_use_symlinks=False, + ) + + print("Converting metadata to Parquet...") + dataset = datasets.load_dataset("json", data_files=local_metadata_path, split="train") + + parquet_path = os.path.join(fold_dir, "llava-video-178k-caption-00000-of-00001.parquet") + dataset.to_parquet(parquet_path) + print(f"Saved parquet to {parquet_path}") + + # 2. Download and extract videos + if args.download_videos: + print("Downloading videos...") + if args.video_subset is not None: + tar_files = [f"{args.fold}/{args.fold}_videos_{i}.tar.gz" for i in range(1, args.video_subset + 1)] + else: + tar_files = [f"{args.fold}/{args.fold}_videos_{i}.tar.gz" for i in range(1, 9)] + + for tar_file in tar_files: + print(f"Downloading {tar_file}...") + local_tar_path = hf_hub_download( + repo_id=repo_id, + filename=tar_file, + repo_type="dataset", + local_dir=local_output_dir, + local_dir_use_symlinks=False, + ) + + print(f"Extracting {local_tar_path} to {fold_dir}...") + with tarfile.open(local_tar_path, "r:gz") as tar: + tar.extractall(path=fold_dir) + + # Optionally remove tar file to save space + # os.remove(local_tar_path) + + print("Videos preparation complete.") + else: + print("Skipping videos download. Use --download_videos to download them.") + + if is_gcs: + print(f"Uploading prepared dataset from {fold_dir} to {args.output_dir}...") + + # Ensure trailing slash for GCS destination to copy directory correctly + gcs_dest = args.output_dir + if not gcs_dest.endswith("/"): + gcs_dest += "/" + + cmd = ["gsutil", "-m", "cp", "-r", fold_dir, gcs_dest] + print(f"Running command: {' '.join(cmd)}") + result = subprocess.run(cmd, capture_output=True, text=True, check=False) + if result.returncode != 0: + print(f"Error uploading to GCS: {result.stderr}") + raise RuntimeError(f"Failed to upload to GCS: {result.stderr}") + print("Upload complete.") + + print(f"Cleaning up local temp directory {fold_dir}...") + shutil.rmtree(fold_dir) + print("Cleanup complete.") + + +if __name__ == "__main__": + main()