Update app.py

app.py

@@ -1,285 +1,327 @@
+	
+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
-import numpy as np
-from PIL import Image, ImageFilter
-import cv2
-import os
+10	
 import torch
-import torch.nn.functional as F
-from torchvision import transforms
-import warnings
-from video_processing import process_video
+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	
 
-warnings.filterwarnings("ignore")
+18	
 
-# ZeroGPU decorator (if available)
-import spaces
-@spaces.GPU()
-# VAAPI acceleration check
-def check_vaapi_support():
-    """Check if VAAPI is available for hardware acceleration"""
-    try:
-        # Check if VAAPI devices are available
-        vaapi_devices = [f for f in os.listdir('/dev/dri') if f.startswith('render')]
-        return len(vaapi_devices) > 0
-    except:
-        return False
-HAS_VAAPI = check_vaapi_support()
+19	
 
-class TorchUpscaler:
-    """PyTorch-based upscaler that can use GPU acceleration"""
+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	
     
-    def __init__(self, device='auto'):
-        if device == 'auto':
-            self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-        else:
-            self.device = torch.device(device)
+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	
         
-        print(f"Using device: {self.device}")
-    
-    def bicubic_torch(self, image_tensor, scale_factor):
-        """GPU-accelerated bicubic upscaling using PyTorch"""
-        return F.interpolate(
-            image_tensor, 
-            scale_factor=scale_factor, 
-            mode='bicubic', 
-            align_corners=False,
-            antialias=True
-        )
-    
-    def lanczos_torch(self, image_tensor, scale_factor):
-        """GPU-accelerated Lanczos-style upscaling"""
-        return F.interpolate(
-            image_tensor,
-            scale_factor=scale_factor,
-            mode='bicubic',
-            align_corners=False,
-            antialias=True
-        )
-    
-    def esrgan_style_upscale(self, image_tensor, scale_factor):
-        """Simple ESRGAN-style upscaling using convolutions"""
-        b, c, h, w = image_tensor.shape
-        upscaled = F.interpolate(image_tensor, scale_factor=scale_factor, mode='bicubic', align_corners=False)
-        kernel = torch.tensor([[[[-1, -1, -1],
-                                [-1,  9, -1],
-                                [-1, -1, -1]]]], dtype=torch.float32, device=self.device)
-        kernel = kernel.repeat(c, 1, 1, 1)
-        sharpened = F.conv2d(upscaled, kernel, padding=1, groups=c)
-        result = 0.8 * upscaled + 0.2 * sharpened
-        return torch.clamp(result, 0, 1)
+60	
+        inputs = ImageLoader.load_image(pil_frame)
+61	
+        preds = model(inputs)
+62	
+        output_frame = ImageLoader.save_image(preds, mode='RGB').convert("RGB")
+63	
 
-class VAAPIUpscaler:
-    """VAAPI hardware-accelerated upscaler"""
-    
-    def __init__(self):
-        self.vaapi_available = HAS_VAAPI
-        if self.vaapi_available:
-            print("VAAPI hardware acceleration available")
-        else:
-            print("VAAPI hardware acceleration not available")
+64	
+        video_writer.write(cv2.cvtColor(np.array(output_frame), cv2.COLOR_RGB2BGR))
+65	
+
+66	
+    video_capture.release()
+67	
+    video_writer.release()
+68	
     
-    def upscale_vaapi(self, image_array, scale_factor, method):
-        """Use VAAPI for hardware-accelerated upscaling"""
-        if not self.vaapi_available:
-            return None
-        try:
-            h, w = image_array.shape[:2]
-            new_h, new_w = int(h * scale_factor), int(w * scale_factor)
-            if method == "VAAPI_BICUBIC":
-                return cv2.resize(image_array, (new_w, new_h), interpolation=cv2.INTER_CUBIC)
-            elif method == "VAAPI_LANCZOS":
-                return cv2.resize(image_array, (new_w, new_h), interpolation=cv2.INTER_LANCZOS4)
-        except Exception as e:
-            print(f"VAAPI upscaling failed: {e}")
-            return None
+69	
+    return output_path
+70	
 
-torch_upscaler = TorchUpscaler()
-vaapi_upscaler = VAAPIUpscaler()
+71	
+from RIFE import Model as RIFEModel
+72	
+from safetensors.torch import load_file
+73	
 
