From 17de6ca853ecfe628480b85aea767eaec2c8b788 Mon Sep 17 00:00:00 2001 From: "Salvador E. Tropea" Date: Wed, 23 Jul 2025 13:15:25 -0300 Subject: [PATCH] [Torch] Moved to SeCoNoHe --- src/nodes/inference/demixer.py | 6 +- src/nodes/inference/stft.py | 4 +- src/nodes/nodes.py | 2 +- src/nodes/utils/torch.py | 128 --------------------------------- 4 files changed, 6 insertions(+), 134 deletions(-) delete mode 100644 src/nodes/utils/torch.py diff --git a/src/nodes/inference/demixer.py b/src/nodes/inference/demixer.py index 1d6e808..1d3fe82 100644 --- a/src/nodes/inference/demixer.py +++ b/src/nodes/inference/demixer.py @@ -6,6 +6,7 @@ # Wrappers for the model and inference import logging import math +from seconohe.torch import model_to_target, get_offload_device import torch import tqdm # ComfyUI imports @@ -18,7 +19,6 @@ except Exception: from .stft import stft_chunk_process, stft_get_chunks from ..db.load_model import load_model from .. import NODES_NAME -from ..utils.torch import model_to_target, get_offload_device from .demucs_api import apply_model, BagOfModels from torchaudio.transforms import Fade @@ -259,7 +259,7 @@ class DemixerDemucs(DemixerGeneric): # Using their example they aren't produced model.target_device = self.device logger.debug("Using PyTorch Audio chunking for old model") - with model_to_target(model): + with model_to_target(logger, model): separated_tensors = separate_sources( model, input_tensor_on_device, @@ -273,7 +273,7 @@ class DemixerDemucs(DemixerGeneric): if with_comfy: comfy_progress_bar = comfy.utils.ProgressBar(self.get_steps(waveform_tensor, forced_segment, shifts, overlap)) - with model_to_target(model): + with model_to_target(logger, model): separated_tensors = apply_model( model, input_tensor_on_device, diff --git a/src/nodes/inference/stft.py b/src/nodes/inference/stft.py index 63f9f0d..950cf42 100644 --- a/src/nodes/inference/stft.py +++ b/src/nodes/inference/stft.py @@ -1,11 +1,11 @@ # Short-Time Fourier Transform (STFT). import logging import numpy as np +from seconohe.torch import model_to_target import torch from tqdm import tqdm # Local imports from .. import NODES_NAME -from ..utils.torch import model_to_target logger = logging.getLogger(f"{NODES_NAME}.stft") @@ -127,7 +127,7 @@ def stft_chunk_process(waveform, d, model_run, device, segment_size=256, hop_len total_chunks = 1 + (mixture.shape[1] - chunk_size + step - 1) // step logger.info(f"⚙️ Processing {total_chunks} chunks...") - with model_to_target(model_run): + with model_to_target(logger, model_run): for i in tqdm(range(0, mixture.shape[1] - chunk_size + 1, step)): start = i end = i + chunk_size diff --git a/src/nodes/nodes.py b/src/nodes/nodes.py index 123d703..b1ec802 100644 --- a/src/nodes/nodes.py +++ b/src/nodes/nodes.py @@ -5,13 +5,13 @@ import os from typing import Dict from seconohe.comfy_node_action import send_node_action +from seconohe.torch import get_torch_device_options, get_canonical_device # ComfyUI imports import folder_paths # ComfyUI's way to access model paths # Local imports # We are the main source, so we use the main_logger from . import main_logger from .utils.load_audio import audio_get_channels, force_stereo, force_sample_rate -from .utils.torch import get_torch_device_options, get_canonical_device from .db.models_db import ModelsDB from .inference.demixer import get_demixer diff --git a/src/nodes/utils/torch.py b/src/nodes/utils/torch.py deleted file mode 100644 index 22409aa..0000000 --- a/src/nodes/utils/torch.py +++ /dev/null @@ -1,128 +0,0 @@ -# -*- coding: utf-8 -*- -# Copyright (c) 2025 Salvador E. Tropea -# Copyright (c) 2025 Instituto Nacional de Tecnología Industrial -# License: CC BY-NC-SA 4.0 -# Project: ComfyUI-Float_Optimized -import contextlib # For context manager -import logging -import torch -try: - import comfy.model_management as mm - with_comfy = True -except Exception: - with_comfy = False -from .. import NODES_NAME - -logger = logging.getLogger(f"{NODES_NAME}.torch") - - -def get_torch_device_options(): - # We always have CPU - default = "cpu" - options = [default] - # Do we have CUDA? - if torch.cuda.is_available(): - default = "cuda" - options.append(default) - for i in range(torch.cuda.device_count()): - options.append(f"cuda:{i}") # Specific CUDA devices - # Is this a Mac? - if torch.backends.mps.is_available() and torch.backends.mps.is_built(): - options.append("mps") - if default == "cpu": - default = "mps" - return options, default - - -def get_offload_device(): - return mm.unet_offload_device() if with_comfy else torch.device("cpu") - - -def get_canonical_device(device: str | torch.device) -> torch.device: - """Converts a device string or object into a canonical torch.device object with an explicit index.""" - if not isinstance(device, torch.device): - device = torch.device(device) - - # If it's a CUDA device and no index is specified, get the default one. - if device.type == 'cuda' and device.index is None: - # NOTE: This adds a dependency on torch.cuda.current_device() - # The first solution is often better as it doesn't need this. - return torch.device(f'cuda:{torch.cuda.current_device()}') - return device - - -# ################################################################################## -# # Helper for inference (Target device, offload, eval, no_grad and cuDNN Benchmark) -# ################################################################################## - -@contextlib.contextmanager -def model_to_target(model): - """ - Consolidated context manager for model device placement and inference state. - - - Moves the model to its designated `model.target_device`. - - Sets `torch.backends.cudnn.benchmark` based on `model.cudnn_benchmark_setting` if available. - - Sets the model to `eval()` mode. - - Wraps the operation in a `torch.no_grad()` context. - - Offloads the model to the CPU (`mm.unet_offload_device()`) afterwards. - """ - if not isinstance(model, torch.nn.Module): - with torch.no_grad(): - yield # The code inside the 'with' statement runs here - return - - # 1. Determine target device from the model object - try: - target_device = model.target_device - assert isinstance(target_device, torch.device) - except Exception as e: - logger.warning(f"model_to_target: Could not get 'target_device' from model ({e}). " - "Defaulting to model's current device.") - target_device = next(model.parameters()).device - - # 2. Get CUDNN benchmark setting from the model object (optional) - # Use hasattr as this is an optional setting that not all models might have. - cudnn_benchmark_enabled = None # Default is to keep the current setting - if hasattr(model, 'cudnn_benchmark_setting'): - cudnn_benchmark_enabled = model.cudnn_benchmark_setting - - original_device = next(model.parameters()).device - original_cudnn_benchmark_state = None - is_cuda_target = target_device.type == 'cuda' - - try: - # 3. Manage cuDNN benchmark state - if (cudnn_benchmark_enabled is not None and is_cuda_target and hasattr(torch.backends, 'cudnn') and - torch.backends.cudnn.is_available()): - if torch.backends.cudnn.benchmark != cudnn_benchmark_enabled: - original_cudnn_benchmark_state = torch.backends.cudnn.benchmark - torch.backends.cudnn.benchmark = cudnn_benchmark_enabled - logger.debug(f"Temporarily set cuDNN benchmark to {torch.backends.cudnn.benchmark}") - - # 4. Move model to target device if not already there - if original_device != target_device: - logger.debug(f"Moving model from `{original_device}` to target device `{target_device}` for inference.") - model.to(target_device) - - # 5. Set to eval mode and disable gradients for the operation - model.eval() - with torch.no_grad(): - yield # The code inside the 'with' statement runs here - - finally: - # 6. Restore original cuDNN benchmark state - if original_cudnn_benchmark_state is not None: - # This check is sufficient because it will only be not None if we set it inside the try block - torch.backends.cudnn.benchmark = original_cudnn_benchmark_state - logger.debug(f"Restored cuDNN benchmark to {original_cudnn_benchmark_state}") - - # 7. Offload model back to CPU - if with_comfy: - offload_device = get_offload_device() - current_device_after_yield = next(model.parameters()).device - if current_device_after_yield != offload_device: - logger.debug(f"Offloading model from `{current_device_after_yield}` to offload device `{offload_device}`.") - model.to(offload_device) - # Clear cache if we were on a CUDA device - if 'cuda' in str(current_device_after_yield): - torch.cuda.empty_cache()