Skip to content

Commit f0f2192

Browse files
hooke007hooke007
committed
vs： k7sfunc 更新 0.4.6
检查部分模块所需的模型是否存在。 模块： - 修复 MVT_STD 的内部参数错误 #397 - 限制 RIFE_STD 的参数 stat_th 超过 60.0 - UAI_DML 和 UAI_NV_TRT 现在支持调用外部目录的模型文件
1 parent 50bc4d3 commit f0f2192

File tree

1 file changed

+72
-8
lines changed

1 file changed

+72
-8
lines changed

‎k7sfunc.py

Lines changed: 72 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22
### 文档： https://github.com/hooke007/MPV_lazy/wiki/3_K7sfunc
33
##################################################
44

5-
__version__ = "0.4.3"
5+
__version__ = "0.4.6"
66

77
__all__ = [
88
"FMT_CHANGE", "FMT_CTRL", "FPS_CHANGE", "FPS_CTRL",
@@ -537,6 +537,12 @@ def CUGAN_NV(
537537
if not hasattr(core, "trt") :
538538
raise ModuleNotFoundError(f"模块 {func_name} 依赖错误：缺失插件，检查项目 trt")
539539

540+
plg_dir = os.path.dirname(core.trt.Version()["path"]).decode()
541+
mdl_fname = ["pro-no-denoise3x-up2x", "pro-conservative-up2x", "pro-denoise3x-up2x"][[-1, 0, 3].index(nr_lv)]
542+
mdl_pth = plg_dir + "/models/cugan/" + mdl_fname + ".onnx"
543+
if not os.path.exists(mdl_pth) :
544+
raise vs.Error(f"模块 {func_name} 所请求的模型缺失")
545+
540546
global vsmlrt
541547
if vsmlrt is None :
542548
try :
@@ -669,6 +675,12 @@ def ESRGAN_DML(
669675
if not hasattr(core, "ort") :
670676
raise ModuleNotFoundError(f"模块 {func_name} 依赖错误：缺失插件，检查项目 ort")
671677

678+
plg_dir = os.path.dirname(core.ort.Version()["path"]).decode()
679+
mdl_fname = ["RealESRGANv2-animevideo-xsx2", "realesr-animevideov3", "animejanaiV2L1", "animejanaiV2L2", "animejanaiV2L3"][[0, 2, 5005, 5006, 5007].index(model)]
680+
mdl_pth = plg_dir + "/models/RealESRGANv2/" + mdl_fname + ".onnx"
681+
if not os.path.exists(mdl_pth) :
682+
raise vs.Error(f"模块 {func_name} 所请求的模型缺失")
683+
672684
global vsmlrt
673685
if vsmlrt is None :
674686
try :
@@ -730,6 +742,12 @@ def ESRGAN_NV(
730742
if not hasattr(core, "trt") :
731743
raise ModuleNotFoundError(f"模块 {func_name} 依赖错误：缺失插件，检查项目 trt")
732744

745+
plg_dir = os.path.dirname(core.trt.Version()["path"]).decode()
746+
mdl_fname = ["RealESRGANv2-animevideo-xsx2", "realesr-animevideov3", "animejanaiV2L1", "animejanaiV2L2", "animejanaiV2L3"][[0, 2, 5005, 5006, 5007].index(model)]
747+
mdl_pth = plg_dir + "/models/RealESRGANv2/" + mdl_fname + ".onnx"
748+
if not os.path.exists(mdl_pth) :
749+
raise vs.Error(f"模块 {func_name} 所请求的模型缺失")
750+
733751
global vsmlrt
734752
if vsmlrt is None :
735753
try :
@@ -830,6 +848,13 @@ def WAIFU_DML(
830848
if not hasattr(core, "ort") :
831849
raise ModuleNotFoundError(f"模块 {func_name} 依赖错误：缺失插件，检查项目 ort")
832850

851+
plg_dir = os.path.dirname(core.ort.Version()["path"]).decode()
852+
mdl_pname = {3:"upconv_7_anime_style_art_rgb/", 5:"upresnet10/", 6:"cunet/"}.get(model)
853+
mdl_fname = ["scale2.0x_model", "noise0_scale2.0x_model", "noise1_scale2.0x_model", "noise2_scale2.0x_model", "noise3_scale2.0x_model"][[-1, 0, 1, 2, 3].index(nr_lv)]
854+
mdl_pth = plg_dir + "/models/waifu2x/" + mdl_pname + mdl_fname + ".onnx"
855+
if not os.path.exists(mdl_pth) :
856+
raise vs.Error(f"模块 {func_name} 所请求的模型缺失")
857+
833858
global vsmlrt
834859
if vsmlrt is None :
835860
try :
@@ -894,6 +919,13 @@ def WAIFU_NV(
894919
if not hasattr(core, "trt") :
895920
raise ModuleNotFoundError(f"模块 {func_name} 依赖错误：缺失插件，检查项目 trt")
896921

922+
plg_dir = os.path.dirname(core.trt.Version()["path"]).decode()
923+
mdl_pname = {3:"upconv_7_anime_style_art_rgb/", 5:"upresnet10/", 6:"cunet/"}.get(model)
924+
mdl_fname = ["scale2.0x_model", "noise0_scale2.0x_model", "noise1_scale2.0x_model", "noise2_scale2.0x_model", "noise3_scale2.0x_model"][[-1, 0, 1, 2, 3].index(nr_lv)]
925+
mdl_pth = plg_dir + "/models/waifu2x/" + mdl_pname + mdl_fname + ".onnx"
926+
if not os.path.exists(mdl_pth) :
927+
raise vs.Error(f"模块 {func_name} 所请求的模型缺失")
928+
897929
global vsmlrt
898930
if vsmlrt is None :
899931
try :
@@ -1032,8 +1064,8 @@ def _ffps(fps) :
10321064

10331065
cut1 = core.std.AssumeFPS(clip=cut0, fpsnum=int(vfps), fpsden=vden)
10341066
cut_s = core.mv.Super(clip=cut1, sharp=1, rfilter=4)
1035-
cut_b = core.mv.Analyse(super=cut_s, blksize=64, searchparam=0, pelsearch=3, isb=True, _lambda=0, lsad=10000, overlapv=16, badrange=0, search_coarse=4)
1036-
cut_f = core.mv.Analyse(super=cut_s, blksize=64, searchparam=0, pelsearch=3, _lambda=0, lsad=10000, overlapv=16, badrange=0, search_coarse=4)
1067+
cut_b = core.mv.Analyse(super=cut_s, blksize=64, searchparam=0, pelsearch=3, isb=True, lambda_=0, lsad=10000, overlapv=16, badrange=0, search_coarse=4)
1068+
cut_f = core.mv.Analyse(super=cut_s, blksize=64, searchparam=0, pelsearch=3, lambda_=0, lsad=10000, overlapv=16, badrange=0, search_coarse=4)
10371069

10381070
output = core.mv.BlockFPS(clip=cut1, super=cut_s, mvbw=cut_b, mvfw=cut_f, num=fps_out * 1000, den=vden, mode=2, thscd1=970, thscd2=255, blend=False)
10391071
if w_tmp + h_tmp > 0 :
@@ -1185,7 +1217,7 @@ def RIFE_STD(
11851217
raise vs.Error(f"模块 {func_name} 的子参数 sc_mode 的值无效")
11861218
if not isinstance(skip, bool) :
11871219
raise vs.Error(f"模块 {func_name} 的子参数 skip 的值无效")
1188-
if not isinstance(stat_th, (int, float)) or stat_th <= 0.0 :
1220+
if not isinstance(stat_th, (int, float)) or stat_th <= 0.0 or stat_th > 60.0 :
11891221
raise vs.Error(f"模块 {func_name} 的子参数 stat_th 的值无效")
11901222
if gpu not in [0, 1, 2] :
11911223
raise vs.Error(f"模块 {func_name} 的子参数 gpu 的值无效")
@@ -1288,6 +1320,16 @@ def RIFE_NV(
12881320
if not hasattr(core, "akarin") :
12891321
raise ModuleNotFoundError(f"模块 {func_name} 依赖错误：缺失插件，检查项目 akarin")
12901322

1323+
plg_dir = os.path.dirname(core.trt.Version()["path"]).decode()
1324+
mdl_pname = "rife/" if ext_proc else "rife_v2/"
1325+
if t_tta :
1326+
mdl_fname = ["rife_v4.6_ensemble", "rife_v4.13_ensemble", "rife_v4.13_lite_ensemble"][[46, 413, 4131].index(model)]
1327+
else :
1328+
mdl_fname = ["rife_v4.6", "rife_v4.13", "rife_v4.13_lite"][[46, 413, 4131].index(model)]
1329+
mdl_pth = plg_dir + "/models/" + mdl_pname + mdl_fname + ".onnx"
1330+
if not os.path.exists(mdl_pth) :
1331+
raise vs.Error(f"模块 {func_name} 所请求的模型缺失")
1332+
12911333
global vsmlrt
12921334
if vsmlrt is None :
12931335
try :
@@ -1911,6 +1953,12 @@ def DPIR_NR_NV(
19111953
if not hasattr(core, "trt") :
19121954
raise ModuleNotFoundError(f"模块 {func_name} 依赖错误：缺失插件，检查项目 trt")
19131955

1956+
plg_dir = os.path.dirname(core.trt.Version()["path"]).decode()
1957+
mdl_fname = ["drunet_gray", "drunet_color"][[0, 1].index(model)]
1958+
mdl_pth = plg_dir + "/models/dpir/" + mdl_fname + ".onnx"
1959+
if not os.path.exists(mdl_pth) :
1960+
raise vs.Error(f"模块 {func_name} 所请求的模型缺失")
1961+
19141962
global vsmlrt
19151963
if vsmlrt is None :
19161964
try :
@@ -2508,6 +2556,12 @@ def DPIR_DBLK_NV(
25082556
if not hasattr(core, "trt") :
25092557
raise ModuleNotFoundError(f"模块 {func_name} 依赖错误：缺失插件，检查项目 trt")
25102558

2559+
plg_dir = os.path.dirname(core.trt.Version()["path"]).decode()
2560+
mdl_fname = ["drunet_deblocking_grayscale", "drunet_deblocking_color"][[2, 3].index(model)]
2561+
mdl_pth = plg_dir + "/models/dpir/" + mdl_fname + ".onnx"
2562+
if not os.path.exists(mdl_pth) :
2563+
raise vs.Error(f"模块 {func_name} 所请求的模型缺失")
2564+
25112565
global vsmlrt
25122566
if vsmlrt is None :
25132567
try :
@@ -2801,7 +2855,7 @@ def UAI_DML(
28012855
raise vs.Error(f"模块 {func_name} 的子参数 input 的值无效")
28022856
if not isinstance(clamp, bool) :
28032857
raise vs.Error(f"模块 {func_name} 的子参数 clamp 的值无效")
2804-
if len(model_pth) == 0 :
2858+
if len(model_pth) <= 5 :
28052859
raise vs.Error(f"模块 {func_name} 的子参数 model_pth 的值无效")
28062860
if gpu not in [0, 1, 2] :
28072861
raise vs.Error(f"模块 {func_name} 的子参数 gpu 的值无效")
@@ -2816,6 +2870,11 @@ def UAI_DML(
28162870
if not hasattr(core, "akarin") :
28172871
raise ModuleNotFoundError(f"模块 {func_name} 依赖错误：缺失插件，检查项目 akarin")
28182872

2873+
mdl_pth_rel = os.path.join(vsmlrt.models_path, model_pth)
2874+
if not os.path.exists(mdl_pth_rel) and not os.path.exists(model_pth) :
2875+
raise vs.Error(f"模块 {func_name} 所请求的模型缺失")
2876+
mdl_pth = mdl_pth_rel if os.path.exists(mdl_pth_rel) else model_pth
2877+
28192878
global vsmlrt
28202879
if vsmlrt is None :
28212880
try :
@@ -2833,7 +2892,7 @@ def UAI_DML(
28332892
if clamp :
28342893
clip = core.akarin.Expr(clips=clip, expr="x 0 1 clamp")
28352894
be_param = vsmlrt.BackendV2.ORT_DML(device_id=gpu, num_streams=gpu_t, fp16=True)
2836-
infer = vsmlrt.inference(clips=clip, network_path=os.path.join(vsmlrt.models_path, model_pth), backend=be_param)
2895+
infer = vsmlrt.inference(clips=clip, network_path=mdl_pth, backend=be_param)
28372896
output = core.resize.Bilinear(clip=infer, format=fmt_in, matrix_s="709", range=1 if colorlv==0 else None)
28382897

28392898
return output
@@ -2863,7 +2922,7 @@ def UAI_NV_TRT(
28632922
raise vs.Error(f"模块 {func_name} 的子参数 input 的值无效")
28642923
if not isinstance(clamp, bool) :
28652924
raise vs.Error(f"模块 {func_name} 的子参数 clamp 的值无效")
2866-
if len(model_pth) == 0 :
2925+
if len(model_pth) <= 5 :
28672926
raise vs.Error(f"模块 {func_name} 的子参数 model_pth 的值无效")
28682927
if opt_lv not in [0, 1, 2, 3, 4, 5] :
28692928
raise vs.Error(f"模块 {func_name} 的子参数 opt_lv 的值无效")
@@ -2896,6 +2955,11 @@ def UAI_NV_TRT(
28962955
if not hasattr(core, "akarin") :
28972956
raise ModuleNotFoundError(f"模块 {func_name} 依赖错误：缺失插件，检查项目 akarin")
28982957

2958+
mdl_pth_rel = os.path.join(vsmlrt.models_path, model_pth)
2959+
if not os.path.exists(mdl_pth_rel) and not os.path.exists(model_pth) :
2960+
raise vs.Error(f"模块 {func_name} 所请求的模型缺失")
2961+
mdl_pth = mdl_pth_rel if os.path.exists(mdl_pth_rel) else model_pth
2962+
28992963
global vsmlrt
29002964
if vsmlrt is None :
29012965
try :
@@ -2918,7 +2982,7 @@ def UAI_NV_TRT(
29182982
num_streams=gpu_t, use_cuda_graph=nv1, use_cublas=nv2, use_cudnn=nv3,
29192983
fp16=fp16, force_fp16=False, tf32=True, output_format=1 if fp16 else 0, workspace=None if ws_size < 128 else (ws_size if st_eng else ws_size * 2),
29202984
static_shape=st_eng, min_shapes=[0, 0] if st_eng else [64, 64], opt_shapes=None if st_eng else res_opt, max_shapes=None if st_eng else res_max)
2921-
infer = vsmlrt.inference(clips=clip, network_path=os.path.join(vsmlrt.models_path, model_pth), backend=be_param)
2985+
infer = vsmlrt.inference(clips=clip, network_path=mdl_pth, backend=be_param)
29222986
output = core.resize.Bilinear(clip=infer, format=fmt_in, matrix_s="709", range=1 if colorlv==0 else None)
29232987

29242988
return output

0 commit comments

Comments
 (0)