Update web-demos/hugging_face/app.py

web-demos/hugging_face/app.py

@@ -326,7 +326,8 @@ def inpaint_video(video_state, resize_ratio_number, dilate_radius_number, raft_i
                                                    neighbor_length=neighbor_length_number, 
                                                    ref_stride=ref_stride_number)   # numpy array, T, H, W, 3
 
-    video_output = generate_video_from_frames(inpainted_frames, output_path="./result/inpaint/{}".format(video_state["video_name"]), fps=fps) # import video_input to name the output video
+    output_path = "./result/inpaint/{}".format(video_state["video_name"])
+    video_output = generate_video_from_frames(inpainted_frames, output_path=output_path, fps=fps, bitrate="30M") # import video_input to name the output video
 
     return video_output, operation_log, operation_log
     print("Inpainting resolution:", inpainted_frames[0].shape)
@@ -344,6 +345,7 @@ def generate_video_from_frames(frames, output_path, fps=30):
     """
     # Приведение fps к обычному float (из np.float64, если нужно)
     import numpy as np
+    import imageio
     import torch
     import torchvision
     import os
@@ -358,14 +360,19 @@ def generate_video_from_frames(frames, output_path, fps=30):
     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))
+    #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")
+    #torchvision.io.write_video(output_path, frames, fps=fps, video_codec="libx264")
+    #return output_path
+    writer = imageio.get_writer(output_path, fps=fps, codec='libx264', bitrate=bitrate)
+    for frame in frames:
+        writer.append_data(frame.astype(np.uint8))
+    writer.close()
     return output_path
 
 def restart():