fix that gpu watcher is not completely clear after generation

app.py

@@ -9,11 +9,10 @@ from typing import List, Tuple
 import gradio as gr
 import spaces
 
-from gpu_info import watch_gpu_memory
+from gpu_info import stop_watcher, watch_gpu_memory
 
 PWD = os.path.dirname(__file__)
 CHECKPOINTS_PATH = "/data/checkpoints"
-# CHECKPOINTS_PATH = os.path.join(PWD, "checkpoints")
 LOG_DIR = os.path.join(PWD, "logs")
 os.makedirs(LOG_DIR, exist_ok=True)
 
@@ -346,7 +345,7 @@ def generate_video(
     log.info(f"Time taken: {end_time - start_time} s")
 
     # stop the watcher
-    watcher.cancel()
+    stop_watcher()
 
     video = videos[0]
 

gpu_info.py

@@ -47,6 +47,16 @@ def watch_gpu_memory(interval: int = 1, callback: Callable[[List[int]], None] =
     return gpu_memory_watcher
 
 
+def stop_watcher():
+    global gpu_memory_watcher
+    if gpu_memory_watcher is None:
+        return
+
+    gpu_memory_watcher.cancel()
+    del gpu_memory_watcher
+    gpu_memory_watcher = None
+
+
 if __name__ == "__main__":
     from time import sleep