From 6d5e3613a0eb4fa37352fc610565dcebf29edcb5 Mon Sep 17 00:00:00 2001
From: ZHANG Hao <zhanghao@4paradigm.com>
Date: Wed, 3 Sep 2025 10:51:05 +0800
Subject: [PATCH] patch for torch >= 2.6

---
 iic.py  | 12 ++++++++++++
 main.py | 11 +++++++++++
 2 files changed, 23 insertions(+)

diff --git a/iic.py b/iic.py
index 04e30fd..905fff9 100644
--- a/iic.py
+++ b/iic.py
@@ -1,4 +1,16 @@
 import os
+import torch
+from functools import wraps
+
+_orig_load = torch.load
+
+@wraps(_orig_load)
+def _load_patch(*args, **kwargs):
+    kwargs.setdefault("weights_only", False)
+    return _orig_load(*args, **kwargs)
+
+torch.load = _load_patch
+
 
 from modelscope.pipelines import pipeline
 from modelscope.outputs import OutputKeys
diff --git a/main.py b/main.py
index 7a4d8e6..5e77c0f 100644
--- a/main.py
+++ b/main.py
@@ -10,6 +10,17 @@ from datetime import datetime
 from pathlib import Path
 
 import torch
+from functools import wraps
+
+_orig_load = torch.load
+
+@wraps(_orig_load)
+def _load_patch(*args, **kwargs):
+    kwargs.setdefault("weights_only", False)
+    return _orig_load(*args, **kwargs)
+
+torch.load = _load_patch
+
 from modelscope.pipelines import pipeline
 from modelscope.outputs import OutputKeys