mirror of
https://github.com/snapotter-hq/SnapOtter.git
synced 2026-08-03 07:46:42 +02:00
ashim-hq
39078f8e7f
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.
268 lines
9.5 KiB
Python
268 lines
9.5 KiB
Python
"""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):
|
|
"""Emit structured progress to stderr for bridge.ts to capture."""
|
|
print(json.dumps({"progress": percent, "stage": stage}), file=sys.stderr, flush=True)
|
|
|
|
|
|
_MODELS_BASE = os.environ.get("MODELS_PATH", "/opt/models")
|
|
|
|
REALESRGAN_MODEL_PATH = os.environ.get(
|
|
"REALESRGAN_MODEL_PATH",
|
|
os.path.join(_MODELS_BASE, "realesrgan", "RealESRGAN_x4plus.pth"),
|
|
)
|
|
|
|
GFPGAN_MODEL_PATH = os.environ.get(
|
|
"GFPGAN_MODEL_PATH",
|
|
os.path.join(_MODELS_BASE, "gfpgan", "GFPGANv1.3.pth"),
|
|
)
|
|
|
|
|
|
def apply_denoise(img, strength):
|
|
"""Apply denoising to a PIL image. Uses OpenCV when available, falls back to PIL."""
|
|
if strength <= 0:
|
|
return img
|
|
try:
|
|
import numpy as np
|
|
import cv2
|
|
from PIL import Image
|
|
|
|
arr = np.array(img)
|
|
# Map 0-1 strength to filter parameter (3-15 range)
|
|
h = int(3 + strength * 12)
|
|
if len(arr.shape) == 3 and arr.shape[2] >= 3:
|
|
denoised = cv2.fastNlMeansDenoisingColored(arr, None, h, h, 7, 21)
|
|
else:
|
|
denoised = cv2.fastNlMeansDenoising(arr, None, h, 7, 21)
|
|
return Image.fromarray(denoised)
|
|
except ImportError:
|
|
from PIL import ImageFilter
|
|
|
|
radius = max(0.5, strength * 1.5)
|
|
return img.filter(ImageFilter.GaussianBlur(radius=radius))
|
|
|
|
|
|
def main():
|
|
input_path = sys.argv[1]
|
|
output_path = sys.argv[2]
|
|
settings = json.loads(sys.argv[3]) if len(sys.argv) > 3 else {}
|
|
|
|
scale = settings.get("scale", 2)
|
|
model_choice = settings.get("model", "auto")
|
|
face_enhance = settings.get("faceEnhance", False)
|
|
denoise_strength = float(settings.get("denoise", 0))
|
|
output_format = settings.get("format", "png")
|
|
quality = int(settings.get("quality", 95))
|
|
|
|
try:
|
|
emit_progress(5, "Opening image")
|
|
from PIL import Image
|
|
|
|
img = Image.open(input_path)
|
|
new_size = (img.width * scale, img.height * scale)
|
|
|
|
method = "lanczos"
|
|
result = None
|
|
|
|
# Try Real-ESRGAN if requested
|
|
if model_choice in ("auto", "realesrgan"):
|
|
try:
|
|
emit_progress(10, "Loading AI model")
|
|
|
|
# Redirect stdout to stderr for the ENTIRE AI pipeline.
|
|
# Libraries like basicsr, realesrgan, gfpgan, and torch print
|
|
# download progress and init messages to stdout which would
|
|
# corrupt our JSON result.
|
|
stdout_fd = None
|
|
try:
|
|
stdout_fd = os.dup(1)
|
|
os.dup2(2, 1)
|
|
except OSError:
|
|
# os.dup may fail on Windows with piped stdio
|
|
stdout_fd = None
|
|
sys.stdout = sys.stderr
|
|
|
|
try:
|
|
from basicsr.archs.rrdbnet_arch import RRDBNet
|
|
from realesrgan import RealESRGANer
|
|
from gpu import gpu_available
|
|
import numpy as np
|
|
import torch
|
|
|
|
if not os.path.exists(REALESRGAN_MODEL_PATH):
|
|
raise FileNotFoundError(
|
|
f"RealESRGAN model not found: {REALESRGAN_MODEL_PATH}"
|
|
)
|
|
|
|
use_gpu = gpu_available()
|
|
device = torch.device("cuda" if use_gpu else "cpu")
|
|
|
|
# RealESRGAN_x4plus is a 4x model internally
|
|
ai_model = RRDBNet(
|
|
num_in_ch=3,
|
|
num_out_ch=3,
|
|
num_feat=64,
|
|
num_block=23,
|
|
num_grow_ch=32,
|
|
scale=4,
|
|
)
|
|
upsampler = RealESRGANer(
|
|
scale=4,
|
|
model_path=REALESRGAN_MODEL_PATH,
|
|
model=ai_model,
|
|
half=use_gpu,
|
|
device=device,
|
|
)
|
|
emit_progress(20, "AI model loaded")
|
|
|
|
img_array = np.array(img.convert("RGB"))
|
|
emit_progress(30, "Enhancing image with AI")
|
|
output_array, _ = upsampler.enhance(img_array, outscale=scale)
|
|
emit_progress(80, "AI enhancement complete")
|
|
result = Image.fromarray(output_array)
|
|
method = "realesrgan"
|
|
|
|
if face_enhance:
|
|
emit_progress(82, "Enhancing faces")
|
|
from gfpgan import GFPGANer
|
|
|
|
if not os.path.exists(GFPGAN_MODEL_PATH):
|
|
raise FileNotFoundError(
|
|
f"GFPGAN model not found at {GFPGAN_MODEL_PATH}. "
|
|
"Install the upscale-enhance feature or disable faceEnhance."
|
|
)
|
|
face_enhancer = GFPGANer(
|
|
model_path=GFPGAN_MODEL_PATH,
|
|
upscale=scale,
|
|
arch="clean",
|
|
channel_multiplier=2,
|
|
bg_upsampler=upsampler,
|
|
)
|
|
_, _, face_output = face_enhancer.enhance(
|
|
img_array,
|
|
has_aligned=False,
|
|
only_center_face=False,
|
|
paste_back=True,
|
|
)
|
|
result = Image.fromarray(face_output)
|
|
emit_progress(88, "Face enhancement complete")
|
|
|
|
finally:
|
|
# Restore stdout after ALL AI processing
|
|
if stdout_fd is not None:
|
|
os.dup2(stdout_fd, 1)
|
|
os.close(stdout_fd)
|
|
sys.stdout = sys.__stdout__
|
|
|
|
except (ImportError, FileNotFoundError, RuntimeError, OSError) as e:
|
|
import traceback
|
|
print(f"[upscale] Real-ESRGAN failed: {e}", file=sys.stderr, flush=True)
|
|
traceback.print_exc(file=sys.stderr)
|
|
print(json.dumps({
|
|
"success": False,
|
|
"error": (
|
|
f"Real-ESRGAN is not available: {e}. "
|
|
"Install the upscale-enhance feature or use model=lanczos for basic upscaling."
|
|
),
|
|
}))
|
|
sys.exit(1)
|
|
|
|
if result is None and model_choice == "lanczos":
|
|
emit_progress(50, "Upscaling with Lanczos")
|
|
result = img.resize(new_size, Image.LANCZOS)
|
|
method = "lanczos"
|
|
|
|
if result is None:
|
|
raise RuntimeError(f"Requested model '{model_choice}' is not available")
|
|
|
|
# Denoise
|
|
if denoise_strength > 0:
|
|
emit_progress(90, "Reducing noise")
|
|
result = apply_denoise(result, denoise_strength)
|
|
|
|
# Determine final output path based on format
|
|
base_path = output_path.rsplit(".", 1)[0]
|
|
EXT_MAP = {
|
|
"jpeg": ".jpg",
|
|
"jpg": ".jpg",
|
|
"png": ".png",
|
|
"webp": ".webp",
|
|
"tiff": ".tiff",
|
|
"gif": ".gif",
|
|
}
|
|
final_path = base_path + EXT_MAP.get(output_format, ".png")
|
|
|
|
# Save with format-specific options
|
|
emit_progress(95, "Saving result")
|
|
save_kwargs = {}
|
|
if output_format in ("jpeg", "jpg"):
|
|
result = result.convert("RGB") # Strip alpha for JPEG
|
|
save_kwargs["quality"] = quality
|
|
save_kwargs["optimize"] = True
|
|
elif output_format == "webp":
|
|
save_kwargs["quality"] = quality
|
|
elif output_format == "tiff":
|
|
save_kwargs["compression"] = "tiff_lzw"
|
|
elif output_format == "gif":
|
|
result = result.convert("P", palette=Image.ADAPTIVE, colors=256)
|
|
|
|
result.save(final_path, **save_kwargs)
|
|
|
|
# Get actual dimensions of the saved result
|
|
actual_w, actual_h = result.size
|
|
|
|
print(
|
|
json.dumps(
|
|
{
|
|
"success": True,
|
|
"scale": scale,
|
|
"width": actual_w,
|
|
"height": actual_h,
|
|
"method": method,
|
|
"output_path": final_path,
|
|
"format": output_format,
|
|
}
|
|
)
|
|
)
|
|
|
|
except ImportError:
|
|
print(
|
|
json.dumps(
|
|
{
|
|
"success": False,
|
|
"error": "Pillow is not installed. Install with: pip install Pillow",
|
|
}
|
|
)
|
|
)
|
|
sys.exit(1)
|
|
except Exception as e:
|
|
print(json.dumps({"success": False, "error": str(e)}))
|
|
sys.exit(1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|