mirror of
https://github.com/pewdiepie-archdaemon/odysseus.git
synced 2026-06-17 10:15:27 -04:00
fix(image): patch realesrgan torchvision compatibility (#4110)
This commit is contained in:
@@ -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"]
|
||||
Reference in New Issue
Block a user