waifu2x on cpu

nihui · Jan 2, 2021 · c564778 · c564778
1 parent faa34ed
commit c564778
Show file tree

Hide file tree

Showing 3 changed files with 396 additions and 22 deletions.
diff --git a/src/main.cpp b/src/main.cpp
@@ -109,7 +109,7 @@ static void print_usage()
     fprintf(stdout, "  -s scale             upscale ratio (1/2, default=2)\n");
     fprintf(stdout, "  -t tile-size         tile size (>=32/0=auto, default=0) can be 0,0,0 for multi-gpu\n");
     fprintf(stdout, "  -m model-path        waifu2x model path (default=models-cunet)\n");
-    fprintf(stdout, "  -g gpu-id            gpu device to use (default=auto) can be 0,1,2 for multi-gpu\n");
+    fprintf(stdout, "  -g gpu-id            gpu device to use (-1=cpu, default=auto) can be 0,1,2 for multi-gpu\n");
     fprintf(stdout, "  -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(stdout, "  -x                   enable tta mode\n");
     fprintf(stdout, "  -f format            output image format (jpg/png/webp, default=ext/png)\n");
@@ -775,7 +775,7 @@ int main(int argc, char** argv)
     int gpu_count = ncnn::get_gpu_count();
     for (int i=0; i<use_gpu_count; i++)
     {
-        if (gpuid[i] < 0 || gpuid[i] >= gpu_count)
+        if (gpuid[i] < -1 || gpuid[i] >= gpu_count)
         {
             fprintf(stderr, "invalid gpu device\n");
 
@@ -787,16 +787,31 @@ int main(int argc, char** argv)
     int total_jobs_proc = 0;
     for (int i=0; i<use_gpu_count; i++)
     {
-        int gpu_queue_count = ncnn::get_gpu_info(gpuid[i]).compute_queue_count();
-        jobs_proc[i] = std::min(jobs_proc[i], gpu_queue_count);
-        total_jobs_proc += jobs_proc[i];
+        if (gpuid[i] == -1)
+        {
+            jobs_proc[i] = std::min(jobs_proc[i], cpu_count);
+            total_jobs_proc += 1;
+        }
+        else
+        {
+            int gpu_queue_count = ncnn::get_gpu_info(gpuid[i]).compute_queue_count();
+            jobs_proc[i] = std::min(jobs_proc[i], gpu_queue_count);
+            total_jobs_proc += jobs_proc[i];
+        }
     }
 
     for (int i=0; i<use_gpu_count; i++)
     {
         if (tilesize[i] != 0)
             continue;
 
+        if (gpuid[i] == -1)
+        {
+            // cpu only
+            tilesize[i] = 4000;
+            continue;
+        }
+
         uint32_t heap_budget = ncnn::get_gpu_device(gpuid[i])->get_heap_budget();
 
         // more fine-grained tilesize policy here
@@ -830,7 +845,9 @@ int main(int argc, char** argv)
 
         for (int i=0; i<use_gpu_count; i++)
         {
-            waifu2x[i] = new Waifu2x(gpuid[i], tta_mode);
+            int num_threads = gpuid[i] == -1 ? jobs_proc[i] : 1;
+
+            waifu2x[i] = new Waifu2x(gpuid[i], tta_mode, num_threads);
 
             waifu2x[i]->load(paramfullpath, modelfullpath);
 
@@ -863,10 +880,17 @@ int main(int argc, char** argv)
                 int total_jobs_proc_id = 0;
                 for (int i=0; i<use_gpu_count; i++)
                 {
-                    for (int j=0; j<jobs_proc[i]; j++)
+                    if (gpuid[i] == -1)
                     {
                         proc_threads[total_jobs_proc_id++] = new ncnn::Thread(proc, (void*)&ptp[i]);
                     }
+                    else
+                    {
+                        for (int j=0; j<jobs_proc[i]; j++)
+                        {
+                            proc_threads[total_jobs_proc_id++] = new ncnn::Thread(proc, (void*)&ptp[i]);
+                        }
+                    }
                 }
             }