-@spaces.GPU if HAS_ZEROGPU else lambda x: x
-def upscale_image_accelerated(image, scale_factor, method, enhance_quality, use_gpu_acceleration):
-    if image is None:
-        return None
-    
-    original_width, original_height = image.size
-    new_width = int(original_width * scale_factor)
-    new_height = int(original_height * scale_factor)
-    
-    try:
-        if use_gpu_acceleration and torch.cuda.is_available():
-            print("Using GPU acceleration")
-            transform = transforms.Compose([transforms.ToTensor()])
-            image_tensor = transform(image).unsqueeze(0).to(torch_upscaler.device)
-            
-            if method == "GPU_Bicubic":
-                upscaled_tensor = torch_upscaler.bicubic_torch(image_tensor, scale_factor)
-            elif method == "GPU_Lanczos":
-                upscaled_tensor = torch_upscaler.lanczos_torch(image_tensor, scale_factor)
-            elif method == "GPU_ESRGAN_Style":
-                upscaled_tensor = torch_upscaler.esrgan_style_upscale(image_tensor, scale_factor)
-            else:
-                upscaled_tensor = torch_upscaler.bicubic_torch(image_tensor, scale_factor)
+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	
             
-            upscaled = transforms.ToPILImage()(upscaled_tensor.squeeze(0).cpu())
+115	
+            # Run inference
+116	
+            with torch.no_grad():
+117	
+                interpolated_frame = model.inference(img0, img1)[0].cpu().numpy().transpose(1, 2, 0) * 255
+118	
             
-        elif method.startswith("VAAPI_") and HAS_VAAPI:
-            print("Using VAAPI acceleration")
-            img_array = np.array(image)
-            upscaled_array = vaapi_upscaler.upscale_vaapi(img_array, scale_factor, method)
-            upscaled = Image.fromarray(upscaled_array) if upscaled_array is not None else image.resize((new_width, new_height), Image.BICUBIC)
-                
-        else:
-            print("Using CPU methods")
-            if method == "Bicubic":
-                upscaled = image.resize((new_width, new_height), Image.BICUBIC)
-            elif method == "Lanczos":
-                upscaled = image.resize((new_width, new_height), Image.LANCZOS)
-            else:
-                upscaled = image.resize((new_width, new_height), Image.BICUBIC)
-        
-        if enhance_quality:
-            upscaled = upscaled.filter(ImageFilter.UnsharpMask(radius=1, percent=120, threshold=3))
-        
-        return upscaled
-        
-    except Exception as e:
-        print(f"Error during upscaling: {e}")
-        return image
+119	
+            video_writer.write(interpolated_frame.astype(np.uint8))
+120	
 
-def get_available_methods():
-    methods = ["Bicubic", "Lanczos"]
-    if torch.cuda.is_available():
-        methods.extend(["GPU_Bicubic", "GPU_Lanczos", "GPU_ESRGAN_Style"])
-    if HAS_VAAPI:
-        methods.extend(["VAAPI_BICUBIC", "VAAPI_LANCZOS"])
-    return methods
+121	
+        video_writer.write(frame)
+122	
+        prev_frame = frame
+123	
 
-def get_system_info():
-    info = []
-    if torch.cuda.is_available():
-        gpu_name = torch.cuda.get_device_name(0)
-        gpu_memory = torch.cuda.get_device_properties(0).total_memory / 1024**3
-        info.append(f"🚀 CUDA GPU: {gpu_name} ({gpu_memory:.1f} GB)")
-    else:
-        info.append("❌ CUDA not available")
-    if HAS_ZEROGPU:
-        info.append("✅ ZeroGPU support enabled")
-    if HAS_VAAPI:
-        info.append("✅ VAAPI hardware acceleration available")
-    return "\n".join(info)
+124	
+    video_capture.release()
+125	
+    video_writer.release()
+126	
 
-def process_and_info_accelerated(image, scale_factor, method, enhance_quality, use_gpu_acceleration):
-    if image is None:
-        return None, "Please upload an image first"
-    
-    original_info = f"Original: {image.size[0]} × {image.size[1]} pixels"
-    result = upscale_image_accelerated(image, scale_factor, method, enhance_quality, use_gpu_acceleration)
-    if result is None:
-        return None, "Error processing image"
-    
-    result_info = f"Upscaled: {result.size[0]} × {result.size[1]} pixels"
-    accel_info = "GPU/Hardware" if use_gpu_acceleration else "CPU"
-    
-    combined_info = f"""
-    ## Processing Details
-    {original_info}  
-    {result_info}  
-    **Scale Factor:** {scale_factor}x  
-    **Method:** {method}  
-    **Acceleration:** {accel_info}  
-    **Quality Enhancement:** {'✅' if enhance_quality else '❌'}
-    
-    ## System Status
-    {get_system_info()}
-    """
-    return result, combined_info
+127	
+    return output_path
+128	
 
