diff --git a/src/h3_blackwell_runtime/checkpoint.py b/src/h3_blackwell_runtime/checkpoint.py index 4ddd308..11b9d5c 100644 --- a/src/h3_blackwell_runtime/checkpoint.py +++ b/src/h3_blackwell_runtime/checkpoint.py @@ -32,7 +32,7 @@ class H3Checkpoint: with fastsafe_open(filenames=[self.path], nogds=True, device=self._fast_device()) as checkpoint: self._no_mmap_tensors = { name: checkpoint.get_tensor(name).clone().detach() - for name in checkpoint.get_keys() + for name in checkpoint.keys() } return self._no_mmap_tensors diff --git a/src/h3_blackwell_runtime/vae_decoder.py b/src/h3_blackwell_runtime/vae_decoder.py index 39dba9c..67763e3 100644 --- a/src/h3_blackwell_runtime/vae_decoder.py +++ b/src/h3_blackwell_runtime/vae_decoder.py @@ -172,7 +172,7 @@ class MiniMaxH3VideoVAE(nn.Module): with fastsafe_open(filenames=[str(path)], nogds=True, device=fast_device) as checkpoint: available_weights = { name: checkpoint.get_tensor(name).clone().detach() - for name in checkpoint.get_keys() + for name in checkpoint.keys() } available = set(available_weights) missing = sorted(set(expected) - available) diff --git a/tools/direct_t2v_preview.py b/tools/direct_t2v_preview.py index e62b0ca..af60f7c 100644 --- a/tools/direct_t2v_preview.py +++ b/tools/direct_t2v_preview.py @@ -4,6 +4,8 @@ import argparse import os from pathlib import Path import subprocess +import time +from datetime import datetime, timezone import warnings warnings.filterwarnings("ignore", message="Found GPU0 NVIDIA GB10 which is of cuda capability 12.1.*", category=UserWarning) @@ -32,11 +34,15 @@ parser.add_argument("--attention", choices=("sage2", "sdpa", "sage3"), default=" parser.add_argument("--model-timesteps-capture", type=Path, help="Directory containing captured input_XX.pt H3 timesteps for strict parity checks.") parser.add_argument("--profile-memory", action="store_true") args = parser.parse_args() +started = time.perf_counter() +last_report = started def report_memory(stage: str) -> None: + global last_report if not args.profile_memory: return + now = time.perf_counter() rss_kb = 0 try: with open("/proc/self/status", encoding="utf-8") as file: @@ -48,7 +54,8 @@ def report_memory(stage: str) -> None: pass cuda_alloc = torch.cuda.memory_allocated() / 1024**3 if torch.cuda.is_available() else 0.0 cuda_reserved = torch.cuda.memory_reserved() / 1024**3 if torch.cuda.is_available() else 0.0 - print({"stage": stage, "rss_gb": round(rss_kb / 1024**2, 3), "cuda_alloc_gb": round(cuda_alloc, 3), "cuda_reserved_gb": round(cuda_reserved, 3), "fast_safetensors": os.getenv("H3_FAST_SAFETENSORS", ""), "disable_mmap": os.getenv("H3_DISABLE_MMAP", "")}, flush=True) + print({"stage": stage, "ts": datetime.now(timezone.utc).isoformat(), "epoch_s": round(time.time(), 3), "elapsed_s": round(now - started, 3), "delta_s": round(now - last_report, 3), "rss_gb": round(rss_kb / 1024**2, 3), "cuda_alloc_gb": round(cuda_alloc, 3), "cuda_reserved_gb": round(cuda_reserved, 3), "fast_safetensors": os.getenv("H3_FAST_SAFETENSORS", ""), "disable_mmap": os.getenv("H3_DISABLE_MMAP", "")}, flush=True) + last_report = now report_memory("start") checkpoint = H3Checkpoint("/models/minimax_h3_fl2va_pruned_nvfp4.safetensors")