Upload 4 files

README.md

@@ -1,6 +1,6 @@
 ---
-title: Test Mcp1
-emoji: 🌍
+title: FPS_Enhancer
+emoji: 🎞️⏫
 colorFrom: pink
 colorTo: red
 sdk: gradio

app.py

@@ -2,127 +2,72 @@ import gradio as gr
 import subprocess
 import os
 import spaces
+import inference_video_w
+import torch
 
 # Download the file
 subprocess.run([
     "wget",
-    "--no-check-certificate",
-    "https://drive.google.com/uc?id=1mj9lH6Be7ztYtHAr1xUUGT3hRtWJBy_5",
+    "https://huggingface.co/r3gm/RIFE/resolve/main/RIFEv4.26_0921.zip",
     "-O",
-    "RIFE_trained_model_v4.13.2.zip"
+    "RIFEv4.26_0921.zip"
 ], check=True)
 
 # Unzip the downloaded file
 subprocess.run([
     "unzip",
-    "RIFE_trained_model_v4.13.2.zip"
+    "-o", 
+    "RIFEv4.26_0921.zip"
 ], check=True)
 
-# The name of your script
-SCRIPT_NAME = "inference_video.py"
-
-@spaces.GPU()
+@spaces.GPU(duration=120)
 def run_rife(
     input_video, 
-    model_dir, 
-    multi, 
-    exp, 
-    fps, 
-    scale, 
-    uhd, 
-    fp16, 
-    skip, 
-    montage, 
-    png_mode, 
-    ext
+    frame_multiplier, 
+    time_exponent, 
+    fixed_fps, 
+    video_scale, 
+    remove_duplicate_frames, 
+    create_montage,
+    progress=gr.Progress(track_tqdm=True),
 ):
-    """
-    Constructs the command line arguments based on Gradio inputs 
-    and runs the inference_video.py script via subprocess.
-    """
-    
     if input_video is None:
         raise gr.Error("Please upload a video first.")
 
-    # 1. Define Output Filename
-    output_name = f"output_{multi}X.{ext}"
-    
-    # 2. Build the Command
-    cmd = ["python3", SCRIPT_NAME]
-    
-    # --video
-    cmd.extend(["--video", input_video])
-    
-    # --output
-    cmd.extend(["--output", output_name])
+    ext = "mp4"
+    model_dir = "train_log"
     
-    # --multi (Multiplier)
-    cmd.extend(["--multi", str(int(multi))])
+    # Construct output filename pattern to match what inference_video.py expects/generates
+    video_path_wo_ext = os.path.splitext(os.path.basename(input_video))[0]
+    # We pass the desired output name, though the function logic tries to stick to this pattern anyway
+    output_base_name = "{}_{}X_fps.{}".format(video_path_wo_ext, int(frame_multiplier), ext)
     
-    # --exp 
-    # Only add exp if it is not default, or if specific logic requires it. 
-    # Usually multi overrides exp in RIFE logic, but we pass it if set.
-    if exp != 1:
-        cmd.extend(["--exp", str(int(exp))])
-
-    # --fps (Target FPS)
-    if fps > 0:
-        cmd.extend(["--fps", str(int(fps))])
-        
-    # --scale (Resolution scale)
-    # Check against float 1.0
-    if scale != 1.0:
-        cmd.extend(["--scale", str(scale)])
-
-    # --ext (Extension)
-    cmd.extend(["--ext", ext])
-
-    # --model (Model directory)
-    if model_dir and model_dir.strip() != "":
-        cmd.extend(["--model", model_dir])
-
-    # --- Boolean Flags ---
-    
-    if uhd:
-        cmd.append("--UHD")
-        
-    if fp16:
-        cmd.append("--fp16")
-        
-    if skip:
-        cmd.append("--skip")
-        
-    if montage:
-        cmd.append("--montage")
-        
-    if png_mode:
-        cmd.append("--png")
-
-    print(f"Executing command: {' '.join(cmd)}")
-
-    # 3. Run the Subprocess
+    if fixed_fps > 0:
+        gr.Warning("Will not merge audio because using fps flag!")
+
+    print(f"Starting Inference for: {input_video}")
+
     try:
