SpatialDiffusion / compat_patch.py
zimhe
compat patch
a8c0d06
raw
history blame contribute delete
829 Bytes
import sys
from pathlib import Path
import importlib.util
import torchvision
# 检查是否有 functional_tensor 模块
if not hasattr(torchvision.transforms, 'functional_tensor'):
# 创建一个虚拟模块
module_name = 'torchvision.transforms.functional_tensor'
spec = importlib.machinery.ModuleSpec(module_name, None)
functional_tensor = importlib.util.module_from_spec(spec)
# 将 functional 中的函数添加到 functional_tensor
from torchvision.transforms import functional
functional_tensor.rgb_to_grayscale = functional.rgb_to_grayscale
# 将模块添加到 sys.modules
sys.modules[module_name] = functional_tensor
# 将模块添加到 torchvision.transforms
torchvision.transforms.functional_tensor = functional_tensor
# 然后导入 basicsr
import basicsr