Unload latent upscale model on OOM

dumas_h3_latent_upscale.py

@@ -1,4 +1,5 @@
 from functools import lru_cache
+import gc
 import glob
 import math
 import os
@@ -489,6 +490,7 @@ def _upscale_video_model_core(video, param):
     model = load_upscale_model(model_name, dev, precision)
     norm_mean, norm_std = _make_norm_tensors(dev, compute_dtype)
 
+    try:
         with torch.inference_mode():
             s = s.sub(norm_mean).div(norm_std)
             out = model(s, scale=eff, target_size=(t, h_out, w_out))
@@ -496,8 +498,13 @@ def _upscale_video_model_core(video, param):
             out = out.mul(norm_std).add(norm_mean)
 
         out = out.to(device="cpu", dtype=orig_dtype)
-    unload_upscale_model(model_name, dev, precision)
         return out, h_out, w_out
+    finally:
+        unload_upscale_model(model_name, dev, precision)
+        try:
+            gc.collect()
+        except Exception:
+            pass
 
 
 def upscale_video_model(video, param):
@@ -509,6 +516,15 @@ def upscale_video_model(video, param):
         smaller = _shrink_model_tile_param(param)
         if smaller is None:
             raise
+        try:
+            gc.collect()
+        except Exception:
+            pass
+        if torch.cuda.is_available():
+            try:
+                torch.cuda.empty_cache()
+            except Exception:
+                pass
         return upscale_video_model(video, smaller)