-        # We use a large timeout because video processing takes time
-        process = subprocess.run(cmd, capture_output=True, text=True)
-        
-        # Log stdout/stderr
-        if process.stdout:
-            print("STDOUT:", process.stdout)
-        if process.stderr:
-            print("STDERR:", process.stderr)
-
-        if process.returncode != 0:
-            raise gr.Error(f"Inference failed. Error: {process.stderr}")
-
-        # 4. Return Result
-        if png_mode:
-            gr.Info("Processing complete. Output is a folder of PNGs (Video preview unavailable for PNG mode).")
-            return None 
-        
-        if os.path.exists(output_name):
-            return output_name
+        # Call the imported function directly
+        result_path = inference_video_w.inference(
+            video=input_video,
+            output=output_base_name,
+            modelDir=model_dir,
+            fp16=(True if torch.cuda.is_available() else False),
+            UHD=False,
+            scale=video_scale,
+            skip=remove_duplicate_frames,
+            fps=int(fixed_fps) if fixed_fps > 0 else None,
+            ext=ext,
+            exp=int(time_exponent),
+            multi=int(frame_multiplier),
+            montage=create_montage
+        )
+
+        if result_path and os.path.exists(result_path):
+            return result_path
         else:
-            raise gr.Error("Output file was not found. Check console for details.")
+            raise gr.Error(f"Output file not found. Expected: {result_path}")
 
     except Exception as e:
         raise gr.Error(f"An error occurred: {str(e)}")
