Update app.py

app.py

@@ -1,6 +1,5 @@
 import spaces
 import os
-import sys
 import shutil
 import subprocess
 import glob
@@ -12,174 +11,155 @@ BASE_DIR = os.getcwd()
 RIFE_DIR = os.path.join(BASE_DIR, "Practical-RIFE")
 MODEL_URL = "https://huggingface.co/hzwer/RIFE/resolve/main/RIFEv4.26_0921.zip"
 
+# --- Setup Functions ---
 def run_command(command):
-    """Helper to run shell commands"""
     subprocess.run(command, shell=True, check=True)
 
 def setup_environment():
-    """
-    1. Clones the RIFE repo.
-    2. Downloads weights.
-    3. Re-organizes folders (User's Fix).
-    4. Patches scikit-video.
-    """
+    """Sets up RIFE, downloads weights, and patches libraries."""
     print("--- Starting Environment Setup ---")
-
-    # 1. Clone Repo if not exists
     if not os.path.exists(RIFE_DIR):
         print("Cloning Practical-RIFE...")
         run_command(f"git clone https://github.com/hzwer/Practical-RIFE {RIFE_DIR}")
     
-    # 2. Download and Organize Weights
-    # We check if 'HDv3' exists to avoid re-downloading on restart
     if not os.path.exists(os.path.join(RIFE_DIR, "HDv3")):
         print("Downloading RIFE v4.26 model weights...")
         zip_path = os.path.join(RIFE_DIR, "RIFEv4.26_0921.zip")
-        
-        # Download
         run_command(f"wget -O {zip_path} {MODEL_URL}")
-        
-        # Unzip
-        print("Unzipping...")
         run_command(f"unzip -o {zip_path} -d {RIFE_DIR}")
         
-        # --- APPLY USER'S DIRECTORY FIX ---
-        print("Applying directory structure fix...")
-        
-        # Create train_log
+        # Fix directory structure
         train_log_dir = os.path.join(RIFE_DIR, "train_log")
         os.makedirs(train_log_dir, exist_ok=True)
-        
-        # Move Python files
         extract_folder = os.path.join(RIFE_DIR, "RIFEv4.26_0921")
         shutil.move(os.path.join(extract_folder, "RIFE_HDv3.py"), train_log_dir)
         shutil.move(os.path.join(extract_folder, "IFNet_HDv3.py"), train_log_dir)
-        
-        # Create __init__.py
-        with open(os.path.join(train_log_dir, "__init__.py"), 'w') as f:
-            pass
-            
-        # Create HDv3 and move weights
+        with open(os.path.join(train_log_dir, "__init__.py"), 'w') as f: pass
         hdv3_dir = os.path.join(RIFE_DIR, "HDv3")
         os.makedirs(hdv3_dir, exist_ok=True)
         shutil.move(os.path.join(extract_folder, "flownet.pkl"), hdv3_dir)
-        
-        # Cleanup
         shutil.rmtree(extract_folder)
-        if os.path.exists(zip_path):
-            os.remove(zip_path)
-            
-        print("Files moved and directories created.")
+        if os.path.exists(zip_path): os.remove(zip_path)
 
-    # 3. Patch skvideo (The numpy fix)
-    print("Checking skvideo for necessary patches...")
+    # Patch skvideo
     try:
         spec = importlib.util.find_spec('skvideo.io.abstract')
         if spec and spec.origin:
-            abstract_path = spec.origin
-            
-            # Read file to see if patching is needed
-            with open(abstract_path, 'r') as f:
-                content = f.read()
-            
-            if 'np.float' in content or 'np.int' in content:
-                print(f"Patching {abstract_path}...")
-                run_command(f"sed -i 's/np.float/float/g' {abstract_path}")
-                run_command(f"sed -i 's/np.int/int/g' {abstract_path}")
-                print("skvideo patched.")
-            else:
-                print("skvideo already patched or compatible.")
+            with open(spec.origin, 'r') as f: content = f.read()
+            if 'np.float' in content:
+                run_command(f"sed -i 's/np.float/float/g' {spec.origin}")
+                run_command(f"sed -i 's/np.int/int/g' {spec.origin}")
     except Exception as e:
