Skip to content

Commit

Permalink
speed up ncnn fallback
Browse files Browse the repository at this point in the history
  • Loading branch information
TNTwise committed Oct 8, 2024
1 parent d377fbd commit 7641f69
Show file tree
Hide file tree
Showing 2 changed files with 35 additions and 17 deletions.
50 changes: 34 additions & 16 deletions backend/src/UpscaleNCNN.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,9 +6,10 @@
try:
from upscale_ncnn_py import UPSCALE
method = "upscale_ncnn_py"
s
except:
from ncnn_vulkan import ncnn
method = "ncnn_vulkan"
import ncnn
method = "ncnn_vulkan"
class NCNNParam:
"""
Puts the last time an op shows up in a param in a dict
Expand Down Expand Up @@ -109,22 +110,39 @@ def hotUnload(self):

def hotReload(self):
self._load()

def NCNNImageMatFromNP(self, npArray: np.array):
return ncnn.Mat.from_pixels(
npArray,
ncnn.Mat.PixelType.PIXEL_BGR,
self.width,
self.height,
)

def NormalizeImage(self, mat, norm_vals):
mean_vals = []
mat.substract_mean_normalize(mean_vals, norm_vals)

def ClampNPArray(self, nparray: np.array) -> np.array:

return nparray.clip(0, 255)


def procNCNNVk(self, imageChunk):
ex = self.net.create_extractor()
mat_in = ncnn.Mat.from_pixels(imageChunk, ncnn.Mat.PixelType.PIXEL_RGB, self.width, self.height)

mat_in.substract_mean_normalize(self.mean_vals, self.norm_vals)
try:
# Make sure the input and output names match the param file
ex.input("data", mat_in)
ret, mat_out = ex.extract("output")
out = np.array(mat_out)

# Transpose the output from `c, h, w` to `h, w, c` and put it back in 0-255 range
return out.clip(0,1).__mul__(255.0).astype(np.uint8).transpose(1, 2, 0).tobytes()
except:
ncnn.destroy_gpu_instance()
frame = self.NCNNImageMatFromNP(imageChunk)
# norm
self.NormalizeImage(mat=frame, norm_vals=[1 / 255.0, 1 / 255.0, 1 / 255.0])
# render frame
ex.input("data", frame)
ret, frame = ex.extract("output")

# norm
self.NormalizeImage(mat=frame, norm_vals=[255.0, 255.0, 255.0])

frame = np.ascontiguousarray(frame)
frame = self.ClampNPArray(frame)
frame = frame.transpose(1, 2, 0)
return np.ascontiguousarray(frame, dtype=np.uint8)

def Upscale(self, imageChunk):
while self.net is None:
Expand Down
2 changes: 1 addition & 1 deletion src/DownloadDeps.py
Original file line number Diff line number Diff line change
Expand Up @@ -264,7 +264,7 @@ def downloadNCNNDeps(self):
ncnnDeps = [
"rife-ncnn-vulkan-python-tntwise==1.4.2",
"upscale_ncnn_py==1.2.0",
"ncnn-vulkan==2023.6.18",
"ncnn==1.0.20240820",
"numpy==1.26.4",
"opencv-python-headless",
] + self.getPlatformIndependentDeps()
Expand Down

0 comments on commit 7641f69

Please sign in to comment.