fix(image): patch realesrgan torchvision compatibility (#4110)

tests/test_realesrgan_torchvision_compat.py

@@ -0,0 +1,47 @@
+import sys
+import types
+
+from src.optional_deps import (
+    patch_realesrgan_torchvision_compat,
+    prepare_optional_dependency_import,
+)
+
+
+def test_realesrgan_patch_restores_removed_functional_tensor_module(monkeypatch):
+    for name in list(sys.modules):
+        if name.startswith("torchvision"):
+            monkeypatch.delitem(sys.modules, name, raising=False)
+
+    sentinel = object()
+    torchvision = types.ModuleType("torchvision")
+    transforms = types.ModuleType("torchvision.transforms")
+    functional = types.ModuleType("torchvision.transforms.functional")
+    functional.rgb_to_grayscale = sentinel
+    transforms.functional = functional
+    torchvision.transforms = transforms
+    monkeypatch.setitem(sys.modules, "torchvision", torchvision)
+    monkeypatch.setitem(sys.modules, "torchvision.transforms", transforms)
+    monkeypatch.setitem(sys.modules, "torchvision.transforms.functional", functional)
+
+    patch_realesrgan_torchvision_compat()
+
+    shim = sys.modules["torchvision.transforms.functional_tensor"]
+    assert shim.rgb_to_grayscale is sentinel
+    assert shim.rgb_to_grayscale is functional.rgb_to_grayscale
+
+
+def test_prepare_optional_dependency_import_scopes_patch_to_realesrgan(monkeypatch):
+    import src.optional_deps as optional_deps
+
+    calls = []
+    monkeypatch.setattr(
+        optional_deps,
+        "patch_realesrgan_torchvision_compat",
+        lambda: calls.append("patched"),
+    )
+
+    prepare_optional_dependency_import("diffusers")
+    assert calls == []
+
+    prepare_optional_dependency_import("realesrgan")
+    assert calls == ["patched"]

tests/test_shell_routes.py

@@ -13,6 +13,7 @@ import pytest
 
 from routes.shell_routes import (
     _find_line_break,
+    _import_optional_dependency_for_status,
     _running_in_container,
     _docker_row_status,
     _package_installed_from_probe,
@@ -376,6 +377,26 @@ class TestPackageProbeStatus:
         assert "add_user_install_bins_to_path()" in script
         assert "shutil.which(b)" in script
 
+    def test_status_import_prepares_optional_dependency(self, monkeypatch):
+        import routes.shell_routes as shell_routes
+
+        calls = []
+        monkeypatch.setattr(
+            shell_routes,
+            "prepare_optional_dependency_import",
+            lambda name: calls.append(name),
+        )
+        monkeypatch.setattr(
+            shell_routes.importlib,
+            "import_module",
+            lambda name: SimpleNamespace(__name__=name),
+        )
+
+        module = _import_optional_dependency_for_status("realesrgan")
+
+        assert module.__name__ == "realesrgan"
+        assert calls == ["realesrgan"]
+
 
 class TestSshBaseArgv:
     def test_basic_host_no_port(self):