diff --git a/pyproject.toml b/pyproject.toml index 318e4f1..ea8a901 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "seedvr2_videoupscaler" description = "SeedVR2 official ComfyUI integration: ByteDance-Seed's one-step diffusion-based video/image upscaling with memory-efficient inference" -version = "2.5.1" +version = "2.5.2" authors = [ {name = "numz"}, {name = "adrientoupet"} diff --git a/src/common/distributed/basic.py b/src/common/distributed/basic.py index 195a714..8b0360f 100644 --- a/src/common/distributed/basic.py +++ b/src/common/distributed/basic.py @@ -47,7 +47,7 @@ def get_device() -> torch.device: """ Get current rank device. """ - if torch.mps.is_available(): + if hasattr(torch, 'mps') and callable(getattr(torch.mps, 'is_available', None)) and torch.mps.is_available(): return torch.device("mps") return torch.device("cuda", get_local_rank()) diff --git a/src/data/image/transforms/area_resize.py b/src/data/image/transforms/area_resize.py index a360c1f..2b6dc22 100644 --- a/src/data/image/transforms/area_resize.py +++ b/src/data/image/transforms/area_resize.py @@ -31,7 +31,7 @@ class AreaResize: self.max_area = max_area self.downsample_only = downsample_only self.interpolation = interpolation - if torch.mps.is_available(): + if hasattr(torch, 'mps') and callable(getattr(torch.mps, 'is_available', None)) and torch.mps.is_available(): self.interpolation = InterpolationMode.BILINEAR def __call__(self, image: Union[torch.Tensor, Image.Image]): diff --git a/src/data/image/transforms/na_resize.py b/src/data/image/transforms/na_resize.py index 14ddb5b..ecb49d5 100644 --- a/src/data/image/transforms/na_resize.py +++ b/src/data/image/transforms/na_resize.py @@ -26,7 +26,7 @@ def NaResize( max_resolution: int = 0, interpolation: InterpolationMode = InterpolationMode.BICUBIC, ): - Interpolation = InterpolationMode.BILINEAR if torch.mps.is_available() else interpolation + Interpolation = InterpolationMode.BILINEAR if (hasattr(torch, 'mps') and callable(getattr(torch.mps, 'is_available', None)) and torch.mps.is_available()) else interpolation if mode == "area": return AreaResize( max_area=resolution**2, diff --git a/src/data/image/transforms/side_resize.py b/src/data/image/transforms/side_resize.py index 4b6c664..5334e98 100644 --- a/src/data/image/transforms/side_resize.py +++ b/src/data/image/transforms/side_resize.py @@ -30,7 +30,7 @@ class SideResize: self.max_size = max_size self.downsample_only = downsample_only self.interpolation = interpolation - if torch.mps.is_available(): + if hasattr(torch, 'mps') and callable(getattr(torch.mps, 'is_available', None)) and torch.mps.is_available(): self.interpolation = InterpolationMode.BILINEAR def __call__(self, image: Union[torch.Tensor, Image.Image]): diff --git a/src/optimization/compatibility.py b/src/optimization/compatibility.py index f270c17..f640f02 100644 --- a/src/optimization/compatibility.py +++ b/src/optimization/compatibility.py @@ -207,7 +207,7 @@ class FP8CompatibleDiT(torch.nn.Module): self._convert_rope_freqs(target_dtype=self.compute_dtype) self.debug.end_timer("_convert_rope_freqs", "RoPE freqs conversion") - if torch.mps.is_available(): + if hasattr(torch, 'mps') and callable(getattr(torch.mps, 'is_available', None)) and torch.mps.is_available(): self.debug.log(f"Also converting NaDiT parameters/buffers for MPS backend", category="setup", force=True) self.debug.start_timer("_force_nadit_precision") self._force_nadit_precision(target_dtype=self.compute_dtype) @@ -510,7 +510,7 @@ class FP8CompatibleDiT(torch.nn.Module): k = k.view(batch_size, seq_len, num_heads, head_dim).transpose(1, 2) v = v.view(batch_size, seq_len, num_heads, head_dim).transpose(1, 2) - if torch.mps.is_available(): + if hasattr(torch, 'mps') and callable(getattr(torch.mps, 'is_available', None)) and torch.mps.is_available(): attn_output = torch.nn.functional.scaled_dot_product_attention( q, k, v, dropout_p=0.0, diff --git a/src/optimization/memory_manager.py b/src/optimization/memory_manager.py index 71fa501..19715b0 100644 --- a/src/optimization/memory_manager.py +++ b/src/optimization/memory_manager.py @@ -80,7 +80,7 @@ def get_basic_vram_info(device: Optional[torch.device] = None) -> Dict[str, Any] elif not isinstance(device, torch.device): device = torch.device(device) free_memory, total_memory = torch.cuda.mem_get_info(device) - elif torch.mps.is_available(): + elif hasattr(torch, 'mps') and callable(getattr(torch.mps, 'is_available', None)) and torch.mps.is_available(): # MPS doesn't support per-device queries or mem_get_info # Use system memory as proxy mem = psutil.virtual_memory() @@ -100,7 +100,7 @@ def get_basic_vram_info(device: Optional[torch.device] = None) -> Dict[str, Any] # Initial VRAM check at module load vram_info = get_basic_vram_info(device=None) if "error" not in vram_info: - backend = "MPS" if torch.mps.is_available() else "CUDA" + backend = "MPS" if (hasattr(torch, 'mps') and callable(getattr(torch.mps, 'is_available', None)) and torch.mps.is_available()) else "CUDA" print(f"📊 Initial {backend} memory: {vram_info['free_gb']:.2f}GB free / {vram_info['total_gb']:.2f}GB total") else: print(f"⚠️ Memory check failed: {vram_info['error']} - No available backend!") @@ -129,7 +129,7 @@ def get_vram_usage(device: Optional[torch.device] = None, debug: Optional['Debug reserved = torch.cuda.memory_reserved(device) / (1024**3) max_allocated = torch.cuda.max_memory_allocated(device) / (1024**3) return allocated, reserved, max_allocated - elif torch.mps.is_available(): + elif hasattr(torch, 'mps') and callable(getattr(torch.mps, 'is_available', None)) and torch.mps.is_available(): # MPS doesn't support per-device queries - uses global memory tracking allocated = torch.mps.current_allocated_memory() / (1024**3) reserved = torch.mps.driver_allocated_memory() / (1024**3) @@ -235,11 +235,11 @@ def clear_memory(debug: Optional['Debug'] = None, deep: bool = False, force: boo if free_ratio < 0.05: should_clear = True if debug: - backend = "MPS" if torch.mps.is_available() else "VRAM" + backend = "MPS" if (hasattr(torch, 'mps') and callable(getattr(torch.mps, 'is_available', None)) and torch.mps.is_available()) else "VRAM" debug.log(f"{backend} pressure: {mem_info['free_gb']:.2f}GB free of {mem_info['total_gb']:.2f}GB", category="memory") # For non-MPS systems, also check system RAM separately - if not should_clear and not torch.mps.is_available(): + if not should_clear and not (hasattr(torch, 'mps') and callable(getattr(torch.mps, 'is_available', None)) and torch.mps.is_available()): mem = psutil.virtual_memory() if mem.available < mem.total * 0.05: should_clear = True @@ -265,7 +265,7 @@ def clear_memory(debug: Optional['Debug'] = None, deep: bool = False, force: boo if torch.cuda.is_available(): torch.cuda.empty_cache() torch.cuda.ipc_collect() - elif torch.mps.is_available(): + elif hasattr(torch, 'mps') and callable(getattr(torch.mps, 'is_available', None)) and torch.mps.is_available(): torch.mps.empty_cache() if debug: @@ -302,7 +302,7 @@ def clear_memory(debug: Optional['Debug'] = None, deep: bool = False, force: boo handle = _os_memory_lib.GetCurrentProcess() _os_memory_lib.SetProcessWorkingSetSize(handle, -1, -1) - elif torch.mps.is_available(): + elif hasattr(torch, 'mps') and callable(getattr(torch.mps, 'is_available', None)) and torch.mps.is_available(): # macOS with MPS import ctypes # Import only when needed import ctypes.util diff --git a/src/utils/debug.py b/src/utils/debug.py index e1040f6..24d0739 100644 --- a/src/utils/debug.py +++ b/src/utils/debug.py @@ -333,7 +333,7 @@ class Debug: } # VRAM metrics - if torch.cuda.is_available() or torch.mps.is_available(): + if torch.cuda.is_available() or (hasattr(torch, 'mps') and callable(getattr(torch.mps, 'is_available', None)) and torch.mps.is_available()): metrics['vram_allocated'], metrics['vram_reserved'], current_global_peak = get_vram_usage(device=None, debug=self) # Calculate peak since last log_memory_state @@ -346,7 +346,7 @@ class Debug: metrics['vram_free'] = vram_info["free_gb"] metrics['vram_total'] = vram_info["total_gb"] - backend = "MPS" if torch.mps.is_available() else "VRAM" + backend = "MPS" if (hasattr(torch, 'mps') and callable(getattr(torch.mps, 'is_available', None)) and torch.mps.is_available()) else "VRAM" metrics['summary_vram'] = (f" [{backend}] {metrics['vram_allocated']:.2f}GB allocated / " f"{metrics['vram_reserved']:.2f}GB reserved / " f"Peak: {metrics['vram_peak_since_last']:.2f}GB / " @@ -369,7 +369,7 @@ class Debug: metrics['summary_ram'] = "" # Update VRAM history for tracking - if torch.cuda.is_available() or torch.mps.is_available(): + if torch.cuda.is_available() or (hasattr(torch, 'mps') and callable(getattr(torch.mps, 'is_available', None)) and torch.mps.is_available()): self.vram_history.append(metrics['vram_allocated']) return metrics