Save the prompt in the metadata

app.py

@@ -4,6 +4,11 @@
 from __future__ import annotations
 
 import os
+import imageio_ffmpeg
+import tempfile
+import shutil
+import subprocess
+from datetime import datetime
 import pathlib
 import random
 import sys
@@ -1074,6 +1079,52 @@ def _make_handler(mode_name: str, h: dict):
 
     return handler
 
+def set_mp4_comments_imageio_ffmpeg(input_file, comments):
+    try:
+        # Get the path to the bundled FFmpeg binary from imageio-ffmpeg
+        ffmpeg_path = imageio_ffmpeg.get_ffmpeg_exe()
+
+        # Check if input file exists
+        if not os.path.exists(input_file):
+            #print(f"Error: Input file {input_file} does not exist")
+            return False
+
+        # Create a temporary file path
+        temp_file = tempfile.NamedTemporaryFile(suffix='.mp4', delete=False).name
+
+        # FFmpeg command using the bundled binary
+        command = [
+            ffmpeg_path,                   # Use imageio-ffmpeg's FFmpeg
+            '-i', input_file,              # input file
+            '-metadata', f'comment={comments}',  # set comment metadata
+            '-c:v', 'copy',                # copy video stream without re-encoding
+            '-c:a', 'copy',                # copy audio stream without re-encoding
+            '-y',                          # overwrite output file if it exists
+            temp_file                      # temporary output file
+        ]
+
+        # Run the FFmpeg command
+        result = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
+
+        if result.returncode == 0:
+            # Replace the original file with the modified one
+            shutil.move(temp_file, input_file)
+            #print(f"Successfully added comments to {input_file}")
+            return True
+        else:
+            # Clean up temp file if FFmpeg fails
+            if os.path.exists(temp_file):
+                os.remove(temp_file)
+            #print(f"Error: FFmpeg failed with message:\n{result.stderr}")
+            return False
+
+    except Exception as e:
+        # Clean up temp file in case of other errors
+        if 'temp_file' in locals() and os.path.exists(temp_file):
+            os.remove(temp_file)
+        print(f"Error saving prompt to video metadata, ffmpeg may be required: "+str(e))
+        return False
+
 async def run_example(dummy_image):
     example_inputs = {
         "prompt": prompt_debug_value[0],
@@ -1127,6 +1178,8 @@ async def run_example(dummy_image):
             else:
                 count = int(count) - 1
 
+    set_mp4_comments_imageio_ffmpeg(final_output, f"Prompt: {example_inputs["prompt"]} | Negative Prompt: {example_inputs["negative_prompt"]}");
+    print("Video exported: " + final_output)
     return final_output[1]
 
 if __name__ == "__main__":