From 39078f8e7f3824705a6210ab04ae84d11d96a9d0 Mon Sep 17 00:00:00 2001
From: ashim-hq <ashim.hq@gmail.com>
Date: Tue, 21 Apr 2026 00:02:31 +0800
Subject: [PATCH] fix: add torchvision shim to upscale.py and enhance_faces.py

When the Python dispatcher crashes and bridge.ts retries via per-request
spawning, the shim from dispatcher.py isn't loaded. basicsr then fails
importing torchvision.transforms.functional_tensor (removed in v0.17).

Adding the shim directly to both scripts ensures they work regardless
of whether they run through the dispatcher or standalone.
---
 packages/ai/python/enhance_faces.py | 19 +++++++++++++++++++
 packages/ai/python/upscale.py       | 21 ++++++++++++++++++++-
 2 files changed, 39 insertions(+), 1 deletion(-)

diff --git a/packages/ai/python/enhance_faces.py b/packages/ai/python/enhance_faces.py
index 6f4c37d2..7639b386 100644
--- a/packages/ai/python/enhance_faces.py
+++ b/packages/ai/python/enhance_faces.py
@@ -2,6 +2,25 @@
 import sys
 import json
 import os
+import types
+
+# basicsr imports torchvision.transforms.functional_tensor which was removed
+# in torchvision >= 0.17. This shim must exist before basicsr is imported.
+try:
+    import torchvision.transforms.functional_tensor  # noqa: F401
+except (ImportError, ModuleNotFoundError):
+    try:
+        import torchvision.transforms.functional as _F
+        import torchvision.transforms
+
+        _shim = types.ModuleType("torchvision.transforms.functional_tensor")
+        for _attr in dir(_F):
+            if not _attr.startswith("_"):
+                setattr(_shim, _attr, getattr(_F, _attr))
+        sys.modules["torchvision.transforms.functional_tensor"] = _shim
+        torchvision.transforms.functional_tensor = _shim
+    except ImportError:
+        pass
 
 
 def emit_progress(percent, stage):
diff --git a/packages/ai/python/upscale.py b/packages/ai/python/upscale.py
index 503a45a7..36ee3777 100644
--- a/packages/ai/python/upscale.py
+++ b/packages/ai/python/upscale.py
@@ -1,7 +1,26 @@
-"""Image upscaling with Real-ESRGAN fallback to Lanczos."""
+"""Image upscaling with Real-ESRGAN."""
 import sys
 import json
 import os
+import types
+
+# basicsr imports torchvision.transforms.functional_tensor which was removed
+# in torchvision >= 0.17. This shim must exist before basicsr is imported.
+try:
+    import torchvision.transforms.functional_tensor  # noqa: F401
+except (ImportError, ModuleNotFoundError):
+    try:
+        import torchvision.transforms.functional as _F
+        import torchvision.transforms
+
+        _shim = types.ModuleType("torchvision.transforms.functional_tensor")
+        for _attr in dir(_F):
+            if not _attr.startswith("_"):
+                setattr(_shim, _attr, getattr(_F, _attr))
+        sys.modules["torchvision.transforms.functional_tensor"] = _shim
+        torchvision.transforms.functional_tensor = _shim
+    except ImportError:
+        pass
 
 
 def emit_progress(percent, stage):