Spaces:
Running
Running
| """ | |
| File upload handler for Video Model Studio. | |
| Processes uploaded files including videos, images, ZIPs, and WebDataset archives. | |
| """ | |
| import os | |
| import shutil | |
| import zipfile | |
| import tarfile | |
| import tempfile | |
| import gradio as gr | |
| from pathlib import Path | |
| from typing import List, Dict, Optional, Tuple, Any, Union | |
| import logging | |
| import traceback | |
| from vms.config import NORMALIZE_IMAGES_TO, TRAINING_VIDEOS_PATH, VIDEOS_TO_SPLIT_PATH, STAGING_PATH, DEFAULT_PROMPT_PREFIX | |
| from vms.utils import normalize_image, is_image_file, is_video_file, add_prefix_to_caption, webdataset_handler | |
| logger = logging.getLogger(__name__) | |
| class FileUploadHandler: | |
| """Handles processing of uploaded files""" | |
| def process_uploaded_files(self, file_paths: List[str], enable_splitting: bool) -> str: | |
| """Process uploaded file (ZIP, TAR, MP4, or image) | |
| Args: | |
| file_paths: File paths to the uploaded files from Gradio | |
| enable_splitting: Whether to enable automatic video splitting | |
| Returns: | |
| Status message string | |
| """ | |
| print(f"process_uploaded_files called with enable_splitting={enable_splitting} and file_paths = {str(file_paths)}") | |
| if not file_paths or len(file_paths) == 0: | |
| logger.warning("No files provided to process_uploaded_files") | |
| return "No files provided" | |
| for file_path in file_paths: | |
| print(f" - {str(file_path)}") | |
| file_path = Path(file_path) | |
| try: | |
| original_name = file_path.name | |
| logger.info(f"Processing uploaded file: {original_name}") | |
| # Determine file type from name | |
| file_ext = file_path.suffix.lower() | |
| if file_ext == '.zip': | |
| return self.process_zip_file(file_path, enable_splitting) | |
| elif file_ext == '.tar': | |
| return self.process_tar_file(file_path, enable_splitting) | |
| elif file_ext == '.mp4' or file_ext == '.webm': | |
| return self.process_mp4_file(file_path, original_name, enable_splitting) | |
| elif is_image_file(file_path): | |
| return self.process_image_file(file_path, original_name) | |
| else: | |
| logger.error(f"Unsupported file type: {file_ext}") | |
| raise gr.Error(f"Unsupported file type: {file_ext}") | |
| except Exception as e: | |
| logger.error(f"Error processing file {file_path}: {str(e)}", exc_info=True) | |
| raise gr.Error(f"Error processing file: {str(e)}") | |
| def process_zip_file(self, file_path: Path, enable_splitting: bool) -> str: | |
| """Process uploaded ZIP file containing media files or WebDataset tar files | |
| Args: | |
| file_path: Path to the uploaded ZIP file | |
| enable_splitting: Whether to enable automatic video splitting | |
| Returns: | |
| Status message string | |
| """ | |
| try: | |
| video_count = 0 | |
| image_count = 0 | |
| tar_count = 0 | |
| logger.info(f"Processing ZIP file: {file_path}") | |
| # Create temporary directory | |
| with tempfile.TemporaryDirectory() as temp_dir: | |
| # Extract ZIP | |
| extract_dir = Path(temp_dir) / "extracted" | |
| extract_dir.mkdir() | |
| with zipfile.ZipFile(file_path, 'r') as zip_ref: | |
| zip_ref.extractall(extract_dir) | |
| # Process each file | |
| for root, _, files in os.walk(extract_dir): | |
| for file in files: | |
| if file.startswith('._'): # Skip Mac metadata | |
| continue | |
| file_path = Path(root) / file | |
| try: | |
| # Check if it's a WebDataset tar file | |
| if file.lower().endswith('.tar'): | |
| logger.info(f"Processing WebDataset archive from ZIP: {file}") | |
| # Process WebDataset shard | |
| vid_count, img_count = webdataset_handler.process_webdataset_shard( | |
| file_path, VIDEOS_TO_SPLIT_PATH if enable_splitting else STAGING_PATH, STAGING_PATH | |
| ) | |
| video_count += vid_count | |
| image_count += img_count | |
| tar_count += 1 | |
| elif is_video_file(file_path): | |
| # Choose target directory based on auto-splitting setting | |
| target_dir = VIDEOS_TO_SPLIT_PATH if enable_splitting else STAGING_PATH | |
| target_path = target_dir / file_path.name | |
| counter = 1 | |
| while target_path.exists(): | |
| target_path = target_dir / f"{file_path.stem}___{counter}{file_path.suffix}" | |
| counter += 1 | |
| shutil.copy2(file_path, target_path) | |
| logger.info(f"Extracted video from ZIP: {file} -> {target_path.name}") | |
| video_count += 1 | |
| elif is_image_file(file_path): | |
| # Convert image and save to staging | |
| target_path = STAGING_PATH / f"{file_path.stem}.{NORMALIZE_IMAGES_TO}" | |
| counter = 1 | |
| while target_path.exists(): | |
| target_path = STAGING_PATH / f"{file_path.stem}___{counter}.{NORMALIZE_IMAGES_TO}" | |
| counter += 1 | |
| if normalize_image(file_path, target_path): | |
| logger.info(f"Extracted image from ZIP: {file} -> {target_path.name}") | |
| image_count += 1 | |
| # Copy associated caption file if it exists | |
| txt_path = file_path.with_suffix('.txt') | |
| if txt_path.exists() and not file.lower().endswith('.tar'): | |
| if is_video_file(file_path): | |
| shutil.copy2(txt_path, target_path.with_suffix('.txt')) | |
| logger.info(f"Copied caption file for {file}") | |
| elif is_image_file(file_path): | |
| caption = txt_path.read_text() | |
| caption = add_prefix_to_caption(caption, DEFAULT_PROMPT_PREFIX) | |
| target_path.with_suffix('.txt').write_text(caption) | |
| logger.info(f"Processed caption for {file}") | |
| except Exception as e: | |
| logger.error(f"Error processing {file_path.name} from ZIP: {str(e)}", exc_info=True) | |
| continue | |
| # Generate status message | |
| parts = [] | |
| if tar_count > 0: | |
| parts.append(f"{tar_count} WebDataset shard{'s' if tar_count != 1 else ''}") | |
| if video_count > 0: | |
| parts.append(f"{video_count} video{'s' if video_count != 1 else ''}") | |
| if image_count > 0: | |
| parts.append(f"{image_count} image{'s' if image_count != 1 else ''}") | |
| if not parts: | |
| logger.warning("No supported media files found in ZIP") | |
| return "No supported media files found in ZIP" | |
| status = f"Successfully stored {', '.join(parts)}" | |
| logger.info(status) | |
| gr.Info(status) | |
| return status | |
| except Exception as e: | |
| logger.error(f"Error processing ZIP: {str(e)}", exc_info=True) | |
| raise gr.Error(f"Error processing ZIP: {str(e)}") | |
| def process_tar_file(self, file_path: Path, enable_splitting: bool) -> str: | |
| """Process a WebDataset tar file | |
| Args: | |
| file_path: Path to the uploaded tar file | |
| enable_splitting: Whether to enable automatic video splitting | |
| Returns: | |
| Status message string | |
| """ | |
| try: | |
| logger.info(f"Processing WebDataset TAR file: {file_path}") | |
| video_count, image_count = webdataset_handler.process_webdataset_shard( | |
| file_path, VIDEOS_TO_SPLIT_PATH if enable_splitting else STAGING_PATH, STAGING_PATH | |
| ) | |
| # Generate status message | |
| parts = [] | |
| if video_count > 0: | |
| parts.append(f"{video_count} video{'s' if video_count != 1 else ''}") | |
| if image_count > 0: | |
| parts.append(f"{image_count} image{'s' if image_count != 1 else ''}") | |
| if not parts: | |
| logger.warning("No supported media files found in WebDataset") | |
| return "No supported media files found in WebDataset" | |
| status = f"Successfully extracted {' and '.join(parts)} from WebDataset" | |
| logger.info(status) | |
| gr.Info(status) | |
| return status | |
| except Exception as e: | |
| logger.error(f"Error processing WebDataset tar file: {str(e)}", exc_info=True) | |
| raise gr.Error(f"Error processing WebDataset tar file: {str(e)}") | |
| def process_mp4_file(self, file_path: Path, original_name: str, enable_splitting: bool) -> str: | |
| """Process a single video file | |
| Args: | |
| file_path: Path to the file | |
| original_name: Original filename | |
| enable_splitting: Whether to enable automatic video splitting | |
| Returns: | |
| Status message string | |
| """ | |
| print(f"process_mp4_file(self, file_path={str(file_path)}, original_name={str(original_name)}, enable_splitting={enable_splitting})") | |
| try: | |
| # Choose target directory based on auto-splitting setting | |
| target_dir = VIDEOS_TO_SPLIT_PATH if enable_splitting else STAGING_PATH | |
| print(f"target_dir = {target_dir}") | |
| # Create a unique filename | |
| target_path = target_dir / original_name | |
| # If file already exists, add number suffix | |
| counter = 1 | |
| while target_path.exists(): | |
| stem = Path(original_name).stem | |
| target_path = target_dir / f"{stem}___{counter}.mp4" | |
| counter += 1 | |
| logger.info(f"Processing video file: {original_name} -> {target_path}") | |
| # Copy the file to the target location | |
| shutil.copy2(file_path, target_path) | |
| logger.info(f"Successfully stored video: {target_path.name}") | |
| gr.Info(f"Successfully stored video: {target_path.name}") | |
| return f"Successfully stored video: {target_path.name}" | |
| except Exception as e: | |
| logger.error(f"Error processing video file: {str(e)}", exc_info=True) | |
| raise gr.Error(f"Error processing video file: {str(e)}") |