Test

app.py

@@ -11,125 +11,139 @@ import numpy as np
 import os
 import tempfile
 import uuid
+from concurrent.futures import ThreadPoolExecutor
+import torch.nn as nn
+import torch.cuda.amp  # for mixed precision training
 
-torch.set_float32_matmul_precision(["high", "highest"][0])
+# Enable tensor cores for faster computation
+torch.set_float32_matmul_precision("high")
+torch.backends.cudnn.benchmark = True  # Enable cudnn autotuner
 
+# Initialize model with optimization flags
 birefnet = AutoModelForImageSegmentation.from_pretrained(
     "ZhengPeng7/BiRefNet", trust_remote_code=True
 )
-birefnet.to("cuda")
-transform_image = transforms.Compose(
-    [
-        transforms.Resize((1024, 1024)),
-        transforms.ToTensor(),
-        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
-    ]
-)
-
-BATCH_SIZE = 3
+birefnet.to("cuda").eval()  # Ensure model is in eval mode
+birefnet = torch.jit.script(birefnet)  # JIT compilation for faster inference
+
+# Pre-compile transforms for better performance
+transform_image = transforms.Compose([
+    transforms.Resize((1024, 1024), antialias=True),
+    transforms.ToTensor(),
+    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+])
+
+# Increased batch size for better GPU utilization
+BATCH_SIZE = 8  # Increased from 3
+NUM_WORKERS = 4  # For parallel processing
+
+# Create a thread pool for parallel processing
+executor = ThreadPoolExecutor(max_workers=NUM_WORKERS)
+
+def process_batch(batch_data):
+    """Process a batch of frames in parallel"""
+    images, backgrounds, image_sizes = zip(*batch_data)
+    
+    # Stack images for batch processing
+    input_tensor = torch.stack(images).to("cuda")
+    
+    # Use automatic mixed precision for faster computation
+    with torch.cuda.amp.autocast():
+        with torch.no_grad():
+            preds = birefnet(input_tensor)[-1].sigmoid().cpu()
+    
+    processed_frames = []
+    for pred, bg, size in zip(preds, backgrounds, image_sizes):
+        mask = transforms.ToPILImage()(pred.squeeze()).resize(size)
+        
+        if isinstance(bg, str) and bg.startswith("#"):
+            color_rgb = tuple(int(bg[i:i+2], 16) for i in (1, 3, 5))
+            background = Image.new("RGBA", size, color_rgb + (255,))
+        elif isinstance(bg, Image.Image):
+            background = bg.convert("RGBA").resize(size)
+        else:
+            background = Image.open(bg).convert("RGBA").resize(size)
+        
+        # Use PIL's faster composite operation
+        image = Image.composite(images[0].resize(size), background, mask)
+        processed_frames.append(np.array(image))
+    
+    return processed_frames
 
 @spaces.GPU
 def fn(vid, bg_type="Color", bg_image=None, bg_video=None, color="#00FF00", fps=0, video_handling="slow_down"):
     try:
-        video = mp.VideoFileClip(vid)
+        # Load video more efficiently
+        video = mp.VideoFileClip(vid, audio_buffersize=2000)
         if fps == 0:
             fps = video.fps
         audio = video.audio
-        frames = video.iter_frames(fps=fps)
-        processed_frames = []
-        yield gr.update(visible=True), gr.update(visible=False)
-
+        
+        # Pre-calculate video parameters
+        total_frames = int(video.fps * video.duration)
+        frames = list(video.iter_frames(fps=fps))  # Load all frames at once
+        
+        # Pre-process background if using video
         if bg_type == "Video":
-            background_video = mp.VideoFileClip(bg_video)
-            if background_video.duration < video.duration:
+            bg_video_clip = mp.VideoFileClip(bg_video)
+            if bg_video_clip.duration < video.duration:
                 if video_handling == "slow_down":
-                    background_video = background_video.fx(mp.vfx.speedx, factor=video.duration / background_video.duration)
+                    bg_video_clip = bg_video_clip.fx(mp.vfx.speedx, 
+                                                   factor=video.duration / bg_video_clip.duration)
                 else:
