Update web-demos/hugging_face/app.py

web-demos/hugging_face/app.py

@@ -133,7 +133,6 @@ def get_frames_from_video(video_input, video_state):
                         gr.update(visible=status_ok), gr.update(visible=status_ok, choices=[], value=[]), \
                         gr.update(visible=True, value=operation_log), gr.update(visible=status_ok, value=operation_log)
 
-# get the select frame from gradio slider
 def select_template(image_selection_slider, video_state, interactive_state, mask_dropdown):
 
     # images = video_state[1]
@@ -365,50 +364,49 @@ def inpaint_video(video_state, *_args):
 
     return video_output, operation_log, operation_log
 
-
 # generate video after vos inference
-def generate_video_from_frames(frames, output_path, fps=30):
+def generate_video_from_frames(frames, output_path, fps=30, bitrate=None):
     """
-    Save high-quality video using ffmpeg and CRF.
+    Generates a video from a list of frames.
+    
+    Args:
+        frames (list of numpy arrays): The frames to include in the video.
+        output_path (str): The path to save the generated video.
+        fps (int, optional): The frame rate of the output video. Defaults to 30.
     """
-    import os
-    import cv2
+    # Приведение fps к обычному float (из np.float64, если нужно)
     import numpy as np
-    import subprocess
-    import tempfile
-
-    os.makedirs(os.path.dirname(output_path), exist_ok=True)
-
-    # Убедимся, что все кадры — одинакового размера
-    h, w, _ = frames[0].shape
-    h, w = h // 2 * 2, w // 2 * 2
-    frames = [cv2.resize(f, (w, h)) for f in frames]
-
-    # Временная папка
-    with tempfile.TemporaryDirectory() as tmpdir:
-        frame_dir = os.path.join(tmpdir, "frames")
-        os.makedirs(frame_dir, exist_ok=True)
-
-        # Сохраняем каждый кадр как PNG
-        for idx, frame in enumerate(frames):
-            cv2.imwrite(os.path.join(frame_dir, f"{idx:05d}.png"), frame)
-
-        # Генерируем видео через ffmpeg
-        ffmpeg_cmd = [
-            "ffmpeg",
-            "-y",
-            "-framerate", str(fps),
-            "-i", os.path.join(frame_dir, "%05d.png"),
-            "-c:v", "libx264",
-            "-preset", "slow",             # качество > скорость
-            "-crf", "18",                  # высокое качество (0 — без потерь)
-            "-pix_fmt", "yuv420p",
-            output_path
-        ]
-
-        print("Running ffmpeg:", " ".join(ffmpeg_cmd))
-        subprocess.run(ffmpeg_cmd, check=True)
+    import imageio
+    import torch
+    import torchvision
+    import os
+    from fractions import Fraction
 
+    # Convert fps to a clean format
+    if isinstance(fps, np.generic):
+        fps = fps.item()
+    fps = float(fps)
+
+    # Ensure all frames are the same shape
+    assert all(f.shape == frames[0].shape for f in frames), "All frames must have the same shape"
+
+    # Convert to tensor (T, H, W, C)
+    #frames = torch.from_numpy(np.asarray(frames))
+
+    # Ensure output directory exists
+    if not os.path.exists(os.path.dirname(output_path)):
+        os.makedirs(os.path.dirname(output_path))
+
+    # Write the video
+    #torchvision.io.write_video(output_path, frames, fps=fps, video_codec="libx264")
+    #return output_path
+    if bitrate is not None:
+        writer = imageio.get_writer(output_path, fps=fps, codec='libx264', bitrate=bitrate)
+    else:
+        writer = imageio.get_writer(output_path, fps=fps, codec='libx264')
+    for frame in frames:
+        writer.append_data(frame.astype(np.uint8))
+    writer.close()
     return output_path
 
 def restart():