diff --git a/dumas_h3_latent_upscale.py b/dumas_h3_latent_upscale.py index 81ecf85..eef86b5 100644 --- a/dumas_h3_latent_upscale.py +++ b/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)