-                    background_video = mp.concatenate_videoclips([background_video] * int(video.duration / background_video.duration + 1))
-            background_frames = list(background_video.iter_frames(fps=fps))
-        else:
-            background_frames = None
-
-        bg_frame_index = 0
-        frame_batch = []
-
-        for i, frame in enumerate(frames):
-            frame_batch.append(frame)
-            if len(frame_batch) == BATCH_SIZE or i == video.fps * video.duration -1:  # Process batch or last frames
-                pil_images = [Image.fromarray(f) for f in frame_batch]
-
-
+                    multiplier = int(video.duration / bg_video_clip.duration + 1)
+                    bg_video_clip = mp.concatenate_videoclips([bg_video_clip] * multiplier)
+            background_frames = list(bg_video_clip.iter_frames(fps=fps))
+        
+        # Process frames in batches
+        processed_frames = []
+        for i in range(0, len(frames), BATCH_SIZE):
+            batch_frames = frames[i:i + BATCH_SIZE]
+            batch_data = []
+            
+            for j, frame in enumerate(batch_frames):
+                pil_image = Image.fromarray(frame)
+                image_size = pil_image.size
+                transformed_image = transform_image(pil_image)
+                
                 if bg_type == "Color":
-                    processed_images = [process(img, color) for img in pil_images]
+                    bg = color
                 elif bg_type == "Image":
-                    processed_images = [process(img, bg_image) for img in pil_images]
-                elif bg_type == "Video":
-                    processed_images = []
-                    for _ in range(len(frame_batch)):
-                        if video_handling == "slow_down":
-                            background_frame = background_frames[bg_frame_index % len(background_frames)]
-                            bg_frame_index += 1
-                            background_image = Image.fromarray(background_frame)
-                        else:  # video_handling == "loop"
-                            background_frame = background_frames[bg_frame_index % len(background_frames)]
-                            bg_frame_index += 1
-                            background_image = Image.fromarray(background_frame)
-
-                        processed_images.append(process(pil_images[_],background_image))
-
-
-                else:
-                    processed_images = pil_images
-
-                for processed_image in processed_images:
-                    processed_frames.append(np.array(processed_image))
-                    yield processed_image, None
-                frame_batch = []  # Clear the batch
-
-
+                    bg = bg_image
+                else:  # Video
+                    frame_idx = (i + j) % len(background_frames)
+                    bg = Image.fromarray(background_frames[frame_idx])
+                
+                batch_data.append((transformed_image, bg, image_size))
+            
+            # Process batch
+            batch_results = process_batch(batch_data)
+            processed_frames.extend(batch_results)
+            
+            # Yield progress updates
+            if len(batch_results) > 0:
+                yield batch_results[-1], None
+        
+        # Create output video
         processed_video = mp.ImageSequenceClip(processed_frames, fps=fps)
-        processed_video = processed_video.set_audio(audio)
-
-        temp_dir = "temp"
-        os.makedirs(temp_dir, exist_ok=True)
-        unique_filename = str(uuid.uuid4()) + ".mp4"
-        temp_filepath = os.path.join(temp_dir, unique_filename)
-
-        processed_video.write_videofile(temp_filepath, codec="libx264", logger=None)
-
+        if audio is not None:
+            processed_video = processed_video.set_audio(audio)
+        
+        # Use temporary file
+        with tempfile.NamedTemporaryFile(suffix='.mp4', delete=False) as tmp_file:
+            output_path = tmp_file.name
+            processed_video.write_videofile(output_path, codec="libx264", 
+                                         preset='ultrafast', threads=NUM_WORKERS)
+        
         yield gr.update(visible=False), gr.update(visible=True)
-        yield processed_image, temp_filepath
-
+        yield processed_frames[-1], output_path
+        
     except Exception as e:
         print(f"Error: {e}")
         yield gr.update(visible=False), gr.update(visible=True)
         yield None, f"Error processing video: {e}"
 
 
-def process(image, bg):
-    image_size = image.size
-    input_images = transform_image(image).unsqueeze(0).to("cuda")
-    # Prediction
-    with torch.no_grad():
-        preds = birefnet(input_images)[-1].sigmoid().cpu()
-    pred = preds[0].squeeze()
-    pred_pil = transforms.ToPILImage()(pred)
-    mask = pred_pil.resize(image_size)
-
-    if isinstance(bg, str) and bg.startswith("#"):
-        color_rgb = tuple(int(bg[i:i+2], 16) for i in (1, 3, 5))
-        background = Image.new("RGBA", image_size, color_rgb + (255,))
-    elif isinstance(bg, Image.Image):
-        background = bg.convert("RGBA").resize(image_size)
-    else:
-        background = Image.open(bg).convert("RGBA").resize(image_size)
-
-    # Composite the image onto the background using the mask
-    image = Image.composite(image, background, mask)
-
-    return image
-
-
 with gr.Blocks(theme=gr.themes.Ocean()) as demo:
     with gr.Row():
         in_video = gr.Video(label="Input Video", interactive=True)