clear gpu vram for film

libs/film/eval/interpolator.py

@@ -149,6 +149,10 @@ class Interpolator:
     self._align = align or None
     self._block_shape = block_shape or None
 
+  def __del__(self):
+    del self._model
+    tf.keras.backend.clear_session()
+  
   def interpolate(self, x0: np.ndarray, x1: np.ndarray,
                   dt: np.ndarray) -> np.ndarray:
     """Generates an interpolated frame between given two batches of frames.
@@ -171,6 +175,8 @@ class Interpolator:
     result = self._model(inputs, training=False)
     image = result['image']
 
+    del self._model
+
     if self._align is not None:
       image = tf.image.crop_to_bounding_box(image, **bbox_to_crop)
     return image.numpy()

libs/film/predict.py

@@ -68,4 +68,7 @@ class Predictor(cog.BasePredictor):
                 input_frames, times_to_interpolate, self.interpolator))
         print('Interpolated frames generated.')
 
+        del self.interpolator
+        tf.keras.backend.clear_session()
+
         return frames

main.py

@@ -64,7 +64,8 @@ import requests
 import json
 
 from huggingface_hub import hf_hub_download, HfApi
-
+from numba import cuda
+from multiprocessing import Pool
 
 # Inputs ===================================================================================================
 
@@ -1220,7 +1221,7 @@ def run_app(images, video_path, train_steps=100, inference_steps=10, fps=12, bg_
     return out_vid+'.webm', results
 
 
-def run_interpolate_frames(frame1, frame2, times_to_interp):
+def interpolate_frames(frame1, frame2, times_to_interp):
     film = Predictor()
     film.setup()
 
@@ -1238,8 +1239,16 @@ def run_interpolate_frames(frame1, frame2, times_to_interp):
     for r in results:
         r.thumbnail((width, height))
 
+    del film
+
     return results, getThumbnails(results)
 
+def run_interpolate_frames(frame1, frame2, times_to_interp):
+    with Pool() as pool:
+        results = pool.starmap(interpolate_frames, [(frame1, frame2, times_to_interp)])
+    print(results[0])
+    return results[0]
+
 
 """
 train_steps = 100