MogensR commited on
Commit
a9f51ee
·
1 Parent(s): bd64f57

agent 1.6

Browse files
Files changed (1) hide show
  1. models/matanyone_loader.py +212 -40
models/matanyone_loader.py CHANGED
@@ -417,51 +417,223 @@ def process_stream(
417
  log.info(f"[MATANY] Processing {N} frames ({W}x{H} @ {fps:.1f}fps) from {video_path}")
418
  _emit_progress(progress_cb, 0.05, f"Processing {N} frames ({W}x{H} @ {fps:.1f}fps)")
419
 
420
- if self._api_mode == "process_video":
421
- # --- PATH-BASED CALL (this wheel expects a video path, not tensors) ---
422
- _emit_progress(progress_cb, 0.1, "Using MatAnyone video mode (GPU-accelerated)")
423
-
424
- # Log before starting video processing
425
- if torch.cuda.is_available():
426
- mem_alloc, _ = self._log_gpu_memory()
427
- _emit_progress(progress_cb, 0.12, f"GPU memory before processing: {mem_alloc:.1f}MB")
428
-
429
- # Some builds accept (video_path, seed_mask_path), others just (video_path)
430
- try:
431
- _emit_progress(progress_cb, 0.15, "Starting video processing with mask...")
432
- res = self._core.process_video(
433
- str(video_path),
434
- str(seed_mask_path) if seed_mask_path is not None else None
435
- )
436
- except TypeError as e:
437
- if "takes 2 positional arguments but 3 were given" in str(e):
438
- _emit_progress(progress_cb, 0.15, "Starting video processing without mask...")
439
- res = self._core.process_video(str(video_path))
440
- else:
441
- raise
442
 
443
- # Log after processing
444
  if torch.cuda.is_available():
445
- _emit_progress(progress_cb, 0.9, f"Processing complete. GPU memory used: {torch.cuda.memory_allocated()/1024**2:.1f}MB")
446
- else:
447
- _emit_progress(progress_cb, 0.9, "Processing complete.")
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
448
 
449
- # Normalize output files
450
- _emit_progress(progress_cb, 0.95, "Finalizing output files...")
451
- alpha_path, fg_path = self._harvest_process_video_output(res, out_dir, base=video_path.stem)
452
- _validate_nonempty(alpha_path)
453
- _validate_nonempty(fg_path)
454
 
455
- _emit_progress(progress_cb, 1.0, "Processing complete!")
456
- return alpha_path, fg_path
 
 
 
 
 
 
 
 
 
 
 
 
 
 
457
 
458
- except Exception as e:
459
- error_msg = f"Error during video processing: {str(e)}"
460
- log.error(error_msg, exc_info=True)
461
- if torch.cuda.is_available():
462
- error_msg += f"\nGPU Memory: {torch.cuda.memory_allocated()/1024**2:.1f}MB allocated"
463
- _emit_progress(progress_cb, -1, error_msg)
464
- raise MatAnyError(error_msg) from e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
465
  else:
466
  # Frame-by-frame (preferred)
467
  log.info(f"[MATANY] Using frame-by-frame mode: {self._api_mode}")
 
417
  log.info(f"[MATANY] Processing {N} frames ({W}x{H} @ {fps:.1f}fps) from {video_path}")
418
  _emit_progress(progress_cb, 0.05, f"Processing {N} frames ({W}x{H} @ {fps:.1f}fps)")
419
 
420
+ try:
421
+ if self._api_mode == "process_video":
422
+ # --- PATH-BASED CALL (this wheel expects a video path, not tensors) ---
423
+ _emit_progress(progress_cb, 0.1, "Using MatAnyone video mode (GPU-accelerated)")
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
424
 
425
+ # Log before starting video processing
426
  if torch.cuda.is_available():
427
+ mem_alloc, _ = self._log_gpu_memory()
428
+ _emit_progress(progress_cb, 0.12, f"GPU memory before processing: {mem_alloc:.1f}MB")
429
+
430
+ # Some builds accept (video_path, seed_mask_path), others just (video_path)
431
+ try:
432
+ _emit_progress(progress_cb, 0.15, "Starting video processing with mask...")
433
+ res = self._core.process_video(
434
+ str(video_path),
435
+ str(seed_mask_path) if seed_mask_path is not None else None
436
+ )
437
+ except TypeError as e:
438
+ if "takes 2 positional arguments but 3 were given" in str(e):
439
+ _emit_progress(progress_cb, 0.15, "Starting video processing without mask...")
440
+ res = self._core.process_video(str(video_path))
441
+ else:
442
+ raise
443
+
444
+ # Log after processing
445
+ if torch.cuda.is_available():
446
+ _emit_progress(progress_cb, 0.9, f"Processing complete. GPU memory used: {torch.cuda.memory_allocated()/1024**2:.1f}MB")
447
+ else:
448
+ _emit_progress(progress_cb, 0.9, "Processing complete.")
449
+
450
+ # Normalize output files
451
+ _emit_progress(progress_cb, 0.95, "Finalizing output files...")
452
+ alpha_path, fg_path = self._harvest_process_video_output(res, out_dir, base=video_path.stem)
453
+ _validate_nonempty(alpha_path)
454
+ _validate_nonempty(fg_path)
455
+
456
+ _emit_progress(progress_cb, 1.0, "Processing complete!")
457
+ return alpha_path, fg_path
458
+
459
+ else:
460
+ # Frame-by-frame (preferred)
461
+ log.info(f"[MATANY] Using frame-by-frame mode: {self._api_mode}")
462
+ _emit_progress(progress_cb, 0.1, f"Using {self._api_mode} mode (frame-by-frame)")
463
 
464
+ cap = cv2.VideoCapture(str(video_path))
465
+ alpha_path = out_dir / "alpha.mp4"
466
+ fg_path = out_dir / "fg.mp4"
 
 
467
 
468
+ # Initialize video writers
469
+ _emit_progress(progress_cb, 0.12, "Initializing video writers...")
470
+ alpha_writer = cv2.VideoWriter(
471
+ str(alpha_path),
472
+ cv2.VideoWriter_fourcc(*'mp4v'),
473
+ fps,
474
+ (W, H),
475
+ isColor=False
476
+ )
477
+ fg_writer = cv2.VideoWriter(
478
+ str(fg_path),
479
+ cv2.VideoWriter_fourcc(*'mp4v'),
480
+ fps,
481
+ (W, H),
482
+ isColor=True
483
+ )
484
 
485
+ if not alpha_writer.isOpened() or not fg_writer.isOpened():
486
+ raise MatAnyError("Failed to initialize video writers")
487
+
488
+ try:
489
+ # Load seed mask if provided
490
+ seed_1hw = None
491
+ if seed_mask_path is not None:
492
+ seed_1hw = _read_mask_hw(seed_mask_path, (H, W))
493
+
494
+ idx = 0
495
+ last_progress_update = 0
496
+ frame_times = []
497
+ start_time = time.time()
498
+
499
+ while True:
500
+ ret, frame = cap.read()
501
+ if not ret:
502
+ break
503
+
504
+ frame_start_time = time.time()
505
+
506
+ # Update progress more frequently (every 1% or 5 frames, whichever is more frequent)
507
+ current_progress = (idx / N) if N > 0 else 0.0
508
+ if idx % max(5, N//100) == 0 or time.time() - last_progress_update > 2.0:
509
+ # Calculate progress metrics
510
+ elapsed = time.time() - start_time
511
+ if idx > 0 and current_progress > 0:
512
+ # Calculate ETA
513
+ eta_seconds = (elapsed / current_progress) * (1 - current_progress)
514
+ if eta_seconds > 3600:
515
+ eta_str = f"{eta_seconds/3600:.1f} hours"
516
+ elif eta_seconds > 60:
517
+ eta_str = f"{eta_seconds/60:.1f} minutes"
518
+ else:
519
+ eta_str = f"{eta_seconds:.0f} seconds"
520
+
521
+ # Calculate processing speed
522
+ fps = idx / elapsed if elapsed > 0 else 0
523
+
524
+ # Add GPU memory info if available
525
+ gpu_info = ""
526
+ if torch.cuda.is_available():
527
+ mem_alloc = torch.cuda.memory_allocated() / 1024**2
528
+ mem_cached = torch.cuda.memory_reserved() / 1024**2
529
+ gpu_info = f" | GPU: {mem_alloc:.1f}/{mem_cached:.1f}MB"
530
+
531
+ status = (f"Processing frame {idx+1}/{N} (ETA: {eta_str}, "
532
+ f"{fps:.1f} FPS{gpu_info}")
533
+ _emit_progress(progress_cb, min(0.99, current_progress), status)
534
+ last_progress_update = time.time()
535
+
536
+ # Process frame
537
+ log.debug(f"[MATANY] Processing frame {idx+1}/{N}")
538
+ # Only pass seed mask on first frame
539
+ current_mask = seed_1hw if idx == 0 else None
540
+ alpha_hw = self._run_frame(frame, current_mask, is_first=(idx == 0))
541
+
542
+ # Calculate frame processing time
543
+ frame_time = time.time() - frame_start_time
544
+ frame_times.append(frame_time)
545
+ if len(frame_times) > 10: # Keep last 10 frame times for average
546
+ frame_times.pop(0)
547
+
548
+ # Log GPU memory usage occasionally
549
+ if idx % 50 == 0 and torch.cuda.is_available():
550
+ log.info(f"[GPU] Memory allocated: {torch.cuda.memory_allocated()/1024**2:.1f}MB, "
551
+ f"Cached: {torch.cuda.memory_reserved()/1024**2:.1f}MB, "
552
+ f"Avg frame time: {sum(frame_times)/len(frame_times)*1000:.1f}ms")
553
+
554
+ # Compose output frames
555
+ alpha_u8 = (alpha_hw * 255.0 + 0.5).astype(np.uint8)
556
+ alpha_rgb = cv2.cvtColor(alpha_u8, cv2.COLOR_GRAY2BGR)
557
+ fg_bgr = (frame.astype(np.float32) * (alpha_hw[..., None] / 255.0)).astype(np.uint8)
558
+
559
+ # Write outputs
560
+ alpha_writer.write(alpha_rgb)
561
+ fg_writer.write(fg_bgr)
562
+ idx += 1
563
+
564
+ except Exception as e:
565
+ # Log detailed error information
566
+ error_msg = f"Error processing frame {idx+1}/{N}: {str(e)}"
567
+ log.error(error_msg, exc_info=True)
568
+
569
+ # Add GPU memory info if available
570
+ if torch.cuda.is_available():
571
+ mem_alloc = torch.cuda.memory_allocated() / 1024**2
572
+ mem_cached = torch.cuda.memory_reserved() / 1024**2
573
+ error_msg += (f"\nGPU Memory - Allocated: {mem_alloc:.1f}MB, "
574
+ f"Cached: {mem_cached:.1f}MB")
575
+
576
+ # Add frame processing stats
577
+ if frame_times:
578
+ avg_time = sum(frame_times) / len(frame_times)
579
+ error_msg += f"\nAvg frame time: {avg_time*1000:.1f}ms"
580
+
581
+ _emit_progress(progress_cb, -1, f"ERROR: {error_msg}")
582
+ raise MatAnyError(error_msg) from e
583
+
584
+ finally:
585
+ # Cleanup resources
586
+ try:
587
+ if 'cap' in locals() and hasattr(cap, 'isOpened') and cap.isOpened():
588
+ cap.release()
589
+ if 'alpha_writer' in locals() and alpha_writer is not None:
590
+ if hasattr(alpha_writer, 'isOpened') and alpha_writer.isOpened():
591
+ alpha_writer.release()
592
+ if 'fg_writer' in locals() and fg_writer is not None:
593
+ if hasattr(fg_writer, 'isOpened') and fg_writer.isOpened():
594
+ fg_writer.release()
595
+
596
+ # Log final stats
597
+ total_time = time.time() - start_time
598
+ fps = idx / total_time if total_time > 0 else 0
599
+
600
+ # Log GPU memory info if available
601
+ gpu_info = ""
602
+ if torch.cuda.is_available():
603
+ mem_alloc = torch.cuda.memory_allocated() / 1024**2
604
+ mem_cached = torch.cuda.memory_reserved() / 1024**2
605
+ gpu_info = f"\nGPU Memory - Allocated: {mem_alloc:.1f}MB, Cached: {mem_cached:.1f}MB"
606
+
607
+ log.info(
608
+ f"[MATANY] Processed {idx} frames in {total_time:.1f}s ({fps:.1f} FPS){gpu_info}"
609
+ )
610
+
611
+ # Validate outputs
612
+ _validate_nonempty(alpha_path)
613
+ _validate_nonempty(fg_path)
614
+
615
+ # Final progress update
616
+ _emit_progress(
617
+ progress_cb,
618
+ 1.0,
619
+ f"Complete! Processed {idx} frames at {fps:.1f} FPS{gpu_info}"
620
+ )
621
+
622
+ return alpha_path, fg_path
623
+
624
+ except Exception as e:
625
+ error_msg = f"Error during cleanup: {str(e)}"
626
+ log.error(error_msg, exc_info=True)
627
+ _emit_progress(progress_cb, -1, f"CLEANUP ERROR: {error_msg}")
628
+ raise MatAnyError(error_msg) from e
629
+
630
+ except Exception as e:
631
+ error_msg = f"Error during video processing: {str(e)}"
632
+ log.error(error_msg, exc_info=True)
633
+ if torch.cuda.is_available():
634
+ error_msg += f"\nGPU Memory: {torch.cuda.memory_allocated()/1024**2:.1f}MB allocated"
635
+ _emit_progress(progress_cb, -1, error_msg)
636
+ raise MatAnyError(error_msg) from e
637
  else:
638
  # Frame-by-frame (preferred)
639
  log.info(f"[MATANY] Using frame-by-frame mode: {self._api_mode}")