agent 1.8

models/matanyone_loader.py

@@ -821,71 +821,56 @@ def _flush_chunk(self, frames_bgr, seed_1hw, alpha_writer, fg_writer):
         if not frames_bgr:
             return
 
-        try:
-            # Prepare inputs
-            frames_chw = [_to_chw01(f) for f in frames_bgr]  # list of CHW
-            frames_t = torch.from_numpy(np.stack(frames_chw)).to(self.device)  # T,C,H,W
-            mask_t = None
-            if seed_1hw is not None:
-                mask_t = torch.from_numpy(seed_1hw).to(self.device)
+        # Prepare inputs
+        frames_chw = [_to_chw01(f) for f in frames_bgr]  # list of CHW
+        frames_t = torch.from_numpy(np.stack(frames_chw)).to(self.device)  # T,C,H,W
+        mask_t = None
+        if seed_1hw is not None:
+            mask_t = torch.from_numpy(seed_1hw).to(self.device)
 
-            try:
-                with torch.no_grad(), self._maybe_amp():
-                    # Process frames in batch
-                    if self._api_mode == "process_frame":
-                        alphas = []
-                        for i in range(len(frames_t)):
-                            # Only use mask on first frame if provided
-                            current_mask = mask_t if (i == 0 and mask_t is not None) else None
-                            alpha = self._core.process_frame(frames_t[i].unsqueeze(0), current_mask)
-                            alphas.append(alpha.squeeze(0))
-                        alphas = torch.stack(alphas)
-                    elif hasattr(self._core, '_process_tensor_video'):
-                        # Try direct tensor processing (newer versions)
-                        alphas = self._core._process_tensor_video(frames_t, mask_t)
-                    else:  # step mode
-                        alphas = self._core.step(frames_t, mask_t)
-
-                    # Convert to numpy and write frames
-                    alphas_np = alphas.cpu().numpy()
-                    for i, alpha in enumerate(alphas_np):
-                        # Convert alpha to uint8 and write
-                        alpha_uint8 = (alpha * 255).astype(np.uint8)
-                        if len(alpha_uint8.shape) == 2:  # If single channel, convert to 3 channels
-                            alpha_uint8 = cv2.cvtColor(alpha_uint8, cv2.COLOR_GRAY2BGR)
-                        alpha_writer.write(alpha_uint8)
-                        
-                        # Write foreground (frame * alpha)
-                        fg = frames_bgr[i] * (alpha[..., None] if alpha.ndim == 2 else alpha[0:1].permute(1, 2, 0))
-                        fg = fg.astype(np.uint8)
-                        fg_writer.write(fg)
+        try:
+            with torch.no_grad(), self._maybe_amp():
+                # Process frames in batch
+                if self._api_mode == "process_frame":
+                    alphas = []
+                    for i in range(len(frames_t)):
+                        # Only use mask on first frame if provided
+                        current_mask = mask_t if (i == 0 and mask_t is not None) else None
+                        alpha = self._core.process_frame(frames_t[i].unsqueeze(0), current_mask)
+                        alphas.append(alpha.squeeze(0))
+                    alphas = torch.stack(alphas)
+                elif hasattr(self._core, '_process_tensor_video'):
+                    # Try direct tensor processing (newer versions)
+                    alphas = self._core._process_tensor_video(frames_t, mask_t)
+                else:  # step mode
+                    alphas = self._core.step(frames_t, mask_t)
+
+                # Convert to numpy and write frames
+                alphas_np = alphas.cpu().numpy()
+                for i, alpha in enumerate(alphas_np):
+                    # Convert alpha to uint8 and write
+                    alpha_uint8 = (alpha * 255).astype(np.uint8)
+                    if len(alpha_uint8.shape) == 2:  # If single channel, convert to 3 channels
+                        alpha_uint8 = cv2.cvtColor(alpha_uint8, cv2.COLOR_GRAY2BGR)
+                    alpha_writer.write(alpha_uint8)
+                    
+                    # Write foreground (frame * alpha)
+                    fg = frames_bgr[i] * (alpha[..., None] if alpha.ndim == 2 else alpha[0:1].permute(1, 2, 0))
+                    fg = fg.astype(np.uint8)
+                    fg_writer.write(fg)
+            
+        except RuntimeError as e:
+            if "out of memory" in str(e).lower():
+                # Clear CUDA cache and retry once
+                torch.cuda.empty_cache()
+                log.warning("CUDA out of memory, retrying after cache clear")
+                return self._flush_chunk(frames_bgr, seed_1hw, alpha_writer, fg_writer)
+            raise
             
-            except RuntimeError as e:
-                if "out of memory" in str(e).lower():
-                    # Clear CUDA cache and retry once
-                    torch.cuda.empty_cache()
-                    log.warning("CUDA out of memory, retrying after cache clear")
-                    return self._flush_chunk(frames_bgr, seed_1hw, alpha_writer, fg_writer)
-                raise
-                
         except Exception as e:
             error_msg = f"Error processing frame chunk: {str(e)}"
             log.error(error_msg, exc_info=True)
             raise MatAnyError(error_msg) from e
-                            cv2.imwrite(path, frame)
-                            frame_paths.append(path)
-                        
-                        # Process video from frames
-                        alphas = self._core.process_video(tmpdir, 
-                                                         mask_path=seed_1hw_path if seed_1hw is not None else None)
-                        
-                        # Ensure alphas is a tensor
-                        if not isinstance(alphas, torch.Tensor):
-                            alphas = torch.from_numpy(alphas).to(self.device)
-                            
-            except Exception as e:
-                log.error(f"Error in _flush_chunk: {str(e)}")
-                raise
 
         # Normalize to numpy list of HW float32 [0,1]
         if isinstance(alphas, torch.Tensor):