diff --git a/src/main.cpp b/src/main.cpp index ebe0e62..ddfb742 100644 --- a/src/main.cpp +++ b/src/main.cpp @@ -109,8 +109,7 @@ static void print_usage() fprintf(stderr, " -o output-path output image path (jpg/png/webp) or directory\n"); fprintf(stderr, " -s scale upscale ratio (can be 2, 3, 4. default=4)\n"); fprintf(stderr, " -t tile-size tile size (>=32/0=auto, default=0) can be 0,0,0 for multi-gpu\n"); - fprintf(stderr, " -m model-path folder path to the pre-trained models. default=models\n"); - fprintf(stderr, " -n model-name model name (default=realesr-animevideov3, can be realesr-animevideov3 | realesrgan-x4plus | realesrgan-x4plus-anime | realesrnet-x4plus)\n"); + fprintf(stderr, " -m model-path model and parameter file name (sans .bin and .param extension)\n"); fprintf(stderr, " -g gpu-id gpu device to use (default=auto) can be 0,1,2 for multi-gpu\n"); fprintf(stderr, " -j load:proc:save thread count for load/proc/save (default=1:2:2) can be 1:2,2,2:2 for multi-gpu\n"); fprintf(stderr, " -x enable tta mode\n"); @@ -438,8 +437,7 @@ int main(int argc, char** argv) path_t outputpath; int scale = 4; std::vector tilesize; - path_t model = PATHSTR("models"); - path_t modelname = PATHSTR("realesr-animevideov3"); + path_t model = PATHSTR(""); std::vector gpuid; int jobs_load = 1; std::vector jobs_proc; @@ -451,7 +449,7 @@ int main(int argc, char** argv) #if _WIN32 setlocale(LC_ALL, ""); wchar_t opt; - while ((opt = getopt(argc, argv, L"i:o:s:t:m:n:g:j:f:vxh")) != (wchar_t)-1) + while ((opt = getopt(argc, argv, L"i:o:t:m:g:j:f:vxh")) != (wchar_t)-1) { switch (opt) { @@ -461,18 +459,12 @@ int main(int argc, char** argv) case L'o': outputpath = optarg; break; - case L's': - scale = _wtoi(optarg); - break; case L't': tilesize = parse_optarg_int_array(optarg); break; case L'm': model = optarg; break; - case L'n': - modelname = optarg; - break; case L'g': gpuid = parse_optarg_int_array(optarg); break; @@ -497,7 +489,7 @@ int main(int argc, char** argv) } #else // _WIN32 int opt; - while ((opt = getopt(argc, argv, "i:o:s:t:m:n:g:j:f:vxh")) != -1) + while ((opt = getopt(argc, argv, "i:o:t:m:g:j:f:vxh")) != -1) { switch (opt) { @@ -507,18 +499,12 @@ int main(int argc, char** argv) case 'o': outputpath = optarg; break; - case 's': - scale = atoi(optarg); - break; case 't': tilesize = parse_optarg_int_array(optarg); break; case 'm': model = optarg; break; - case 'n': - modelname = optarg; - break; case 'g': gpuid = parse_optarg_int_array(optarg); break; @@ -549,6 +535,12 @@ int main(int argc, char** argv) return -1; } + if (model.empty()) + { + fprintf(stderr, "no model given\n"); + return -1; + } + if (tilesize.size() != (gpuid.empty() ? 1 : gpuid.size()) && !tilesize.empty()) { fprintf(stderr, "invalid tilesize argument\n"); @@ -671,61 +663,17 @@ int main(int argc, char** argv) } } - int prepadding = 0; - - if (model.find(PATHSTR("models")) != path_t::npos - || model.find(PATHSTR("models2")) != path_t::npos) - { - prepadding = 10; - } - else - { - fprintf(stderr, "unknown model dir type\n"); - return -1; - } + int prepadding = 10; - // if (modelname.find(PATHSTR("realesrgan-x4plus")) != path_t::npos - // || modelname.find(PATHSTR("realesrnet-x4plus")) != path_t::npos - // || modelname.find(PATHSTR("esrgan-x4")) != path_t::npos) - // {} - // else - // { - // fprintf(stderr, "unknown model name\n"); - // return -1; - // } #if _WIN32 - wchar_t parampath[256]; - wchar_t modelpath[256]; - - if (modelname == PATHSTR("realesr-animevideov3")) - { - swprintf(parampath, 256, L"%s/%s-x%s.param", model.c_str(), modelname.c_str(), std::to_string(scale)); - swprintf(modelpath, 256, L"%s/%s-x%s.bin", model.c_str(), modelname.c_str(), std::to_string(scale)); - } - else{ - swprintf(parampath, 256, L"%s/%s.param", model.c_str(), modelname.c_str()); - swprintf(modelpath, 256, L"%s/%s.bin", model.c_str(), modelname.c_str()); - } - + path_t parampath = model + L".param"; + path_t modelpath = model + L".bin"; #else - char parampath[256]; - char modelpath[256]; - - if (modelname == PATHSTR("realesr-animevideov3")) - { - sprintf(parampath, "%s/%s-x%s.param", model.c_str(), modelname.c_str(), std::to_string(scale).c_str()); - sprintf(modelpath, "%s/%s-x%s.bin", model.c_str(), modelname.c_str(), std::to_string(scale).c_str()); - } - else{ - sprintf(parampath, "%s/%s.param", model.c_str(), modelname.c_str()); - sprintf(modelpath, "%s/%s.bin", model.c_str(), modelname.c_str()); - } + path_t parampath = model + ".param"; + path_t modelpath = model + ".bin"; #endif - path_t paramfullpath = sanitize_filepath(parampath); - path_t modelfullpath = sanitize_filepath(modelpath); - #if _WIN32 CoInitializeEx(NULL, COINIT_MULTITHREADED); #endif @@ -781,17 +729,14 @@ int main(int argc, char** argv) uint32_t heap_budget = ncnn::get_gpu_device(gpuid[i])->get_heap_budget(); // more fine-grained tilesize policy here - if (model.find(PATHSTR("models")) != path_t::npos) - { - if (heap_budget > 1900) - tilesize[i] = 200; - else if (heap_budget > 550) - tilesize[i] = 100; - else if (heap_budget > 190) - tilesize[i] = 64; - else - tilesize[i] = 32; - } + if (heap_budget > 1900) + tilesize[i] = 200; + else if (heap_budget > 550) + tilesize[i] = 100; + else if (heap_budget > 190) + tilesize[i] = 64; + else + tilesize[i] = 32; } { @@ -801,7 +746,7 @@ int main(int argc, char** argv) { realesrgan[i] = new RealESRGAN(gpuid[i], tta_mode); - realesrgan[i]->load(paramfullpath, modelfullpath); + realesrgan[i]->load(parampath, modelpath); realesrgan[i]->scale = scale; realesrgan[i]->tilesize = tilesize[i];