Update app.py

app.py

@@ -18,30 +18,42 @@ def run_command(command):
 def setup_environment():
     """Sets up RIFE, downloads weights, and patches libraries."""
     print("--- Starting Environment Setup ---")
+    
+    # 1. Clone Repo
     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 Weights
     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")
         run_command(f"wget -O {zip_path} {MODEL_URL}")
         run_command(f"unzip -o {zip_path} -d {RIFE_DIR}")
         
-        # Fix directory structure
+        # Fix directory structure (User Fix)
         train_log_dir = os.path.join(RIFE_DIR, "train_log")
         os.makedirs(train_log_dir, exist_ok=True)
         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)
+        
+        # Move files safely
+        if os.path.exists(os.path.join(extract_folder, "RIFE_HDv3.py")):
+            shutil.move(os.path.join(extract_folder, "RIFE_HDv3.py"), train_log_dir)
+        if os.path.exists(os.path.join(extract_folder, "IFNet_HDv3.py")):
+            shutil.move(os.path.join(extract_folder, "IFNet_HDv3.py"), train_log_dir)
+            
         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)
+        
+        if os.path.exists(os.path.join(extract_folder, "flownet.pkl")):
+            shutil.move(os.path.join(extract_folder, "flownet.pkl"), hdv3_dir)
+            
         shutil.rmtree(extract_folder)
         if os.path.exists(zip_path): os.remove(zip_path)
 
-    # Patch skvideo
+    # 3. Patch skvideo (Numpy Fix)
     try:
         spec = importlib.util.find_spec('skvideo.io.abstract')
         if spec and spec.origin:
@@ -52,13 +64,14 @@ def setup_environment():
     except Exception as e:
         print(f"Warning: skvideo patch failed: {e}")
 
-    # Patch RIFE inference for libx264
+    # 4. Patch RIFE inference for libx264
     inference_script = os.path.join(RIFE_DIR, "inference_video.py")
     if os.path.exists(inference_script):
         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 ---")
 
 setup_environment()
@@ -66,24 +79,49 @@ setup_environment()
 @spaces.GPU(required=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_interpolated.mp4")
 
+    # Clean previous runs
     if os.path.exists(output_path): os.remove(output_path)
     if os.path.exists(final_output_path): os.remove(final_output_path)
 
+    # RIFE script creates a specific output name if audio transfer fails
+    # It often appends _noaudio, so we must watch for that.
+    expected_no_audio = output_path.replace(".mp4", "_noaudio.mp4")
+    if os.path.exists(expected_no_audio): os.remove(expected_no_audio)
+
     os.chdir(RIFE_DIR)
     try:
+        print(f"Running RIFE with {factor}x on {input_video_path}")
         cmd = ['python3', 'inference_video.py', '--video', input_video_path, '--output', output_path, '--multi', factor, '--model', 'HDv3']
         subprocess.run(cmd, check=True)
         
+        # --- Logic to handle the file finding ---
+        # If moviepy fails, RIFE creates 'output_rife.mp4' OR 'output_rife_noaudio.mp4'
+        # depending on where it crashed. We check for both.
+        
+        source_to_encode = None
+        
         if os.path.exists(output_path):
-            # 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
+            source_to_encode = output_path
+        elif os.path.exists(expected_no_audio):
+            source_to_encode = expected_no_audio
+            print("Audio transfer failed inside RIFE, using no-audio version.")
+        else:
+            print("Error: Output video file not found after inference.")
+            return None
+
+        # Re-encode for web compatibility
+        print(f"Re-encoding {source_to_encode} to {final_output_path}...")
+        subprocess.run(['ffmpeg', '-i', source_to_encode, '-c:v', 'libx264', '-pix_fmt', 'yuv420p', '-movflags', '+faststart', '-y', final_output_path], check=True)
+        
+        return final_output_path
+
     except Exception as e:
-        print(f"Error: {e}")
+        print(f"Error during interpolation: {e}")
         return None
     finally:
         os.chdir(BASE_DIR)
@@ -92,13 +130,11 @@ def interpolate_video(input_video_path, multi_factor):
 def stitch_videos(video_files, resolution_choice):
     if not video_files: return None
     
-    # Parse resolution (e.g., "1024x1024" -> w=1024, h=1024)
+    # Parse resolution
     try:
         target_w, target_h = resolution_choice.split("x")
-        target_w = target_w.strip()
-        target_h = target_h.strip()
+        target_w, target_h = target_w.strip(), target_h.strip()
     except:
-        # Fallback default
         target_w, target_h = "1920", "1080"
 
     print(f"Stitching {len(video_files)} videos into {target_w}x{target_h}...")
@@ -110,30 +146,27 @@ def stitch_videos(video_files, resolution_choice):
     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 User-Selected Scale
+    # 1. Normalize
     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}...")
         
