Update app.py

app.py

@@ -91,10 +91,7 @@ torch.cuda.current_device = lambda: 0
 torch.cuda.device_count = lambda: 0
 torch.cuda.set_device = lambda *a, **kw: None
 torch.cuda.init = lambda *a, **kw: None
-# memory_stats возвращаем пустым — GPUMemoryManager посчитает фрагментацию = 0.
 torch.cuda.memory_stats = lambda *a, **kw: {}
-# mem_get_info возвращаем маленькими числами, чтобы defragment_memory не пытался
-# выделить огромный тензор на "device 0" (это CPU) и не упал по OOM.
 torch.cuda.mem_get_info = lambda *a, **kw: (100, 100)
 torch.cuda.empty_cache = lambda *a, **kw: None
 torch.cuda.synchronize = lambda *a, **kw: None
@@ -104,6 +101,42 @@ torch.cuda.manual_seed_all = lambda *a, **kw: None
 torch.cuda.get_device_properties = lambda *a, **kw: _DummyDeviceProperties()
 
 
+# =====================================================================
+# ПАТЧ 3.5. Принудительно грузим чекпоинты на CPU.
+# Чекпоинты PlayDiffusion (last_250k_fixed.pkl, xlsr2_1b_v2_custom.pt и др.)
+# сохранены с CUDA-тензорами. Без map_location="cpu" torch.load падает с
+# "Attempting to deserialize object on a CUDA device but ... is False".
+# Патчим torch.load так, чтобы он всегда использовал map_location="cpu".
+# =====================================================================
+_orig_torch_load = torch.load
+
+
+def _remap_map_location(ml):
+    """Любое CUDA-устройство (или None) превращаем в 'cpu'."""
+    if ml is None:
+        return "cpu"
+    if isinstance(ml, str) and ml.startswith("cuda"):
+        return "cpu"
+    if isinstance(ml, torch.device) and ml.type == "cuda":
+        return "cpu"
+    return ml
+
+
+def _patched_torch_load(*args, **kwargs):
+    # torch.load(f, map_location=None, ...) — второй позиционный аргумент
+    # это map_location. Обрабатываем оба варианта: позиционный и ключевой.
+    if len(args) >= 2:
+        args = (args[0], _remap_map_location(args[1])) + args[2:]
+    else:
+        kwargs["map_location"] = _remap_map_location(
+            kwargs.get("map_location")
+        )
+    return _orig_torch_load(*args, **kwargs)
+
+
+torch.load = _patched_torch_load
+
+
 # =====================================================================
 # Обычные импорты (все патчи уже применены).
 # =====================================================================
@@ -123,9 +156,7 @@ _orig_load_preset = PlayDiffusion.load_preset
 def _patched_load_preset(self, **kwargs):
     preset = _orig_load_preset(self, **kwargs)
     if str(self.device).startswith("cpu"):
-        # Речевой токенизатор (XLSR-1B) — в float32
         preset["speech_tokenizer"]["dtype"] = torch.float32
-        # Вокодер BigVGAN — в float32
         preset["vocoder"]["dtype"] = torch.float32
     return preset