From a56bc058dc7fa5ef4e41aa607fe232059f697780 Mon Sep 17 00:00:00 2001
From: neuregex <sanightblack@gmail.com>
Date: Fri, 5 Jun 2026 17:55:01 -0600
Subject: [PATCH] fix(offload): aparcar experto inactivo en disco (mmap) en vez
 de copiar 14.6GB a RAM (v0.3.14)

---
 bernini/loader.py |  5 ++++-
 bernini/model.py  | 21 +++++++++++++++------
 pyproject.toml    |  2 +-
 3 files changed, 20 insertions(+), 8 deletions(-)

diff --git a/bernini/loader.py b/bernini/loader.py
index 9ee0003..5ed33fb 100644
--- a/bernini/loader.py
+++ b/bernini/loader.py
@@ -211,7 +211,10 @@ def _load_prefp8(model_dir, subfolder):
     for shard in sorted(glob.glob(os.path.join(model_dir, subfolder, "*.safetensors"))):
         sd.update(_safe_load(shard, keep))
     missing, unexpected = m.load_state_dict(sd, strict=False, assign=True)
-    m._bernini_keep = keep   # mantiene vivos mmap+file mientras viva el modelo
+    m._bernini_keep = keep      # mantiene vivos mmap+file mientras viva el modelo
+    m._bernini_cpu_sd = sd      # pesos read-only mmap (CPU, perezosos): para volver a
+    #                             "aparcar" el experto en disco al hacer offload (to_idle)
+    #                             sin copiar 14.6GB a RAM.
     if missing:
         print(f"[BerniniR] fp8 load {subfolder}: {len(missing)} missing (ej: {missing[:2]})")
     if unexpected:
diff --git a/bernini/model.py b/bernini/model.py
index cff3de3..0a0205c 100644
--- a/bernini/model.py
+++ b/bernini/model.py
@@ -119,12 +119,21 @@ class BerniniExpert(torch.nn.Module):
             self.t.scale_shift_table.data = self.t.scale_shift_table.data.to(device)
 
     def to_idle(self):
-        """Manda TODO el experto (residentes + swap + no-bloque) a CPU (inactivo). Si YA
-        está en CPU no hace nada: llamar .to('cpu') sobre los tensores fp8 de SOLO-LECTURA
-        (respaldados por mmap) dispara un 'access violation' -> torch recrea el Parameter
-        sobre memoria no-escribible. Solo movemos de verdad cuando viene de GPU (ahí
-        .to('cpu') crea copias escribibles nuevas, sin tocar el mmap original)."""
-        if self.device.type != "cpu":
+        """Aparca el experto FUERA de la GPU sin meterlo en RAM. Si ya está en CPU, nada.
+        Si está en GPU, RE-ASIGNA los pesos read-only mmap originales (CPU, perezosos,
+        ~0 RAM) y libera la VRAM -> el experto inactivo vuelve a 'vivir en disco'.
+
+        Antes hacía `self.t.to('cpu')`, que copia ~14.6GB a RAM escribible y, con la RAM
+        llena, revienta (access violation / WinError 1450/1455). Re-asignar el state_dict
+        mmap (assign=True) NO copia: solo reapunta los params al mmap y suelta los tensores
+        GPU (se liberan con empty_cache en el sampler). Al reactivarse, to_active vuelve a
+        streamear los pesos a GPU desde el mmap (file cache, no commit de RAM)."""
+        if self.device.type == "cpu":
+            return
+        cpu_sd = getattr(self.t, "_bernini_cpu_sd", None)
+        if cpu_sd is not None:
+            self.t.load_state_dict(cpu_sd, strict=False, assign=True)   # re-apunta a mmap; libera GPU
+        else:
             self.t.to("cpu")
 
     @property
diff --git a/pyproject.toml b/pyproject.toml
index 0793caf..0e76fdb 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,7 +1,7 @@
 [project]
 name = "comfyui-berninir"
 description = "ByteDance Bernini-R (Wan2.2-T2V-A14B + source-id RoPE + APG multi-condition guidance) for ComfyUI: t2v/t2i, image & video editing, reference-to-video."
-version = "0.3.13"
+version = "0.3.14"
 license = { file = "LICENSE" }
 requires-python = ">=3.10"
 dependencies = [