-        # FFmpeg Filter Breakdown:
-        # scale=W:H:force_original_aspect_ratio=decrease -> Resize video to FIT inside the box
-        # pad=W:H:(ow-iw)/2:(oh-ih)/2 -> Create the box and center the video (adding black bars)
+        # Scale and Pad Filter
         scale_filter = f"scale={target_w}:{target_h}:force_original_aspect_ratio=decrease,pad={target_w}:{target_h}:(ow-iw)/2:(oh-ih)/2"
         
         cmd = [
             'ffmpeg', '-i', vid_path,
-            '-r', '60',                 # Force Framerate to 60
-            '-vf', scale_filter,        # Apply Dynamic Scaling
+            '-r', '60',                 
+            '-vf', scale_filter,        
             '-c:v', 'libx264', '-crf', '23',
-            '-c:a', 'aac',              # Force Audio Codec
-            '-ar', '44100',             # Force Audio Rate
+            '-c:a', 'aac',              
+            '-ar', '44100',             
             '-y', temp_filepath
         ]
         subprocess.run(cmd, check=True)
         normalized_files.append(temp_filepath)
 
-    # 2. Create List File
+    # 2. Create List
     with open(stitch_list_path, 'w') as f:
         for path in normalized_files:
             f.write(f"file '{path}'\n")
@@ -142,7 +175,6 @@ def stitch_videos(video_files, resolution_choice):
     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
 
@@ -153,11 +185,11 @@ with gr.Blocks() as demo:
     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).")
+            gr.Markdown("Upload a single video to increase its framerate.")
             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.)")
+                    multi_select = gr.Dropdown(["2", "4", "8"], value="2", label="Multiplier")
                     interp_btn = gr.Button("Interpolate", variant="primary")
                 with gr.Column():
                     video_output = gr.Video(label="Smoothed Output")
@@ -165,23 +197,15 @@ with gr.Blocks() as demo:
 
         # TAB 2: Stitch
         with gr.TabItem("2. Stitch Videos"):
-            gr.Markdown("Upload multiple videos. They will be normalized to **60fps** and your selected resolution.")
+            gr.Markdown("Upload multiple videos. They will be normalized to **60fps**.")
             with gr.Row():
                 with gr.Column():
                     stitch_inputs = gr.File(label="Upload Clips", file_count="multiple")
-                    
-                    # New Resolution Dropdown
-                    res_select = gr.Dropdown(
-                        choices=["1920x1080", "1280x1280", "1024x1024"], 
-                        value="1920x1080", 
-                        label="Output Resolution"
-                    )
-                    
+                    res_select = gr.Dropdown(choices=["1920x1080", "1280x1280", "1024x1024"], value="1920x1080", label="Resolution")
                     stitch_btn = gr.Button("Stitch Videos", variant="primary")
                 with gr.Column():
                     stitch_output = gr.Video(label="Stitched Result")
             
-            # Updated click event passing the resolution choice
             stitch_btn.click(stitch_videos, inputs=[stitch_inputs, res_select], outputs=stitch_output)
 
 demo.launch()