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,15 +490,21 @@ 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)
 
-    with torch.inference_mode():
-        s = s.sub(norm_mean).div(norm_std)
-        out = model(s, scale=eff, target_size=(t, h_out, w_out))
-        del s
-        out = out.mul(norm_std).add(norm_mean)
+    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))
+            del s
+            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
+        out = out.to(device="cpu", dtype=orig_dtype)
+        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)