Spaces:
Runtime error
Runtime error
File size: 2,989 Bytes
c19ca42 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 |
from __future__ import annotations
import gc
import numpy as np
try:
from ncnn_vulkan import ncnn
use_gpu = True
except ImportError:
from ncnn import ncnn
use_gpu = False
from nodes.log import logger
from ...utils.utils import get_h_w_c
from ..image_utils import to_uint8
from ..upscale.auto_split import Split, Tiler, auto_split
def ncnn_auto_split(
img: np.ndarray,
net,
input_name: str,
output_name: str,
blob_vkallocator,
staging_vkallocator,
tiler: Tiler,
) -> np.ndarray:
def upscale(img: np.ndarray, _):
ex = net.create_extractor()
if use_gpu:
ex.set_blob_vkallocator(blob_vkallocator)
ex.set_workspace_vkallocator(blob_vkallocator)
ex.set_staging_vkallocator(staging_vkallocator)
# ex.set_light_mode(True)
try:
lr_c = get_h_w_c(img)[2]
lr_img_fix = to_uint8(img)
if lr_c == 1:
pixel_type = ncnn.Mat.PixelType.PIXEL_GRAY
elif lr_c == 3:
pixel_type = ncnn.Mat.PixelType.PIXEL_RGB
else:
pixel_type = ncnn.Mat.PixelType.PIXEL_RGBA
mat_in = ncnn.Mat.from_pixels(
lr_img_fix,
pixel_type,
lr_img_fix.shape[1],
lr_img_fix.shape[0],
)
mean_vals = []
norm_vals = [1 / 255.0] * lr_c
mat_in.substract_mean_normalize(mean_vals, norm_vals)
ex.input(input_name, mat_in)
_, mat_out = ex.extract(output_name)
result = np.array(mat_out).transpose(1, 2, 0).astype(np.float32)
del ex, mat_in, mat_out
gc.collect()
if use_gpu:
# Clear VRAM
blob_vkallocator.clear()
staging_vkallocator.clear()
return result
except Exception as e:
if "vkQueueSubmit" in str(e):
ex = None
del ex
gc.collect()
if use_gpu:
blob_vkallocator.clear()
staging_vkallocator.clear()
# ncnn.destroy_gpu_instance()
raise RuntimeError(
"A critical error has occurred. You may need to restart chaiNNer in order for NCNN upscaling to start working again."
) from e
# Check to see if its actually the NCNN out of memory error
if "failed" in str(e):
# clear VRAM
logger.debug("NCNN out of VRAM, clearing VRAM and splitting.")
ex = None
del ex
gc.collect()
if use_gpu:
blob_vkallocator.clear()
staging_vkallocator.clear()
return Split()
else:
# Re-raise the exception if not an OOM error
raise
return auto_split(img, upscale, tiler)
|