Update app.py

app.py

@@ -11,133 +11,96 @@ 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
 
-# Enable tensor cores for faster computation
-torch.set_float32_matmul_precision("high")
-torch.backends.cudnn.benchmark = True  # Enable cudnn autotuner
+torch.set_float32_matmul_precision("highest")
 
-# Initialize model with optimization flags
 birefnet = AutoModelForImageSegmentation.from_pretrained(
     "ZhengPeng7/BiRefNet", trust_remote_code=True
+).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]),
+    ]
 )
-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:
-        # Load video more efficiently
-        video = mp.VideoFileClip(vid, audio_buffersize=2000)
+        video = mp.VideoFileClip(vid)
         if fps == 0:
             fps = video.fps
         audio = video.audio
-        
-        # 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
+        frames = video.iter_frames(fps=fps)
+        processed_frames = []
+        yield gr.update(visible=True), gr.update(visible=False)
+
         if bg_type == "Video":
-            bg_video_clip = mp.VideoFileClip(bg_video)
-            if bg_video_clip.duration < video.duration:
+            background_video = mp.VideoFileClip(bg_video)
+            if background_video.duration < video.duration:
                 if video_handling == "slow_down":
-                    bg_video_clip = bg_video_clip.fx(mp.vfx.speedx, 
-                                                   factor=video.duration / bg_video_clip.duration)
+                    background_video = background_video.fx(mp.vfx.speedx, factor=video.duration / background_video.duration)
                 else:
-                    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":
-                    bg = color
-                elif bg_type == "Image":
-                    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
+                    background_video = mp.concatenate_videoclips([background_video] * int(video.duration / background_video.duration + 1))
+            background_frames = list(background_video.iter_frames(fps=fps))
+        elif bg_type in ["Color", "Image"]:
+             # Prepare background once if it's a static image or color
+            if bg_type == "Color":
+                color_rgb = tuple(int(color[i:i+2], 16) for i in (1, 3, 5))
+                background_pil = Image.new("RGBA", (1024, 1024), color_rgb + (255,))
+            else: # bg_type == "Image":
+                background_pil = Image.open(bg_image).convert("RGBA").resize((1024, 1024))
+            background_tensor = transforms.ToTensor(background_pil).to("cuda")
+        else:
+            background_tensor = None
+
+
+        bg_frame_index = 0
+        frame_batch = []
+        for i, frame in enumerate(frames):
+            frame = Image.fromarray(frame)
+            frame = transforms.ToTensor(frame).to('cuda')
+            frame_batch.append(frame)
+
+            if len(frame_batch) >= 3 or i == int(video.fps * video.duration) - 1 :
+                input_images = torch.stack(frame_batch).to("cuda")
+                with torch.no_grad():
+                    preds = birefnet(input_images)[-1].sigmoid()
+                for j, pred in enumerate(preds):
+                    if bg_type == "Video":
+                        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).resize((1024, 1024))
+                            background_tensor = transforms.ToTensor(background_image).to("cuda")
+                        else:  # video_handling == "loop"
+                            background_frame = background_frames[bg_frame_index % len(background_frames)]
+                            bg_frame_index += 1
+                            background_image = Image.fromarray(background_frame).resize((1024, 1024))
+                            background_tensor = transforms.ToTensor(background_image).to("cuda")
+                    mask = transforms.ToPILImage()(pred.cpu().squeeze())
+                    processed_image = Image.composite(transforms.ToPILImage()(frame_batch[j].cpu()), transforms.ToPILImage()(background_tensor.cpu()), mask).resize(video.size)
+
+                    processed_frames.append(np.array(processed_image))
+                    yield processed_image, None
+
+                frame_batch = []
+
+
         processed_video = mp.ImageSequenceClip(processed_frames, fps=fps)
-        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)
-        
+        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)
+
         yield gr.update(visible=False), gr.update(visible=True)
-        yield processed_frames[-1], output_path
-        
+        yield processed_image, temp_filepath
+
     except Exception as e:
         print(f"Error: {e}")
         yield gr.update(visible=False), gr.update(visible=True)