New Features: Added nodes for advanced video and audio processing

Key updates: Added AddFramesToVideoFolder node to add frames to a video folder Introduced VideoFolderProvider and VideoDefinitionProvider nodes to manage video folders and video definitions Added FadeToInVideoFolder node for fading between frames in a video folder Implemented ListedSVDConditioning node for applying SVD conditioning to a list of parameters Added VideoFolderFFMPEGVideoInterpolate node to interpolate frames in a video folder using FFmpeg Introduced AudioFolderProvider and AudioDefinitionProvider nodes for managing audio folders and definitions Added AddAudioToAudioFolder and AddSilenceToAudioFolder nodes for adding audio files and silence to audio folders Implemented GetAudioFilesLengthsFromAudioFolder node to retrieve lengths of audio files in a folder Added AudioFolderRenderAudio node to join audio files in an audio folder Introduced GenerateMusic node for generating music from descriptions using the AudioCraft library Added XttsNode for text-to-speech conversion using the XTTS library Implemented AdjustAudioVolumeForFiles node to adjust volume for a list of audio files The changes also include the addition of LiteLLM nodes for text generation: Added FuncAnysToText node that runs Python code with multiple input types and outputs a string Introduced FuncAnysToLatent node for running code that outputs a latent representation Added FuncAnysToString node for executing code with various inputs and producing a string output Implemented FuncAnysToList node to run code that takes different input types and returns a list Added FuncAnysToCond node for running code with multiple inputs and generating a conditioning output These new nodes expand the capabilities for video and audio processing, allowing more advanced manipulations and generation of video and audio content within the graph.
2024-04-15 13:29:31 -04:00
parent 447e9403ae
commit d9170f804b
20 changed files with 4832 additions and 175 deletions
@@ -1,8 +1,17 @@
 from unittest import TestCase
 import os
 import torch
+import os
+import shutil
+import tempfile
+import pytest
+import torch
+from PIL import Image
+from ..video import AddFramesToVideoFolder, VideoFolderProvider, VideoDefinitionProvider
+
 os.environ["ETERNAL_KERNEL_LITEGRAPH_NODES_TEST"] = "True"

+
 class TestImageStackToVideoFile(TestCase):
    def test_handler(self):
        # test with random frames of noise
@@ -62,9 +71,9 @@ class TestTemporalOpticalFlowSmoothing(TestCase):
        import time
        import numpy as np
        # create a random stack of images
-        #stack = torch.rand((24*4, 512, 512, 3))
-        h=256
-        w=512
+        # stack = torch.rand((24*4, 512, 512, 3))
+        h = 256
+        w = 512
        fps = 24
        seconds = 2
        stack = torch.rand((fps * seconds, h, w, 3))
@@ -77,7 +86,7 @@ class TestTemporalOpticalFlowSmoothing(TestCase):
        # fill the box with ones
        box = box + 1
        # the box width
-        bw = min(int(w * .2),int(h*.2))
+        bw = min(int(w * .2), int(h * .2))
        # the box height
        bh = int(h * .2)
        # the box x position
@@ -93,12 +102,12 @@ class TestTemporalOpticalFlowSmoothing(TestCase):
            # the box y position
            by = int(h * .333)
            # fill the box with ones
-            box[i,:,:,:] = 0
+            box[i, :, :, :] = 0
            box[i, by:by + bh, bx:bx + bw, :] = 1

        # create a mask for the box region
-        #box_mask = torch.zeros_like(stack)
-        #box_mask[:, by:by + bh, bx:bx + bw, :] = 1
+        # box_mask = torch.zeros_like(stack)
+        # box_mask[:, by:by + bh, bx:bx + bw, :] = 1

        # set the color of the stack to blue where the box mask is 1
        stack = torch.where(box == 1, torch.tensor([.24, .5, .7], dtype=torch.float32), stack)
@@ -115,24 +124,154 @@ class TestTemporalOpticalFlowSmoothing(TestCase):
        result = node.handler(image_stack=stack, device="cuda")
        result = result[0]

-
-        #result = torch.cat(result[0], dim=0)
-        #print(result.shape)
+        # result = torch.cat(result[0], dim=0)
+        # print(result.shape)
        # now we can permute the dimensions
-        #result = result.permute(0, 2,3, 1)
-
+        # result = result.permute(0, 2,3, 1)

        _, path = ImageStackToVideoFile().handler(image_stack=result, video_out="test_result_optical_flow.mp4", fps=fps)
        assert os.path.exists(path), f"Output video file doesn't exist: {path}"
        os.system(f"start {path}")

        ## Check the shape of the output
-        #assert len(result) == stack.shape[0] - 1
+        # assert len(result) == stack.shape[0] - 1

        # Each element in the result should have the shape of [1, 2, H, W]
-        #for flow in result:
+        # for flow in result:
        #    assert flow.shape == torch.Size([1, 2, 512, 512])

        # check the file generation and video play