-def create_accelerated_upscaler_ui():
-    available_methods = get_available_methods()
-    
-    gr.Markdown("## 🚀 Accelerated Image Upscaler")
-    with gr.Row():
-        with gr.Column(scale=1):
-            input_image = gr.Image(type="pil", label="Upload Image", sources=["upload", "clipboard"])
-            scale_factor = gr.Slider(minimum=1.5, maximum=4.0, step=0.5, value=2.0, label="Scale Factor")
-            method = gr.Dropdown(choices=available_methods, value=available_methods[0], label="Upscaling Method")
-            use_gpu_acceleration = gr.Checkbox(label="Use GPU Acceleration", value=torch.cuda.is_available())
-            enhance_quality = gr.Checkbox(label="Apply Quality Enhancement", value=True)
-            process_btn = gr.Button("🚀 Upscale Image", variant="primary")
-        
-        with gr.Column(scale=2):
-            output_image = gr.Image(label="Upscaled Image", type="pil")
-            image_info = gr.Markdown(value=f"## System Status\n{get_system_info()}", label="Processing Information")
-    
-    process_btn.click(
-        fn=process_and_info_accelerated,
-        inputs=[input_image, scale_factor, method, enhance_quality, use_gpu_acceleration],
-        outputs=[output_image, image_info]
-    )
+129	
 
-def create_video_interface_ui():
-    gr.Markdown("## 🚀 Video Upscaler and Frame Interpolator")
-    with gr.Row():
-        with gr.Column(scale=1):
-            input_video = gr.Video(label="Upload Video", sources=["upload"])
-            scale_factor = gr.Slider(minimum=1.5, maximum=4.0, step=0.5, value=2.0, label="Scale Factor")
-            multi = gr.Slider(minimum=2, maximum=8, step=1, value=2, label="Frame Multiplier")
-            use_gpu_acceleration = gr.Checkbox(label="Use GPU Acceleration", value=torch.cuda.is_available())
-            process_btn = gr.Button("🚀 Process Video", variant="primary")
-        
-        with gr.Column(scale=2):
-            output_video = gr.Video(label="Processed Video")
-            processing_info = gr.Markdown(value=f"## System Status\n{get_system_info()}", label="Processing Information")
-    
-    process_btn.click(
-        fn=process_video_wrapper,
-        inputs=[input_video, scale_factor, multi, use_gpu_acceleration],
-        outputs=[output_video, processing_info]
-    )
+130	
 
-def process_video_wrapper(video_path, scale_factor, multi, use_gpu):
-    if video_path is None:
-        return None, "Please upload a video first"
-    
-    output_path = "temp_output.mp4"
-    modelDir = 'rife/train_log'
-    
-    processed_video_path = process_video(
-        video=video_path, 
-        output=output_path, 
-        modelDir=modelDir, 
-        fp16=use_gpu, 
-        UHD=False, 
-        scale=scale_factor, 
-        skip=False, 
-        fps=None, 
-        png=False, 
-        ext='mp4', 
-        exp=1, 
-        multi=multi
+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	
     
-    info = f"""
-    ## Processing Details
-    **Scale Factor:** {scale_factor}x  
-    **Frame Multiplier:** {multi}x
-    **Acceleration:** {'GPU' if use_gpu else 'CPU'}
-    
-    ## System Status
-    {get_system_info()}
-    """
-    return processed_video_path, info
-
-with gr.Blocks(title="Accelerated Media Processor", theme=gr.themes.Soft()) as demo:
-    with gr.Tab("Image Upscaler"):
-        create_accelerated_upscaler_ui()
-    with gr.Tab("Video Processing"):
-        create_video_interface_ui()
+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__":
-    demo.launch(
-        server_name="0.0.0.0",
-        server_port=7860,
-        share=False,
-        debug=True
-    )
+163	
+    demo.launch(share=True)
+164