Update app.py

app.py

@@ -1,327 +1,163 @@
-	
 try:
-2	
     import spaces
-3	
 except ImportError:
-4	
     # Create a dummy decorator if spaces is not available
-5	
     def spaces_gpu(func):
-6	
         return func
-7	
     spaces = type('spaces', (), {'GPU': spaces_gpu})()
-8	
 
-9	
 import gradio as gr
-10	
 import torch
-11	
 from torchvision.transforms import functional as F
-12	
 from PIL import Image
-13	
 import os
-14	
 import cv2
-15	
 import numpy as np
-16	
 from super_image import EdsrModel, ImageLoader
-17	
 
-18	
 
-19	
 
-20	
 @spaces.GPU
-21	
 def upscale_video(video_path, scale_factor, progress=gr.Progress()):
-22	
     """
-23	
     Upscales a video using EDSR model.
-24	
     This function is decorated with @spaces.GPU to run on ZeroGPU.
-25	
     """
-26	
     # Load models inside the function for ZeroGPU compatibility
-27	
     if scale_factor == 2:
-28	
         model = EdsrModel.from_pretrained('eugenesiow/edsr-base', scale=2)
-29	
     elif scale_factor == 4:
-30	
         model = EdsrModel.from_pretrained('eugenesiow/edsr-base', scale=4)
-31	
     else:
-32	
         raise gr.Error("Invalid scale factor. Choose 2 or 4.")
-33	
 
-34	
     if not os.path.exists(video_path):
-35	
         raise gr.Error(f"Input file not found at {video_path}")
-36	
 
-37	
     video_capture = cv2.VideoCapture(video_path)
-38	
     if not video_capture.isOpened():
-39	
         raise gr.Error(f"Could not open video file {video_path}")
-40	
 
-41	
     fourcc = cv2.VideoWriter_fourcc(*'mp4v')
-42	
     fps = video_capture.get(cv2.CAP_PROP_FPS)
-43	
     width = int(video_capture.get(cv2.CAP_PROP_FRAME_WIDTH))
-44	
     height = int(video_capture.get(cv2.CAP_PROP_FRAME_HEIGHT))
-45	
     frame_count = int(video_capture.get(cv2.CAP_PROP_FRAME_COUNT))
-46	
 
-47	
     output_width = width * scale_factor
-48	
     output_height = height * scale_factor
-49	
     
-50	
     output_path = f"upscaled_{scale_factor}x_{os.path.basename(video_path)}"
-51	
     video_writer = cv2.VideoWriter(output_path, fourcc, fps, (output_width, output_height))
-52	
 
-53	
     for i in progress.tqdm(range(frame_count), desc=f"Upscaling {scale_factor}x"):
-54	
         ret, frame = video_capture.read()
-55	
         if not ret:
-56	
             break
-57	
 
-58	
         pil_frame = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
-59	
         
-60	
         inputs = ImageLoader.load_image(pil_frame)
-61	
         preds = model(inputs)
-62	
         output_frame = ImageLoader.save_image(preds, mode='RGB').convert("RGB")
-63	
 
-64	
         video_writer.write(cv2.cvtColor(np.array(output_frame), cv2.COLOR_RGB2BGR))
-65	
 
-66	
     video_capture.release()
-67	
     video_writer.release()
-68	
     
-69	
     return output_path
-70	
 
-71	
 from RIFE import Model as RIFEModel
-72	
 from safetensors.torch import load_file
-73	
 
-74	
 # ... (existing code)
-75	
 
-76	
 @spaces.GPU
-77	
 def rife_interpolate_video(video_path, progress=gr.Progress()):
-78	
     """
-79	
     Interpolates a video using the RIFE model.
-80	
     This function is decorated with @spaces.GPU to run on ZeroGPU.
-81	
     """
-82	
     if not os.path.exists(video_path):
-83	
         raise gr.Error(f"Input file not found at {video_path}")