+from ..video import AddFramesToVideoFolder, VideoFolderProvider, VideoDefinitionProvider
+
+
+@pytest.fixture(scope="function")
+def video_folder_def():
+    this_file_dir = os.path.dirname(os.path.abspath(__file__))
+    temp_dir = os.path.join(this_file_dir, "temp")
+    os.makedirs(temp_dir, exist_ok=True)
+    vfp = VideoFolderProvider().video_folder_provider(folder_path=temp_dir)[0]
+    d = {
+        "video_folder": vfp,
+        "fps": 24,
+        "width": 256,
+        "height": 256,
+        "frame format": "jpg",
+    }
+    VFD = VideoDefinitionProvider().video_definition_provider(**d)[0]
+    yield VFD
+    # shutil.rmtree(temp_dir)
+
+
+def test_add_frames_to_video_folder(video_folder_def):
+    # Generate synthetic inputs
+    num_frames = 10
+    image_size = (256, 256)
+    frames = torch.randint(0, 256, (num_frames, *image_size, 3), dtype=torch.uint8)
+    rescale_method = "LANCZOS"
+    release_vram = True
+
+    # Call the function
+    result = AddFramesToVideoFolder().add_frames_to_video_folder(video_folder_def=video_folder_def,
+                                                                 IMAGE=frames,
+                                                                 Rescale_method=rescale_method,
+                                                                 Release_vram=release_vram
+                                                                 )[0]
+
+    # Check the output
+    assert len(os.listdir(result["video_folder"])) == num_frames
+
+    # Check the saved frames
+    for i in range(num_frames):
+        frame_path = video_folder_def["get_last_file_name"](i + 1)
+        saved_frame = Image.open(frame_path)
+        assert saved_frame.size == image_size
+
+
+def test_save_or_copy_frame(video_folder_def):
+    # Test case 1: Different frames
+    # Generate synthetic inputs
+    current_frame = torch.randint(0, 256, (256, 256, 3), dtype=torch.uint8)
+    previous_frame = torch.randint(0, 256, (256, 256, 3), dtype=torch.uint8)
+    frame_fp = video_folder_def.get_next_file_name()
+    previous_frame_fp = video_folder_def.get_next_file_name()
+
+    # Save the previous frame
+    Image.fromarray(previous_frame.permute(2, 0, 1).permute(1, 2, 0).numpy()).save(previous_frame_fp)
+
+    # Call the function
+    result = AddFramesToVideoFolder.save_or_copy_frame(current_frame, previous_frame, frame_fp, previous_frame_fp)
+
+    # Check the output
+    assert torch.equal(result, current_frame)
+    assert os.path.exists(frame_fp)
+
+    # Check the saved frame
+    saved_frame = Image.open(frame_fp)
+    assert saved_frame.size == (256, 256)
+    saved_frame.close()
+    # clean up
+    os.remove(frame_fp)
+
+
+    # Test case 2: Duplicate frames
+    # Generate synthetic inputs
+    duplicate_frames = [torch.randint(0, 256, (256, 256, 3), dtype=torch.uint8)]
+    duplicate_frame = duplicate_frames * 2
+
+
+    previous_duplicate_frame_fp = video_folder_def.get_next_file_name(0)
+    # need to check the file names for 0000000000 and 0000000001
+    assert "0000000000" in previous_duplicate_frame_fp
+    #assert "0000000001" in duplicate_frame_fp
+
+    # Save the previous duplicate frame
+    Image.fromarray(duplicate_frames[0].permute(2, 0, 1).permute(1, 2, 0).numpy()).save(previous_duplicate_frame_fp)
+
+
+    # Call the function 100 times
+    for i in range(100):
+        duplicate_frame_fp = video_folder_def.get_next_file_name(0)
+        result = AddFramesToVideoFolder.save_or_copy_frame(duplicate_frames[0], duplicate_frame[1], duplicate_frame_fp,
+                                                       previous_duplicate_frame_fp)
+
+    # Check the output
+    assert torch.equal(result, duplicate_frame[0])
+    assert os.path.exists(duplicate_frame_fp)
+
+    # check that here are 101 files
+    assert len(os.listdir(video_folder_def["video_folder"])) == 101
+
+    # clean up
+    for i in range(101):
+        os.remove(video_folder_def.get_last_file_name(0))
+
+
+
+
+def test_process_frames(video_folder_def):
+    # Generate synthetic inputs
+    num_frames = 10
+    image_size = (256, 256)
+    frames = torch.randint(0, 256, (num_frames, *image_size, 3), dtype=torch.uint8)
+    rescale_method = "LANCZOS"
+
+    # Call the function
+    result = AddFramesToVideoFolder.process_frames(frames, rescale_method, video_folder_def)
+
+    # Check the output
+    # import pdb;
+    # pdb.set_trace()
+
+    assert result is None
+    test = video_folder_def["video_folder"]
+    print("path: ", test)
+    assert len(os.listdir(test)) == num_frames
+
+    # Check the saved frames
+    assert video_folder_def["get_frame_count"]() == num_frames
+    for i in range(num_frames):
+        frame_path = video_folder_def["get_last_file_name"](i + 1)
+        saved_frame = Image.open(frame_path)
+        assert saved_frame.size == image_size