Merge branch 'free_vram'

Author: nagadomi

iw3/gui.py

@@ -3,7 +3,6 @@
 import sys
 import os
 from os import path
-import gc
 import traceback
 import functools
 from time import time
@@ -16,6 +15,7 @@
     create_parser, set_state_args, iw3_main,
     is_text, is_video, is_output_dir, is_yaml, make_output_filename,
     has_rembg_model)
+from nunif.initializer import gc_collect
 from nunif.device import mps_is_available, xpu_is_available
 from nunif.utils.image_loader import IMG_EXTENSIONS as LOADER_SUPPORTED_EXTENSIONS
 from nunif.utils.video import VIDEO_EXTENSIONS as KNOWN_VIDEO_EXTENSIONS, has_nvenc
@@ -1096,9 +1096,7 @@ def parse_args(self):
             self.depth_model = None
             self.depth_model_type = None
             self.depth_model_device_id = None
-            gc.collect()
-            if torch.cuda.is_available():
-                torch.cuda.empty_cache()
+            gc_collect()
 
         remove_bg = self.chk_rembg.GetValue()
         bg_model_type = self.cbo_bg_model.GetValue()
@@ -1241,9 +1239,7 @@ def on_exit_worker(self, result):
         self.update_start_button_state()
 
         # free vram
-        gc.collect()
-        if torch.cuda.is_available:
-            torch.cuda.empty_cache()
+        gc_collect()
 
     def on_click_btn_cancel(self, event):
         self.suspend_event.set()

iw3/utils.py

@@ -10,6 +10,7 @@
 import math
 from tqdm import tqdm
 from PIL import ImageDraw, Image
+from nunif.initializer import gc_collect
 from nunif.utils.image_loader import ImageLoader
 from nunif.utils.pil_io import load_image_simple
 from nunif.models import load_model  # , compile_model
@@ -1673,10 +1674,12 @@ def iw3_main(args):
         if not args.recursive:
             image_files = ImageLoader.listdir(args.input)
             process_images(image_files, args.output, args, depth_model, side_model, title="Images")
+            gc_collect()
             for video_file in VU.list_videos(args.input):
                 if args.state["stop_event"] is not None and args.state["stop_event"].is_set():
                     return args
                 process_video(video_file, args.output, args, depth_model, side_model)
+                gc_collect()
         else:
             subdirs = list_subdir(args.input, include_root=True, excludes=args.output)
             for input_dir in subdirs:
@@ -1685,10 +1688,12 @@ def iw3_main(args):
                 if image_files:
                     process_images(image_files, output_dir, args, depth_model, side_model,
                                    title=path.relpath(input_dir, args.input))
+                    gc_collect()
                 for video_file in VU.list_videos(input_dir):
                     if args.state["stop_event"] is not None and args.state["stop_event"].is_set():
                         return args
                     process_video(video_file, output_dir, args, depth_model, side_model)
+                    gc_collect()
 
     elif is_yaml(args.input):
         config = export_config.ExportConfig.load(args.input)
@@ -1712,6 +1717,7 @@ def iw3_main(args):
             if args.state["stop_event"] is not None and args.state["stop_event"].is_set():
                 return args
             process_video(video_file, args.output, args, depth_model, side_model)
+            gc_collect()
     elif is_video(args.input):
         process_video(args.input, args.output, args, depth_model, side_model)
     elif is_image(args.input):

nunif/initializer.py

@@ -3,6 +3,7 @@
 import random
 import numpy as np
 import secrets
+import gc
 
 
 def disable_image_lib_threads():
@@ -35,3 +36,16 @@ def set_seed(seed):
         torch.cuda.manual_seed_all(seed)
     random.seed(seed)
     np.random.seed(seed)
+
+
+def gc_collect():
+    gc.collect()
+
+    if hasattr(torch, "_dynamo") and hasattr(torch._dynamo, "reset"):
+        torch._dynamo.reset()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+    if torch.backends.mps.is_available():
+        torch.mps.empty_cache()
+    if hasattr(torch, "xpu") and torch.xpu.is_available():
+        torch.xpu.empty_cache()