@@ -130,111 +75,75 @@ def run_rife(
 
 # --- Gradio UI Layout ---
 
-with gr.Blocks(title="RIFE Video Interpolation") as app:
-    gr.Markdown("# 🚀 RIFE: Real-Time Intermediate Flow Estimation")
-    gr.Markdown("Upload a video to increase its frame rate (smoothness) using AI.")
+with gr.Blocks(title="Frame Rate Enhancer") as app:
+    gr.Markdown("# ⚡ RIFE: Frame Rate Enhancer")
+    gr.Markdown("Creates extra frames between the original ones to make motion in your videos smoother and more fluid.")
+    gr.Markdown("⚠️ **Notice:** Keep input videos under 60 seconds for frame interpolation to prevent GPU task aborts.")
 
     with gr.Row():
         # --- Left Column: Inputs & Settings ---
         with gr.Column(scale=1):
-            input_vid = gr.Video(label="Input Video", sources=["upload"])
+            input_vid = gr.Video(label="🎬 Input Source Video", sources=["upload"])
             
-            with gr.Group():
-                gr.Markdown("### 🎯 Core Parameters")
-                
-                with gr.Row():
-                    multi_param = gr.Dropdown(
-                        choices=["2", "4", "8", "16", "32"], 
-                        value="2", 
-                        label="Interpolation Multiplier (--multi)",
-                        info="How many times to multiply the frames. 2X doubles the FPS (e.g., 30fps -> 60fps). 4X quadruples it."
-                    )
-                    ext_param = gr.Dropdown(
-                        choices=["mp4", "avi", "mov", "mkv"], 
-                        value="mp4", 
-                        label="Output Format (--ext)",
-                        info="The file extension for the generated video."
-                    )
-                
-                model_param = gr.Textbox(
-                    value="train_log", 
-                    label="Model Directory (--model)", 
-                    placeholder="train_log",
-                    info="Path to the folder containing the trained model files (e.g., 'train_log' or 'rife-v4.6')."
+            with gr.Group():                
+                multi_param = gr.Dropdown(
+                    choices=["2", "3", "4", "5", "6"], 
+                    value="2", 
+                    label="🗃️ Frame Multiplier",
+                    info="2X = Double FPS (e.g. 30 -> 60). Higher multipliers create more intermediate frames."
                 )
 
-            with gr.Accordion("⚙️ Advanced Settings", open=False):
-                gr.Markdown("Fine-tune the inference process.")
+            with gr.Accordion("🛠️ Advanced Configuration", open=False):
+                gr.Markdown("Control rendering parameters.")
                 
                 with gr.Row():
-                    scale_param = gr.Slider(
-                        minimum=0.1, maximum=1.0, value=1.0, step=0.1, 
-                        label="Input Scale (--scale)", 
-                        info="1.0 = Original resolution. Set to 0.5 to reduce memory usage for 4K video inputs."
+                    scale_param = gr.Dropdown(
+                        choices=[0.25, 0.5, 1.0, 2.0, 4.0],
+                        value=1.0,
+                        label="📉 Render Scale",
+                        info="1.0 = Original Resolution. Reduce to 0.5 for faster processing on 4K content."
                     )
                     fps_param = gr.Number(
                         value=0, 
-                        label="Force Target FPS (--fps)", 
-                        info="0 = Auto-calculate based on multiplier. Enter a number (e.g., 60) to force a specific output frame rate."
+                        label="🎯 Force Output FPS", 
+                        info="0 = Auto-calculate. Set to 30 or 60 to lock the framerate. Audio will be removed when forcing FPS"
                     )
                     exp_param = gr.Number(
                         value=1, 
-                        label="Exponent Power (--exp)", 
-                        info="Alternative to Multiplier. Sets multiplier to 2^exp. (Usually left at 1 if Multiplier is set)."
+                        label="🔢 Exponent Power", 
+                        info="Alternative multiplier calculation (2^exp)."
                     )
 
-                with gr.Row():
-                    uhd_chk = gr.Checkbox(
-                        label="UHD Mode (--UHD)", 
-                        value=False,
-                        info="Optimized for 4K video. Equivalent to setting scale=0.5 manually."
-                    )
-                    fp16_chk = gr.Checkbox(
-                        label="FP16 Mode (--fp16)", 
-                        value=True, 
-                        info="Uses half-precision floating point. Faster and uses less VRAM with minimal quality loss."
-                    )
-                
                 with gr.Row():
                     skip_chk = gr.Checkbox(
-                        label="Skip Static Frames (--skip)", 
+                        label="⏩ Skip Static Frames", 
                         value=False,
-                        info="If the video has frames that don't move, skip processing them to save time."
+                        info="Bypass processing for static frames to save time."
                     )
                     montage_chk = gr.Checkbox(
-                        label="Montage (--montage)", 
-                        value=False,
-                        info="Creates a video with the Original on the Left and Interpolated on the Right for comparison."
-                    )
-                    png_chk = gr.Checkbox(
-                        label="Output as PNGs (--png)", 
+                        label="🆚 Split-Screen Comparison", 
                         value=False,
-                        info="Outputs a sequence of images instead of a video file. (Video Preview will be disabled)."
+                        info="Output video showing Original vs. Processed."
                     )
 
-            btn_run = gr.Button("✨ Start Interpolation", variant="primary", size="lg")
+            btn_run = gr.Button("GENERATE INTERMEDIATE FRAMES", variant="primary", size="lg")
 
         # --- Right Column: Output ---
         with gr.Column(scale=1):
-            output_vid = gr.Video(label="Interpolated Result")
-            gr.Markdown("**Note:** Processing time depends on video length, resolution, and your GPU speed.")
+            output_vid = gr.Video(label="INTERPOLATED RESULT")
+            gr.Markdown("**Status:** Rendering time depends on input resolution and duration.")
 
     # --- Bind Logic ---
     btn_run.click(
         fn=run_rife,
         inputs=[
             input_vid, 
-            model_param, 
             multi_param, 
             exp_param, 
             fps_param, 
             scale_param, 
-            uhd_chk, 
-            fp16_chk, 
             skip_chk, 
-            montage_chk, 
-            png_chk, 
-            ext_param
+            montage_chk
         ],
         outputs=output_vid
     )
@@ -242,4 +151,5 @@ with gr.Blocks(title="RIFE Video Interpolation") as app:
 if __name__ == "__main__":
     app.launch(
         theme=gr.themes.Soft(),
+        mcp_server=True,
     )

inference_video_w.py

@@ -0,0 +1,316 @@
+import os
+import cv2
+import torch
+import numpy as np
+from tqdm import tqdm
+from torch.nn import functional as F
+import warnings
+import _thread
+import skvideo.io
+from queue import Queue, Empty
+from model.pytorch_msssim import ssim_matlab
+import shutil
+import tempfile
+import time
+
+warnings.filterwarnings("ignore")
+
+# Utility class to mimic argparse object
+class Args:
+    def __init__(self, **kwargs):
+        self.__dict__.update(kwargs)
+
+def transferAudio(sourceVideo, targetVideo):
+    # generate a unique temp directory for this user
+    unique_temp_dir = tempfile.mkdtemp()
+    tempAudioFileName = os.path.join(unique_temp_dir, "audio.mkv")
+
+    # extract audio from video
+    os.system('ffmpeg -hide_banner -loglevel error -y -i "{}" -c:a copy -vn {}'.format(sourceVideo, tempAudioFileName))
+
+    targetNoAudio = os.path.splitext(targetVideo)[0] + "_noaudio" + os.path.splitext(targetVideo)[1]
+    os.rename(targetVideo, targetNoAudio)
+    # combine audio file and new video file
+    os.system('ffmpeg -hide_banner -loglevel error -y -i "{}" -i {} -c copy "{}"'.format(targetNoAudio, tempAudioFileName, targetVideo))
+
+    if os.path.getsize(targetVideo) == 0: # if ffmpeg failed to merge the video and audio together try converting the audio to aac
+        tempAudioFileName = os.path.join(unique_temp_dir, "audio.m4a")
+        os.system('ffmpeg -hide_banner -loglevel error -y -i "{}" -c:a aac -b:a 160k -vn {}'.format(sourceVideo, tempAudioFileName))
+        os.system('ffmpeg -hide_banner -loglevel error -y -i "{}" -i {} -c copy "{}"'.format(targetNoAudio, tempAudioFileName, targetVideo))
+        if (os.path.getsize(targetVideo) == 0): # if aac is not supported by selected format
+            os.rename(targetNoAudio, targetVideo)
+            print("Audio transfer failed. Interpolated video will have no audio")
+        else:
+            print("Lossless audio transfer failed. Audio was transcoded to AAC (M4A) instead.")
+            # remove audio-less video
+            os.remove(targetNoAudio)
+    else:
+        os.remove(targetNoAudio)
+
+    # remove temp directory
+    shutil.rmtree(unique_temp_dir)
+
+def inference(
+    video=None,
+    output=None,
+    img=None,
+    montage=False,
+    modelDir='train_log',
+    fp16=False,
+    UHD=False,
+    scale=1.0,
+    skip=False,
+    fps=None,
+    png=False,
+    ext='mp4',
+    exp=1,
+    multi=2
+):
+    # Initialize Arguments Object
+    args = Args(
+        video=video, output=output, img=img, montage=montage,
+        modelDir=modelDir, fp16=fp16, UHD=UHD, scale=scale,
+        skip=skip, fps=fps, png=png, ext=ext, exp=exp, multi=multi
+    )
+
+    # Argument Logic Adjustment
+    if args.exp != 1:
+        args.multi = (2 ** args.exp)
+    
+    # Assertions
+    assert (not args.video is None or not args.img is None)
+    if args.skip:
+        print("skip flag is abandoned, please refer to issue #207.")
+    if args.UHD and args.scale==1.0:
+        args.scale = 0.5
+    assert args.scale in [0.25, 0.5, 1.0, 2.0, 4.0]
+    if not args.img is None:
+        args.png = True
+
+    # Device Setup
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    torch.set_grad_enabled(False)
+    if torch.cuda.is_available():
+        torch.backends.cudnn.enabled = True
+        torch.backends.cudnn.benchmark = True
+        if(args.fp16):
+            torch.set_default_tensor_type(torch.cuda.HalfTensor)
+
+    # Load Model
+    from train_log.RIFE_HDv3 import Model
+    model = Model()
+    if not hasattr(model, 'version'):
+        model.version = 0
+    model.load_model(args.modelDir, -1)
+    print("Loaded 3.x/4.x HD model.")
+    model.eval()
+    model.device()
+
+    # Video/Image Setup
+    if not args.video is None:
+        videoCapture = cv2.VideoCapture(args.video)
+        original_fps = videoCapture.get(cv2.CAP_PROP_FPS)
+        tot_frame = videoCapture.get(cv2.CAP_PROP_FRAME_COUNT)
+        videoCapture.release()
+        
+        if args.fps is None or args.fps == 0:
+            fpsNotAssigned = True
+            args.fps = original_fps * args.multi
+        else:
+            fpsNotAssigned = False
+            
+        videogen = skvideo.io.vreader(args.video)
+        lastframe = next(videogen)
+        # fourcc = cv2.VideoWriter_fourcc('m', 'p', '4', 'v') # Unused in original logic for skvideo
+        video_path_wo_ext, ext = os.path.splitext(args.video)
+        print('{}.{}, {} frames in total, {}FPS to {}FPS'.format(video_path_wo_ext, args.ext, tot_frame, original_fps, args.fps))
+        
+        if args.png == False and fpsNotAssigned == True:
+            print("The audio will be merged after interpolation process")
+        else:
+            print("Will not merge audio because using png or fps flag!")
+    else:
+        videogen = []
+        for f in os.listdir(args.img):
+            if 'png' in f:
+                videogen.append(f)
+        tot_frame = len(videogen)
+        videogen.sort(key= lambda x:int(x[:-4]))
+        lastframe = cv2.imread(os.path.join(args.img, videogen[0]), cv2.IMREAD_UNCHANGED)[:, :, ::-1].copy()
+        videogen = videogen[1:]
+    
+    h, w, _ = lastframe.shape
+    vid_out_name = None
+    vid_out = None
+    
+    if args.png:
+        if not os.path.exists('vid_out'):
+            os.mkdir('vid_out')
+    else:
+        if args.output is not None:
+            vid_out_name = args.output
+        else:
+            vid_out_name = '{}_{}X_{}fps.{}'.format(video_path_wo_ext, args.multi, int(np.round(args.fps)), args.ext)
+        
+        outputdict = {
+            '-c:v': 'libx264', 
+            '-crf': '17', 
+            '-preset': 'slow',
+            '-pix_fmt': 'yuv420p'
+        }
+        vid_out = skvideo.io.FFmpegWriter(vid_out_name, inputdict={'-r': str(args.fps)}, outputdict=outputdict)
+
+    # --- Nested Helper Functions to capture 'args', 'model', 'vid_out' scope ---
+
+    def clear_write_buffer(write_buffer):
+        cnt = 0
+        while True:
+            item = write_buffer.get()
+            if item is None:
+                break
+            if args.png:
+                cv2.imwrite('vid_out/{:0>7d}.png'.format(cnt), item[:, :, ::-1])
+                cnt += 1
+            else:
+                vid_out.writeFrame(item)
+
+    def build_read_buffer(read_buffer, videogen):
+        try:
+            for frame in videogen:
+                if not args.img is None:
+                    frame = cv2.imread(os.path.join(args.img, frame), cv2.IMREAD_UNCHANGED)[:, :, ::-1].copy()
+                if args.montage:
+                    frame = frame[:, left: left + w]
+                read_buffer.put(frame)
+        except:
+            pass
+        read_buffer.put(None)
+
+    def make_inference(I0, I1, n):    
+        if model.version >= 3.9:
+            res = []
+            for i in range(n):
+                res.append(model.inference(I0, I1, (i+1) * 1. / (n+1), args.scale))
+            return res
+        else:
+            middle = model.inference(I0, I1, args.scale)
+            if n == 1:
+                return [middle]
+            first_half = make_inference(I0, middle, n=n//2)
+            second_half = make_inference(middle, I1, n=n//2)
+            if n%2:
+                return [*first_half, middle, *second_half]
+            else:
+                return [*first_half, *second_half]
+
+    def pad_image(img):
+        if(args.fp16):
+            return F.pad(img, padding).half()
+        else:
+            return F.pad(img, padding)
+
+    # --- Pre-Loop Setup ---
+
+    left = 0 # Define default
+    if args.montage:
+        left = w // 4
+        w = w // 2
+    
+    tmp = max(128, int(128 / args.scale))
+    ph = ((h - 1) // tmp + 1) * tmp
+    pw = ((w - 1) // tmp + 1) * tmp
+    padding = (0, pw - w, 0, ph - h)
+    
+    pbar = tqdm(total=tot_frame)
+    if args.montage:
+        lastframe = lastframe[:, left: left + w]
+        
+    write_buffer = Queue(maxsize=500)
+    read_buffer = Queue(maxsize=500)
+    
+    # Start threads
+    _thread.start_new_thread(build_read_buffer, (read_buffer, videogen))
+    _thread.start_new_thread(clear_write_buffer, (write_buffer,))
+
+    I1 = torch.from_numpy(np.transpose(lastframe, (2,0,1))).to(device, non_blocking=True).unsqueeze(0).float() / 255.
+    I1 = pad_image(I1)
+    temp = None 
+
+    # --- Main Loop ---
+    
+    while True:
+        if temp is not None:
+            frame = temp
+            temp = None
+        else:
+            frame = read_buffer.get()
+        if frame is None:
+            break
+        I0 = I1
+        I1 = torch.from_numpy(np.transpose(frame, (2,0,1))).to(device, non_blocking=True).unsqueeze(0).float() / 255.
+        I1 = pad_image(I1)
+        I0_small = F.interpolate(I0, (32, 32), mode='bilinear', align_corners=False)
+        I1_small = F.interpolate(I1, (32, 32), mode='bilinear', align_corners=False)
+        ssim = ssim_matlab(I0_small[:, :3], I1_small[:, :3])
+
+        break_flag = False
+        if ssim > 0.996:
+            frame = read_buffer.get() # read a new frame
+            if frame is None:
+                break_flag = True
+                frame = lastframe
+            else:
+                temp = frame
+            I1 = torch.from_numpy(np.transpose(frame, (2,0,1))).to(device, non_blocking=True).unsqueeze(0).float() / 255.
+            I1 = pad_image(I1)
+            I1 = model.inference(I0, I1, scale=args.scale)
+            I1_small = F.interpolate(I1, (32, 32), mode='bilinear', align_corners=False)
+            ssim = ssim_matlab(I0_small[:, :3], I1_small[:, :3])
+            frame = (I1[0] * 255).byte().cpu().numpy().transpose(1, 2, 0)[:h, :w]
+            
+        if ssim < 0.2:
+            output_frames = []
+            for i in range(args.multi - 1):
+                output_frames.append(I0)
+        else:
+            output_frames = make_inference(I0, I1, args.multi - 1)
+
+        if args.montage:
+            write_buffer.put(np.concatenate((lastframe, lastframe), 1))
+            for mid in output_frames:
+                mid = (((mid[0] * 255.).byte().cpu().numpy().transpose(1, 2, 0)))
+                write_buffer.put(np.concatenate((lastframe, mid[:h, :w]), 1))
+        else:
+            write_buffer.put(lastframe)
+            for mid in output_frames:
+                mid = (((mid[0] * 255.).byte().cpu().numpy().transpose(1, 2, 0)))
+                write_buffer.put(mid[:h, :w])
+        pbar.update(1)
+        lastframe = frame
+        if break_flag:
+            break
+
+    if args.montage:
+        write_buffer.put(np.concatenate((lastframe, lastframe), 1))
+    else:
+        write_buffer.put(lastframe)
+        
+    write_buffer.put(None)
+
+    while(not write_buffer.empty()):
+        time.sleep(0.1)
+    pbar.close()
+    
+    if not vid_out is None:
+        vid_out.close()
+
+    # Audio Transfer Logic
+    if args.png == False and fpsNotAssigned == True and not args.video is None:
+        try:
+            transferAudio(args.video, vid_out_name)
+        except:
+            print("Audio transfer failed. Interpolated video will have no audio")
+            targetNoAudio = os.path.splitext(vid_out_name)[0] + "_noaudio" + os.path.splitext(vid_out_name)[1]
+            os.rename(targetNoAudio, vid_out_name)
+            
+    return vid_out_name

packages.txt

@@ -0,0 +1 @@
+ffmpeg