-        print(f"Warning: Could not patch skvideo: {e}")
+        print(f"Warning: skvideo patch failed: {e}")
 
-    # 4. Patch inference_video.py for libx264
+    # Patch RIFE inference for libx264
     inference_script = os.path.join(RIFE_DIR, "inference_video.py")
     if os.path.exists(inference_script):
-        # Read to check if already patched
-        with open(inference_script, 'r') as f:
-            script_content = f.read()
-        
-        if "libx264" not in script_content:
-            print("Patching inference_video.py for libx264...")
-            # Replicating the sed command using python replace for safety
-            new_content = script_content.replace(
-                "-c:v', 'mpeg4', '-qscale:v', '1'", 
-                "-c:v', 'libx264', '-preset', 'medium', '-crf', '23'"
-            )
-            with open(inference_script, 'w') as f:
-                f.write(new_content)
-
+        with open(inference_script, 'r') as f: content = f.read()
+        if "libx264" not in content:
+            new_content = content.replace("-c:v', 'mpeg4', '-qscale:v', '1'", "-c:v', 'libx264', '-preset', 'medium', '-crf', '23'")
+            with open(inference_script, 'w') as f: f.write(new_content)
     print("--- Setup Complete ---")
 
-# Run Setup immediately upon app launch
 setup_environment()
 
-@spaces.GPU(required=True)
-def interpolate_video(input_video_path):
-    if input_video_path is None:
-        return None
-
-    print(f"Processing video: {input_video_path}")
-    
+@spaces.GPU(requires=True)
+def interpolate_video(input_video_path, multi_factor):
+    if input_video_path is None: return None
+    factor = str(multi_factor).replace("x", "").strip()
     output_path = os.path.join(BASE_DIR, "output_rife.mp4")
-    final_output_path = os.path.join(BASE_DIR, "final_output.mp4")
+    final_output_path = os.path.join(BASE_DIR, "final_interpolated.mp4")
 
-    # Clean up previous runs
     if os.path.exists(output_path): os.remove(output_path)
     if os.path.exists(final_output_path): os.remove(final_output_path)
 
-    # Change dir to RIFE for execution
     os.chdir(RIFE_DIR)
-
     try:
-        # Run RIFE inference
-        # Note: --multi 4 means 4x interpolation (e.g. 30fps -> 120fps)
-        cmd = [
-            'python3', 'inference_video.py',
-            '--video', input_video_path,
-            '--output', output_path,
-            '--multi', '4',
-            '--model', 'HDv3'
-        ]
+        cmd = ['python3', 'inference_video.py', '--video', input_video_path, '--output', output_path, '--multi', factor, '--model', 'HDv3']
         subprocess.run(cmd, check=True)
-
-        # Post-process for Web Compatibility (Faststart + h264)
+        
         if os.path.exists(output_path):
-            print("Re-encoding for web compatibility...")
-            ffmpeg_cmd = [
-                'ffmpeg', '-i', output_path,
-                '-c:v', 'libx264',
-                '-pix_fmt', 'yuv420p',
-                '-preset', 'medium',
-                '-crf', '23',
-                '-movflags', '+faststart',
-                '-y', final_output_path
-            ]
-            subprocess.run(ffmpeg_cmd, check=True)
-            
+            # Re-encode for web compatibility
+            subprocess.run(['ffmpeg', '-i', output_path, '-c:v', 'libx264', '-pix_fmt', 'yuv420p', '-movflags', '+faststart', '-y', final_output_path], check=True)
             return final_output_path
-        else:
-            return None
-            
     except Exception as e:
-        print(f"Error during inference: {e}")
+        print(f"Error: {e}")
         return None
     finally:
-        # Reset directory
         os.chdir(BASE_DIR)
 
