app.py

import gradio as gr
import torch
import os
import tempfile
import shutil
import imageio
import logging
from pathlib import Path

# Import from our modules
from model_loader import ModelLoader, MODELS_ROOT_DIR
from video_processor import VideoProcessor
from config import CAMERA_TRANSFORMATIONS, TEST_DATA_DIR

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Global model loader instance
model_loader = ModelLoader()
video_processor = None

def init_video_processor():
    """Initialize video processor"""
    global video_processor
    if model_loader.is_loaded and video_processor is None:
        video_processor = VideoProcessor(model_loader.pipe)
    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 max_frames frames"""
    os.makedirs(output_dir, exist_ok=True)
    
    reader = imageio.get_reader(video_path)
    fps = reader.get_meta_data()['fps']
    total_frames = reader.count_frames()
    
    frames = []
    for i, frame in enumerate(reader):
        frames.append(frame)
    reader.close()
    
    # If we have fewer than required frames, repeat the last frame
    if len(frames) < max_frames:
        logger.info(f"Video has {len(frames)} frames, padding to {max_frames} frames")
        last_frame = frames[-1]
        while len(frames) < max_frames:
            frames.append(last_frame)
    
    # Save frames
    for i, frame in enumerate(frames[:max_frames]):
        frame_path = os.path.join(output_dir, f"frame_{i:04d}.png")
        imageio.imwrite(frame_path, frame)
    
    return len(frames[:max_frames]), fps

def generate_recammaster_video(
    video_file,
    text_prompt,
    camera_type,
    num_frames,
    resolution,
    progress=gr.Progress()
):
    """Main function to generate video with ReCamMaster"""
    
    if not model_loader.is_loaded:
        return None, "Error: Models not loaded! Please load models first."
    
    if not init_video_processor():
        return None, "Error: Failed to initialize video processor."
    
    if video_file is None:
        return None, "Please upload a video file."
    
    try:
        # Create temporary directory for processing
        with tempfile.TemporaryDirectory() as temp_dir:
            progress(0.1, desc="Processing input video...")
            
            # Copy uploaded video to temp directory
            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...")
            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,
                num_frames=num_frames,
                height=height,
                width=width
            )
            
            # Save output video
            progress(0.9, desc="Saving output video...")
            output_path = os.path.join(temp_dir, "output.mp4")
            from diffsynth import save_video
            save_video(output_video, output_path, fps=30, quality=5)
            
            # Copy to persistent location
            final_output_path = tempfile.NamedTemporaryFile(suffix='.mp4', delete=False).name
            shutil.copy(output_path, final_output_path)
            
            progress(1.0, desc="Done!")
            
            transformation_name = CAMERA_TRANSFORMATIONS.get(str(camera_type), "Unknown")
            status_msg = f"Successfully generated video with '{transformation_name}' camera movement!"
            
            return final_output_path, status_msg
    
    except Exception as e:
        logger.error(f"Error generating video: {str(e)}")
        return None, f"Error: {str(e)}"

# Create Gradio interface
with gr.Blocks(title="ReCamMaster") as demo:

    gr.Markdown(f"""
    # ReCamMaster 🎥

    This is a demo of [ReCamMaster](https://jianhongbai.github.io/ReCamMaster/), an amazing model that allows you to reshoot any video!
    
    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():
        with gr.Column():
            # Video input section
            with gr.Group():
                gr.Markdown("### Step 1: Upload Video")
                video_input = gr.Video(label="Input Video")
                text_prompt = gr.Textbox(
                    label="Text Prompt (describe your video)",
                    placeholder="A person walking in the street",
                    value="A dynamic scene"
                )
            
            # Camera selection
            with gr.Group():
                gr.Markdown("### Step 2: Select Camera Movement")
                camera_type = gr.Radio(
                    choices=[(v, k) for k, v in CAMERA_TRANSFORMATIONS.items()],
                    label="Camera Transformation",
                    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 (~10 min)", variant="primary")
        
        with gr.Column():
            # Output section
            output_video = gr.Video(label="Output Video")
            status_output = gr.Textbox(label="Status", interactive=False)
   
    # Event handlers
    generate_btn.click(
        fn=generate_recammaster_video,
        inputs=[video_input, text_prompt, camera_type, num_frames, resolution],
        outputs=[output_video, status_output]
    )

if __name__ == "__main__":
    model_loader.load_models()
    demo.launch(share=True)