working on dynamic number of frames and multiple resolutions

app.py

@@ -27,7 +27,7 @@ def init_video_processor():
     return video_processor is not None
 
 def extract_frames_from_video(video_path, output_dir, max_frames=81):
-    """Extract frames from video and ensure we have at least 81 frames"""
+    """Extract frames from video and ensure we have at least max_frames frames"""
     os.makedirs(output_dir, exist_ok=True)
     
     reader = imageio.get_reader(video_path)
@@ -57,6 +57,8 @@ def generate_recammaster_video(
     video_file,
     text_prompt,
     camera_type,
+    num_frames,
+    resolution,
     progress=gr.Progress()
 ):
     """Main function to generate video with ReCamMaster"""
@@ -79,17 +81,27 @@ def generate_recammaster_video(
             input_video_path = os.path.join(temp_dir, "input.mp4")
             shutil.copy(video_file, input_video_path)
             
+            # Parse resolution
+            width, height = map(int, resolution.split('x'))
+            
             # Extract frames
             progress(0.2, desc="Extracting video frames...")
-            num_frames, fps = extract_frames_from_video(input_video_path, os.path.join(temp_dir, "frames"))
-            logger.info(f"Extracted {num_frames} frames at {fps} fps")
+            extracted_frames, fps = extract_frames_from_video(
+                input_video_path, 
+                os.path.join(temp_dir, "frames"),
+                max_frames=num_frames
+            )
+            logger.info(f"Extracted {extracted_frames} frames at {fps} fps")
             
             # Process with ReCamMaster
             progress(0.3, desc="Processing with ReCamMaster...")
             output_video = video_processor.process_video(
                 input_video_path,
                 text_prompt,
-                camera_type
+                camera_type,
+                num_frames=num_frames,
+                height=height,
+                width=width
             )
             
             # Save output video
@@ -114,13 +126,14 @@ def generate_recammaster_video(
         return None, f"Error: {str(e)}"
 
 # Create Gradio interface
-with gr.Blocks(title="ReCamMaster Demo") as demo:
+with gr.Blocks(title="ReCamMaster") as demo:
 
     gr.Markdown(f"""
-    # 🎥 ReCamMaster
+    # ReCamMaster 🎥
+
+    This is a demo of [ReCamMaster](https://jianhongbai.github.io/ReCamMaster/), an amazing model that allows you to reshoot any video!
     
-    ReCamMaster allows you to re-capture videos with novel camera trajectories.
-    Upload a video and select a camera transformation to see the magic!
+    Due to the long generation times (~ 10 min) this space [should be duplicated](https://huggingface.co/spaces/jbilcke-hf/ReCamMaster?duplicate=true) to your own account for the best experience (please select at least a Nvidia L40S).
     """)
     
     with gr.Row():
@@ -144,28 +157,36 @@ with gr.Blocks(title="ReCamMaster Demo") as demo:
                     value="1"
                 )
             
+            # Video settings
+            with gr.Group():
+                gr.Markdown("### Step 3: Video Settings")
+                num_frames = gr.Slider(
+                    minimum=17,
+                    maximum=81,
+                    value=81,
+                    step=16,
+                    label="Number of Frames",
+                    info="Must be 16n+1 (17, 33, 49, 65, 81)"
+                )
+                resolution = gr.Dropdown(
+                    choices=["832x480", "480x480", "480x832", "576x320", "320x576"],
+                    value="832x480",
+                    label="Resolution",
+                    info="Output video resolution"
+                )
+            
             # Generate button
-            generate_btn = gr.Button("Generate Video", variant="primary")
+            generate_btn = gr.Button("Generate Video (~10 min)", variant="primary")
         
         with gr.Column():
             # Output section
             output_video = gr.Video(label="Output Video")
-            status_output = gr.Textbox(label="Generation Status", interactive=False)
-    
-    # Example videos
-    gr.Markdown("### Example Videos")
-    gr.Examples(
-        examples=[
-            [f"{TEST_DATA_DIR}/videos/case0.mp4", "A person dancing", "1"],
-            [f"{TEST_DATA_DIR}/videos/case1.mp4", "A scenic view", "5"],
-        ],
-        inputs=[video_input, text_prompt, camera_type],
-    )
-    
+            status_output = gr.Textbox(label="Status", interactive=False)
+   
     # Event handlers
     generate_btn.click(
         fn=generate_recammaster_video,
-        inputs=[video_input, text_prompt, camera_type],
+        inputs=[video_input, text_prompt, camera_type, num_frames, resolution],
         outputs=[output_video, status_output]
     )
 

video_processor.py

@@ -15,21 +15,13 @@ logger = logging.getLogger(__name__)
 class VideoProcessor:
     def __init__(self, pipe):
         self.pipe = pipe
-        self.height = 480
-        self.width = 832
-        
-        # Create frame processor
-        self.frame_process = v2.Compose([
-            v2.CenterCrop(size=(self.height, self.width)),
-            v2.Resize(size=(self.height, self.width), antialias=True),
-            v2.ToTensor(),
-            v2.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
-        ])
+        self.default_height = 480
+        self.default_width = 832
     
-    def crop_and_resize(self, image):
+    def crop_and_resize(self, image, height, width):
         """Crop and resize image to match target dimensions"""
         width_img, height_img = image.size
-        scale = max(self.width / width_img, self.height / height_img)
+        scale = max(width / width_img, height / height_img)
         image = torchvision.transforms.functional.resize(
             image,
             (round(height_img*scale), round(width_img*scale)),
@@ -37,17 +29,25 @@ class VideoProcessor:
         )
         return image
     
-    def load_video_frames(self, video_path):
+    def load_video_frames(self, video_path, num_frames=81, height=480, width=832):
         """Load and process video frames"""
         reader = imageio.get_reader(video_path)
         frames = []
         
-        for i in range(81):  # ReCamMaster needs exactly 81 frames
+        # Create frame processor with specified dimensions
+        frame_process = v2.Compose([
+            v2.CenterCrop(size=(height, width)),
+            v2.Resize(size=(height, width), antialias=True),
+            v2.ToTensor(),
+            v2.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
+        ])
+        
+        for i in range(num_frames):
             try:
                 frame = reader.get_data(i)
                 frame = Image.fromarray(frame)
-                frame = self.crop_and_resize(frame)
-                frame = self.frame_process(frame)
+                frame = self.crop_and_resize(frame, height, width)
+                frame = frame_process(frame)
                 frames.append(frame)
             except:
                 # If we run out of frames, repeat the last one
@@ -64,14 +64,14 @@ class VideoProcessor:
         
         return video_tensor
     
-    def load_camera_trajectory(self, cam_type):
+    def load_camera_trajectory(self, cam_type, num_frames=81):
         """Load camera trajectory for the selected type"""
         tgt_camera_path = "./camera_trajectories/camera_extrinsics.json"
         with open(tgt_camera_path, 'r') as file:
             cam_data = json.load(file)
         
         # Get camera trajectory for selected type
-        cam_idx = list(range(81))[::4]  # Sample every 4 frames
+        cam_idx = list(range(num_frames))[::4]  # Sample every 4 frames
         traj = [parse_matrix(cam_data[f"frame{idx}"][f"cam{int(cam_type):02d}"]) for idx in cam_idx]
         traj = np.stack(traj).transpose(0, 2, 1)
         
@@ -94,14 +94,14 @@ class VideoProcessor:
         
         return camera_tensor
     
-    def process_video(self, video_path, text_prompt, cam_type):
+    def process_video(self, video_path, text_prompt, cam_type, num_frames=81, height=480, width=832):
         """Process video through ReCamMaster model"""
         
         # Load video frames
-        video_tensor = self.load_video_frames(video_path)
+        video_tensor = self.load_video_frames(video_path, num_frames, height, width)
         
         # Load camera trajectory
-        camera_tensor = self.load_camera_trajectory(cam_type)
+        camera_tensor = self.load_camera_trajectory(cam_type, num_frames)
         
         # Generate video with ReCamMaster
         video = self.pipe(
@@ -109,6 +109,9 @@ class VideoProcessor:
             negative_prompt=["worst quality, low quality, blurry, jittery, distorted"],
             source_video=video_tensor,
             target_camera=camera_tensor,
+            height=height,
+            width=width,
+            num_frames=num_frames,
             cfg_scale=5.0,
             num_inference_steps=50,
             seed=0,