File size: 74,364 Bytes
84669a3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 450 451 452 453 454 455 456 457 458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495 496 497 498 499 500 501 502 503 504 505 506 507 508 509 510 511 512 513 514 515 516 517 518 519 520 521 522 523 524 525 526 527 528 529 530 531 532 533 534 535 536 537 538 539 540 541 542 543 544 545 546 547 548 549 550 551 552 553 554 555 556 557 558 559 560 561 562 563 564 565 566 567 568 569 570 571 572 573 574 575 576 577 578 579 580 581 582 583 584 585 586 587 588 589 590 591 592 593 594 595 596 597 598 599 600 601 602 603 604 605 606 607 608 609 610 611 612 613 614 615 616 617 618 619 620 621 622 623 624 625 626 627 628 629 630 631 632 633 634 635 636 637 638 639 640 641 642 643 644 645 646 647 648 649 650 651 652 653 654 655 656 657 658 659 660 661 662 663 664 665 666 667 668 669 670 671 672 673 674 675 676 677 678 679 680 681 682 683 684 685 686 687 688 689 690 691 692 693 694 695 696 697 698 699 700 701 702 703 704 705 706 707 708 709 710 711 712 713 714 715 716 717 718 719 720 721 722 723 724 725 726 727 728 729 730 731 732 733 734 735 736 737 738 739 740 741 742 743 744 745 746 747 748 749 750 751 752 753 754 755 756 757 758 759 760 761 762 763 764 765 766 767 768 769 770 771 772 773 774 775 776 777 778 779 780 781 782 783 784 785 786 787 788 789 790 791 792 793 794 795 796 797 798 799 800 801 802 803 804 805 806 807 808 809 810 811 812 813 814 815 816 817 818 819 820 821 822 823 824 825 826 827 828 829 830 831 832 833 834 835 836 837 838 839 840 841 842 843 844 845 846 847 848 849 850 851 852 853 854 855 856 857 858 859 860 861 862 863 864 865 866 867 868 869 870 871 872 873 874 875 876 877 878 879 880 881 882 883 884 885 886 887 888 889 890 891 892 893 894 895 896 897 898 899 900 901 902 903 904 905 906 907 908 909 910 911 912 913 914 915 916 917 918 919 920 921 922 923 924 925 926 927 928 929 930 931 932 933 934 935 936 937 938 939 940 941 942 943 944 945 946 947 948 949 950 951 952 953 954 955 956 957 958 959 960 961 962 963 964 965 966 967 968 969 970 971 972 973 974 975 976 977 978 979 980 981 982 983 984 985 986 987 988 989 990 991 992 993 994 995 996 997 998 999 1000 1001 1002 1003 1004 1005 1006 1007 1008 1009 1010 1011 1012 1013 1014 1015 1016 1017 1018 1019 1020 1021 1022 1023 1024 1025 1026 1027 1028 1029 1030 1031 1032 1033 1034 1035 1036 1037 1038 1039 1040 1041 1042 1043 1044 1045 1046 1047 1048 1049 1050 1051 1052 1053 1054 1055 1056 1057 1058 1059 1060 1061 1062 1063 1064 1065 1066 1067 1068 1069 1070 1071 1072 1073 1074 1075 1076 1077 1078 1079 1080 1081 1082 1083 1084 1085 1086 1087 1088 1089 1090 1091 1092 1093 1094 1095 1096 1097 1098 1099 1100 1101 1102 1103 1104 1105 1106 1107 1108 1109 1110 1111 1112 1113 1114 1115 1116 1117 1118 1119 1120 1121 1122 1123 1124 1125 1126 1127 1128 1129 1130 1131 1132 1133 1134 1135 1136 1137 1138 1139 1140 1141 1142 1143 1144 1145 1146 1147 1148 1149 1150 1151 1152 1153 1154 1155 1156 1157 1158 1159 1160 1161 1162 1163 1164 1165 1166 1167 1168 1169 1170 1171 1172 1173 1174 1175 1176 1177 1178 1179 1180 1181 1182 1183 1184 1185 1186 1187 1188 1189 1190 1191 1192 1193 1194 1195 1196 1197 1198 1199 1200 1201 1202 1203 1204 1205 1206 1207 1208 1209 1210 1211 1212 1213 1214 1215 1216 1217 1218 1219 1220 1221 1222 1223 1224 1225 1226 1227 1228 1229 1230 1231 1232 1233 1234 1235 1236 1237 1238 1239 1240 1241 1242 1243 1244 1245 1246 1247 1248 1249 1250 1251 1252 1253 1254 1255 1256 1257 1258 1259 1260 |
import os
import gc
import sys
import re
import numpy as np
import torch
import imageio
import gradio as gr
import subprocess
import devicetorch
import json
import math
import shutil # For moving files
from datetime import datetime
from pathlib import Path
from huggingface_hub import snapshot_download
from tqdm.auto import tqdm
from torchvision.transforms.functional import to_tensor, to_pil_image
from modules.toolbox.rife_core import RIFEHandler
from modules.toolbox.esrgan_core import ESRGANUpscaler
from modules.toolbox.message_manager import MessageManager
device_name_str = devicetorch.get(torch)
VIDEO_QUALITY = 8 # Used by imageio.mimwrite quality/quantizer
class VideoProcessor:
def __init__(self, message_manager: MessageManager, settings):
self.message_manager = message_manager
self.rife_handler = RIFEHandler(message_manager)
self.device_obj = torch.device(device_name_str) # Store device_obj
self.esrgan_upscaler = ESRGANUpscaler(message_manager, self.device_obj)
self.settings = settings
# FFmpeg/FFprobe paths and status flags
self.ffmpeg_exe = None
self.ffprobe_exe = None
self.has_ffmpeg = False
self.has_ffprobe = False
# --- NEW: Add source tracking ---
self.ffmpeg_source = None
self.ffprobe_source = None
self._tb_initialize_ffmpeg() # Finds executables and sets flags
studio_output_dir = Path(self.settings.get("output_dir"))
self.postprocessed_output_root_dir = studio_output_dir / "postprocessed_output"
self._base_temp_output_dir = self.postprocessed_output_root_dir / "temp_processing"
self._base_permanent_save_dir = self.postprocessed_output_root_dir / "saved_videos"
self.toolbox_video_output_dir = self._base_temp_output_dir
self.toolbox_permanent_save_dir = self._base_permanent_save_dir
# Ensure all necessary directories exist
os.makedirs(self.postprocessed_output_root_dir, exist_ok=True)
os.makedirs(self._base_temp_output_dir, exist_ok=True)
os.makedirs(self._base_permanent_save_dir, exist_ok=True)
self.extracted_frames_target_path = self.postprocessed_output_root_dir / "frames" / "extracted_frames"
os.makedirs(self.extracted_frames_target_path, exist_ok=True)
self.reassembled_video_target_path = self.postprocessed_output_root_dir / "frames" / "reassembled_videos"
os.makedirs(self.reassembled_video_target_path, exist_ok=True)
def _tb_initialize_ffmpeg(self):
"""Finds FFmpeg/FFprobe and sets status flags and sources."""
(
self.ffmpeg_exe,
self.ffmpeg_source,
self.ffprobe_exe,
self.ffprobe_source,
) = self._tb_find_ffmpeg_executables()
self.has_ffmpeg = bool(self.ffmpeg_exe)
self.has_ffprobe = bool(self.ffprobe_exe)
self._report_ffmpeg_status()
def _tb_find_ffmpeg_executables(self):
"""
Finds ffmpeg and ffprobe with a priority system.
Priority: 1. Bundled -> 2. System PATH -> 3. imageio-ffmpeg
Returns (ffmpeg_path, ffmpeg_source, ffprobe_path, ffprobe_source)
"""
ffmpeg_path, ffprobe_path = None, None
ffmpeg_source, ffprobe_source = None, None
ffmpeg_name = "ffmpeg.exe" if sys.platform == "win32" else "ffmpeg"
ffprobe_name = "ffprobe.exe" if sys.platform == "win32" else "ffprobe"
# --- Priority 1: Bundled ---
try:
script_dir = os.path.dirname(os.path.abspath(__file__))
bin_dir = os.path.join(script_dir, 'bin')
bundled_ffmpeg = os.path.join(bin_dir, ffmpeg_name)
bundled_ffprobe = os.path.join(bin_dir, ffprobe_name)
if os.path.exists(bundled_ffmpeg):
ffmpeg_path = bundled_ffmpeg
ffmpeg_source = "Bundled"
if os.path.exists(bundled_ffprobe):
ffprobe_path = bundled_ffprobe
ffprobe_source = "Bundled"
except Exception:
pass # Silently fail and move to next priority
# --- Priority 2: System PATH ---
# Use shutil.which to find executables in the system's PATH
if not ffmpeg_path:
path_from_env = shutil.which(ffmpeg_name)
if path_from_env:
ffmpeg_path = path_from_env
ffmpeg_source = "System PATH"
if not ffprobe_path:
path_from_env = shutil.which(ffprobe_name)
if path_from_env:
ffprobe_path = path_from_env
ffprobe_source = "System PATH"
# --- Priority 3: imageio-ffmpeg ---
# This will only provide ffmpeg, not ffprobe.
if not ffmpeg_path:
try:
imageio_ffmpeg_exe = imageio.plugins.ffmpeg.get_exe()
if os.path.isfile(imageio_ffmpeg_exe):
ffmpeg_path = imageio_ffmpeg_exe
ffmpeg_source = "imageio-ffmpeg"
except Exception:
pass # Silently fail
return ffmpeg_path, ffmpeg_source, ffprobe_path, ffprobe_source
def _report_ffmpeg_status(self):
"""Provides a summary of FFmpeg/FFprobe status based on what was found."""
# Ideal case: Bundled version is used
if self.ffmpeg_source == "Bundled" and self.ffprobe_source == "Bundled":
self.message_manager.add_message(f"Bundled FFmpeg found: {self.ffmpeg_exe}", "SUCCESS")
self.message_manager.add_message(f"Bundled FFprobe found: {self.ffprobe_exe}", "SUCCESS")
self.message_manager.add_message("All video and audio features are enabled.", "SUCCESS")
return
# Fallback cases: Report what was found and where
if self.has_ffmpeg:
self.message_manager.add_message(f"FFmpeg found via {self.ffmpeg_source}: {self.ffmpeg_exe}", "SUCCESS")
else:
self.message_manager.add_error(
"Critical: FFmpeg executable could not be found. "
"Most video processing operations will fail. Please try running the setup script."
)
if self.has_ffprobe:
self.message_manager.add_message(f"FFprobe found via {self.ffprobe_source}: {self.ffprobe_exe}", "SUCCESS")
else:
self.message_manager.add_warning(
"FFprobe not found. Audio detection and full video analysis will be limited."
)
# Add a specific nag if the bundled version should exist but doesn't
if self.ffmpeg_source != "Bundled":
self.message_manager.add_warning(
"For full functionality, please run the 'setup_ffmpeg.py' script."
)
def set_autosave_mode(self, autosave_enabled: bool):
if autosave_enabled:
self.toolbox_video_output_dir = self._base_permanent_save_dir
self.message_manager.add_message("Autosave ENABLED: Processed videos will be saved to the permanent folder.", "SUCCESS")
else:
self.toolbox_video_output_dir = self._base_temp_output_dir
self.message_manager.add_message("Autosave DISABLED: Processed videos will be saved to the temporary folder.", "INFO")
def _tb_log_ffmpeg_error(self, e_ffmpeg: subprocess.CalledProcessError, operation_description: str):
self.message_manager.add_error(f"FFmpeg failed during {operation_description}.")
ffmpeg_stderr_str = e_ffmpeg.stderr.strip() if e_ffmpeg.stderr else ""
ffmpeg_stdout_str = e_ffmpeg.stdout.strip() if e_ffmpeg.stdout else ""
details_log = []
if ffmpeg_stderr_str: details_log.append(f"FFmpeg Stderr: {ffmpeg_stderr_str}")
if ffmpeg_stdout_str: details_log.append(f"FFmpeg Stdout: {ffmpeg_stdout_str}")
if details_log:
self.message_manager.add_message("FFmpeg Output:\n" + "\n".join(details_log), "INFO")
else:
self.message_manager.add_message(f"No specific output from FFmpeg. (Return code: {e_ffmpeg.returncode}, Command: '{e_ffmpeg.cmd}')", "INFO")
def tb_extract_frames(self, video_path, extraction_rate, progress=gr.Progress()):
if video_path is None:
self.message_manager.add_warning("No input video for frame extraction.")
return None
if not isinstance(extraction_rate, int) or extraction_rate < 1:
self.message_manager.add_error("Extraction rate must be a positive integer (1 for all frames, N for every Nth frame).")
return None
resolved_video_path = str(Path(video_path).resolve())
output_folder_name = self._tb_generate_output_folder_path(
resolved_video_path,
suffix=f"extracted_every_{extraction_rate}")
os.makedirs(output_folder_name, exist_ok=True)
self.message_manager.add_message(
f"Starting frame extraction for {os.path.basename(resolved_video_path)} (every {extraction_rate} frame(s))."
)
self.message_manager.add_message(f"Outputting to: {output_folder_name}")
progress(0, desc="Initializing frame extraction...")
reader = None
try:
reader = imageio.get_reader(resolved_video_path) # Default 'ffmpeg' plugin if available
total_frames = None
try:
# Try to get nframes metadata first, as count_frames can be slow or Inf for streams
meta_nframes = reader.get_meta_data().get('nframes')
if meta_nframes and meta_nframes != float('inf'):
total_frames = int(meta_nframes)
elif hasattr(reader, 'count_frames'): # Fallback to count_frames if available and nframes is not
total_frames_counted = reader.count_frames()
if total_frames_counted != float('inf'):
total_frames = total_frames_counted
except Exception:
self.message_manager.add_warning("Could not accurately determine total frames. Progress might be approximate.")
total_frames = None
extracted_count = 0
frame_iterable = reader
if total_frames:
frame_iterable = progress.tqdm(reader, total=total_frames, desc="Extracting frames")
else:
self.message_manager.add_message("Processing frames (total unknown)...")
for i, frame in enumerate(frame_iterable):
if not total_frames and i % 100 == 0:
progress(i / (i + 1000.0), desc=f"Extracting frame {i+1}...")
if i % extraction_rate == 0:
frame_filename = f"frame_{extracted_count:06d}.png"
output_frame_path = os.path.join(output_folder_name, frame_filename)
imageio.imwrite(output_frame_path, frame, format='PNG')
extracted_count += 1
progress(1.0, desc="Extraction complete.")
self.message_manager.add_success(f"Successfully extracted {extracted_count} frames to: {output_folder_name}")
return output_folder_name
except Exception as e:
self.message_manager.add_error(f"Error during frame extraction: {e}")
import traceback
self.message_manager.add_error(traceback.format_exc())
if "Could not find a backend" in str(e) or "No such file or directory: 'ffmpeg'" in str(e).lower():
self.message_manager.add_error("This might indicate an issue with FFmpeg backend for imageio. Ensure 'imageio-ffmpeg' is installed or FFmpeg is in PATH.")
progress(1.0, desc="Error during extraction.")
return None
finally:
if reader:
reader.close()
gc.collect()
def tb_get_extracted_frame_folders(self) -> list:
if not os.path.exists(self.extracted_frames_target_path):
self.message_manager.add_warning(f"Extracted frames directory not found: {self.extracted_frames_target_path}")
return []
try:
folders = [
d for d in os.listdir(self.extracted_frames_target_path)
if os.path.isdir(os.path.join(self.extracted_frames_target_path, d))
]
folders.sort()
# self.message_manager.add_message(f"Found {len(folders)} extracted frame folders.") # Can be noisy
return folders
except Exception as e:
self.message_manager.add_error(f"Error scanning for extracted frame folders: {e}")
return []
def tb_delete_extracted_frames_folder(self, folder_name_to_delete: str) -> bool:
if not folder_name_to_delete:
self.message_manager.add_warning("No folder selected for deletion.")
return False
folder_path_to_delete = os.path.join(self.extracted_frames_target_path, folder_name_to_delete)
if not os.path.exists(folder_path_to_delete) or not os.path.isdir(folder_path_to_delete):
self.message_manager.add_error(f"Folder not found or is not a directory: {folder_path_to_delete}")
return False
try:
shutil.rmtree(folder_path_to_delete)
self.message_manager.add_success(f"Successfully deleted folder: {folder_name_to_delete}")
return True
except Exception as e:
self.message_manager.add_error(f"Error deleting folder '{folder_name_to_delete}': {e}")
self.message_manager.add_error(traceback.format_exc() if 'traceback' in sys.modules else str(e))
return False
def tb_reassemble_frames_to_video(self, frames_source, output_fps, output_base_name_override=None, progress=gr.Progress()):
if not frames_source:
self.message_manager.add_warning("No frames source (folder or files) provided for reassembly.")
return None
# This operation primarily uses imageio.
# FFmpeg dependency is indirect via imageio-ffmpeg for mimwrite.
try:
output_fps = int(output_fps)
if output_fps <= 0:
self.message_manager.add_error("Output FPS must be a positive number.")
return None
except ValueError:
self.message_manager.add_error("Invalid FPS value for reassembly.")
return None
self.message_manager.add_message(f"Starting frame reassembly to video at {output_fps} FPS.")
frame_info_list = []
frames_data_prepared = False # To track if frames_data list was populated for cleanup
try:
if isinstance(frames_source, str) and os.path.isdir(frames_source):
self.message_manager.add_message(f"Processing frames from directory: {frames_source}")
for filename in os.listdir(frames_source):
if filename.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.webp')):
full_path = os.path.join(frames_source, filename)
frame_info_list.append({
'original_like_filename': filename,
'temp_path': full_path
})
elif isinstance(frames_source, list): # List of Gradio FileData objects
self.message_manager.add_message(f"Processing {len(frames_source)} uploaded files for reassembly.")
for temp_file_wrapper in frames_source:
# Gradio temp files might have generic names, try to use original if available
original_like_filename = getattr(temp_file_wrapper, 'orig_name', None) or os.path.basename(temp_file_wrapper.name)
frame_info_list.append({
'original_like_filename': original_like_filename,
'temp_path': temp_file_wrapper.name
})
else:
self.message_manager.add_error("Invalid frames_source type for reassembly.")
return None
if not frame_info_list:
self.message_manager.add_warning("No valid image files found in the provided source to reassemble.")
return None
def natural_sort_key_for_dict(item):
filename = item['original_like_filename']
return [int(text) if text.isdigit() else text.lower() for text in re.split('([0-9]+)', filename)]
frame_info_list.sort(key=natural_sort_key_for_dict)
self.message_manager.add_message(f"Sorted {len(frame_info_list)} frames based on their filenames.")
# For debugging, log first few sorted names
# if frame_info_list:
# debug_sorted_names = [info['original_like_filename'] for info in frame_info_list[:min(5, len(frame_info_list))]]
# self.message_manager.add_message(f"DEBUG: First {len(debug_sorted_names)} sorted filenames: {debug_sorted_names}", "DEBUG")
output_file_basename = "reassembled_video"
if output_base_name_override and isinstance(output_base_name_override, str) and output_base_name_override.strip():
sanitized_name = "".join(c if c.isalnum() or c in (' ', '_', '-') else '_' for c in output_base_name_override.strip())
output_file_basename = Path(sanitized_name).stem
if not output_file_basename: output_file_basename = "reassembled_video"
self.message_manager.add_message(f"Using custom output video base name: {output_file_basename}")
output_video_path = self._tb_generate_output_path(
input_material_name=output_file_basename,
suffix=f"{output_fps}fps_reassembled",
target_dir=self.reassembled_video_target_path, # Specific target for reassembled
ext=".mp4"
)
frames_data = []
frames_data_prepared = True
self.message_manager.add_message("Reading frame images (in sorted order)...")
frame_iterator = frame_info_list
if frame_info_list and progress is not None and hasattr(progress, 'tqdm'):
frame_iterator = progress.tqdm(frame_info_list, desc="Reading frames")
for frame_info in frame_iterator:
frame_actual_path = frame_info['temp_path']
filename_for_log = frame_info['original_like_filename']
try:
if not filename_for_log.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.webp')):
self.message_manager.add_warning(f"Skipping non-standard image file: {filename_for_log}.")
continue
frames_data.append(imageio.imread(frame_actual_path))
except Exception as e_read_frame:
self.message_manager.add_warning(f"Could not read frame ({filename_for_log}): {e_read_frame}. Skipping.")
if not frames_data:
self.message_manager.add_error("No valid frames could be successfully read for reassembly.")
return None
self.message_manager.add_message(f"Writing {len(frames_data)} frames to video: {output_video_path}")
# Ensure macro_block_size is None if not multiple of 16, or handle fps issues
imageio.mimwrite(output_video_path, frames_data, fps=output_fps, quality=VIDEO_QUALITY, macro_block_size=None) # macro_block_size often problematic
self.message_manager.add_success(f"Successfully reassembled {len(frames_data)} frames into: {output_video_path}")
return output_video_path
except Exception as e:
self.message_manager.add_error(f"Error during frame reassembly: {e}")
import traceback
self.message_manager.add_error(traceback.format_exc())
if "Could not find a backend" in str(e) or "No such file or directory: 'ffmpeg'" in str(e).lower():
self.message_manager.add_error("This might indicate an issue with FFmpeg backend for imageio. Ensure 'imageio-ffmpeg' is installed or FFmpeg is in PATH.")
return None
finally:
if frames_data_prepared and 'frames_data' in locals():
del frames_data # Explicitly delete large list of frames
gc.collect()
def _tb_clean_filename(self, filename):
filename = re.sub(r'_\d{6}_\d{6}', '', filename) # Example timestamp pattern
filename = re.sub(r'_\d{6}_\d{4}', '', filename) # Another example
return filename.strip('_')
def _tb_generate_output_path(self, input_material_name, suffix, target_dir, ext=".mp4"):
base_name = Path(input_material_name).stem
if not base_name: base_name = "untitled_video"
cleaned_name = self._tb_clean_filename(base_name)
timestamp = datetime.now().strftime("%y%m%d_%H%M%S")
filename = f"{cleaned_name}_{suffix}_{timestamp}{ext}"
return os.path.join(target_dir, filename)
def _tb_generate_output_folder_path(self, input_video_path, suffix):
base_name = Path(input_video_path).stem
if not base_name: base_name = "untitled_video_frames"
cleaned_name = self._tb_clean_filename(base_name)
timestamp = datetime.now().strftime("%y%m%d_%H%M%S")
folder_name = f"{cleaned_name}_{suffix}_{timestamp}"
return os.path.join(self.extracted_frames_target_path, folder_name)
def tb_copy_video_to_permanent_storage(self, temp_video_path):
if not temp_video_path or not os.path.exists(temp_video_path):
self.message_manager.add_error("No video file provided or file does not exist to save.")
return temp_video_path
try:
video_filename = Path(temp_video_path).name
permanent_video_path = os.path.join(self.toolbox_permanent_save_dir, video_filename)
os.makedirs(self.toolbox_permanent_save_dir, exist_ok=True)
self.message_manager.add_message(f"Copying '{video_filename}' to permanent storage: '{permanent_video_path}'")
shutil.copy2(temp_video_path, permanent_video_path)
self.message_manager.add_success(f"Video saved to: {permanent_video_path}")
return permanent_video_path
except Exception as e:
self.message_manager.add_error(f"Error saving video to permanent storage: {e}")
self.message_manager.add_error(traceback.format_exc())
return temp_video_path
def tb_analyze_video_input(self, video_path):
if video_path is None:
self.message_manager.add_warning("No video provided for analysis.")
return "Please upload a video."
resolved_video_path = str(Path(video_path).resolve())
analysis_report_lines = [] # Use a list to build the report string
# Variables to hold parsed info, initialized to defaults
video_width, video_height = 0, 0
num_frames_value = None # For the upscale warning
duration_display, fps_display, resolution_display, nframes_display, has_audio_str = "N/A", "N/A", "N/A", "N/A", "No"
analysis_source = "imageio" # Default analysis source
if self.has_ffprobe:
self.message_manager.add_message(f"Analyzing video with ffprobe: {os.path.basename(video_path)}")
try:
probe_cmd = [
self.ffprobe_exe, "-v", "error", "-show_format", "-show_streams",
"-of", "json", resolved_video_path
]
result = subprocess.run(probe_cmd, capture_output=True, text=True, check=True, errors='ignore')
probe_data = json.loads(result.stdout)
video_stream = next((s for s in probe_data.get("streams", []) if s.get("codec_type") == "video"), None)
audio_stream = next((s for s in probe_data.get("streams", []) if s.get("codec_type") == "audio"), None)
if not video_stream:
self.message_manager.add_error("No video stream found in the file (ffprobe).")
# Fall through to imageio or return error, depending on desired strictness
# For now, let's allow imageio to try
else:
analysis_source = "ffprobe"
duration_str = probe_data.get("format", {}).get("duration", "0")
duration = float(duration_str) if duration_str and duration_str.replace('.', '', 1).isdigit() else 0.0
duration_display = f"{duration:.2f} seconds"
r_frame_rate_str = video_stream.get("r_frame_rate", "0/0")
avg_frame_rate_str = video_stream.get("avg_frame_rate", "0/0")
calculated_fps = 0.0
def parse_fps(fps_s):
if isinstance(fps_s, (int, float)): return float(fps_s)
if isinstance(fps_s, str) and "/" in fps_s:
try: num, den = map(float, fps_s.split('/')); return num / den if den != 0 else 0.0
except ValueError: return 0.0
try: return float(fps_s)
except ValueError: return 0.0
r_fps_val = parse_fps(r_frame_rate_str); avg_fps_val = parse_fps(avg_frame_rate_str)
if r_fps_val > 0: calculated_fps = r_fps_val; fps_display = f"{r_fps_val:.2f} FPS"
if avg_fps_val > 0 and abs(r_fps_val - avg_fps_val) > 0.01 : # Only show average if meaningfully different
calculated_fps = avg_fps_val # Prefer average if it's different and valid
fps_display = f"{avg_fps_val:.2f} FPS (Avg, r: {r_fps_val:.2f})"
elif avg_fps_val > 0 and r_fps_val <=0:
calculated_fps = avg_fps_val; fps_display = f"{avg_fps_val:.2f} FPS (Average)"
video_width = video_stream.get("width", 0)
video_height = video_stream.get("height", 0)
resolution_display = f"{video_width}x{video_height}" if video_width and video_height else "N/A"
nframes_str_probe = video_stream.get("nb_frames")
if nframes_str_probe and nframes_str_probe.isdigit():
num_frames_value = int(nframes_str_probe)
nframes_display = str(num_frames_value)
elif duration > 0 and calculated_fps > 0:
num_frames_value = int(duration * calculated_fps)
nframes_display = f"{num_frames_value} (Calculated)"
if audio_stream:
has_audio_str = (f"Yes (Codec: {audio_stream.get('codec_name', 'N/A')}, "
f"Channels: {audio_stream.get('channels', 'N/A')}, "
f"Rate: {audio_stream.get('sample_rate', 'N/A')} Hz)")
self.message_manager.add_success("Video analysis complete (using ffprobe).")
except (subprocess.CalledProcessError, json.JSONDecodeError, Exception) as e_ffprobe:
self.message_manager.add_warning(f"ffprobe analysis failed ({type(e_ffprobe).__name__}). Trying imageio fallback.")
if isinstance(e_ffprobe, subprocess.CalledProcessError):
self._tb_log_ffmpeg_error(e_ffprobe, "video analysis with ffprobe")
analysis_source = "imageio" # Ensure fallback if ffprobe fails midway
if analysis_source == "imageio": # Either ffprobe not available, or it failed
self.message_manager.add_message(f"Analyzing video with imageio: {os.path.basename(video_path)}")
reader = None
try:
reader = imageio.get_reader(resolved_video_path)
meta = reader.get_meta_data()
duration_imgio_val = meta.get('duration')
duration_display = f"{float(duration_imgio_val):.2f} seconds" if duration_imgio_val is not None else "N/A"
fps_val_imgio = meta.get('fps')
fps_display = f"{float(fps_val_imgio):.2f} FPS" if fps_val_imgio is not None else "N/A"
size_imgio = meta.get('size')
if isinstance(size_imgio, tuple) and len(size_imgio) == 2:
video_width, video_height = int(size_imgio[0]), int(size_imgio[1])
resolution_display = f"{video_width}x{video_height}"
else:
resolution_display = "N/A"
nframes_val_imgio_meta = meta.get('nframes')
if nframes_val_imgio_meta not in [float('inf'), "N/A", None] and isinstance(nframes_val_imgio_meta, (int,float)):
num_frames_value = int(nframes_val_imgio_meta)
nframes_display = str(num_frames_value)
elif hasattr(reader, 'count_frames'):
try:
nframes_val_imgio_count = reader.count_frames()
if nframes_val_imgio_count != float('inf'):
num_frames_value = int(nframes_val_imgio_count)
nframes_display = f"{num_frames_value} (Counted)"
else: nframes_display = "Unknown (Stream or very long)"
except Exception: nframes_display = "Unknown (Frame count failed)"
has_audio_str = "(Audio info not available via imageio)"
self.message_manager.add_success("Video analysis complete (using imageio).")
except Exception as e_imgio:
self.message_manager.add_error(f"Error analyzing video with imageio: {e_imgio}")
import traceback
self.message_manager.add_error(traceback.format_exc())
return f"Error analyzing video: Both ffprobe (if attempted) and imageio failed."
finally:
if reader: reader.close()
# --- Construct Main Analysis Report ---
analysis_report_lines.append(f"Video Analysis ({analysis_source}):")
analysis_report_lines.append(f"File: {os.path.basename(video_path)}")
analysis_report_lines.append("------------------------------------")
analysis_report_lines.append(f"Duration: {duration_display}")
analysis_report_lines.append(f"Frame Rate: {fps_display}")
analysis_report_lines.append(f"Resolution: {resolution_display}")
analysis_report_lines.append(f"Frames: {nframes_display}")
analysis_report_lines.append(f"Audio: {has_audio_str}")
analysis_report_lines.append(f"Source: {video_path}")
# --- Append UPSCALE ADVISORY Conditionally ---
if video_width > 0 and video_height > 0: # Ensure we have dimensions
HD_WIDTH_THRESHOLD = 1920
FOUR_K_WIDTH_THRESHOLD = 3800
is_hd_or_larger = (video_width >= HD_WIDTH_THRESHOLD or video_height >= (HD_WIDTH_THRESHOLD * 9/16 * 0.95)) # Adjusted height for aspect ratios
is_4k_or_larger = (video_width >= FOUR_K_WIDTH_THRESHOLD or video_height >= (FOUR_K_WIDTH_THRESHOLD * 9/16 * 0.95))
upscale_warnings = []
if is_4k_or_larger:
upscale_warnings.append(
"This video is 4K resolution or higher. Upscaling (e.g., to 8K+) will be very "
"slow, memory-intensive, and may cause issues. Proceed with caution."
)
elif is_hd_or_larger:
upscale_warnings.append(
"This video is HD or larger. Upscaling (e.g., to 4K+) will be resource-intensive "
"and slow. Ensure your system is prepared."
)
if num_frames_value and num_frames_value > 900: # e.g., > 30 seconds at 30fps
upscale_warnings.append(
f"With {num_frames_value} frames, upscaling will also be very time-consuming."
)
if upscale_warnings:
analysis_report_lines.append("\n--- UPSCALE ADVISORY ---")
for warning_msg in upscale_warnings:
analysis_report_lines.append(f"β οΈ {warning_msg}")
# analysis_report_lines.append("------------------------") # Optional closing separator
return "\n".join(analysis_report_lines)
def _tb_has_audio_stream(self, video_path_to_check):
if not self.has_ffprobe: # Critical check
self.message_manager.add_warning(
"FFprobe not available. Cannot reliably determine if video has audio. "
"Assuming no audio for operations requiring this check. "
"Install FFmpeg with ffprobe for full audio support."
)
return False
try:
resolved_path = str(Path(video_path_to_check).resolve())
ffprobe_cmd = [
self.ffprobe_exe, "-v", "error", "-select_streams", "a:0",
"-show_entries", "stream=codec_type", "-of", "csv=p=0", resolved_path
]
# check=False because a non-zero return often means no audio stream, which is a valid outcome here.
audio_check_result = subprocess.run(ffprobe_cmd, capture_output=True, text=True, check=False, errors='ignore')
if audio_check_result.returncode == 0 and "audio" in audio_check_result.stdout.strip().lower():
return True
else:
# Optionally log if ffprobe ran but found no audio, or if it errored for other reasons
# if audio_check_result.returncode != 0 and audio_check_result.stderr:
# self.message_manager.add_message(f"FFprobe check for audio stream in {os.path.basename(video_path_to_check)} completed. Stderr: {audio_check_result.stderr.strip()}", "DEBUG")
return False
except FileNotFoundError:
self.message_manager.add_warning("FFprobe executable not found during audio stream check (should have been caught by self.has_ffprobe). Assuming no audio.")
return False # Should ideally not happen if self.has_ffprobe is true and self.ffprobe_exe is set
except Exception as e:
self.message_manager.add_warning(f"Error checking for audio stream in {os.path.basename(video_path_to_check)}: {e}. Assuming no audio.")
return False
def tb_process_frames(self, video_path, target_fps_mode, speed_factor, progress=gr.Progress()):
if video_path is None: self.message_manager.add_warning("No input video for frame processing."); return None
# Core video processing relies on imageio for reading/writing frames, RIFE for interpolation.
# FFmpeg is primarily for audio handling here.
final_output_path = None
try:
self.message_manager.add_message(
f"Starting frame processing for {os.path.basename(video_path)}: "
f"FPS Mode: {target_fps_mode}, Speed: {speed_factor}x"
)
progress(0, desc="Initializing...")
resolved_video_path = str(Path(video_path).resolve())
self.message_manager.add_message("Reading video frames...")
progress(0.05, desc="Reading video...")
reader = imageio.get_reader(resolved_video_path)
original_fps = reader.get_meta_data().get('fps', 30.0) # Default if not found
video_frames = [frame for frame in reader]
reader.close()
self.message_manager.add_message(f"Read {len(video_frames)} frames at {original_fps} FPS.")
processed_frames = video_frames
current_fps = original_fps # This will be the FPS for the *output video stream*
if speed_factor != 1.0:
self.message_manager.add_message(f"Adjusting speed by {speed_factor}x (frame sampling/duplication)...")
progress(0.2, desc="Adjusting speed...")
if speed_factor > 1.0:
indices = np.arange(0, len(video_frames), speed_factor).astype(int)
processed_frames = [video_frames[i] for i in indices if i < len(video_frames)]
else:
new_len = int(len(video_frames) / speed_factor)
indices = np.linspace(0, len(video_frames) - 1, new_len).astype(int)
processed_frames = [video_frames[i] for i in indices]
self.message_manager.add_message(f"Speed adjustment (sampling) resulted in {len(processed_frames)} frames.")
should_interpolate = (target_fps_mode == "2x RIFE Interpolation")
if should_interpolate and len(processed_frames) > 1:
self.message_manager.add_message("Attempting to load RIFE model for 2x interpolation...")
if not self.rife_handler._ensure_model_downloaded_and_loaded():
self.message_manager.add_error("RIFE model could not be loaded. Skipping interpolation.")
else:
self.message_manager.add_message("RIFE model loaded. Starting RIFE 2x interpolation...")
interpolated_video_frames = []
num_pairs = len(processed_frames) - 1
for i in progress.tqdm(range(num_pairs), desc="RIFE Interpolating (2x)"):
frame1_np, frame2_np = processed_frames[i], processed_frames[i+1]
interpolated_video_frames.append(frame1_np)
middle_frame_np = self.rife_handler.interpolate_between_frames(frame1_np, frame2_np)
if middle_frame_np is not None: interpolated_video_frames.append(middle_frame_np)
else: interpolated_video_frames.append(frame1_np) # Duplicate on failure
interpolated_video_frames.append(processed_frames[-1])
processed_frames = interpolated_video_frames
# The video stream FPS itself doesn't change due to RIFE; it just has more frames.
# If RIFE is used, the perceived playback smoothness increases as if FPS doubled.
# The container FPS (current_fps) should reflect the intended playback rate of these frames.
# If original FPS was 30, and we RIFE, we now have 2x frames intended to still play over
# the same original duration segment, effectively meaning playback at 2*original_fps.
current_fps = original_fps * 2
self.message_manager.add_message(f"RIFE 2x interpolation resulted in {len(processed_frames)} frames. Effective FPS: {current_fps:.2f}")
elif should_interpolate and len(processed_frames) <= 1:
self.message_manager.add_warning("Not enough frames for RIFE interpolation. Skipping.")
op_suffix_parts = []
if speed_factor != 1.0: op_suffix_parts.append(f"speed{speed_factor:.2f}x".replace('.',','))
if should_interpolate and self.rife_handler.rife_model is not None: op_suffix_parts.append("RIFE2x")
if not op_suffix_parts: op_suffix_parts.append("processed")
op_suffix = "_".join(op_suffix_parts)
temp_video_suffix = f"{op_suffix}_temp_video"
video_stream_output_path = self._tb_generate_output_path(
resolved_video_path, suffix=temp_video_suffix, target_dir=self.toolbox_video_output_dir
)
final_muxed_output_path = video_stream_output_path.replace("_temp_video", "")
self.message_manager.add_message(f"Saving video stream to {video_stream_output_path} at {current_fps:.2f} FPS...")
progress(0.85, desc="Saving video stream...")
imageio.mimwrite(video_stream_output_path, processed_frames, fps=current_fps, quality=VIDEO_QUALITY, macro_block_size=None)
final_output_path = final_muxed_output_path
can_process_audio = self.has_ffmpeg
original_video_has_audio = self._tb_has_audio_stream(resolved_video_path) if can_process_audio else False
if can_process_audio and original_video_has_audio:
self.message_manager.add_message("Original video has audio. Processing audio with FFmpeg...")
progress(0.9, desc="Processing audio...")
ffmpeg_mux_cmd = [self.ffmpeg_exe, "-y", "-loglevel", "error", "-i", video_stream_output_path]
audio_filters = []
if speed_factor != 1.0: # Only apply atempo if speed actually changed
# Complex atempo for large speed changes (FFmpeg's atempo is 0.5-100.0)
# This simplified version handles common cases. For extreme speed_factor, might need more atempo stages.
if 0.5 <= speed_factor <= 100.0:
audio_filters.append(f"atempo={speed_factor:.4f}")
elif speed_factor < 0.5: # Needs multiple 0.5 steps
num_half_steps = int(np.ceil(np.log(speed_factor) / np.log(0.5)))
for _ in range(num_half_steps): audio_filters.append("atempo=0.5")
final_factor = speed_factor / (0.5**num_half_steps)
if abs(final_factor - 1.0) > 1e-4 and 0.5 <= final_factor <= 100.0: # Add final adjustment if needed
audio_filters.append(f"atempo={final_factor:.4f}")
elif speed_factor > 100.0: # Needs multiple 2.0 (or higher, like 100.0) steps
num_double_steps = int(np.ceil(np.log(speed_factor / 100.0) / np.log(2.0))) # Example for steps of 2 after 100
audio_filters.append("atempo=100.0") # Max one step
remaining_factor = speed_factor / 100.0
if abs(remaining_factor - 1.0) > 1e-4 and 0.5 <= remaining_factor <= 100.0:
audio_filters.append(f"atempo={remaining_factor:.4f}")
self.message_manager.add_message(f"Applying audio speed adjustment with atempo: {','.join(audio_filters) if audio_filters else 'None (speed_factor out of simple atempo range or 1.0)'}")
ffmpeg_mux_cmd.extend(["-i", resolved_video_path]) # Input for audio
ffmpeg_mux_cmd.extend(["-c:v", "copy"])
if audio_filters:
ffmpeg_mux_cmd.extend(["-filter:a", ",".join(audio_filters)])
# Always re-encode audio to AAC for MP4 compatibility, even if no speed change,
# as original audio might not be AAC.
ffmpeg_mux_cmd.extend(["-c:a", "aac", "-b:a", "192k"])
ffmpeg_mux_cmd.extend(["-map", "0:v:0", "-map", "1:a:0?", "-shortest", final_muxed_output_path])
try:
subprocess.run(ffmpeg_mux_cmd, check=True, capture_output=True, text=True)
self.message_manager.add_success(f"Video saved with processed audio: {final_muxed_output_path}")
except subprocess.CalledProcessError as e_mux:
self._tb_log_ffmpeg_error(e_mux, "audio processing/muxing")
self.message_manager.add_message("Saving video without audio as fallback.")
if os.path.exists(final_muxed_output_path): os.remove(final_muxed_output_path)
os.rename(video_stream_output_path, final_muxed_output_path)
except FileNotFoundError: # Should not happen if self.has_ffmpeg is true
self.message_manager.add_error(f"FFmpeg not found during muxing. This is unexpected if has_ffmpeg was True.")
if os.path.exists(final_muxed_output_path): os.remove(final_muxed_output_path)
os.rename(video_stream_output_path, final_muxed_output_path)
else:
if original_video_has_audio and not can_process_audio:
self.message_manager.add_warning("Original video has audio, but FFmpeg is not available to process it. Output will be silent. Install FFmpeg for audio support.")
elif not original_video_has_audio:
self.message_manager.add_message("No audio in original or audio processing skipped (e.g. FFprobe missing for detection). Saving video-only.")
if os.path.exists(final_muxed_output_path) and final_muxed_output_path != video_stream_output_path :
os.remove(final_muxed_output_path)
os.rename(video_stream_output_path, final_muxed_output_path)
if os.path.exists(video_stream_output_path) and video_stream_output_path != final_muxed_output_path:
try: os.remove(video_stream_output_path)
except Exception as e_clean: self.message_manager.add_warning(f"Could not remove temp video file {video_stream_output_path}: {e_clean}")
progress(1.0, desc="Complete.")
self.message_manager.add_success(f"Frame processing complete: {final_output_path}")
return final_output_path
except Exception as e:
self.message_manager.add_error(f"Error during frame processing: {e}")
import traceback; self.message_manager.add_error(traceback.format_exc())
progress(1.0, desc="Error.")
return None
finally:
if self.rife_handler: self.rife_handler.unload_model()
devicetorch.empty_cache(torch); gc.collect()
def tb_create_loop(self, video_path, loop_type, num_loops, progress=gr.Progress()):
if video_path is None: self.message_manager.add_warning("No input video for loop creation."); return None
if not self.has_ffmpeg: # FFmpeg is essential for this function's stream_loop and complex filter
self.message_manager.add_error("FFmpeg is required for creating video loops. This operation cannot proceed.")
return video_path # Return original video path
if loop_type == "none": self.message_manager.add_message("Loop type 'none'. No action."); return video_path
progress(0, desc="Initializing loop creation...")
resolved_video_path = str(Path(video_path).resolve())
output_path = self._tb_generate_output_path(
resolved_video_path,
suffix=f"looped_{loop_type}_{num_loops}x",
target_dir=self.toolbox_video_output_dir
)
self.message_manager.add_message(f"Creating {loop_type} ({num_loops}x) for {os.path.basename(resolved_video_path)}...")
ping_pong_unit_path = None
original_video_has_audio = self._tb_has_audio_stream(resolved_video_path) # Check once
try:
progress(0.2, desc=f"Preparing {loop_type} loop...")
if loop_type == "ping-pong":
ping_pong_unit_path = self._tb_generate_output_path(
resolved_video_path,
suffix="pingpong_unit_temp",
target_dir=self.toolbox_video_output_dir
)
# Create video-only ping-pong unit first
ffmpeg_pp_unit_cmd = [
self.ffmpeg_exe, "-y", "-loglevel", "error",
"-i", resolved_video_path,
"-vf", "split[main][tmp];[tmp]reverse[rev];[main][rev]concat=n=2:v=1:a=0", # Video only
"-an", ping_pong_unit_path
]
subprocess.run(ffmpeg_pp_unit_cmd, check=True, capture_output=True, text=True)
self.message_manager.add_message(f"Created ping-pong unit (video-only): {ping_pong_unit_path}")
ffmpeg_cmd = [
self.ffmpeg_exe, "-y", "-loglevel", "error",
"-stream_loop", str(num_loops - 1), # Loop the video unit N-1 times (total N plays)
"-i", ping_pong_unit_path,
"-c:v", "copy"
]
if original_video_has_audio:
self.message_manager.add_message("Original video has audio. Will loop audio for ping-pong.")
# Audio duration needs to match 2T * num_loops. FFmpeg aloop is complex.
# A simpler approach for ping-pong audio might be to create a 2T audio segment (original + reversed original)
# then loop that audio segment N times and mux with the N-times looped ping-pong video.
# Current FFmpeg command for aloop:
audio_loop_count_for_ffmpeg = (num_loops * 2) -1 # Total plays of original audio needed, minus one for initial play by -i
ffmpeg_cmd.extend([
"-i", resolved_video_path, # Original video for audio source
"-filter_complex", f"[1:a]areverse[areva];[1:a][areva]concat=n=2:v=0:a=1[ppa];[ppa]aloop=loop={num_loops-1}:size=2147483647[a_looped]",
"-map", "0:v:0", "-map", "[a_looped]",
"-c:a", "aac", "-b:a", "192k", "-shortest"
])
else:
self.message_manager.add_message("No audio in original or detection issue. Creating video-only ping-pong loop.")
ffmpeg_cmd.extend(["-an"])
ffmpeg_cmd.append(output_path)
else: # Regular 'loop'
# Hi Gemini/Claude!
# The 'num_loops' parameter comes from the UI slider labeled "Number of Loops/Repeats".
# The user's intent is that this slider value (num_loops) represents the number of
# *additional* times the video should be repeated after its initial playthrough.
# For example, if num_loops = 1 (from slider), the video should play twice (original + 1 repeat).
#
# FFmpeg's -stream_loop option takes a value (let's call it X_ffmpeg),
# meaning the input is looped X_ffmpeg times *in addition* to the first play.
# So, X_ffmpeg should be equal to the slider value 'num_loops'.
ffmpeg_stream_loop_value = num_loops
# Ensure ffmpeg_stream_loop_value is non-negative.
# Given the UI slider minimum is typically 1, num_loops should always be >= 1.
# This check is for robustness if the input num_loops could ever be less than 0
# (e.g., if UI constraints change or input comes from elsewhere).
if ffmpeg_stream_loop_value < 0:
ffmpeg_stream_loop_value = 0 # Should ideally not be hit if slider min is 1.
# Total plays will be the original play + ffmpeg_stream_loop_value additional plays.
total_plays = ffmpeg_stream_loop_value + 1
self.message_manager.add_message(
f"Regular loop: original video + {ffmpeg_stream_loop_value} additional repeat(s). Total {total_plays} plays."
)
ffmpeg_cmd = [
self.ffmpeg_exe, "-y", "-loglevel", "error",
"-stream_loop", str(ffmpeg_stream_loop_value), # This now uses num_loops directly
"-i", resolved_video_path,
"-c:v", "copy"
]
if original_video_has_audio:
self.message_manager.add_message("Original video has audio. Re-encoding to AAC for looped MP4 (if not already AAC).")
ffmpeg_cmd.extend(["-c:a", "aac", "-b:a", "192k", "-map", "0:v:0", "-map", "0:a:0?"])
else:
self.message_manager.add_message("No audio in original or detection issue. Looped video will be silent.")
ffmpeg_cmd.extend(["-an", "-map", "0:v:0"])
ffmpeg_cmd.append(output_path)
self.message_manager.add_message(f"Processing video {loop_type} with FFmpeg...")
progress(0.5, desc=f"Running FFmpeg for {loop_type}...")
subprocess.run(ffmpeg_cmd, check=True, capture_output=True, text=True, errors='ignore')
progress(1.0, desc=f"{loop_type.capitalize()} loop created successfully.")
self.message_manager.add_success(f"Loop creation complete: {output_path}")
return output_path
except subprocess.CalledProcessError as e_loop:
self._tb_log_ffmpeg_error(e_loop, f"{loop_type} creation")
progress(1.0, desc=f"Error creating {loop_type}.")
return None
except Exception as e:
self.message_manager.add_error(f"Error creating loop: {e}")
import traceback; self.message_manager.add_error(traceback.format_exc())
progress(1.0, desc="Error creating loop.")
return None
finally:
if ping_pong_unit_path and os.path.exists(ping_pong_unit_path):
try: os.remove(ping_pong_unit_path)
except Exception as e_clean_pp: self.message_manager.add_warning(f"Could not remove temp ping-pong unit: {e_clean_pp}")
gc.collect()
def _tb_get_video_dimensions(self, video_path):
video_width, video_height = 0, 0
# Prefer ffprobe if available for dimensions
if self.has_ffprobe:
try:
probe_cmd = [self.ffprobe_exe, "-v", "error", "-select_streams", "v:0",
"-show_entries", "stream=width,height", "-of", "csv=s=x:p=0", video_path]
result = subprocess.run(probe_cmd, capture_output=True, text=True, check=True, errors='ignore')
w_str, h_str = result.stdout.strip().split('x')
video_width, video_height = int(w_str), int(h_str)
if video_width > 0 and video_height > 0: return video_width, video_height
except Exception as e_probe_dim:
self.message_manager.add_warning(f"ffprobe failed to get dimensions ({e_probe_dim}), trying imageio.")
# Fallback to imageio
reader = None
try:
reader = imageio.get_reader(video_path)
meta = reader.get_meta_data()
size_imgio = meta.get('size')
if size_imgio and isinstance(size_imgio, tuple) and len(size_imgio) == 2:
video_width, video_height = int(size_imgio[0]), int(size_imgio[1])
except Exception as e_meta:
self.message_manager.add_warning(f"Error getting video dimensions for vignette (imageio): {e_meta}. Defaulting aspect to 1/1.")
finally:
if reader: reader.close()
return video_width, video_height # Might be 0,0 if all failed
def _tb_create_vignette_filter(self, strength_percent, width, height):
min_angle_rad = math.pi / 3.5; max_angle_rad = math.pi / 2
normalized_strength = strength_percent / 100.0
angle_rad = min_angle_rad + normalized_strength * (max_angle_rad - min_angle_rad)
vignette_aspect_ratio_val = "1/1"
if width > 0 and height > 0: vignette_aspect_ratio_val = f"{width/height:.4f}"
return f"vignette=angle={angle_rad:.4f}:mode=forward:eval=init:aspect={vignette_aspect_ratio_val}"
def tb_apply_filters(self, video_path, brightness, contrast, saturation, temperature,
sharpen, blur, denoise, vignette, s_curve_contrast, film_grain_strength,
progress=gr.Progress()):
if video_path is None: self.message_manager.add_warning("No input video for filters."); return None
if not self.has_ffmpeg: # FFmpeg is essential for this function
self.message_manager.add_error("FFmpeg is required for applying video filters. This operation cannot proceed.")
return video_path
progress(0, desc="Initializing filter application...")
resolved_video_path = str(Path(video_path).resolve())
output_path = self._tb_generate_output_path(resolved_video_path, "filtered", self.toolbox_video_output_dir)
self.message_manager.add_message(f"π¨ Applying filters to {os.path.basename(resolved_video_path)}...")
video_width, video_height = 0,0
if vignette > 0: # Only get dimensions if vignette is used
video_width, video_height = self._tb_get_video_dimensions(resolved_video_path)
if video_width > 0 and video_height > 0: self.message_manager.add_message(f"Video dimensions for vignette: {video_width}x{video_height}", "DEBUG")
filters, applied_filter_descriptions = [], []
# Filter definitions (unchanged, assuming they are correct)
if denoise > 0: filters.append(f"hqdn3d={denoise*0.8:.1f}:{denoise*0.6:.1f}:{denoise*0.7:.1f}:{denoise*0.5:.1f}"); applied_filter_descriptions.append(f"Denoise (hqdn3d)")
if temperature != 0: mid_shift = (temperature/100.0)*0.3; filters.append(f"colorbalance=rm={mid_shift:.2f}:bm={-mid_shift:.2f}"); applied_filter_descriptions.append(f"Color Temp")
eq_parts = []; desc_eq = []
if brightness != 0: eq_parts.append(f"brightness={brightness/100.0:.2f}"); desc_eq.append(f"Brightness")
if contrast != 1: eq_parts.append(f"contrast={contrast:.2f}"); desc_eq.append(f"Contrast (Linear)")
if saturation != 1: eq_parts.append(f"saturation={saturation:.2f}"); desc_eq.append(f"Saturation")
if eq_parts: filters.append(f"eq={':'.join(eq_parts)}"); applied_filter_descriptions.append(" & ".join(desc_eq))
if s_curve_contrast > 0: s = s_curve_contrast/100.0; y1 = 0.25-s*(0.25-0.10); y2 = 0.75+s*(0.90-0.75); filters.append(f"curves=all='0/0 0.25/{y1:.2f} 0.75/{y2:.2f} 1/1'"); applied_filter_descriptions.append(f"S-Curve Contrast")
if blur > 0: filters.append(f"gblur=sigma={blur*0.4:.1f}"); applied_filter_descriptions.append(f"Blur")
if sharpen > 0: filters.append(f"unsharp=luma_msize_x=5:luma_msize_y=5:luma_amount={sharpen*0.3:.2f}"); applied_filter_descriptions.append(f"Sharpen")
if film_grain_strength > 0: filters.append(f"noise=alls={film_grain_strength*0.5:.1f}:allf=t+u"); applied_filter_descriptions.append(f"Film Grain")
if vignette > 0: filters.append(self._tb_create_vignette_filter(vignette, video_width, video_height)); applied_filter_descriptions.append(f"Vignette")
if not filters: self.message_manager.add_message("βΉοΈ No filters selected."); progress(1.0); return video_path
if applied_filter_descriptions: self.message_manager.add_message("π§ Applying FFmpeg filters: " + ", ".join(applied_filter_descriptions))
progress(0.2, desc="Preparing filter command...")
original_video_has_audio = self._tb_has_audio_stream(resolved_video_path)
try:
ffmpeg_cmd = [
self.ffmpeg_exe, "-y", "-loglevel", "error", "-i", resolved_video_path,
"-vf", ",".join(filters),
"-c:v", "libx264", "-preset", "medium", "-crf", "20",
"-pix_fmt", "yuv420p",
"-map", "0:v:0"
]
if original_video_has_audio:
self.message_manager.add_message("Original video has audio. Re-encoding to AAC for filtered video.", "INFO")
ffmpeg_cmd.extend(["-c:a", "aac", "-b:a", "192k", "-map", "0:a:0?"])
else:
self.message_manager.add_message("No audio in original or detection issue. Filtered video will be silent.", "INFO")
ffmpeg_cmd.extend(["-an"])
ffmpeg_cmd.append(output_path)
self.message_manager.add_message("π Processing filters with FFmpeg...")
progress(0.5, desc="Running FFmpeg for filters...")
subprocess.run(ffmpeg_cmd, check=True, capture_output=True, text=True, errors='ignore')
progress(1.0, desc="Filters applied successfully.")
self.message_manager.add_success(f"β
Filters applied! Output: {output_path}")
return output_path
except subprocess.CalledProcessError as e_filters:
self._tb_log_ffmpeg_error(e_filters, "filter application")
progress(1.0, desc="Error applying filters."); return None
except Exception as e:
self.message_manager.add_error(f"β An unexpected error occurred: {e}")
import traceback; self.message_manager.add_error(traceback.format_exc())
progress(1.0, desc="Error applying filters."); return None
finally: gc.collect()
def tb_upscale_video(self, video_path, model_key: str, output_scale_factor_ui: float,
tile_size: int, enhance_face: bool,
denoise_strength_ui: float | None, # New parameter
progress=gr.Progress()):
if video_path is None: self.message_manager.add_warning("No input video for upscaling."); return None
final_output_path = None; reader = None
try:
if model_key not in self.esrgan_upscaler.supported_models:
self.message_manager.add_error(f"Upscale model key '{model_key}' not found in supported models.")
return None
model_native_scale = self.esrgan_upscaler.supported_models[model_key].get('scale', 0)
tile_size_str_for_log = str(tile_size) if tile_size > 0 else "Auto"
face_enhance_str_for_log = "+FaceEnhance" if enhance_face else ""
denoise_str_for_log = ""
if model_key == "RealESR-general-x4v3" and denoise_strength_ui is not None:
denoise_str_for_log = f", DNI: {denoise_strength_ui:.2f}"
self.message_manager.add_message(
f"Preparing to load ESRGAN model '{model_key}' for {output_scale_factor_ui:.2f}x target upscale "
f"(Native: {model_native_scale}x, Tile: {tile_size_str_for_log}{face_enhance_str_for_log}{denoise_str_for_log})."
)
progress(0.05, desc=f"Loading ESRGAN model '{model_key}' (Tile: {tile_size_str_for_log}{denoise_str_for_log})...")
# Pass denoise_strength_ui to load_model
upsampler_instance = self.esrgan_upscaler.load_model(
model_key=model_key,
tile_size=tile_size,
denoise_strength=denoise_strength_ui if model_key == "RealESR-general-x4v3" else None
)
if not upsampler_instance:
self.message_manager.add_error(f"Could not load ESRGAN model '{model_key}'. Aborting."); return None
if enhance_face: # Face enhancer loading logic
if not self.esrgan_upscaler._load_face_enhancer(bg_upsampler=upsampler_instance):
self.message_manager.add_warning("Failed to load GFPGAN for face enhancement. Proceeding without it.")
enhance_face = False
face_enhance_str_for_log = "" # Update log string if face enhance fails
self.message_manager.add_message(
f"ESRGAN model '{model_key}' (Native: {model_native_scale}x, Tile: {tile_size_str_for_log}{denoise_str_for_log}) "
f"{'and GFPGAN ' if enhance_face else ''}loaded for target {output_scale_factor_ui:.2f}x output."
)
progress(0.1, desc=f"Initializing {output_scale_factor_ui:.2f}x upscaling{face_enhance_str_for_log}{denoise_str_for_log} process...")
resolved_video_path = str(Path(video_path).resolve())
upscaled_frames = []
progress(0.12, desc="Reading video info...")
reader = imageio.get_reader(resolved_video_path)
meta_data = reader.get_meta_data(); original_fps = meta_data.get('fps', 30.0)
n_frames = meta_data.get('nframes')
if n_frames is None or n_frames == float('inf'):
try: n_frames = reader.count_frames()
except: n_frames = None
if n_frames == float('inf'): n_frames = None
n_frames_display = str(int(n_frames)) if n_frames is not None else "Unknown"
self.message_manager.add_message(f"Original FPS: {original_fps:.2f}. Total frames: {n_frames_display}.")
progress_desc = (
f"Upscaling Frames to {output_scale_factor_ui:.2f}x (Model: {model_key}{face_enhance_str_for_log}{denoise_str_for_log}, "
f"Native: {model_native_scale}x, Tile: {tile_size_str_for_log})"
)
frame_iterator = enumerate(reader)
if n_frames is not None: frame_iterator = progress.tqdm(enumerate(reader), total=int(n_frames), desc=progress_desc)
else: self.message_manager.add_message(f"Total frames unknown, progress per batch ({progress_desc}).")
for i, frame_np in frame_iterator:
if n_frames is None and i % 10 == 0:
current_progress_val = 0.15 + ( (i/(i+500.0)) * 0.65 )
progress(current_progress_val , desc=f"Upscaling frame {i+1} to {output_scale_factor_ui:.2f}x (Tile: {tile_size_str_for_log})...")
upscaled_frame_np = self.esrgan_upscaler.upscale_frame( # DNI is handled by loaded model
frame_np_array=frame_np,
model_key=model_key,
target_outscale_factor=float(output_scale_factor_ui),
enhance_face=enhance_face
)
if upscaled_frame_np is not None: upscaled_frames.append(upscaled_frame_np)
else: # Error handling for frame upscale
self.message_manager.add_error(f"Failed to upscale frame {i+1}. Skipping.")
if "out of memory" in self.message_manager.get_recent_errors_as_str(count=1).lower():
self.message_manager.add_error("CUDA OOM likely. Aborting video upscale."); return None
if (i+1) % 20 == 0: devicetorch.empty_cache(torch); gc.collect()
if reader: reader.close(); reader = None
if not upscaled_frames: self.message_manager.add_error("No frames upscaled."); return None
self.message_manager.add_message(f"Successfully upscaled {len(upscaled_frames)} frames to {output_scale_factor_ui:.2f}x.")
progress(0.80, desc="Saving upscaled video stream...")
temp_video_suffix_base = (
f"upscaled_{model_key}_{output_scale_factor_ui:.2f}x_tile{tile_size_str_for_log}"
f"{face_enhance_str_for_log.replace('+','_')}"
)
if model_key == "RealESR-general-x4v3" and denoise_strength_ui is not None:
temp_video_suffix_base += f"_dni{denoise_strength_ui:.2f}"
temp_video_suffix = temp_video_suffix_base.replace(".","p") + "_temp_video"
video_stream_output_path = self._tb_generate_output_path(resolved_video_path, temp_video_suffix, self.toolbox_video_output_dir)
final_muxed_output_path = video_stream_output_path.replace("_temp_video", "")
imageio.mimwrite(video_stream_output_path, upscaled_frames, fps=original_fps, quality=VIDEO_QUALITY, macro_block_size=None)
del upscaled_frames; devicetorch.empty_cache(torch); gc.collect()
final_output_path = final_muxed_output_path
can_process_audio = self.has_ffmpeg
original_video_has_audio = self._tb_has_audio_stream(resolved_video_path) if can_process_audio else False
if can_process_audio and original_video_has_audio:
progress(0.90, desc="Muxing audio...")
self.message_manager.add_message("Original video has audio. Muxing audio with FFmpeg...")
ffmpeg_mux_cmd = [
self.ffmpeg_exe, "-y", "-loglevel", "error",
"-i", video_stream_output_path, "-i", resolved_video_path,
"-c:v", "copy", "-c:a", "aac", "-b:a", "192k",
"-map", "0:v:0", "-map", "1:a:0?", "-shortest", final_muxed_output_path
]
try:
subprocess.run(ffmpeg_mux_cmd, check=True, capture_output=True, text=True)
self.message_manager.add_success(f"Upscaled video saved with audio: {final_muxed_output_path}")
except subprocess.CalledProcessError as e_mux:
self._tb_log_ffmpeg_error(e_mux, "audio muxing for upscaled video")
if os.path.exists(final_muxed_output_path): os.remove(final_muxed_output_path)
os.rename(video_stream_output_path, final_muxed_output_path)
except FileNotFoundError:
self.message_manager.add_error(f"FFmpeg not found during muxing. Unexpected.")
if os.path.exists(final_muxed_output_path): os.remove(final_muxed_output_path)
os.rename(video_stream_output_path, final_muxed_output_path)
else:
if original_video_has_audio and not can_process_audio:
self.message_manager.add_warning("Original video has audio, but FFmpeg is not available to process it. Upscaled output will be silent.")
elif not original_video_has_audio :
self.message_manager.add_message("No audio in original or detection issue. Saving upscaled video-only.")
if os.path.exists(final_muxed_output_path) and final_muxed_output_path != video_stream_output_path:
os.remove(final_muxed_output_path)
os.rename(video_stream_output_path, final_muxed_output_path)
if os.path.exists(video_stream_output_path) and video_stream_output_path != final_muxed_output_path:
try: os.remove(video_stream_output_path)
except Exception as e_clean: self.message_manager.add_warning(f"Could not remove temp upscaled video: {e_clean}")
progress(1.0, desc="Upscaling complete.")
self.message_manager.add_success(f"Video upscaling to {output_scale_factor_ui:.2f}x complete: {final_output_path}")
return final_output_path
except Exception as e:
self.message_manager.add_error(f"Error during video upscaling: {e}")
import traceback; self.message_manager.add_error(traceback.format_exc())
progress(1.0, desc="Error during upscaling."); return None
finally:
if reader:
try:
if hasattr(reader, 'closed') and not reader.closed: reader.close()
except: pass
if model_key and self.esrgan_upscaler: # Ensure model is unloaded
self.esrgan_upscaler.unload_model(model_key)
if enhance_face and self.esrgan_upscaler and self.esrgan_upscaler.face_enhancer:
self.esrgan_upscaler._unload_face_enhancer()
devicetorch.empty_cache(torch); gc.collect()
def tb_open_output_folder(self):
folder_path = os.path.abspath(self.postprocessed_output_root_dir)
try:
os.makedirs(folder_path, exist_ok=True)
if sys.platform == 'win32': subprocess.run(['explorer', folder_path])
elif sys.platform == 'darwin': subprocess.run(['open', folder_path])
else: subprocess.run(['xdg-open', folder_path])
self.message_manager.add_success(f"Opened postprocessed output folder: {folder_path}")
except Exception as e:
self.message_manager.add_error(f"Error opening folder {folder_path}: {e}")
def tb_clear_temporary_files(self):
temp_dir_path_str = str(self._base_temp_output_dir)
self.message_manager.add_message(f"Attempting to clear temporary files in: {temp_dir_path_str}", "INFO")
cleared_successfully = False
if os.path.exists(temp_dir_path_str):
try:
# Count items for logging
items = os.listdir(temp_dir_path_str)
file_count = sum(1 for item in items if os.path.isfile(os.path.join(temp_dir_path_str, item)))
dir_count = sum(1 for item in items if os.path.isdir(os.path.join(temp_dir_path_str, item)))
shutil.rmtree(temp_dir_path_str)
self.message_manager.add_success(
f"Successfully removed temporary directory and its contents ({file_count} files, {dir_count} subdirectories)."
)
cleared_successfully = True
except Exception as e:
self.message_manager.add_error(f"Error deleting temporary directory '{temp_dir_path_str}': {e}")
self.message_manager.add_error(traceback.format_exc())
else:
self.message_manager.add_message("Temporary directory does not exist. Nothing to clear.", "INFO")
cleared_successfully = True
try:
os.makedirs(temp_dir_path_str, exist_ok=True) # Always recreate
if cleared_successfully: self.message_manager.add_message(f"Recreated temporary directory: {temp_dir_path_str}", "INFO")
except Exception as e_recreate:
self.message_manager.add_error(f"CRITICAL: Failed to recreate temporary directory '{temp_dir_path_str}': {e_recreate}. Processing may fail.")
self.message_manager.add_error(traceback.format_exc())
cleared_successfully = False
return cleared_successfully |