Update app.py

app.py

@@ -6,6 +6,7 @@ import cv2
 from PIL import Image
 from tqdm import tqdm
 import gradio as gr
+from concurrent.futures import ThreadPoolExecutor, as_completed
 
 # Constants
 PADDING_TOKEN = b"<PAD>"
@@ -39,8 +40,10 @@ def encode(file_path):
                 chunk_file.write(chunk)
     # Convert each chunk to an image
     chunk_files = [os.path.join('./parts', f"{i}.part") for i in range(num_chunks)]
-    for i, chunk_file in tqdm(enumerate(chunk_files), desc="Converting chunks to images", unit="chunk"):
-        convert_part_to_image(chunk_file, i)
+    with ThreadPoolExecutor() as executor:
+        futures = [executor.submit(convert_part_to_image, chunk_file, i) for i, chunk_file in enumerate(chunk_files)]
+        for future in tqdm(as_completed(futures), desc="Converting chunks to images", unit="chunk", total=len(chunk_files)):
+            future.result()
     # Create a video from the frames
     create_video_from_frames(FRAMES_DIR, VIDEO_PATH)
     return VIDEO_PATH
@@ -54,18 +57,24 @@ def decode(video_path, output_file_path):
     extracted_frame_files = sorted([os.path.join(extracted_frames_dir, f) for f in os.listdir(extracted_frames_dir) if f.endswith('.png')])
     # Convert frames back to chunk files
     chunk_files = []
-    for i, frame_file in tqdm(enumerate(extracted_frame_files), desc="Converting frames to chunks", unit="frame"):
-        binary_string = hd_rgb_image_to_binary(frame_file)
-        chunk_file_path = os.path.join('./parts', f"{i}.part")
-        with open(chunk_file_path, 'wb') as chunk_file:
-            chunk_file.write(int(binary_string, 2).to_bytes(len(binary_string) // 8, byteorder='big'))
-        chunk_files.append(chunk_file_path)
+    with ThreadPoolExecutor() as executor:
+        futures = [executor.submit(convert_frame_to_chunk, frame_file, i) for i, frame_file in enumerate(extracted_frame_files)]
+        for future in tqdm(as_completed(futures), desc="Converting frames to chunks", unit="frame", total=len(extracted_frame_files)):
+            chunk_file_path = future.result()
+            chunk_files.append(chunk_file_path)
     # Merge the chunks into the original file
     merge_chunks(chunk_files, output_file_path)
     # Verify file integrity
     original_hash = calculate_sha256(output_file_path)
     return f"SHA-256 hash of the decoded file: {original_hash}"
 
+def convert_frame_to_chunk(frame_file, i):
+    binary_string = hd_rgb_image_to_binary(frame_file)
+    chunk_file_path = os.path.join('./parts', f"{i}.part")
+    with open(chunk_file_path, 'wb') as chunk_file:
+        chunk_file.write(int(binary_string, 2).to_bytes(len(binary_string) // 8, byteorder='big'))
+    return chunk_file_path
+
 def binary_string_to_rgb_image(binary_string, width=1920, height=1080, output_path='output_rgb.png'):
     """Convert a binary string to an RGB image and save it."""
     if len(binary_string) != REQUIRED_LENGTH * 3:  # 3 bits per pixel for RGB