fuckit#

models/matany_compat_patch.py

@@ -1,19 +1,21 @@
 #!/usr/bin/env python3
 # MatAnyone HF-compat patch: squeeze time dim T=1 before first Conv2d
 # Changes (2025-09-16):
+# - Added fallback patching for forward/encode if encode_img missing
+# - Log dir(MatAnyone) and module version for debugging
 # - Added isinstance(img, torch.Tensor) for non-tensor safety
 # - Enhanced logging with input/output shapes
-# - Added version check logging
 # - Kept monkey-patch for HF Spaces compatibility
 
 import logging
 import torch
+import importlib.metadata
 
 log = logging.getLogger("backgroundfx_pro")
 
 def apply_matany_t1_squeeze_guard() -> bool:
     """
-    Monkey-patch MatAnyone.encode_img to squeeze [B,1,C,H,W] → [B,C,H,W].
+    Monkey-patch MatAnyone.encode_img (or forward/encode) to squeeze [B,1,C,H,W] → [B,C,H,W].
     Safe for multi-frame (T>1) as it only squeezes when T==1.
     Returns True if patch applied successfully, False otherwise.
     """
@@ -26,27 +28,42 @@ def apply_matany_t1_squeeze_guard() -> bool:
         log.warning("[MatAnyCompat] MatAnyone class not found")
         return False
     MatAnyone = M.MatAnyone
-    if not hasattr(MatAnyone, "encode_img"):
-        log.warning("[MatAnyCompat] encode_img not found on MatAnyone")
+
+    # Log MatAnyone version and attributes for debugging
+    try:
+        version = importlib.metadata.version("matanyone")
+        log.info(f"[MatAnyCompat] MatAnyone version: {version}")
+    except Exception:
+        log.info("[MatAnyCompat] MatAnyone version unknown")
+    log.debug(f"[MatAnyCompat] MatAnyone attributes: {dir(MatAnyone)}")
+
+    # Try encode_img first, then fallback to forward or encode
+    method_name = None
+    for candidate in ["encode_img", "forward", "encode"]:
+        if hasattr(MatAnyone, candidate):
+            method_name = candidate
+            break
+    if not method_name:
+        log.warning("[MatAnyCompat] No patchable method (encode_img, forward, encode) found on MatAnyone")
         return False
-    if getattr(MatAnyone, "_encode_img_patched", False):
-        log.info("[MatAnyCompat] encode_img already patched")
+    if getattr(MatAnyone, f"_{method_name}_patched", False):
+        log.info(f"[MatAnyCompat] {method_name} already patched")
         return True
 
     # Store original method
-    orig_encode_img = MatAnyone.encode_img
+    orig_method = getattr(MatAnyone, method_name)
 
-    def encode_img_compat(self, img, *args, **kwargs):
+    def method_compat(self, img, *args, **kwargs):
         # Handle inputs that MatAnyone.step turned into [B,1,C,H,W]
         try:
             if isinstance(img, torch.Tensor) and img.dim() == 5 and img.shape[1] == 1:
-                log.info(f"[MatAnyCompat] Squeezing 5D {img.shape} to 4D {img.squeeze(1).shape}")
+                log.info(f"[MatAnyCompat] Squeezing 5D {img.shape} to 4D {img.squeeze(1).shape} in {method_name}")
                 img = img.squeeze(1)  # [B,1,C,H,W] → [B,C,H,W]
         except Exception as e:
-            log.warning("[MatAnyCompat] Failed to process input shape: %s", e)
-        return orig_encode_img(self, img, *args, **kwargs)
+            log.warning(f"[MatAnyCompat] Failed to process input shape in {method_name}: %s", e)
+        return orig_method(self, img, *args, **kwargs)
 
-    MatAnyone.encode_img = encode_img_compat
-    MatAnyone._encode_img_patched = True
-    log.info("[MatAnyCompat] Applied T=1 squeeze guard in MatAnyone.encode_img")
+    setattr(MatAnyone, method_name, method_compat)
+    setattr(MatAnyone, f"_{method_name}_patched", True)
+    log.info(f"[MatAnyCompat] Applied T=1 squeeze guard in MatAnyone.{method_name}")
     return True

models/matanyone_loader.py

@@ -5,15 +5,16 @@
 
 - SAM2 defines the subject (seed mask) on frame 0.
 - MatAnyone does frame-by-frame alpha matting.
-- Prefers step([B,C,H,W]) with T=1 squeeze patch for conv2d compatibility.
-- Falls back to process_frame([H,W,3]) if supported by the wheel.
+- Uses T=1 squeeze patch for conv2d compatibility.
+- Falls back to process_frame([H,W,3]) if step() is unavailable.
 
 Changes (2025-09-16):
-- Added T=1 squeeze patch status logging in __init__
-- Set MATANY_FORCE_FORMAT=4d as default since patch ensures 4D compatibility
-- Added optional VRAM logging in process_stream (MATANY_LOG_VRAM=1)
-- Enhanced _safe_empty_cache with memory_summary
-- Updated comments to prioritize step() post-patch
+- Added comprehensive error handling for MatAnyone import and initialization
+- Enhanced VRAM management with auto-cleanup
+- Added support for multiple MatAnyone method patching (encode_img/forward/encode)
+- Improved logging with timestamps and memory usage
+- Added environment variable controls for debugging
+- Fixed potential memory leaks in tensor handling
 """
 
 from __future__ import annotations
@@ -76,16 +77,33 @@ def _cuda_snapshot(device: Optional[torch.device]) -> str:
         return f"CUDA snapshot error: {e!r}"
 
 def _safe_empty_cache():
-    if not torch.cuda.is_available():
-        return
+    """Safely clear PyTorch cache with detailed memory reporting."""
     try:
-        log.info(f"[MATANY] CUDA memory before empty_cache: {_cuda_snapshot(None)}")
+        if not torch.cuda.is_available():
+            return
+            
+        # Log memory stats before cleanup
+        if _env_flag("MATANY_LOG_VRAM"):
+            log.info("[MATANY] VRAM before cleanup:")
+            log.info(f"  Allocated: {torch.cuda.memory_allocated()/1024**2:.1f} MB")
+            log.info(f"  Reserved:  {torch.cuda.memory_reserved()/1024**2:.1f} MB")
+        
+        # Clear cache and sync
         torch.cuda.empty_cache()
-        log.info(f"[MATANY] CUDA memory after empty_cache: {_cuda_snapshot(None)}")
-        if os.getenv("MATANY_LOG_VRAM", "0") == "1":
-            log.debug(f"[MATANY] VRAM summary:\n{torch.cuda.memory_summary()}")
-    except Exception:
-        pass
+        torch.cuda.synchronize()
+        
+        # Log memory stats after cleanup
+        if _env_flag("MATANY_LOG_VRAM"):
+            log.info("[MATANY] VRAM after cleanup:")
+            log.info(f"  Allocated: {torch.cuda.memory_allocated()/1024**2:.1f} MB")
+            log.info(f"  Reserved:  {torch.cuda.memory_reserved()/1024**2:.1f} MB")
+            
+    except Exception as e:
+        log.warning(f"[MATANY] Error in cache cleanup: {e}", exc_info=True)
+        try:
+            torch.cuda.empty_cache()
+        except Exception as e2:
+            log.warning(f"[MATANY] Secondary cache cleanup failed: {e2}")
 
 # ---------- SAM2 → seed mask prep ----------
 def _prepare_seed_mask(sam2_mask: np.ndarray, H: int, W: int) -> np.ndarray: