patch_annotators.py

"""
patch_annotators.py
===================
annotator 内の互換性問題をランタイムでパッチするモジュール。
app.py の冒頭で import するだけで有効になります。

対応する問題:
  1. annotator/zoe/__init__.py: torch.load に strict=False が必要
     (timm のバージョンアップで relative_position_index キーが消えたため)
"""

import importlib
import sys
from unittest.mock import patch


# ── Zoe: torch.load を strict=False でラップ ──────────────────────────────

def _patch_zoe():
    """
    annotator.zoe.__init__ の ZoeDetector.__init__ が
    model.load_state_dict(...) を strict=True (デフォルト) で呼ぶのを
    strict=False に差し替える。
    """
    try:
        import torch
        _original_load_state_dict = torch.nn.Module.load_state_dict

        def _patched_load_state_dict(self, state_dict, strict=True, **kwargs):
            # ZoeDepth モデルのロード時だけ strict=False に緩める
            cls_name = type(self).__name__
            if cls_name == "ZoeDepth":
                strict = False
            return _original_load_state_dict(self, state_dict, strict=strict, **kwargs)

        torch.nn.Module.load_state_dict = _patched_load_state_dict
        print("[patch_annotators] Patched: ZoeDepth.load_state_dict -> strict=False")
    except Exception as e:
        print(f"[patch_annotators] Warning: Could not patch ZoeDepth: {e}")


_patch_zoe()