import os
import shutil
import subprocess
from huggingface_hub import HfApi, create_repo
from pathlib import Path
import json
import re
import logging
from typing import Any, Optional, Dict, List, Union, Tuple
logger = logging.getLogger(__name__)
def make_archive(source: str | Path, destination: str | Path):
source = str(source)
destination = str(destination)
#print(f"make_archive({source}, {destination})")
base = os.path.basename(destination)
name = base.split('.')[0]
format = base.split('.')[1]
archive_from = os.path.dirname(source)
archive_to = os.path.basename(source.strip(os.sep))
shutil.make_archive(name, format, archive_from, archive_to)
shutil.move('%s.%s'%(name,format), destination)
def get_video_fps(video_path: Path) -> Optional[str]:
"""Get FPS information from video file using ffprobe
Args:
video_path: Path to video file
Returns:
FPS string (e.g. "24 FPS, ") or None if unable to determine
"""
try:
cmd = [
'ffprobe',
'-v', 'error',
'-select_streams', 'v:0',
'-show_entries', 'stream=avg_frame_rate',
'-of', 'default=noprint_wrappers=1:nokey=1',
str(video_path)
]
result = subprocess.run(cmd, capture_output=True, text=True)
if result.returncode != 0:
logger.warning(f"Error getting FPS for {video_path}: {result.stderr}")
return None
fps = result.stdout.strip()
if '/' in fps:
# Convert fraction to decimal
num, den = map(int, fps.split('/'))
if den == 0:
return None
fps = str(round(num / den))
return f"{fps} FPS, "
except Exception as e:
logger.warning(f"Failed to get FPS for {video_path}: {e}")
return None
def extract_scene_info(filename: str) -> Tuple[str, Optional[int]]:
"""Extract base name and scene number from filename
Args:
filename: Input filename like "my_cool_video_1___001.mp4"
Returns:
Tuple of (base_name, scene_number)
e.g. ("my_cool_video_1", 1)
"""
# Match numbers at the end of the filename before extension
match = re.search(r'(.+?)___(\d+)$', Path(filename).stem)
if match:
return match.group(1), int(match.group(2))
return Path(filename).stem, None
def is_image_file(file_path: Path) -> bool:
"""Check if file is an image based on extension
Args:
file_path: Path to check
Returns:
bool: True if file has image extension
"""
image_extensions = {'.jpg', '.jpeg', '.png', '.webp', '.avif', '.heic'}
return file_path.suffix.lower() in image_extensions
def is_video_file(file_path: Path) -> bool:
"""Check if file is a video based on extension
Args:
file_path: Path to check
Returns:
bool: True if file has video extension
"""
video_extensions = {'.mp4', '.webm'}
return file_path.suffix.lower() in video_extensions
def parse_bool_env(env_value: Optional[str]) -> bool:
"""Parse environment variable string to boolean
Handles various true/false string representations:
- True: "true", "True", "TRUE", "1", etc
- False: "false", "False", "FALSE", "0", "", None
"""
if not env_value:
return False
return str(env_value).lower() in ('true', '1', 't', 'y', 'yes')
def validate_model_repo(repo_id: str) -> Dict[str, str]:
"""Validate HuggingFace model repository name
Args:
repo_id: Repository ID in format "username/model-name"
Returns:
Dict with error message if invalid, or None if valid
"""
if not repo_id:
return {"error": "Repository name is required"}
if "/" not in repo_id:
return {"error": "Repository name must be in format username/model-name"}
# Check characters
invalid_chars = set('<>:"/\\|?*')
if any(c in repo_id for c in invalid_chars):
return {"error": "Repository name contains invalid characters"}
return {"error": None}
def save_to_hub(model_path: Path, repo_id: str, token: str, commit_message: str = "Update model") -> bool:
"""Save model files to Hugging Face Hub
Args:
model_path: Path to model files
repo_id: Repository ID (username/model-name)
token: HuggingFace API token
commit_message: Optional commit message
Returns:
bool: True if successful, False if failed
"""
try:
api = HfApi(token=token)
# Validate repo_id
validation = validate_model_repo(repo_id)
if validation["error"]:
return False
# Create or get repo
try:
create_repo(repo_id, token=token, repo_type="model", exist_ok=True)
except Exception as e:
print(f"Error creating repo: {e}")
return False
# Upload all files
api.upload_folder(
folder_path=str(model_path),
repo_id=repo_id,
repo_type="model",
commit_message=commit_message
)
return True
except Exception as e:
print(f"Error uploading to hub: {e}")
return False
def parse_training_log(line: str) -> Dict:
"""Parse a training log line for metrics
Args:
line: Log line from training output
Returns:
Dict with parsed metrics (epoch, step, loss, etc)
"""
metrics = {}
try:
# Extract step/epoch info
if "step=" in line:
step = int(line.split("step=")[1].split()[0].strip(","))
metrics["step"] = step
if "epoch=" in line:
epoch = int(line.split("epoch=")[1].split()[0].strip(","))
metrics["epoch"] = epoch
if "loss=" in line:
loss = float(line.split("loss=")[1].split()[0].strip(","))
metrics["loss"] = loss
if "lr=" in line:
lr = float(line.split("lr=")[1].split()[0].strip(","))
metrics["learning_rate"] = lr
except:
pass
return metrics
def format_size(size_bytes: int) -> str:
"""Format bytes into human readable string with appropriate unit
Args:
size_bytes: Size in bytes
Returns:
Formatted string (e.g. "1.5 Gb")
"""
units = ['bytes', 'Kb', 'Mb', 'Gb', 'Tb']
unit_index = 0
size = float(size_bytes)
while size >= 1024 and unit_index < len(units) - 1:
size /= 1024
unit_index += 1
# Special case for bytes - no decimal places
if unit_index == 0:
return f"{int(size)} {units[unit_index]}"
return f"{size:.1f} {units[unit_index]}"
def count_media_files(path: Path) -> Tuple[int, int, int]:
"""Count videos and images in directory
Args:
path: Directory to scan
Returns:
Tuple of (video_count, image_count, total_size)
"""
video_count = 0
image_count = 0
total_size = 0
for file in path.glob("*"):
# Skip hidden files and caption files
if file.name.startswith('.') or file.suffix.lower() == '.txt':
continue
if is_video_file(file):
video_count += 1
total_size += file.stat().st_size
elif is_image_file(file):
image_count += 1
total_size += file.stat().st_size
return video_count, image_count, total_size
def format_media_title(action: str, video_count: int, image_count: int, total_size: int) -> str:
"""Format title with media counts and size
Args:
action: Action (eg "split", "caption")
video_count: Number of videos
image_count: Number of images
total_size: Total size in bytes
Returns:
Formatted title string
"""
parts = []
if image_count > 0:
parts.append(f"{image_count:,} photo{'s' if image_count != 1 else ''}")
if video_count > 0:
parts.append(f"{video_count:,} video{'s' if video_count != 1 else ''}")
if not parts:
return f"## 0 files to {action} (0 bytes)"
return f"## {' and '.join(parts)} to {action} ({format_size(total_size)})"
def add_prefix_to_caption(caption: str, prefix: str) -> str:
"""Add prefix to caption if not already present"""
if not prefix or not caption:
return caption
if caption.startswith(prefix):
return caption
return f"{prefix}{caption}"
def format_time(seconds: float) -> str:
"""Format time duration in seconds to human readable string
Args:
seconds: Time in seconds
Returns:
Formatted string (e.g. "2h 30m 45s")
"""
hours = int(seconds // 3600)
minutes = int((seconds % 3600) // 60)
secs = int(seconds % 60)
parts = []
if hours > 0:
parts.append(f"{hours}h")
if minutes > 0:
parts.append(f"{minutes}m")
if secs > 0 or not parts:
parts.append(f"{secs}s")
return " ".join(parts)