diff --git a/src/processing/models.py b/src/processing/models.py index 3c94938..8613e62 100644 --- a/src/processing/models.py +++ b/src/processing/models.py @@ -1,6 +1,7 @@ """Model management for Real-ESRGAN and GFPGAN.""" import logging +import sys import time from pathlib import Path from threading import Lock @@ -10,6 +11,21 @@ import requests import torch from tqdm import tqdm +# Patch for basicsr compatibility with torchvision 0.18+ +# The functional_tensor module was removed, but basicsr still imports from it +import torchvision.transforms.functional as TF + +if not hasattr(TF, "rgb_to_grayscale"): + TF.rgb_to_grayscale = TF.to_grayscale + +# Create a fake module to satisfy basicsr's import +if "torchvision.transforms.functional_tensor" not in sys.modules: + + class FakeFunctionalTensor: + rgb_to_grayscale = TF.rgb_to_grayscale + + sys.modules["torchvision.transforms.functional_tensor"] = FakeFunctionalTensor() + from src.config import MODEL_REGISTRY, MODELS_DIR, ModelInfo, config logger = logging.getLogger(__name__)