+# --- Stitching Logic ---
+def stitch_videos(video_files):
+    if not video_files: return None
+    
+    print(f"Stitching {len(video_files)} videos...")
+    stitch_list_path = os.path.join(BASE_DIR, "stitch_list.txt")
+    output_stitched = os.path.join(BASE_DIR, "final_stitched.mp4")
+    temp_dir = os.path.join(BASE_DIR, "temp_stitch")
+    
+    if os.path.exists(temp_dir): shutil.rmtree(temp_dir)
+    os.makedirs(temp_dir)
+
+    # 1. Normalize Step: Convert all inputs to same Codec, FPS (60), and Scale
+    # This prevents the "glitch" when combining different framerates
+    normalized_files = []
+    for i, vid_path in enumerate(video_files):
+        temp_filepath = os.path.join(temp_dir, f"norm_{i}.mp4")
+        print(f"Normalizing clip {i+1}...")
+        
+        # We force 60fps (-r 60) and 1080p scale to ensure they match
+        cmd = [
+            'ffmpeg', '-i', vid_path,
+            '-r', '60',                 # Force Framerate
+            '-vf', 'scale=1920:1080:force_original_aspect_ratio=decrease,pad=1920:1080:(ow-iw)/2:(oh-ih)/2', # Fit to 1080p box
+            '-c:v', 'libx264', '-crf', '23',
+            '-c:a', 'aac',              # Force Audio Codec
+            '-ar', '44100',             # Force Audio Rate
+            '-y', temp_filepath
+        ]
+        subprocess.run(cmd, check=True)
+        normalized_files.append(temp_filepath)
+
+    # 2. Create List File
+    with open(stitch_list_path, 'w') as f:
+        for path in normalized_files:
+            f.write(f"file '{path}'\n")
+
+    # 3. Concatenate
+    print("Concatenating...")
+    subprocess.run(['ffmpeg', '-f', 'concat', '-safe', '0', '-i', stitch_list_path, '-c', 'copy', '-y', output_stitched], check=True)
+    
+    # Cleanup
+    shutil.rmtree(temp_dir)
+    return output_stitched
+
 # --- Gradio Interface ---
 with gr.Blocks() as demo:
-    gr.Markdown("# 🎞️ RIFE Video Interpolation (v4.26)")
-    gr.Markdown("Upload a video to smoothen it (4X Frame Interpolation). **Note:** This process is slow on CPU. Using a GPU Space is recommended.")
-    
-    with gr.Row():
-        with gr.Column():
-            video_input = gr.Video(label="Input Video")
-            submit_btn = gr.Button("Interpolate Frames", variant="primary")
-        with gr.Column():
-            video_output = gr.Video(label="Smoothed Output")
+    gr.Markdown("# 🎞️ RIFE: Interpolate & Stitch")
     
-    submit_btn.click(fn=interpolate_video, inputs=video_input, outputs=video_output)
+    with gr.Tabs():
+        # TAB 1: Interpolate
+        with gr.TabItem("1. Smooth Video (Interpolate)"):
+            gr.Markdown("Upload a single video to increase its framerate (make it smoother).")
+            with gr.Row():
+                with gr.Column():
+                    video_input = gr.Video(label="Input Video")
+                    multi_select = gr.Dropdown(["2", "4", "8"], value="2", label="Multiplier (2x, 4x, etc.)")
+                    interp_btn = gr.Button("Interpolate", variant="primary")
+                with gr.Column():
+                    video_output = gr.Video(label="Smoothed Output")
+            interp_btn.click(interpolate_video, inputs=[video_input, multi_select], outputs=video_output)
+
+        # TAB 2: Stitch
+        with gr.TabItem("2. Stitch Videos"):
+            gr.Markdown("Upload multiple videos. They will be **normalized to 60fps 1080p** and joined together.")
+            with gr.Row():
+                with gr.Column():
+                    # file_count="multiple" allows selecting multiple files
+                    stitch_inputs = gr.File(label="Upload Clips", file_count="multiple")
+                    stitch_btn = gr.Button("Stitch Videos", variant="primary")
+                with gr.Column():
+                    stitch_output = gr.Video(label="Stitched Result")
+            stitch_btn.click(stitch_videos, inputs=stitch_inputs, outputs=stitch_output)
 
 demo.launch()