-84	
 
-85	
     # Load the RIFE model
-86	
     model = RIFEModel()
-87	
     model.load_state_dict(load_file("/Users/craigellenwood/Workspace/video_upscaler_rife_interpolator/rife_model_new/rife-flownet-4.13.2.safetensors"))
-88	
     model.eval()
-89	
     model.cuda()
-90	
 
-91	
     video_capture = cv2.VideoCapture(video_path)
-92	
     if not video_capture.isOpened():
-93	
         raise gr.Error(f"Could not open video file {video_path}")
-94	
 
-95	
     fourcc = cv2.VideoWriter_fourcc(*'mp4v')
-96	
     fps = video_capture.get(cv2.CAP_PROP_FPS)
-97	
     width = int(video_capture.get(cv2.CAP_PROP_FRAME_WIDTH))
-98	
     height = int(video_capture.get(cv2.CAP_PROP_FRAME_HEIGHT))
-99	
     frame_count = int(video_capture.get(cv2.CAP_PROP_FRAME_COUNT))
-100	
 
-101	
     output_path = f"interpolated_{os.path.basename(video_path)}"
-102	
     video_writer = cv2.VideoWriter(output_path, fourcc, fps * 2, (width, height))
-103	
 
-104	
     prev_frame = None
-105	
     for i in progress.tqdm(range(frame_count), desc="Interpolating"):
-106	
         ret, frame = video_capture.read()
-107	
         if not ret:
-108	
             break
-109	
 
-110	
         if prev_frame is not None:
-111	
             # Preprocess frames
-112	
             img0 = torch.from_numpy(prev_frame.transpose(2, 0, 1)).float().unsqueeze(0).cuda() / 255.
-113	
             img1 = torch.from_numpy(frame.transpose(2, 0, 1)).float().unsqueeze(0).cuda() / 255.
-114	
             
-115	
             # Run inference
-116	
             with torch.no_grad():
-117	
                 interpolated_frame = model.inference(img0, img1)[0].cpu().numpy().transpose(1, 2, 0) * 255
-118	
             
-119	
             video_writer.write(interpolated_frame.astype(np.uint8))
-120	
 
-121	
         video_writer.write(frame)
-122	
         prev_frame = frame
-123	
 
-124	
     video_capture.release()
-125	
     video_writer.release()
-126	
 
-127	
     return output_path
-128	
 
-129	
 
-130	
 
-131	
 
-132	
 with gr.Blocks() as demo:
-133	
     gr.Markdown("# Video Upscaler and Frame Interpolator")
-134	
     with gr.Tab("Upscale"):
-135	
         with gr.Row():
-136	
             with gr.Column():
-137	
                 video_input_upscale = gr.Video(label="Input Video")
-138	
                 scale_factor = gr.Radio([2, 4], label="Scale Factor", value=2)
-139	
                 upscale_button = gr.Button("Upscale Video")
-140	
             with gr.Column():
-141	
                 video_output_upscale = gr.Video(label="Upscaled Video")
-142	
     with gr.Tab("Interpolate"):
-143	
         with gr.Row():
-144	
             with gr.Column():
-145	
                 video_input_rife = gr.Video(label="Input Video")
-146	
                 rife_button = gr.Button("Interpolate Frames")
-147	
             with gr.Column():
-148	
                 video_output_rife = gr.Video(label="Interpolated Video")
-149	
 
-150	
     upscale_button.click(
-151	
         fn=upscale_video,
-152	
         inputs=[video_input_upscale, scale_factor],
-153	
         outputs=video_output_upscale
-154	
     )
-155	
     
-156	
     rife_button.click(
-157	
         fn=rife_interpolate_video,
-158	
         inputs=[video_input_rife],
-159	
         outputs=video_output_rife
-160	
     )
-161	
 
-162	
 if __name__ == "__main__":
-163	
     demo.launch(share=True)
-164