Upload app.py

app.py

@@ -8,6 +8,7 @@ gradio API for every request; `reference_encoder` stays here, next to the autoen
 from __future__ import annotations
 
 import json
+import gc
 import os
 import random
 import re
@@ -17,6 +18,11 @@ import time
 import traceback
 from functools import cache
 
+# Before torch exists: 72 GiB of resident weights leave the rest of the card in pieces, and a request that needs one
+# more large contiguous block then fails on free memory it cannot use in one piece. Expandable segments let the
+# allocator grow a block instead of hunting for one, which is the single biggest difference on this Space.
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+
 # Before anything that could initialize CUDA: `import spaces` patches `torch.cuda` so the 72 GiB load can happen at
 # startup rather than on GPU time.
 import spaces
@@ -44,7 +50,7 @@ MAX_SEQUENCE = int(os.environ.get("H3_MAX_SEQUENCE", "90000"))
 # An attached adapter adds its own layers and their activations to the same card, so the ceiling above is not the
 # ceiling any more. Refusing a request that is over the reduced one is a sentence on screen; letting it through is a
 # dead worker and a bare "runtime error".
-LORA_SEQUENCE_FRACTION = float(os.environ.get("H3_LORA_SEQUENCE_FRACTION", "0.8"))
+LORA_SEQUENCE_FRACTION = float(os.environ.get("H3_LORA_SEQUENCE_FRACTION", "0.75"))
 
 
 def sequence_ceiling(loras=()) -> int:
@@ -995,6 +1001,13 @@ def apply_loras(transformer, loras) -> list[str]:
 
     for name in list(getattr(transformer, "peft_config", None) or {}):
         transformer.delete_adapters(name)
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
+    base_dtype = next(
+        (param.dtype for key, param in transformer.named_parameters() if ".lora_" not in key), torch.bfloat16
+    )
 
     names, scales = [], []
     for index, (path, scale) in enumerate(loras):
@@ -1011,6 +1024,14 @@ def apply_loras(transformer, loras) -> list[str]:
             else:
                 state_dict = _convert_kohya_lora(state_dict) or state_dict
         state_dict = _fit_to_transformer(transformer, state_dict)
+        # Cast on the host, before anything crosses onto the card. A float32 adapter injected as-is is put on the
+        # card at float32 and cast afterwards, so for a moment both copies are resident - on top of 72 GiB of
+        # weights that is exactly the block the allocator cannot find.
+        state_dict = {
+            key: (value.to(base_dtype) if hasattr(value, "is_floating_point") and value.is_floating_point()
+                  else value)
+            for key, value in state_dict.items()
+        }
         name = f"lora{index}"
         try:
             transformer.load_lora_adapter(state_dict, adapter_name=name, prefix=_lora_prefix(state_dict))
@@ -1021,6 +1042,8 @@ def apply_loras(transformer, loras) -> list[str]:
                 "It is almost always an adapter for another partition or another base model - a `transformer/` "
                 "adapter does not fit the `transformer_ref/` half this Space runs."
             ) from error
+        del state_dict
+        gc.collect()
         names.append(name)
         scales.append(float(scale))
 
@@ -1036,6 +1059,11 @@ def apply_loras(transformer, loras) -> list[str]:
                 param.data = param.data.to(device=base.device, dtype=base.dtype)
 
     transformer.set_adapters(names, scales)
+    # The load and the cast both leave freed blocks behind; handing them back before the denoise starts is what
+    # keeps the first large allocation of the run from failing on a card that has the memory but not in one piece.
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
     return names
 
 
@@ -1222,11 +1250,47 @@ def _generate(prompt_embeds, text_token_tags, references, height, width, num_fra
         state = _run()
     except RuntimeError as error:
         message = str(error).lower()
-        if "no available kernel" not in message and "cudnn" not in message:
+        kernel = "no available kernel" in message or "cudnn" in message
+        # `NVML_SUCCESS == r INTERNAL ASSERT FAILED` out of `CUDACachingAllocator.cpp` is not a bug to report to
+        # PyTorch: it is the allocator failing to grow, i.e. out of memory, with its own error path falling over on
+        # the way to saying so. Same treatment as a plain OOM.
+        starved = (
+            "out of memory" in message
+            or "nvml" in message
+            or "cudacachingallocator" in message
+            or "cuda error" in message
+        )
+        if not (kernel or starved):
             raise
-        print(f"[ref2va] attention backend refused this shape ({error}); retrying on `native`", flush=True)
+
+        gc.collect()
+        try:
+            torch.cuda.empty_cache()
+            torch.cuda.synchronize()
+        except Exception:  # noqa: BLE001
+            pass
         _backend("native")
-        state = _run()
+        print(f"[ref2va] first attempt failed ({type(error).__name__}); retrying once on a cleared card", flush=True)
+
+        try:
+            state = _run()
+        except RuntimeError as second:
+            second_message = str(second).lower()
+            if not (
+                "out of memory" in second_message
+                or "nvml" in second_message
+                or "cudacachingallocator" in second_message
+            ):
+                raise
+            traceback.print_exc()
+            raise gr.Error(
+                f"The card ran out of memory at {width}x{height}, {num_frames / FPS:.1f} s, "
+                f"{len(references)} reference(s)"
+                f"{f' and {len(loras or ())} lora' if loras else ''}. "
+                "A lora takes its share of the same card, so a setting that worked without one can be too much "
+                "with it. Pick a smaller canvas — a 1:1 one is the smallest — shorten the duration, or drop a "
+                "reference, and it goes through."
+            ) from second
 
     return state.get("videos")[0], state.get("audio")[0].cpu(), state.get("sampling_rate")