Bypass learned latent upscaler on low VRAM

dumas_h3_latent_upscale.py

@@ -46,6 +46,7 @@ LOGGER = logging.getLogger(__name__)
 MP_UNIT = 1024 * 1024
 RES_MULTIPLE = 32
 CUDA_MODEL_CHUNK_LENGTH = 17
+CUDA_MODEL_MIN_TOTAL_VRAM_BYTES = 10 * 1024 * 1024 * 1024
 _MODEL_HOLD_DEPTH = 0
 
 
@@ -76,6 +77,42 @@ def _should_retry_temporal_before_spatial(param):
     return _uses_cuda_model_upscale(param) and int(param.get("chunk_length", 0) or 0) > CUDA_MODEL_CHUNK_LENGTH
 
 
+def _cuda_total_memory_bytes():
+    try:
+        if not torch.cuda.is_available():
+            return 0
+        if hasattr(torch.cuda, "mem_get_info"):
+            _free, total = torch.cuda.mem_get_info()
+            return int(total)
+        current_device = torch.cuda.current_device() if hasattr(torch.cuda, "current_device") else 0
+        props = torch.cuda.get_device_properties(current_device)
+        return int(getattr(props, "total_memory", 0) or 0)
+    except Exception:
+        return 0
+
+
+def _should_skip_cuda_model_upscale(param):
+    total = _cuda_total_memory_bytes()
+    return _uses_cuda_model_upscale(param) and 0 < total < CUDA_MODEL_MIN_TOTAL_VRAM_BYTES
+
+
+def _fallback_to_interp(video, param, reason):
+    method = param.get("method", "bilinear")
+    LOGGER.warning("H3 latent upscale model %s; using %s interpolation instead", reason, method)
+    try:
+        gc.collect()
+    except Exception:
+        pass
+    if torch.cuda.is_available():
+        try:
+            torch.cuda.empty_cache()
+        except Exception:
+            pass
+    interp_param = dict(param)
+    interp_param["mode"] = "interp"
+    return upscale_video_interp(video, interp_param)
+
+
 def _retry_with_smaller_temporal(video, param, upscaler, exc):
     if not _is_oom_error(exc):
         raise exc
@@ -902,15 +939,22 @@ def upscale_latent_video(video, param):
     if mode == "off":
         return video, video.shape[-2], video.shape[-1]
     if mode == "model":
+        if _should_skip_cuda_model_upscale(param):
+            return _fallback_to_interp(video, param, "requires more than this card's VRAM")
         model_name = param.get("model_name")
         device = param.get("device", "cuda")
         precision = param.get("precision", "fp16")
         dev = torch.device(device if (device == "cpu" or torch.cuda.is_available()) else "cpu")
-        with _hold_upscale_model_loaded():
-            try:
-                return _upscale_video_temporal_chunks(video, param, upscale_video_model)
-            finally:
-                _unload_upscale_model_now(model_name, dev, precision)
+        try:
+            with _hold_upscale_model_loaded():
+                try:
+                    return _upscale_video_temporal_chunks(video, param, upscale_video_model)
+                finally:
+                    _unload_upscale_model_now(model_name, dev, precision)
+        except RuntimeError as exc:
+            if not _is_oom_error(exc):
+                raise
+            return _fallback_to_interp(video, param, "exhausted GPU memory")
     return _upscale_video_temporal_chunks(video, param, upscale_video_interp)