mirror of
https://github.com/snapotter-hq/SnapOtter.git
synced 2026-08-03 07:46:42 +02:00
SnapOtterandGitHub
36dde9ad87
gpu_available() answers "can ANY framework use a GPU" (torch, then ONNX, then
paddle). But torch tools consumed that shared boolean directly as
device = torch.device("cuda" if gpu_available() else "cpu"). On a GPU host where
gpu_available() is True via paddle or ONNX while torch is a CPU-only build, those
tools would route to a CUDA torch cannot use and crash. Transcription had the
mirror problem: it runs on CTranslate2 (not torch), so on a transcription-only
GPU box gpu_available() returned False and Whisper ran on CPU despite a GPU.
Add per-framework helpers to gpu.py:
- torch_gpu_available(): torch.cuda.is_available(), honoring SNAPOTTER_GPU.
- ctranslate2_gpu_available(): ctranslate2.get_cuda_device_count() > 0.
Point each tool at the helper for its own framework: upscale, noise_removal,
enhance_faces and restore use torch_gpu_available(); transcribe uses
ctranslate2_gpu_available(). ocr.py keeps gpu_available() (paddle-aware) and the
dispatcher keeps it for its startup GPU-status line. The SNAPOTTER_GPU override
check is factored into a shared _override_disables_gpu() helper.
TDD: 7 new tests in tests/test_gpu_detection.py cover both helpers (override,
CPU-only, absent framework), including the crux that torch_gpu_available() stays
False on a CPU-only torch build even when a GPU exists for another framework.
Claude-Session: https://claude.ai/code/session_01NfaRxjek8ex5nawvx3mVMf
312 lines
12 KiB
Python
312 lines
12 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 torch_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 = torch_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,
|
|
)
|
|
tile_size = 512
|
|
upsampler = RealESRGANer(
|
|
scale=4,
|
|
model_path=REALESRGAN_MODEL_PATH,
|
|
model=ai_model,
|
|
tile=tile_size,
|
|
tile_pad=10,
|
|
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")
|
|
|
|
try:
|
|
output_array, _ = upsampler.enhance(img_array, outscale=scale)
|
|
except RuntimeError as oom_err:
|
|
if "out of memory" not in str(oom_err).lower():
|
|
raise
|
|
torch.cuda.empty_cache()
|
|
print(
|
|
f"[upscale] OOM with tile={tile_size}, retrying with tile=256",
|
|
file=sys.stderr,
|
|
flush=True,
|
|
)
|
|
emit_progress(35, "Retrying with smaller tiles")
|
|
upsampler.tile = 256
|
|
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."
|
|
)
|
|
# GFPGANer resolves its facexlib helper weights
|
|
# relative to the cwd and downloads them from GitHub
|
|
# when missing; resolve them from the bundle first so
|
|
# no download is needed (strict offline mode errors
|
|
# instead).
|
|
from offline_guard import prepare_gfpgan_helper_weights
|
|
prepare_gfpgan_helper_weights(_MODELS_BASE)
|
|
face_enhancer = GFPGANer(
|
|
model_path=GFPGAN_MODEL_PATH,
|
|
upscale=scale,
|
|
arch="clean",
|
|
channel_multiplier=2,
|
|
bg_upsampler=upsampler,
|
|
)
|
|
try:
|
|
_, _, face_output = face_enhancer.enhance(
|
|
img_array,
|
|
has_aligned=False,
|
|
only_center_face=False,
|
|
paste_back=True,
|
|
)
|
|
except RuntimeError as oom_err:
|
|
if "out of memory" not in str(oom_err).lower():
|
|
raise
|
|
torch.cuda.empty_cache()
|
|
print(
|
|
"[upscale] OOM during face enhance, retrying with tile=256",
|
|
file=sys.stderr,
|
|
flush=True,
|
|
)
|
|
emit_progress(84, "Retrying face enhancement")
|
|
upsampler.tile = 256
|
|
face_enhancer.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()
|