From 4cc120bbda338324fb093c925f12863fbdd65ed6 Mon Sep 17 00:00:00 2001 From: steve Date: Mon, 6 Jul 2026 12:47:55 -0400 Subject: [PATCH 01/18] feat(inference): add ComponentDeviceMap for multi-GPU placement (PR1) Introduce per-component device mapping so DiT, VAE, text encoder, and LM can be assigned to distinct CUDA devices while preserving single-GPU behavior when no mapping is provided. Co-authored-by: Cursor --- .../handler/init_service_catalog.py | 14 +- .../generation/handler/init_service_loader.py | 5 +- .../handler/init_service_offload_context.py | 24 +- .../handler/init_service_orchestrator.py | 33 ++- .../generation/handler/init_service_setup.py | 56 +++- .../generation/handler/init_service_test.py | 67 +++++ acestep/device_map.py | 278 ++++++++++++++++++ acestep/test_device_map.py | 109 +++++++ 8 files changed, 553 insertions(+), 33 deletions(-) create mode 100644 acestep/device_map.py create mode 100644 acestep/test_device_map.py diff --git a/acestep/core/generation/handler/init_service_catalog.py b/acestep/core/generation/handler/init_service_catalog.py index 9fe339cb0..2fddb76df 100644 --- a/acestep/core/generation/handler/init_service_catalog.py +++ b/acestep/core/generation/handler/init_service_catalog.py @@ -6,6 +6,8 @@ import torch from loguru import logger +from acestep.device_map import cuda_device_index, is_cuda_device + class InitServiceCatalogMixin: """Checkpoint discovery and backend capability helpers.""" @@ -47,16 +49,20 @@ def get_available_acestep_v15_models(self) -> List[str]: def is_flash_attention_available(self, device: Optional[str] = None) -> bool: """Check whether flash attention can be used on the target device.""" - target_device = str(device or self.device or "auto").split(":", 1)[0] + target_device = str(device or self.device or "auto") if target_device == "auto": if not torch.cuda.is_available(): return False - else: - if target_device != "cuda" or not torch.cuda.is_available(): + cuda_index = 0 + elif is_cuda_device(target_device): + if not torch.cuda.is_available(): return False + cuda_index = cuda_device_index(target_device) + else: + return False try: - major, _ = torch.cuda.get_device_capability() + major, _ = torch.cuda.get_device_capability(cuda_index) if major < 8: logger.info( f"[is_flash_attention_available] GPU compute capability {major}.x < 8.0 " diff --git a/acestep/core/generation/handler/init_service_loader.py b/acestep/core/generation/handler/init_service_loader.py index 040c5ce11..921328667 100644 --- a/acestep/core/generation/handler/init_service_loader.py +++ b/acestep/core/generation/handler/init_service_loader.py @@ -8,6 +8,7 @@ from loguru import logger from acestep import gpu_config +from acestep.device_map import cuda_device_index, is_cuda_device from .init_service_loader_components import InitServiceLoaderComponentsMixin @@ -143,7 +144,9 @@ def _load_main_model_from_checkpoint( if use_flash_attention and self.is_flash_attention_available(device): attn_implementation = "flash_attention_2" - elif device == "cuda" and not gpu_config.cuda_supports_bfloat16(): + elif is_cuda_device(device) and not gpu_config.cuda_supports_bfloat16( + cuda_device_index(device) + ): # Pre-Ampere GPUs (compute capability < 8.0) run in float16 which # can overflow in SDPA's fused softmax with longer sequences, # producing NaN/Inf latents (see issues #924, #927). Eager diff --git a/acestep/core/generation/handler/init_service_offload_context.py b/acestep/core/generation/handler/init_service_offload_context.py index 92af32fa7..f148d50f8 100644 --- a/acestep/core/generation/handler/init_service_offload_context.py +++ b/acestep/core/generation/handler/init_service_offload_context.py @@ -16,17 +16,21 @@ def _load_model_context(self, model_name: str): yield return + target_device = self._get_component_device(model_name) + if model_name == "model" and not self.offload_dit_to_cpu: model = getattr(self, model_name, None) if model is not None: try: param = next(model.parameters()) if param.device.type == "cpu": - logger.info(f"[_load_model_context] Moving {model_name} to {self.device} (persistent)") - self._recursive_to_device(model, self.device, self.dtype) + logger.info( + f"[_load_model_context] Moving {model_name} to {target_device} (persistent)" + ) + self._recursive_to_device(model, target_device, self.dtype) self._release_system_memory() if hasattr(self, "silence_latent"): - self.silence_latent = self.silence_latent.to(self.device).to(self.dtype) + self.silence_latent = self.silence_latent.to(target_device).to(self.dtype) except StopIteration: pass yield @@ -38,16 +42,18 @@ def _load_model_context(self, model_name: str): return rss_before = self._get_rss_mb() - logger.info(f"[_load_model_context] Loading {model_name} to {self.device} (RSS: {rss_before:.0f} MB)") + logger.info( + f"[_load_model_context] Loading {model_name} to {target_device} (RSS: {rss_before:.0f} MB)" + ) start_time = time.time() if model_name == "vae": - vae_dtype = self._get_vae_dtype() - self._recursive_to_device(model, self.device, vae_dtype) + vae_dtype = self._get_vae_dtype(target_device) + self._recursive_to_device(model, target_device, vae_dtype) else: - self._recursive_to_device(model, self.device, self.dtype) + self._recursive_to_device(model, target_device, self.dtype) if model_name == "model" and hasattr(self, "silence_latent"): - self.silence_latent = self.silence_latent.to(self.device).to(self.dtype) + self.silence_latent = self.silence_latent.to(target_device).to(self.dtype) load_time = time.time() - start_time self.current_offload_cost += load_time @@ -56,7 +62,7 @@ def _load_model_context(self, model_name: str): self._release_system_memory() rss_after = self._get_rss_mb() logger.info( - f"[_load_model_context] Loaded {model_name} to {self.device} in {load_time:.4f}s " + f"[_load_model_context] Loaded {model_name} to {target_device} in {load_time:.4f}s " f"(RSS: {rss_before:.0f} -> {rss_after:.0f} MB, delta: {rss_after - rss_before:+.0f} MB)" ) diff --git a/acestep/core/generation/handler/init_service_orchestrator.py b/acestep/core/generation/handler/init_service_orchestrator.py index 4b0322470..7192383ee 100644 --- a/acestep/core/generation/handler/init_service_orchestrator.py +++ b/acestep/core/generation/handler/init_service_orchestrator.py @@ -9,6 +9,7 @@ from loguru import logger from acestep import gpu_config +from acestep.device_map import cuda_device_index, device_type, is_cuda_device, set_active_cuda_device _ROCM_DTYPE_MAP = { "float32": torch.float32, @@ -58,6 +59,7 @@ def initialize_service( prefer_source: Optional[str] = None, use_mlx_dit: bool = True, vae_checkpoint: Optional[str] = None, + gpu_mapping: Optional[str] = None, ) -> Tuple[str, bool]: """Initialize model artifacts and runtime backends for generation. @@ -72,24 +74,32 @@ def initialize_service( ) resolved_device = self._resolve_initialize_device(device) - self.device = resolved_device + self.device_map = self._resolve_component_device_map( + resolved_device=resolved_device, + gpu_mapping=gpu_mapping, + ) + self.device = self.device_map.dit self.offload_to_cpu = offload_to_cpu self.offload_dit_to_cpu = offload_dit_to_cpu + if is_cuda_device(self.device): + set_active_cuda_device(self.device) + normalized_compile, normalized_quantization, mlx_compile_requested = self._configure_initialize_runtime( device=resolved_device, compile_model=compile_model, quantization=quantization, ) self.compiled = normalized_compile - if resolved_device == "cuda" and gpu_config.is_rocm_available(): + if is_cuda_device(resolved_device) and gpu_config.is_rocm_available(): self.dtype = _resolve_rocm_dtype() logger.info( f"[initialize_service] ROCm/HIP device detected: using dtype={self.dtype} " "(set ACESTEP_ROCM_DTYPE=bfloat16 or float16 to override)" ) - elif resolved_device == "cuda": - if gpu_config.cuda_supports_bfloat16(): + elif is_cuda_device(resolved_device): + dit_cuda_index = cuda_device_index(self.device) + if gpu_config.cuda_supports_bfloat16(dit_cuda_index): self.dtype = torch.bfloat16 else: self.dtype = torch.float16 @@ -98,7 +108,7 @@ def initialize_service( "using float16 instead of bfloat16." ) else: - self.dtype = torch.bfloat16 if resolved_device == "xpu" else torch.float32 + self.dtype = torch.bfloat16 if device_type(resolved_device) == "xpu" else torch.float32 self.quantization = normalized_quantization try: self._validate_quantization_setup( @@ -155,20 +165,20 @@ def initialize_service( model_path = os.path.join(checkpoint_dir, config_path) self._load_main_model_from_checkpoint( model_checkpoint_path=model_path, - device=resolved_device, + device=self.device_map.dit, use_flash_attention=use_flash_attention, compile_model=normalized_compile, quantization=self.quantization, ) vae_path = self._load_vae_model( checkpoint_dir=checkpoint_dir, - device=resolved_device, + device=self.device_map.vae, compile_model=normalized_compile, vae_variant=resolved_vae_variant, ) text_encoder_path = self._load_text_encoder_and_tokenizer( checkpoint_dir=checkpoint_dir, - device=resolved_device, + device=self.device_map.text_encoder, ) mlx_dit_status, mlx_vae_status = self._initialize_mlx_backends( @@ -178,7 +188,7 @@ def initialize_service( ) status_msg = self._build_initialize_status_message( - device=resolved_device, + device=self.device, model_path=model_path, vae_path=vae_path, text_encoder_path=text_encoder_path, @@ -192,11 +202,14 @@ def initialize_service( mlx_dit_status=mlx_dit_status, mlx_vae_status=mlx_vae_status, ) + if self.device_map.is_multi_device(): + status_msg += f"\nGPU mapping: {self.device_map.summary()}" self.last_init_params = { "project_root": project_root, "config_path": config_path, - "device": resolved_device, + "device": self.device, + "gpu_mapping": self.device_map.summary(), "use_flash_attention": use_flash_attention, "compile_model": normalized_compile, "offload_to_cpu": offload_to_cpu, diff --git a/acestep/core/generation/handler/init_service_setup.py b/acestep/core/generation/handler/init_service_setup.py index c7e7932bf..4f48e96ca 100644 --- a/acestep/core/generation/handler/init_service_setup.py +++ b/acestep/core/generation/handler/init_service_setup.py @@ -6,6 +6,13 @@ from loguru import logger from acestep import gpu_config +from acestep.device_map import ( + ComponentDeviceMap, + is_cuda_device, + log_device_map, + normalize_component_device, + resolve_component_device_map, +) class InitServiceSetupMixin: @@ -23,15 +30,23 @@ def _resolve_initialize_device(self, requested_device: str) -> str: return "xpu" return "cpu" - if device == "cuda" and not gpu_config.is_cuda_available(): - if gpu_config.is_mps_available(): - logger.warning("[initialize_service] CUDA requested but unavailable. Falling back to MPS.") - return "mps" - if gpu_config.is_xpu_available(): - logger.warning("[initialize_service] CUDA requested but unavailable. Falling back to XPU.") - return "xpu" - logger.warning("[initialize_service] CUDA requested but unavailable. Falling back to CPU.") - return "cpu" + if is_cuda_device(device): + if not gpu_config.is_cuda_available(): + if gpu_config.is_mps_available(): + logger.warning( + "[initialize_service] CUDA device requested but unavailable. Falling back to MPS." + ) + return "mps" + if gpu_config.is_xpu_available(): + logger.warning( + "[initialize_service] CUDA device requested but unavailable. Falling back to XPU." + ) + return "xpu" + logger.warning( + "[initialize_service] CUDA device requested but unavailable. Falling back to CPU." + ) + return "cpu" + return normalize_component_device(device) if device == "mps" and not gpu_config.is_mps_available(): if gpu_config.is_cuda_available(): @@ -55,6 +70,29 @@ def _resolve_initialize_device(self, requested_device: str) -> str: return device + def _resolve_component_device_map( + self, + *, + resolved_device: str, + gpu_mapping: Optional[str] = None, + ) -> ComponentDeviceMap: + """Resolve per-component device placement for initialization.""" + device_map = resolve_component_device_map( + requested_device=resolved_device, + gpu_mapping=gpu_mapping, + ) + log_device_map(device_map) + return device_map + + def _get_component_device(self, component: str) -> str: + """Return the device string for a model component.""" + device_map = getattr(self, "device_map", None) + if device_map is None: + return self.device + if component == "model": + component = "dit" + return device_map.device_for(component) + def _configure_initialize_runtime( self, *, diff --git a/acestep/core/generation/handler/init_service_test.py b/acestep/core/generation/handler/init_service_test.py index 7f70bdaba..88f59f045 100644 --- a/acestep/core/generation/handler/init_service_test.py +++ b/acestep/core/generation/handler/init_service_test.py @@ -261,6 +261,33 @@ def test_resolve_initialize_device_auto_prefers_cuda(self): with patch("torch.xpu", new=types.SimpleNamespace(is_available=lambda: False), create=True): self.assertEqual(host._resolve_initialize_device("auto"), "cuda") + def test_resolve_initialize_device_preserves_cuda_index(self): + """It preserves explicit CUDA device indices such as ``cuda:1``.""" + host = _Host(project_root="K:/fake_root", device="auto") + with patch("torch.cuda.is_available", return_value=True): + self.assertEqual(host._resolve_initialize_device("cuda:1"), "cuda:1") + + def test_resolve_component_device_map_uses_explicit_mapping(self): + """It resolves per-component devices from a GPU mapping string.""" + host = _Host(project_root="K:/fake_root", device="cuda:0") + device_map = host._resolve_component_device_map( + resolved_device="cuda:0", + gpu_mapping="dit:0,vae:0,text_encoder:0,lm:1", + ) + self.assertEqual(device_map.dit, "cuda:0") + self.assertEqual(device_map.lm, "cuda:1") + self.assertTrue(device_map.is_multi_device()) + + def test_get_component_device_returns_component_specific_device(self): + """It returns the mapped device for offload/load helpers.""" + host = _Host(project_root="K:/fake_root", device="cuda:0") + host.device_map = host._resolve_component_device_map( + resolved_device="cuda:0", + gpu_mapping="dit:0,vae:2,text_encoder:0,lm:1", + ) + self.assertEqual(host._get_component_device("model"), "cuda:0") + self.assertEqual(host._get_component_device("vae"), "cuda:2") + def test_configure_initialize_runtime_redirects_compile_on_mps(self): """It converts MPS compile intent to MLX compile and disables quantization.""" host = _Host(project_root="K:/fake_root", device="mps") @@ -464,6 +491,7 @@ def _fake_load_main_model(**_kwargs): "project_root", "config_path", "device", + "gpu_mapping", "use_flash_attention", "compile_model", "offload_to_cpu", @@ -476,6 +504,7 @@ def _fake_load_main_model(**_kwargs): self.assertEqual(set(host.last_init_params.keys()), expected_keys) self.assertEqual(host.last_init_params["config_path"], "acestep-v15-turbo") self.assertEqual(host.last_init_params["device"], "cpu") + self.assertEqual(host.last_init_params["gpu_mapping"], "dit:cpu, vae:cpu, text_encoder:cpu, lm:cpu") self.assertEqual(host.last_init_params["vae_checkpoint"], "official") ensure_models.assert_called_once() sync_code.assert_called_once() @@ -918,6 +947,44 @@ def test_resolve_rocm_dtype_unknown_value_falls_back_to_float32(self): result = ORCHESTRATOR_MODULE._resolve_rocm_dtype() self.assertEqual(result, torch.float32) + def test_initialize_service_records_multi_gpu_mapping(self): + """It records an explicit GPU mapping in init params and status output.""" + host = _Host(project_root="K:/fake_root", device="cuda:0") + + def _fake_load_main_model(**_kwargs): + host.config = types.SimpleNamespace(_attn_implementation="sdpa") + host.model = object() + + with patch.object(GPU_CONFIG_MODULE, "is_cuda_available", return_value=True), \ + patch.object(GPU_CONFIG_MODULE, "is_rocm_available", return_value=False), \ + patch.object(GPU_CONFIG_MODULE, "cuda_supports_bfloat16", return_value=True): + with patch.object(host, "_ensure_models_present", return_value=None): + with patch.object(host, "_sync_model_code_if_needed"): + with patch.object(host, "_load_main_model_from_checkpoint", side_effect=_fake_load_main_model) as load_main_mock: + with patch.object(host, "_load_vae_model", return_value="vae") as load_vae_mock: + with patch.object( + host, + "_load_text_encoder_and_tokenizer", + return_value="te", + ) as load_text_mock: + with patch.object( + host, + "_initialize_mlx_backends", + return_value=("Disabled", "Disabled"), + ): + status, ok = host.initialize_service( + project_root="K:/fake_root", + config_path="acestep-v15-turbo", + device="cuda:0", + gpu_mapping="dit:0,vae:0,text_encoder:0,lm:1", + ) + self.assertTrue(ok) + self.assertIn("GPU mapping: dit:0, vae:0, text_encoder:0, lm:1", status) + self.assertEqual(host.last_init_params["gpu_mapping"], "dit:0, vae:0, text_encoder:0, lm:1") + self.assertEqual(load_main_mock.call_args.kwargs["device"], "cuda:0") + self.assertEqual(load_vae_mock.call_args.kwargs["device"], "cuda:0") + self.assertEqual(load_text_mock.call_args.kwargs["device"], "cuda:0") + def test_initialize_service_uses_float32_on_rocm(self): """It sets dtype=float32 during initialization when ROCm is active.""" host = self._make_rocm_host() diff --git a/acestep/device_map.py b/acestep/device_map.py new file mode 100644 index 000000000..6f29e92cc --- /dev/null +++ b/acestep/device_map.py @@ -0,0 +1,278 @@ +""" +Component-level device placement for multi-GPU inference. + +PR1 scope: parse explicit mappings, discover GPUs, and resolve per-component +device strings. Auto-layout across multiple GPUs is deferred to PR2. +""" + +from __future__ import annotations + +import os +import re +from dataclasses import dataclass +from typing import Dict, List, Optional, Tuple + +from loguru import logger + +GPU_MAPPING_ENV = "ACESTEP_GPU_MAPPING" + +_COMPONENT_KEYS = ("dit", "vae", "text_encoder", "lm") +_SINGLE_PATTERN = re.compile(r"^single:(\d+)$") +_PAIR_PATTERN = re.compile(r"^([a-z_]+):(\d+)$") + + +@dataclass(frozen=True) +class GpuInfo: + """Describes one visible CUDA device using logical indices.""" + + logical_index: int + name: str + total_vram_gb: float + free_vram_gb: float + compute_capability: Optional[Tuple[int, int]] = None + + +@dataclass(frozen=True) +class ComponentDeviceMap: + """Per-component runtime device strings for inference.""" + + dit: str + vae: str + text_encoder: str + lm: Optional[str] = None + lm_tensor_parallel: int = 1 + + def device_for(self, component: str) -> str: + """Return the resolved device string for a component key.""" + key = component.strip().lower() + if key == "model": + key = "dit" + if key not in _COMPONENT_KEYS: + raise KeyError(f"Unknown component: {component}") + value = getattr(self, key) + if value is None: + raise ValueError(f"Component '{component}' has no assigned device") + return value + + def is_multi_device(self) -> bool: + """Return True when components target more than one distinct device.""" + devices = {self.dit, self.vae, self.text_encoder} + if self.lm is not None: + devices.add(self.lm) + return len(devices) > 1 + + def summary(self) -> str: + """Return a compact human-readable layout description.""" + parts = [ + f"dit:{device_index_label(self.dit)}", + f"vae:{device_index_label(self.vae)}", + f"text_encoder:{device_index_label(self.text_encoder)}", + ] + if self.lm is not None: + parts.append(f"lm:{device_index_label(self.lm)}") + if self.lm_tensor_parallel > 1: + parts.append(f"lm_tp={self.lm_tensor_parallel}") + return ", ".join(parts) + + @classmethod + def from_single_device(cls, device: str) -> "ComponentDeviceMap": + """Place all inference components on one device.""" + normalized = normalize_component_device(device) + return cls( + dit=normalized, + vae=normalized, + text_encoder=normalized, + lm=normalized, + ) + + +class DeviceMapError(ValueError): + """Raised when a GPU mapping string cannot be parsed or validated.""" + + +def device_type(device: str) -> str: + """Return the backend type token from a device string.""" + return str(device).split(":", 1)[0] + + +def cuda_device_index(device: str) -> int: + """Return the CUDA device index for a device string.""" + normalized = str(device) + if normalized == "cuda": + return 0 + if normalized.startswith("cuda:"): + return int(normalized.split(":", 1)[1]) + raise DeviceMapError(f"Not a CUDA device string: {device!r}") + + +def is_cuda_device(device: str) -> bool: + """Return whether *device* refers to a CUDA backend.""" + return device_type(device) == "cuda" + + +def normalize_component_device(device: str) -> str: + """Normalize device strings used for component placement.""" + value = str(device).strip() + if not value: + raise DeviceMapError("Device string cannot be empty") + if value == "cuda": + return "cuda:0" + return value + + +def device_index_label(device: str) -> str: + """Return a short index label for logs/UI.""" + if is_cuda_device(device): + return str(cuda_device_index(device)) + return device_type(device) + + +def set_active_cuda_device(device: str) -> None: + """Set the active CUDA device when loading or running on a specific GPU.""" + if not is_cuda_device(device): + return + import torch + + if not torch.cuda.is_available(): + return + index = cuda_device_index(device) + torch.cuda.set_device(index) + + +def discover_gpus() -> List[GpuInfo]: + """Enumerate visible CUDA devices and their current VRAM stats.""" + import torch + + if not torch.cuda.is_available(): + return [] + + gpus: List[GpuInfo] = [] + for index in range(torch.cuda.device_count()): + props = torch.cuda.get_device_properties(index) + free_bytes, total_bytes = torch.cuda.mem_get_info(index) + try: + capability = torch.cuda.get_device_capability(index) + except Exception: + capability = None + gpus.append( + GpuInfo( + logical_index=index, + name=props.name, + total_vram_gb=round(total_bytes / (1024**3), 2), + free_vram_gb=round(free_bytes / (1024**3), 2), + compute_capability=capability, + ) + ) + return gpus + + +def _format_device_for_backend(backend: str, index: int) -> str: + if backend == "cuda": + return f"cuda:{index}" + if backend in {"mps", "xpu", "cpu"}: + if index != 0: + raise DeviceMapError( + f"Component index {index} is invalid for backend '{backend}'" + ) + return backend if backend != "xpu" else "xpu:0" + raise DeviceMapError(f"Unsupported backend for component mapping: {backend!r}") + + +def _parse_mapping_pairs(mapping: str) -> Dict[str, int]: + pairs: Dict[str, int] = {} + for raw_part in mapping.split(","): + part = raw_part.strip() + if not part: + continue + match = _PAIR_PATTERN.fullmatch(part) + if match is None: + raise DeviceMapError( + f"Invalid mapping segment {part!r}; expected 'component:index'" + ) + component, index_text = match.group(1), match.group(2) + if component not in _COMPONENT_KEYS: + raise DeviceMapError( + f"Unknown component {component!r}; expected one of {_COMPONENT_KEYS}" + ) + if component in pairs: + raise DeviceMapError(f"Duplicate component entry: {component!r}") + pairs[component] = int(index_text) + if not pairs: + raise DeviceMapError("GPU mapping must include at least one component") + return pairs + + +def parse_gpu_mapping( + mapping: Optional[str], + *, + default_device: str, +) -> Optional[ComponentDeviceMap]: + """ + Parse a GPU mapping string into a component device map. + + Returns None when mapping is unset or set to 'auto' (PR1 legacy behavior). + """ + raw = (mapping or "").strip() + if not raw: + raw = os.environ.get(GPU_MAPPING_ENV, "").strip() + if not raw or raw.lower() == "auto": + return None + + backend = device_type(default_device) + if backend not in {"cuda", "mps", "xpu", "cpu"}: + backend = "cuda" if is_cuda_device(default_device) else device_type(default_device) + + single_match = _SINGLE_PATTERN.fullmatch(raw) + if single_match is not None: + index = int(single_match.group(1)) + device = _format_device_for_backend(backend, index) + return ComponentDeviceMap.from_single_device(device) + + pairs = _parse_mapping_pairs(raw) + dit_index = pairs.get("dit") + if dit_index is None: + raise DeviceMapError("Explicit GPU mapping must include a 'dit' component") + + dit_device = _format_device_for_backend(backend, dit_index) + vae_device = _format_device_for_backend(backend, pairs.get("vae", dit_index)) + text_encoder_device = _format_device_for_backend( + backend, + pairs.get("text_encoder", dit_index), + ) + lm_device = None + if "lm" in pairs: + lm_device = _format_device_for_backend(backend, pairs["lm"]) + + return ComponentDeviceMap( + dit=dit_device, + vae=vae_device, + text_encoder=text_encoder_device, + lm=lm_device, + ) + + +def resolve_component_device_map( + *, + requested_device: str, + gpu_mapping: Optional[str] = None, +) -> ComponentDeviceMap: + """ + Resolve the effective component device map for service initialization. + + When no mapping is provided, all components share ``requested_device``. + """ + normalized_device = normalize_component_device(requested_device) + parsed = parse_gpu_mapping(gpu_mapping, default_device=normalized_device) + if parsed is None: + return ComponentDeviceMap.from_single_device(normalized_device) + return parsed + + +def log_device_map(device_map: ComponentDeviceMap) -> None: + """Emit a startup log line describing the active component layout.""" + logger.info("[device_map] Active layout: {}", device_map.summary()) + if device_map.is_multi_device(): + logger.info( + "[device_map] Multi-device placement enabled; " + "cross-GPU inference routing arrives in PR2" + ) diff --git a/acestep/test_device_map.py b/acestep/test_device_map.py new file mode 100644 index 000000000..884d1eb4a --- /dev/null +++ b/acestep/test_device_map.py @@ -0,0 +1,109 @@ +"""Unit tests for multi-GPU device map parsing and resolution.""" + +import os +import unittest +from unittest.mock import patch + +from acestep.device_map import ( + ComponentDeviceMap, + DeviceMapError, + cuda_device_index, + device_type, + is_cuda_device, + normalize_component_device, + parse_gpu_mapping, + resolve_component_device_map, +) + + +class DeviceMapParsingTests(unittest.TestCase): + """Tests for mapping string parsing and validation.""" + + def test_normalize_component_device_expands_bare_cuda(self): + self.assertEqual(normalize_component_device("cuda"), "cuda:0") + + def test_normalize_component_device_preserves_cuda_index(self): + self.assertEqual(normalize_component_device("cuda:2"), "cuda:2") + + def test_parse_gpu_mapping_returns_none_for_empty_and_auto(self): + self.assertIsNone(parse_gpu_mapping(None, default_device="cuda:0")) + self.assertIsNone(parse_gpu_mapping("", default_device="cuda:0")) + self.assertIsNone(parse_gpu_mapping("auto", default_device="cuda:0")) + + def test_parse_gpu_mapping_single_places_all_components_on_index(self): + device_map = parse_gpu_mapping("single:1", default_device="cuda:0") + self.assertIsNotNone(device_map) + assert device_map is not None + self.assertEqual(device_map.dit, "cuda:1") + self.assertEqual(device_map.vae, "cuda:1") + self.assertEqual(device_map.text_encoder, "cuda:1") + self.assertEqual(device_map.lm, "cuda:1") + + def test_parse_gpu_mapping_explicit_components(self): + device_map = parse_gpu_mapping( + "dit:0,vae:0,text_encoder:0,lm:1", + default_device="cuda:0", + ) + self.assertIsNotNone(device_map) + assert device_map is not None + self.assertEqual(device_map.dit, "cuda:0") + self.assertEqual(device_map.lm, "cuda:1") + self.assertTrue(device_map.is_multi_device()) + + def test_parse_gpu_mapping_defaults_aux_components_to_dit(self): + device_map = parse_gpu_mapping("dit:2", default_device="cuda:0") + self.assertIsNotNone(device_map) + assert device_map is not None + self.assertEqual(device_map.dit, "cuda:2") + self.assertEqual(device_map.vae, "cuda:2") + self.assertEqual(device_map.text_encoder, "cuda:2") + self.assertIsNone(device_map.lm) + + def test_parse_gpu_mapping_rejects_unknown_component(self): + with self.assertRaises(DeviceMapError): + parse_gpu_mapping("decoder:0", default_device="cuda:0") + + def test_parse_gpu_mapping_rejects_missing_dit(self): + with self.assertRaises(DeviceMapError): + parse_gpu_mapping("vae:0,lm:1", default_device="cuda:0") + + def test_parse_gpu_mapping_reads_env_when_argument_missing(self): + with patch.dict(os.environ, {"ACESTEP_GPU_MAPPING": "single:3"}, clear=False): + device_map = parse_gpu_mapping(None, default_device="cuda:0") + self.assertIsNotNone(device_map) + assert device_map is not None + self.assertEqual(device_map.dit, "cuda:3") + + +class DeviceMapResolutionTests(unittest.TestCase): + """Tests for legacy and explicit device map resolution.""" + + def test_resolve_component_device_map_legacy_single_device(self): + device_map = resolve_component_device_map( + requested_device="cuda", + gpu_mapping=None, + ) + self.assertEqual(device_map.dit, "cuda:0") + self.assertFalse(device_map.is_multi_device()) + + def test_resolve_component_device_map_preserves_requested_index(self): + device_map = resolve_component_device_map( + requested_device="cuda:1", + gpu_mapping=None, + ) + self.assertEqual(device_map.dit, "cuda:1") + self.assertEqual(device_map.lm, "cuda:1") + + def test_device_for_maps_model_alias_to_dit(self): + device_map = ComponentDeviceMap.from_single_device("cuda:0") + self.assertEqual(device_map.device_for("model"), "cuda:0") + + def test_cuda_device_index_helpers(self): + self.assertTrue(is_cuda_device("cuda:2")) + self.assertEqual(device_type("cuda:2"), "cuda") + self.assertEqual(cuda_device_index("cuda"), 0) + self.assertEqual(cuda_device_index("cuda:2"), 2) + + +if __name__ == "__main__": + unittest.main() From a38527f6254e85ff41c4f449ab005a82b56916bd Mon Sep 17 00:00:00 2001 From: steve Date: Mon, 6 Jul 2026 12:59:36 -0400 Subject: [PATCH 02/18] feat(inference): add multi-GPU auto-layout and cross-GPU routing (PR2) Compute VRAM-aware device placement when gpu_mapping=auto on multi-GPU systems, route conditioning tensors to the DiT GPU at inference boundaries, and initialize the LM on its mapped GPU with correct memory budgeting. Co-authored-by: Cursor --- acestep/api/startup_llm_init.py | 4 + acestep/api/startup_model_init.py | 2 + .../core/generation/handler/audio_codes.py | 9 +- .../generation/handler/conditioning_embed.py | 20 ++- .../handler/generate_music_decode.py | 15 +- .../handler/init_service_orchestrator.py | 5 + .../generation/handler/init_service_setup.py | 47 ++++- .../generation/handler/init_service_test.py | 18 ++ .../handler/service_generate_execute.py | 1 + acestep/device_map.py | 165 +++++++++++++++++- acestep/gpu_config.py | 14 +- acestep/llm_inference.py | 34 +++- acestep/test_device_map.py | 70 ++++++++ .../gradio/events/generation/service_init.py | 7 +- 14 files changed, 383 insertions(+), 28 deletions(-) diff --git a/acestep/api/startup_llm_init.py b/acestep/api/startup_llm_init.py index dd5199d1c..d0f0c1015 100644 --- a/acestep/api/startup_llm_init.py +++ b/acestep/api/startup_llm_init.py @@ -23,6 +23,7 @@ def initialize_llm_at_startup( get_model_name: Callable[[str], str], ensure_model_downloaded: Callable[[str, str], str], env_bool: Callable[[str, bool], bool], + dit_handler: Any = None, ) -> None: """Initialize LLM model according to GPU config and environment overrides.""" @@ -73,6 +74,9 @@ def initialize_llm_at_startup( lm_backend = resolve_lm_backend(os.getenv("ACESTEP_LM_BACKEND"), gpu_config) lm_device = os.getenv("ACESTEP_LM_DEVICE", device) + device_map = getattr(dit_handler, "device_map", None) if dit_handler is not None else None + if device_map is not None and device_map.lm is not None: + lm_device = device_map.lm lm_offload_env = os.getenv("ACESTEP_LM_OFFLOAD_TO_CPU") lm_offload = env_bool("ACESTEP_LM_OFFLOAD_TO_CPU", False) if lm_offload_env is not None else offload_to_cpu diff --git a/acestep/api/startup_model_init.py b/acestep/api/startup_model_init.py index c8bca850f..563068625 100644 --- a/acestep/api/startup_model_init.py +++ b/acestep/api/startup_model_init.py @@ -85,6 +85,7 @@ def do_model_initialization( compile_model=compile_model, offload_to_cpu=offload_to_cpu, offload_dit_to_cpu=offload_dit_to_cpu, + gpu_mapping=os.getenv("ACESTEP_GPU_MAPPING"), ) if not ok: app.state._init_error = status_msg @@ -157,6 +158,7 @@ def do_model_initialization( get_model_name=get_model_name, ensure_model_downloaded=ensure_model_downloaded, env_bool=env_bool, + dit_handler=handler, ) print("[API Server] All models initialized successfully!") diff --git a/acestep/core/generation/handler/audio_codes.py b/acestep/core/generation/handler/audio_codes.py index 31d6835c9..e0e8664ad 100644 --- a/acestep/core/generation/handler/audio_codes.py +++ b/acestep/core/generation/handler/audio_codes.py @@ -56,7 +56,8 @@ def _decode_audio_codes_to_latents(self, code_str: str) -> Optional[torch.Tensor with self._load_model_context("model"): quantizer = self.model.tokenizer.quantizer detokenizer = self.model.detokenizer - indices = torch.tensor(code_ids, device=self.device, dtype=torch.long) + dit_device = self._get_component_device("model") + indices = torch.tensor(code_ids, device=dit_device, dtype=torch.long) indices = indices.unsqueeze(0).unsqueeze(-1) quantized = quantizer.get_output_from_indices(indices) @@ -83,7 +84,11 @@ def convert_src_audio_to_codes(self, audio_file) -> str: return "❌ Audio file appears to be silent" latents = self._encode_audio_to_latents(processed_audio) - attention_mask = torch.ones(latents.shape[0], dtype=torch.bool, device=self.device) + attention_mask = torch.ones( + latents.shape[0], + dtype=torch.bool, + device=self._get_component_device("model"), + ) with self._load_model_context("model"): hidden_states = latents.unsqueeze(0) _, indices, _ = self.model.tokenize( diff --git a/acestep/core/generation/handler/conditioning_embed.py b/acestep/core/generation/handler/conditioning_embed.py index ef9e4a572..1c0c38292 100644 --- a/acestep/core/generation/handler/conditioning_embed.py +++ b/acestep/core/generation/handler/conditioning_embed.py @@ -56,7 +56,9 @@ def _ensure_latent_3d(z: torch.Tensor) -> torch.Tensor: refer_audio = _normalize_audio_2d(refer_audio) with torch.inference_mode(): refer_audio_latent = self.tiled_encode(refer_audio, offload_latent_to_cpu=True) - refer_audio_latent = refer_audio_latent.to(self.device).to(self.dtype) + refer_audio_latent = refer_audio_latent.to( + self._get_component_device("dit") + ).to(self.dtype) if refer_audio_latent.dim() == 2: refer_audio_latent = refer_audio_latent.unsqueeze(0) refer_audio_latent = _ensure_latent_3d(refer_audio_latent.transpose(1, 2)) @@ -65,7 +67,11 @@ def _ensure_latent_3d(z: torch.Tensor) -> torch.Tensor: refer_audio_order_mask.append(batch_idx) refer_audio_latents = torch.cat(refer_audio_latents, dim=0) - refer_audio_order_mask = torch.tensor(refer_audio_order_mask, device=self.device, dtype=torch.long) + dit_device = self._get_component_device("dit") + refer_audio_latents = refer_audio_latents.to(dit_device).to(self.dtype) + refer_audio_order_mask = torch.tensor( + refer_audio_order_mask, device=dit_device, dtype=torch.long + ) return refer_audio_latents, refer_audio_order_mask def infer_text_embeddings(self, text_token_idss): @@ -124,6 +130,16 @@ def preprocess_batch(self, batch) -> Tuple: repaint_mask = batch.get("repaint_mask", None) + dit_device = self._get_component_device("dit") + text_hidden_states = text_hidden_states.to(dit_device) + lyric_hidden_states = lyric_hidden_states.to(dit_device) + text_attention_mask = text_attention_mask.to(dit_device) + lyric_attention_mask = lyric_attention_mask.to(dit_device) + if non_cover_text_hidden_states is not None: + non_cover_text_hidden_states = non_cover_text_hidden_states.to(dit_device) + if non_cover_text_attention_masks is not None: + non_cover_text_attention_masks = non_cover_text_attention_masks.to(dit_device) + return ( keys, text_inputs, diff --git a/acestep/core/generation/handler/generate_music_decode.py b/acestep/core/generation/handler/generate_music_decode.py index b7330503c..3d58fcc39 100644 --- a/acestep/core/generation/handler/generate_music_decode.py +++ b/acestep/core/generation/handler/generate_music_decode.py @@ -9,6 +9,7 @@ from loguru import logger from acestep.gpu_config import get_effective_free_vram_gb +from acestep.device_map import cuda_device_index, is_cuda_device class GenerateMusicDecodeMixin: @@ -129,7 +130,12 @@ def _decode_generate_music_pred_latents( with torch.inference_mode(): with self._load_model_context("vae"): pred_latents_cpu = pred_latents.detach().cpu() - pred_latents_for_decode = pred_latents.transpose(1, 2).contiguous().to(self.vae.dtype) + vae_device = self._get_component_device("vae") + pred_latents_for_decode = ( + pred_latents.transpose(1, 2) + .contiguous() + .to(device=vae_device, dtype=self.vae.dtype) + ) del pred_latents self._empty_cache() @@ -150,7 +156,12 @@ def _decode_generate_music_pred_latents( "(unified memory), keeping VAE on MPS" ) else: - effective_free = get_effective_free_vram_gb() + vae_cuda_index = ( + cuda_device_index(vae_device) + if is_cuda_device(vae_device) + else 0 + ) + effective_free = get_effective_free_vram_gb(vae_cuda_index) logger.info( "[generate_music] Effective free VRAM before VAE decode: " f"{effective_free:.2f} GB" diff --git a/acestep/core/generation/handler/init_service_orchestrator.py b/acestep/core/generation/handler/init_service_orchestrator.py index 7192383ee..064ba684e 100644 --- a/acestep/core/generation/handler/init_service_orchestrator.py +++ b/acestep/core/generation/handler/init_service_orchestrator.py @@ -74,9 +74,14 @@ def initialize_service( ) resolved_device = self._resolve_initialize_device(device) + if gpu_mapping is None: + gpu_mapping = os.environ.get("ACESTEP_GPU_MAPPING") + lm_model_path = os.environ.get("ACESTEP_LM_MODEL_PATH") self.device_map = self._resolve_component_device_map( resolved_device=resolved_device, gpu_mapping=gpu_mapping, + config_path=config_path, + lm_model_path=lm_model_path, ) self.device = self.device_map.dit self.offload_to_cpu = offload_to_cpu diff --git a/acestep/core/generation/handler/init_service_setup.py b/acestep/core/generation/handler/init_service_setup.py index 4f48e96ca..d7efc6963 100644 --- a/acestep/core/generation/handler/init_service_setup.py +++ b/acestep/core/generation/handler/init_service_setup.py @@ -1,6 +1,6 @@ """Runtime setup helpers for initialization orchestration.""" -from typing import Any, Optional, Tuple +from typing import Any, Dict, Optional, Tuple import torch from loguru import logger @@ -75,11 +75,19 @@ def _resolve_component_device_map( *, resolved_device: str, gpu_mapping: Optional[str] = None, + config_path: Optional[str] = None, + lm_model_path: Optional[str] = None, + use_lm: bool = True, + batch_size: int = 1, ) -> ComponentDeviceMap: """Resolve per-component device placement for initialization.""" device_map = resolve_component_device_map( requested_device=resolved_device, gpu_mapping=gpu_mapping, + config_path=config_path, + lm_model_path=lm_model_path, + use_lm=use_lm, + batch_size=batch_size, ) log_device_map(device_map) return device_map @@ -93,6 +101,43 @@ def _get_component_device(self, component: str) -> str: component = "dit" return device_map.device_for(component) + def _to_component_device(self, value: Any, component: str) -> Any: + """Move a tensor to the device assigned to *component* when needed.""" + if value is None or not isinstance(value, torch.Tensor): + return value + target = self._get_component_device(component) + device = torch.device(target) + if value.device != device: + return value.to(device) + return value + + def _route_service_payload_to_dit(self, payload: Dict[str, Any]) -> Dict[str, Any]: + """Move service-generation tensors onto the DiT device for diffusion.""" + device_map = getattr(self, "device_map", None) + if device_map is None or not device_map.is_multi_device(): + return payload + + routed = dict(payload) + for key in ( + "text_hidden_states", + "text_attention_mask", + "lyric_hidden_states", + "lyric_attention_mask", + "refer_audio_acoustic_hidden_states_packed", + "refer_audio_order_mask", + "src_latents", + "chunk_mask", + "is_covers", + "precomputed_lm_hints_25Hz", + "non_cover_text_hidden_states", + "non_cover_text_attention_masks", + "repaint_mask", + "target_latents", + ): + if key in routed: + routed[key] = self._to_component_device(routed[key], "dit") + return routed + def _configure_initialize_runtime( self, *, diff --git a/acestep/core/generation/handler/init_service_test.py b/acestep/core/generation/handler/init_service_test.py index 88f59f045..69591fecd 100644 --- a/acestep/core/generation/handler/init_service_test.py +++ b/acestep/core/generation/handler/init_service_test.py @@ -288,6 +288,24 @@ def test_get_component_device_returns_component_specific_device(self): self.assertEqual(host._get_component_device("model"), "cuda:0") self.assertEqual(host._get_component_device("vae"), "cuda:2") + def test_route_service_payload_to_dit_moves_tensors(self): + """It routes diffusion tensors onto the DiT device in multi-GPU mode.""" + host = _Host(project_root="K:/fake_root", device="cuda:0") + host.device_map = host._resolve_component_device_map( + resolved_device="cuda:0", + gpu_mapping="dit:0,vae:0,text_encoder:0,lm:1", + ) + tensor = object() + payload = { + "text_hidden_states": tensor, + "ignored": "keep-me", + } + with patch.object(host, "_to_component_device", side_effect=lambda value, component: value) as move_mock: + routed = host._route_service_payload_to_dit(payload) + move_mock.assert_called_once_with(tensor, "dit") + self.assertIs(routed["text_hidden_states"], tensor) + self.assertEqual(routed["ignored"], "keep-me") + def test_configure_initialize_runtime_redirects_compile_on_mps(self): """It converts MPS compile intent to MLX compile and disables quantization.""" host = _Host(project_root="K:/fake_root", device="mps") diff --git a/acestep/core/generation/handler/service_generate_execute.py b/acestep/core/generation/handler/service_generate_execute.py index 46c6b0a4f..1cb8758e9 100644 --- a/acestep/core/generation/handler/service_generate_execute.py +++ b/acestep/core/generation/handler/service_generate_execute.py @@ -155,6 +155,7 @@ def _execute_service_generate_diffusion( self, payload=payload, generate_kwargs=generate_kwargs, seed_param=seed_param, flow_edit_ctx=flow_edit_ctx, ) + payload = self._route_service_payload_to_dit(payload) dit_backend = ( "MLX (native)" if (self.use_mlx_dit and self.mlx_decoder is not None) else f"PyTorch ({self.device})" ) diff --git a/acestep/device_map.py b/acestep/device_map.py index 6f29e92cc..e93c05c40 100644 --- a/acestep/device_map.py +++ b/acestep/device_map.py @@ -1,19 +1,28 @@ """ Component-level device placement for multi-GPU inference. -PR1 scope: parse explicit mappings, discover GPUs, and resolve per-component -device strings. Auto-layout across multiple GPUs is deferred to PR2. +Supports explicit mappings, single-GPU parity, and automatic layout across +multiple CUDA devices when ``gpu_mapping=auto``. """ from __future__ import annotations import os import re -from dataclasses import dataclass -from typing import Dict, List, Optional, Tuple +from dataclasses import dataclass, field +from typing import Dict, List, Optional, Tuple, Union from loguru import logger +from acestep.gpu_config import ( + DIT_INFERENCE_VRAM_PER_BATCH, + LM_VRAM, + MODEL_VRAM, + VRAM_SAFETY_MARGIN_GB, + get_dit_type_from_path, + get_lm_model_size, +) + GPU_MAPPING_ENV = "ACESTEP_GPU_MAPPING" _COMPONENT_KEYS = ("dit", "vae", "text_encoder", "lm") @@ -90,6 +99,110 @@ class DeviceMapError(ValueError): """Raised when a GPU mapping string cannot be parsed or validated.""" +@dataclass(frozen=True) +class LayoutRequest: + """Inputs used to compute an automatic multi-GPU layout.""" + + gpus: List[GpuInfo] + dit_type: str = "turbo" + lm_model_path: Optional[str] = None + use_lm: bool = True + batch_size: int = 1 + + +@dataclass(frozen=True) +class LayoutError: + """Returned when automatic layout cannot place all requested components.""" + + message: str + suggestions: Tuple[str, ...] = field(default_factory=tuple) + + +def _dit_model_vram_key(dit_type: str) -> str: + if dit_type.startswith("dit_"): + return dit_type + return f"dit_{dit_type}" + + +def estimate_dit_peak_gb(dit_type: str, batch_size: int) -> float: + """Estimate DiT GPU memory including co-located VAE and text encoder.""" + model_key = _dit_model_vram_key(dit_type) + weights = MODEL_VRAM.get(model_key, MODEL_VRAM["dit_turbo"]) + per_batch = DIT_INFERENCE_VRAM_PER_BATCH.get( + dit_type, + DIT_INFERENCE_VRAM_PER_BATCH["turbo"], + ) + aux = ( + MODEL_VRAM["vae"] + + MODEL_VRAM["text_encoder"] + + MODEL_VRAM["cuda_context"] + + VRAM_SAFETY_MARGIN_GB + ) + return weights + (per_batch * max(1, batch_size)) + aux + + +def estimate_lm_total_gb(lm_model_path: Optional[str]) -> float: + """Estimate LM weights plus KV cache on the target GPU.""" + model_size = get_lm_model_size(lm_model_path or "acestep-5Hz-lm-1.7B") + lm_info = LM_VRAM.get(model_size, LM_VRAM["0.6B"]) + return lm_info["weights"] + lm_info["kv_cache_4k"] + 0.3 + + +def _sort_gpus_for_layout(gpus: List[GpuInfo]) -> List[GpuInfo]: + return sorted( + gpus, + key=lambda gpu: (-gpu.free_vram_gb, -gpu.total_vram_gb, gpu.logical_index), + ) + + +def compute_auto_device_map(request: LayoutRequest) -> Union[ComponentDeviceMap, LayoutError]: + """Place components across visible CUDA devices based on free VRAM.""" + gpus = _sort_gpus_for_layout(request.gpus) + if not gpus: + return LayoutError( + "No CUDA devices detected for auto layout", + suggestions=("Install CUDA drivers or pass an explicit --gpu-mapping",), + ) + + dit_need_gb = estimate_dit_peak_gb(request.dit_type, request.batch_size) + dit_gpu = next((gpu for gpu in gpus if gpu.free_vram_gb >= dit_need_gb), None) + if dit_gpu is None: + return LayoutError( + f"No GPU has {dit_need_gb:.1f}GB free for DiT ({request.dit_type})", + suggestions=( + "Reduce batch size", + "Use a smaller DiT checkpoint", + "Enable CPU offload", + "Pass an explicit gpu mapping", + ), + ) + + dit_device = f"cuda:{dit_gpu.logical_index}" + lm_device = None + if request.use_lm: + lm_need_gb = estimate_lm_total_gb(request.lm_model_path) + lm_candidates = [gpu for gpu in gpus if gpu.logical_index != dit_gpu.logical_index] + lm_candidates.extend(gpu for gpu in gpus if gpu.logical_index == dit_gpu.logical_index) + lm_gpu = next((gpu for gpu in lm_candidates if gpu.free_vram_gb >= lm_need_gb), None) + if lm_gpu is None: + return LayoutError( + f"No GPU has {lm_need_gb:.1f}GB free for LM ({request.lm_model_path or 'default'})", + suggestions=( + "Use a smaller LM model", + "Pass gpu_mapping with lm on a specific GPU", + "Disable LM initialization", + ), + ) + lm_device = f"cuda:{lm_gpu.logical_index}" + + return ComponentDeviceMap( + dit=dit_device, + vae=dit_device, + text_encoder=dit_device, + lm=lm_device, + ) + + def device_type(device: str) -> str: """Return the backend type token from a device string.""" return str(device).split(":", 1)[0] @@ -210,7 +323,7 @@ def parse_gpu_mapping( """ Parse a GPU mapping string into a component device map. - Returns None when mapping is unset or set to 'auto' (PR1 legacy behavior). + Returns None when mapping is unset or explicitly set to 'auto'. """ raw = (mapping or "").strip() if not raw: @@ -251,17 +364,54 @@ def parse_gpu_mapping( ) +def _raw_gpu_mapping_value(gpu_mapping: Optional[str]) -> str: + raw = (gpu_mapping or "").strip() + if not raw: + raw = os.environ.get(GPU_MAPPING_ENV, "").strip() + return raw + + def resolve_component_device_map( *, requested_device: str, gpu_mapping: Optional[str] = None, + config_path: Optional[str] = None, + lm_model_path: Optional[str] = None, + use_lm: bool = True, + batch_size: int = 1, ) -> ComponentDeviceMap: """ Resolve the effective component device map for service initialization. When no mapping is provided, all components share ``requested_device``. + When mapping is ``auto`` and multiple CUDA devices are visible, compute a + VRAM-aware layout; otherwise fall back to single-device placement. """ normalized_device = normalize_component_device(requested_device) + raw_mapping = _raw_gpu_mapping_value(gpu_mapping) + + if raw_mapping.lower() == "auto" and is_cuda_device(normalized_device): + gpus = discover_gpus() + if len(gpus) >= 2: + dit_type = get_dit_type_from_path(config_path or "") + layout = compute_auto_device_map( + LayoutRequest( + gpus=gpus, + dit_type=dit_type, + lm_model_path=lm_model_path, + use_lm=use_lm, + batch_size=batch_size, + ) + ) + if isinstance(layout, ComponentDeviceMap): + logger.info("[device_map] Auto layout selected from {} GPU(s)", len(gpus)) + return layout + logger.warning( + "[device_map] Auto layout failed: {}. Suggestions: {}", + layout.message, + "; ".join(layout.suggestions), + ) + parsed = parse_gpu_mapping(gpu_mapping, default_device=normalized_device) if parsed is None: return ComponentDeviceMap.from_single_device(normalized_device) @@ -272,7 +422,4 @@ def log_device_map(device_map: ComponentDeviceMap) -> None: """Emit a startup log line describing the active component layout.""" logger.info("[device_map] Active layout: {}", device_map.summary()) if device_map.is_multi_device(): - logger.info( - "[device_map] Multi-device placement enabled; " - "cross-GPU inference routing arrives in PR2" - ) + logger.info("[device_map] Multi-device placement enabled with cross-GPU routing") diff --git a/acestep/gpu_config.py b/acestep/gpu_config.py index 8b1b4b5e5..f835c03ad 100644 --- a/acestep/gpu_config.py +++ b/acestep/gpu_config.py @@ -974,7 +974,11 @@ def find_best_lm_model_on_disk( def get_lm_gpu_memory_ratio( - model_path: str, total_gpu_memory_gb: float + model_path: str, + total_gpu_memory_gb: float, + device_index: int = 0, + *, + reserve_dit_inference_gb: float = 1.5, ) -> Tuple[float, float]: """ Calculate GPU memory utilization ratio for LM model. @@ -987,6 +991,9 @@ def get_lm_gpu_memory_ratio( Args: model_path: LM model path (e.g., "acestep-5Hz-lm-0.6B") total_gpu_memory_gb: Total GPU memory in GB (used as fallback) + device_index: CUDA device index to query for free VRAM + reserve_dit_inference_gb: Headroom to leave on the LM GPU for DiT activations + when DiT is co-located on the same device (set to 0 when split across GPUs) Returns: Tuple of (gpu_memory_utilization_ratio, target_memory_gb) @@ -1008,7 +1015,7 @@ def get_lm_gpu_memory_ratio( import torch if torch.cuda.is_available(): - free_bytes, total_bytes = torch.cuda.mem_get_info() + free_bytes, total_bytes = torch.cuda.mem_get_info(device_index) free_gb = free_bytes / (1024**3) actual_total_gb = total_bytes / (1024**3) @@ -1031,8 +1038,7 @@ def get_lm_gpu_memory_ratio( # The ratio is relative to total GPU memory (nano-vllm convention), # but we compute it so that the LM only claims what's actually free # minus a safety margin for DiT inference activations. - # Reserve at least 1.5 GB for DiT inference activations - dit_reserve_gb = 1.5 + dit_reserve_gb = max(0.0, reserve_dit_inference_gb) usable_for_lm = max(0, free_gb - dit_reserve_gb - VRAM_SAFETY_MARGIN_GB) # Cap to what the LM actually needs diff --git a/acestep/llm_inference.py b/acestep/llm_inference.py index 3690af56d..c8f73db56 100644 --- a/acestep/llm_inference.py +++ b/acestep/llm_inference.py @@ -26,6 +26,12 @@ from acestep.constrained_logits_processor import MetadataConstrainedLogitsProcessor from acestep.constants import DEFAULT_LM_INSTRUCTION, DEFAULT_LM_UNDERSTAND_INSTRUCTION, DEFAULT_LM_INSPIRED_INSTRUCTION, DEFAULT_LM_REWRITE_INSTRUCTION, DURATION_MIN, DURATION_MAX from acestep.gpu_config import get_lm_gpu_memory_ratio, get_gpu_memory_gb, get_lm_model_size, get_global_gpu_config +from acestep.device_map import ( + cuda_device_index, + is_cuda_device, + normalize_component_device, + set_active_cuda_device, +) # Minimum free VRAM (GB) required to attempt vLLM initialization. # vLLM's KV cache allocator adapts to available memory, so we only need a @@ -199,15 +205,23 @@ def get_gpu_memory_utilization(self, model_path: str = None, minimal_gpu: float Tuple of (gpu_memory_utilization_ratio, low_gpu_memory_mode) """ try: - device = torch.device("cuda:0") - total_gpu_mem_bytes = torch.cuda.get_device_properties(device).total_memory + if not is_cuda_device(self.device) or not torch.cuda.is_available(): + raise RuntimeError("CUDA device required for vLLM memory budgeting") + + device_index = cuda_device_index(normalize_component_device(self.device)) + total_gpu_mem_bytes = torch.cuda.get_device_properties(device_index).total_memory total_gpu = total_gpu_mem_bytes / 1024**3 low_gpu_memory_mode = False # Use adaptive GPU memory ratio based on model size if model_path: - ratio, target_memory_gb = get_lm_gpu_memory_ratio(model_path, total_gpu) + ratio, target_memory_gb = get_lm_gpu_memory_ratio( + model_path, + total_gpu, + device_index=device_index, + reserve_dit_inference_gb=0.0, + ) logger.info(f"Adaptive LM memory allocation: model={model_path}, target={target_memory_gb}GB, ratio={ratio:.3f}, total_gpu={total_gpu:.1f}GB") # Enable low memory mode for small GPUs @@ -217,7 +231,8 @@ def get_gpu_memory_utilization(self, model_path: str = None, minimal_gpu: float return ratio, low_gpu_memory_mode # Fallback to original logic if no model_path provided - reserved_mem_bytes = torch.cuda.memory_reserved(device) + device_index = cuda_device_index(normalize_component_device(self.device)) + reserved_mem_bytes = torch.cuda.memory_reserved(device_index) reserved_gpu = reserved_mem_bytes / 1024**3 available_gpu = total_gpu - reserved_gpu @@ -529,7 +544,7 @@ def initialize( device = "xpu" else: device = "cpu" - elif device == "cuda" and not torch.cuda.is_available(): + elif is_cuda_device(device) and not torch.cuda.is_available(): if hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): logger.warning("[initialize] CUDA requested but unavailable. Falling back to MPS.") device = "mps" @@ -539,6 +554,8 @@ def initialize( else: logger.warning("[initialize] CUDA requested but unavailable. Falling back to CPU.") device = "cpu" + elif is_cuda_device(device): + device = normalize_component_device(device) elif device == "mps" and not (hasattr(torch.backends, "mps") and torch.backends.mps.is_available()): if torch.cuda.is_available(): logger.warning("[initialize] MPS requested but unavailable. Falling back to CUDA.") @@ -569,7 +586,7 @@ def initialize( # produce NaN/inf when naively converted to float16 (different exponent range). # The DiT and VAE use float16 on MPS where it actually helps throughput. if dtype is None: - if device in ["cuda", "xpu"]: + if is_cuda_device(device) or device.startswith("xpu"): self.dtype = torch.bfloat16 else: self.dtype = torch.float32 @@ -601,7 +618,8 @@ def initialize( } # Proactive CUDA cleanup before LM load to reduce fragmentation on mode/model switch - if device == "cuda" and torch.cuda.is_available(): + if is_cuda_device(device) and torch.cuda.is_available(): + set_active_cuda_device(device) gc.collect() torch.cuda.empty_cache() torch.cuda.synchronize() @@ -801,6 +819,8 @@ def _initialize_5hz_lm_vllm(self, model_path: str, enforce_eager: bool = False, return "❌ nano-vllm is not installed. Please install it using 'cd acestep/third_parts/nano-vllm && pip install .'" try: + if is_cuda_device(self.device): + set_active_cuda_device(self.device) current_device = torch.cuda.current_device() device_name = torch.cuda.get_device_name(current_device) diff --git a/acestep/test_device_map.py b/acestep/test_device_map.py index 884d1eb4a..a2686aae0 100644 --- a/acestep/test_device_map.py +++ b/acestep/test_device_map.py @@ -7,8 +7,14 @@ from acestep.device_map import ( ComponentDeviceMap, DeviceMapError, + GpuInfo, + LayoutError, + LayoutRequest, + compute_auto_device_map, cuda_device_index, device_type, + estimate_dit_peak_gb, + estimate_lm_total_gb, is_cuda_device, normalize_component_device, parse_gpu_mapping, @@ -94,6 +100,70 @@ def test_resolve_component_device_map_preserves_requested_index(self): self.assertEqual(device_map.dit, "cuda:1") self.assertEqual(device_map.lm, "cuda:1") + def test_resolve_component_device_map_auto_uses_multi_gpu_layout(self): + gpus = [ + GpuInfo(0, "GPU0", 24.0, 22.0), + GpuInfo(1, "GPU1", 24.0, 23.0), + ] + with patch("acestep.device_map.discover_gpus", return_value=gpus): + device_map = resolve_component_device_map( + requested_device="cuda:0", + gpu_mapping="auto", + config_path="acestep-v15-xl-sft", + lm_model_path="acestep-5Hz-lm-4B", + ) + self.assertEqual(device_map.dit, "cuda:1") + self.assertEqual(device_map.lm, "cuda:0") + self.assertTrue(device_map.is_multi_device()) + + def test_resolve_component_device_map_auto_falls_back_on_single_gpu(self): + with patch("acestep.device_map.discover_gpus", return_value=[GpuInfo(0, "GPU0", 24.0, 22.0)]): + device_map = resolve_component_device_map( + requested_device="cuda:0", + gpu_mapping="auto", + ) + self.assertEqual(device_map.dit, "cuda:0") + self.assertFalse(device_map.is_multi_device()) + + +class AutoLayoutTests(unittest.TestCase): + """Tests for VRAM-aware automatic layout selection.""" + + def test_compute_auto_device_map_splits_dit_and_lm(self): + gpus = [ + GpuInfo(0, "GPU0", 24.0, 10.0), + GpuInfo(1, "GPU1", 24.0, 23.0), + ] + layout = compute_auto_device_map( + LayoutRequest( + gpus=gpus, + dit_type="xl_base", + lm_model_path="acestep-5Hz-lm-4B", + ) + ) + self.assertIsInstance(layout, ComponentDeviceMap) + self.assertEqual(layout.dit, "cuda:1") + self.assertEqual(layout.lm, "cuda:0") + + def test_compute_auto_device_map_returns_error_when_dit_does_not_fit(self): + gpus = [GpuInfo(0, "GPU0", 8.0, 2.0)] + layout = compute_auto_device_map( + LayoutRequest(gpus=gpus, dit_type="xl_base") + ) + self.assertIsInstance(layout, LayoutError) + self.assertIn("No GPU has", layout.message) + + def test_estimate_helpers_use_config_profiles(self): + self.assertGreater(estimate_dit_peak_gb("xl_base", batch_size=2), estimate_dit_peak_gb("turbo", 1)) + self.assertGreater( + estimate_lm_total_gb("acestep-5Hz-lm-4B"), + estimate_lm_total_gb("acestep-5Hz-lm-0.6B"), + ) + + +class DeviceAliasTests(unittest.TestCase): + """Tests for CUDA alias helpers.""" + def test_device_for_maps_model_alias_to_dit(self): device_map = ComponentDeviceMap.from_single_device("cuda:0") self.assertEqual(device_map.device_for("model"), "cuda:0") diff --git a/acestep/ui/gradio/events/generation/service_init.py b/acestep/ui/gradio/events/generation/service_init.py index 545aaf1e7..eb7865421 100644 --- a/acestep/ui/gradio/events/generation/service_init.py +++ b/acestep/ui/gradio/events/generation/service_init.py @@ -97,7 +97,11 @@ def init_service_wrapper( ) lm_device = "cpu" else: - lm_device = device + device_map = getattr(dit_handler, "device_map", None) + if device_map is not None and device_map.lm is not None: + lm_device = device_map.lm + else: + lm_device = device if init_llm and lm_model_path and gpu_config.available_lm_models: if not is_lm_model_size_allowed(lm_model_path, gpu_config.available_lm_models): @@ -130,6 +134,7 @@ def init_service_wrapper( offload_to_cpu=offload_to_cpu, offload_dit_to_cpu=offload_dit_to_cpu, quantization=quant_value, use_mlx_dit=mlx_dit, vae_checkpoint=vae_checkpoint, + gpu_mapping=os.environ.get("ACESTEP_GPU_MAPPING"), ) if init_llm: From 83c3c0eb3a2b0a5198bb9d2916a08d629ae79c60 Mon Sep 17 00:00:00 2001 From: steve Date: Mon, 6 Jul 2026 15:04:41 -0400 Subject: [PATCH 03/18] refactor(device_map): split package to satisfy 200-LOC module cap. Move GPU discovery, layout, parsing, resolution, and status helpers into focused submodules under acestep/device_map/ with a stable public facade. Co-authored-by: Cursor --- acestep/device_map.py | 425 -------------------------------- acestep/device_map/__init__.py | 64 +++++ acestep/device_map/constants.py | 12 + acestep/device_map/devices.py | 54 ++++ acestep/device_map/discovery.py | 60 +++++ acestep/device_map/errors.py | 5 + acestep/device_map/layout.py | 100 ++++++++ acestep/device_map/parsing.py | 104 ++++++++ acestep/device_map/resolve.py | 70 ++++++ acestep/device_map/status.py | 81 ++++++ acestep/device_map/types.py | 93 +++++++ acestep/test_device_map.py | 4 +- 12 files changed, 645 insertions(+), 427 deletions(-) delete mode 100644 acestep/device_map.py create mode 100644 acestep/device_map/__init__.py create mode 100644 acestep/device_map/constants.py create mode 100644 acestep/device_map/devices.py create mode 100644 acestep/device_map/discovery.py create mode 100644 acestep/device_map/errors.py create mode 100644 acestep/device_map/layout.py create mode 100644 acestep/device_map/parsing.py create mode 100644 acestep/device_map/resolve.py create mode 100644 acestep/device_map/status.py create mode 100644 acestep/device_map/types.py diff --git a/acestep/device_map.py b/acestep/device_map.py deleted file mode 100644 index e93c05c40..000000000 --- a/acestep/device_map.py +++ /dev/null @@ -1,425 +0,0 @@ -""" -Component-level device placement for multi-GPU inference. - -Supports explicit mappings, single-GPU parity, and automatic layout across -multiple CUDA devices when ``gpu_mapping=auto``. -""" - -from __future__ import annotations - -import os -import re -from dataclasses import dataclass, field -from typing import Dict, List, Optional, Tuple, Union - -from loguru import logger - -from acestep.gpu_config import ( - DIT_INFERENCE_VRAM_PER_BATCH, - LM_VRAM, - MODEL_VRAM, - VRAM_SAFETY_MARGIN_GB, - get_dit_type_from_path, - get_lm_model_size, -) - -GPU_MAPPING_ENV = "ACESTEP_GPU_MAPPING" - -_COMPONENT_KEYS = ("dit", "vae", "text_encoder", "lm") -_SINGLE_PATTERN = re.compile(r"^single:(\d+)$") -_PAIR_PATTERN = re.compile(r"^([a-z_]+):(\d+)$") - - -@dataclass(frozen=True) -class GpuInfo: - """Describes one visible CUDA device using logical indices.""" - - logical_index: int - name: str - total_vram_gb: float - free_vram_gb: float - compute_capability: Optional[Tuple[int, int]] = None - - -@dataclass(frozen=True) -class ComponentDeviceMap: - """Per-component runtime device strings for inference.""" - - dit: str - vae: str - text_encoder: str - lm: Optional[str] = None - lm_tensor_parallel: int = 1 - - def device_for(self, component: str) -> str: - """Return the resolved device string for a component key.""" - key = component.strip().lower() - if key == "model": - key = "dit" - if key not in _COMPONENT_KEYS: - raise KeyError(f"Unknown component: {component}") - value = getattr(self, key) - if value is None: - raise ValueError(f"Component '{component}' has no assigned device") - return value - - def is_multi_device(self) -> bool: - """Return True when components target more than one distinct device.""" - devices = {self.dit, self.vae, self.text_encoder} - if self.lm is not None: - devices.add(self.lm) - return len(devices) > 1 - - def summary(self) -> str: - """Return a compact human-readable layout description.""" - parts = [ - f"dit:{device_index_label(self.dit)}", - f"vae:{device_index_label(self.vae)}", - f"text_encoder:{device_index_label(self.text_encoder)}", - ] - if self.lm is not None: - parts.append(f"lm:{device_index_label(self.lm)}") - if self.lm_tensor_parallel > 1: - parts.append(f"lm_tp={self.lm_tensor_parallel}") - return ", ".join(parts) - - @classmethod - def from_single_device(cls, device: str) -> "ComponentDeviceMap": - """Place all inference components on one device.""" - normalized = normalize_component_device(device) - return cls( - dit=normalized, - vae=normalized, - text_encoder=normalized, - lm=normalized, - ) - - -class DeviceMapError(ValueError): - """Raised when a GPU mapping string cannot be parsed or validated.""" - - -@dataclass(frozen=True) -class LayoutRequest: - """Inputs used to compute an automatic multi-GPU layout.""" - - gpus: List[GpuInfo] - dit_type: str = "turbo" - lm_model_path: Optional[str] = None - use_lm: bool = True - batch_size: int = 1 - - -@dataclass(frozen=True) -class LayoutError: - """Returned when automatic layout cannot place all requested components.""" - - message: str - suggestions: Tuple[str, ...] = field(default_factory=tuple) - - -def _dit_model_vram_key(dit_type: str) -> str: - if dit_type.startswith("dit_"): - return dit_type - return f"dit_{dit_type}" - - -def estimate_dit_peak_gb(dit_type: str, batch_size: int) -> float: - """Estimate DiT GPU memory including co-located VAE and text encoder.""" - model_key = _dit_model_vram_key(dit_type) - weights = MODEL_VRAM.get(model_key, MODEL_VRAM["dit_turbo"]) - per_batch = DIT_INFERENCE_VRAM_PER_BATCH.get( - dit_type, - DIT_INFERENCE_VRAM_PER_BATCH["turbo"], - ) - aux = ( - MODEL_VRAM["vae"] - + MODEL_VRAM["text_encoder"] - + MODEL_VRAM["cuda_context"] - + VRAM_SAFETY_MARGIN_GB - ) - return weights + (per_batch * max(1, batch_size)) + aux - - -def estimate_lm_total_gb(lm_model_path: Optional[str]) -> float: - """Estimate LM weights plus KV cache on the target GPU.""" - model_size = get_lm_model_size(lm_model_path or "acestep-5Hz-lm-1.7B") - lm_info = LM_VRAM.get(model_size, LM_VRAM["0.6B"]) - return lm_info["weights"] + lm_info["kv_cache_4k"] + 0.3 - - -def _sort_gpus_for_layout(gpus: List[GpuInfo]) -> List[GpuInfo]: - return sorted( - gpus, - key=lambda gpu: (-gpu.free_vram_gb, -gpu.total_vram_gb, gpu.logical_index), - ) - - -def compute_auto_device_map(request: LayoutRequest) -> Union[ComponentDeviceMap, LayoutError]: - """Place components across visible CUDA devices based on free VRAM.""" - gpus = _sort_gpus_for_layout(request.gpus) - if not gpus: - return LayoutError( - "No CUDA devices detected for auto layout", - suggestions=("Install CUDA drivers or pass an explicit --gpu-mapping",), - ) - - dit_need_gb = estimate_dit_peak_gb(request.dit_type, request.batch_size) - dit_gpu = next((gpu for gpu in gpus if gpu.free_vram_gb >= dit_need_gb), None) - if dit_gpu is None: - return LayoutError( - f"No GPU has {dit_need_gb:.1f}GB free for DiT ({request.dit_type})", - suggestions=( - "Reduce batch size", - "Use a smaller DiT checkpoint", - "Enable CPU offload", - "Pass an explicit gpu mapping", - ), - ) - - dit_device = f"cuda:{dit_gpu.logical_index}" - lm_device = None - if request.use_lm: - lm_need_gb = estimate_lm_total_gb(request.lm_model_path) - lm_candidates = [gpu for gpu in gpus if gpu.logical_index != dit_gpu.logical_index] - lm_candidates.extend(gpu for gpu in gpus if gpu.logical_index == dit_gpu.logical_index) - lm_gpu = next((gpu for gpu in lm_candidates if gpu.free_vram_gb >= lm_need_gb), None) - if lm_gpu is None: - return LayoutError( - f"No GPU has {lm_need_gb:.1f}GB free for LM ({request.lm_model_path or 'default'})", - suggestions=( - "Use a smaller LM model", - "Pass gpu_mapping with lm on a specific GPU", - "Disable LM initialization", - ), - ) - lm_device = f"cuda:{lm_gpu.logical_index}" - - return ComponentDeviceMap( - dit=dit_device, - vae=dit_device, - text_encoder=dit_device, - lm=lm_device, - ) - - -def device_type(device: str) -> str: - """Return the backend type token from a device string.""" - return str(device).split(":", 1)[0] - - -def cuda_device_index(device: str) -> int: - """Return the CUDA device index for a device string.""" - normalized = str(device) - if normalized == "cuda": - return 0 - if normalized.startswith("cuda:"): - return int(normalized.split(":", 1)[1]) - raise DeviceMapError(f"Not a CUDA device string: {device!r}") - - -def is_cuda_device(device: str) -> bool: - """Return whether *device* refers to a CUDA backend.""" - return device_type(device) == "cuda" - - -def normalize_component_device(device: str) -> str: - """Normalize device strings used for component placement.""" - value = str(device).strip() - if not value: - raise DeviceMapError("Device string cannot be empty") - if value == "cuda": - return "cuda:0" - return value - - -def device_index_label(device: str) -> str: - """Return a short index label for logs/UI.""" - if is_cuda_device(device): - return str(cuda_device_index(device)) - return device_type(device) - - -def set_active_cuda_device(device: str) -> None: - """Set the active CUDA device when loading or running on a specific GPU.""" - if not is_cuda_device(device): - return - import torch - - if not torch.cuda.is_available(): - return - index = cuda_device_index(device) - torch.cuda.set_device(index) - - -def discover_gpus() -> List[GpuInfo]: - """Enumerate visible CUDA devices and their current VRAM stats.""" - import torch - - if not torch.cuda.is_available(): - return [] - - gpus: List[GpuInfo] = [] - for index in range(torch.cuda.device_count()): - props = torch.cuda.get_device_properties(index) - free_bytes, total_bytes = torch.cuda.mem_get_info(index) - try: - capability = torch.cuda.get_device_capability(index) - except Exception: - capability = None - gpus.append( - GpuInfo( - logical_index=index, - name=props.name, - total_vram_gb=round(total_bytes / (1024**3), 2), - free_vram_gb=round(free_bytes / (1024**3), 2), - compute_capability=capability, - ) - ) - return gpus - - -def _format_device_for_backend(backend: str, index: int) -> str: - if backend == "cuda": - return f"cuda:{index}" - if backend in {"mps", "xpu", "cpu"}: - if index != 0: - raise DeviceMapError( - f"Component index {index} is invalid for backend '{backend}'" - ) - return backend if backend != "xpu" else "xpu:0" - raise DeviceMapError(f"Unsupported backend for component mapping: {backend!r}") - - -def _parse_mapping_pairs(mapping: str) -> Dict[str, int]: - pairs: Dict[str, int] = {} - for raw_part in mapping.split(","): - part = raw_part.strip() - if not part: - continue - match = _PAIR_PATTERN.fullmatch(part) - if match is None: - raise DeviceMapError( - f"Invalid mapping segment {part!r}; expected 'component:index'" - ) - component, index_text = match.group(1), match.group(2) - if component not in _COMPONENT_KEYS: - raise DeviceMapError( - f"Unknown component {component!r}; expected one of {_COMPONENT_KEYS}" - ) - if component in pairs: - raise DeviceMapError(f"Duplicate component entry: {component!r}") - pairs[component] = int(index_text) - if not pairs: - raise DeviceMapError("GPU mapping must include at least one component") - return pairs - - -def parse_gpu_mapping( - mapping: Optional[str], - *, - default_device: str, -) -> Optional[ComponentDeviceMap]: - """ - Parse a GPU mapping string into a component device map. - - Returns None when mapping is unset or explicitly set to 'auto'. - """ - raw = (mapping or "").strip() - if not raw: - raw = os.environ.get(GPU_MAPPING_ENV, "").strip() - if not raw or raw.lower() == "auto": - return None - - backend = device_type(default_device) - if backend not in {"cuda", "mps", "xpu", "cpu"}: - backend = "cuda" if is_cuda_device(default_device) else device_type(default_device) - - single_match = _SINGLE_PATTERN.fullmatch(raw) - if single_match is not None: - index = int(single_match.group(1)) - device = _format_device_for_backend(backend, index) - return ComponentDeviceMap.from_single_device(device) - - pairs = _parse_mapping_pairs(raw) - dit_index = pairs.get("dit") - if dit_index is None: - raise DeviceMapError("Explicit GPU mapping must include a 'dit' component") - - dit_device = _format_device_for_backend(backend, dit_index) - vae_device = _format_device_for_backend(backend, pairs.get("vae", dit_index)) - text_encoder_device = _format_device_for_backend( - backend, - pairs.get("text_encoder", dit_index), - ) - lm_device = None - if "lm" in pairs: - lm_device = _format_device_for_backend(backend, pairs["lm"]) - - return ComponentDeviceMap( - dit=dit_device, - vae=vae_device, - text_encoder=text_encoder_device, - lm=lm_device, - ) - - -def _raw_gpu_mapping_value(gpu_mapping: Optional[str]) -> str: - raw = (gpu_mapping or "").strip() - if not raw: - raw = os.environ.get(GPU_MAPPING_ENV, "").strip() - return raw - - -def resolve_component_device_map( - *, - requested_device: str, - gpu_mapping: Optional[str] = None, - config_path: Optional[str] = None, - lm_model_path: Optional[str] = None, - use_lm: bool = True, - batch_size: int = 1, -) -> ComponentDeviceMap: - """ - Resolve the effective component device map for service initialization. - - When no mapping is provided, all components share ``requested_device``. - When mapping is ``auto`` and multiple CUDA devices are visible, compute a - VRAM-aware layout; otherwise fall back to single-device placement. - """ - normalized_device = normalize_component_device(requested_device) - raw_mapping = _raw_gpu_mapping_value(gpu_mapping) - - if raw_mapping.lower() == "auto" and is_cuda_device(normalized_device): - gpus = discover_gpus() - if len(gpus) >= 2: - dit_type = get_dit_type_from_path(config_path or "") - layout = compute_auto_device_map( - LayoutRequest( - gpus=gpus, - dit_type=dit_type, - lm_model_path=lm_model_path, - use_lm=use_lm, - batch_size=batch_size, - ) - ) - if isinstance(layout, ComponentDeviceMap): - logger.info("[device_map] Auto layout selected from {} GPU(s)", len(gpus)) - return layout - logger.warning( - "[device_map] Auto layout failed: {}. Suggestions: {}", - layout.message, - "; ".join(layout.suggestions), - ) - - parsed = parse_gpu_mapping(gpu_mapping, default_device=normalized_device) - if parsed is None: - return ComponentDeviceMap.from_single_device(normalized_device) - return parsed - - -def log_device_map(device_map: ComponentDeviceMap) -> None: - """Emit a startup log line describing the active component layout.""" - logger.info("[device_map] Active layout: {}", device_map.summary()) - if device_map.is_multi_device(): - logger.info("[device_map] Multi-device placement enabled with cross-GPU routing") diff --git a/acestep/device_map/__init__.py b/acestep/device_map/__init__.py new file mode 100644 index 000000000..b992a0c0a --- /dev/null +++ b/acestep/device_map/__init__.py @@ -0,0 +1,64 @@ +""" +Component-level device placement for multi-GPU inference. + +Public API facade — import from ``acestep.device_map`` as before. +""" + +from acestep.device_map.constants import GPU_MAPPING_ENV, LM_DEVICE_ENV +from acestep.device_map.devices import ( + cuda_device_index, + device_index_label, + device_type, + is_cuda_device, + normalize_component_device, + set_active_cuda_device, +) +from acestep.device_map.discovery import discover_gpus, format_gpu_list_text +from acestep.device_map.errors import DeviceMapError +from acestep.device_map.layout import ( + compute_auto_device_map, + estimate_dit_peak_gb, + estimate_lm_total_gb, +) +from acestep.device_map.parsing import parse_gpu_mapping +from acestep.device_map.resolve import log_device_map, resolve_component_device_map +from acestep.device_map.status import ( + collect_gpu_runtime_status, + device_map_to_dict, + gpu_info_to_dict, + log_lm_device_deprecation, +) +from acestep.device_map.types import ( + ComponentDeviceMap, + GpuInfo, + LayoutError, + LayoutRequest, +) + +__all__ = [ + "GPU_MAPPING_ENV", + "LM_DEVICE_ENV", + "ComponentDeviceMap", + "DeviceMapError", + "GpuInfo", + "LayoutError", + "LayoutRequest", + "collect_gpu_runtime_status", + "compute_auto_device_map", + "cuda_device_index", + "device_index_label", + "device_map_to_dict", + "device_type", + "discover_gpus", + "estimate_dit_peak_gb", + "estimate_lm_total_gb", + "format_gpu_list_text", + "gpu_info_to_dict", + "is_cuda_device", + "log_device_map", + "log_lm_device_deprecation", + "normalize_component_device", + "parse_gpu_mapping", + "resolve_component_device_map", + "set_active_cuda_device", +] diff --git a/acestep/device_map/constants.py b/acestep/device_map/constants.py new file mode 100644 index 000000000..c45d91dd9 --- /dev/null +++ b/acestep/device_map/constants.py @@ -0,0 +1,12 @@ +"""Shared constants for multi-GPU device placement.""" + +from __future__ import annotations + +import re + +GPU_MAPPING_ENV = "ACESTEP_GPU_MAPPING" +LM_DEVICE_ENV = "ACESTEP_LM_DEVICE" + +COMPONENT_KEYS = ("dit", "vae", "text_encoder", "lm") +SINGLE_PATTERN = re.compile(r"^single:(\d+)$") +PAIR_PATTERN = re.compile(r"^([a-z_]+):(\d+)$") diff --git a/acestep/device_map/devices.py b/acestep/device_map/devices.py new file mode 100644 index 000000000..fb19c76c7 --- /dev/null +++ b/acestep/device_map/devices.py @@ -0,0 +1,54 @@ +"""Device-string normalization and CUDA helpers.""" + +from __future__ import annotations + +from acestep.device_map.errors import DeviceMapError + + +def device_type(device: str) -> str: + """Return the backend type token from a device string.""" + return str(device).split(":", 1)[0] + + +def cuda_device_index(device: str) -> int: + """Return the CUDA device index for a device string.""" + normalized = str(device) + if normalized == "cuda": + return 0 + if normalized.startswith("cuda:"): + return int(normalized.split(":", 1)[1]) + raise DeviceMapError(f"Not a CUDA device string: {device!r}") + + +def is_cuda_device(device: str) -> bool: + """Return whether *device* refers to a CUDA backend.""" + return device_type(device) == "cuda" + + +def normalize_component_device(device: str) -> str: + """Normalize device strings used for component placement.""" + value = str(device).strip() + if not value: + raise DeviceMapError("Device string cannot be empty") + if value == "cuda": + return "cuda:0" + return value + + +def device_index_label(device: str) -> str: + """Return a short index label for logs/UI.""" + if is_cuda_device(device): + return str(cuda_device_index(device)) + return device_type(device) + + +def set_active_cuda_device(device: str) -> None: + """Set the active CUDA device when loading or running on a specific GPU.""" + if not is_cuda_device(device): + return + import torch + + if not torch.cuda.is_available(): + return + index = cuda_device_index(device) + torch.cuda.set_device(index) diff --git a/acestep/device_map/discovery.py b/acestep/device_map/discovery.py new file mode 100644 index 000000000..5a4d24430 --- /dev/null +++ b/acestep/device_map/discovery.py @@ -0,0 +1,60 @@ +"""CUDA device discovery helpers.""" + +from __future__ import annotations + +from typing import List, Optional + +from acestep.device_map.types import GpuInfo + + +def discover_gpus() -> List[GpuInfo]: + """Enumerate visible CUDA devices and their current VRAM stats.""" + import torch + + if not torch.cuda.is_available(): + return [] + + gpus: List[GpuInfo] = [] + for index in range(torch.cuda.device_count()): + props = torch.cuda.get_device_properties(index) + free_bytes, total_bytes = torch.cuda.mem_get_info(index) + try: + capability = torch.cuda.get_device_capability(index) + except Exception: + capability = None + gpus.append( + GpuInfo( + logical_index=index, + name=props.name, + total_vram_gb=round(total_bytes / (1024**3), 2), + free_vram_gb=round(free_bytes / (1024**3), 2), + compute_capability=capability, + ) + ) + return gpus + + +def format_gpu_list_text(gpus: Optional[List[GpuInfo]] = None) -> str: + """Return a human-readable table of visible CUDA devices.""" + visible = gpus if gpus is not None else discover_gpus() + if not visible: + return "No CUDA devices detected." + + lines = [ + "Visible CUDA devices:", + " idx name total(GB) free(GB) compute", + ] + for gpu in visible: + capability = "" + if gpu.compute_capability is not None: + capability = f"{gpu.compute_capability[0]}.{gpu.compute_capability[1]}" + lines.append( + " {idx:>3} {name:<28} {total:>8.2f} {free:>7.2f} {capability}".format( + idx=gpu.logical_index, + name=gpu.name[:28], + total=gpu.total_vram_gb, + free=gpu.free_vram_gb, + capability=capability or "-", + ) + ) + return "\n".join(lines) diff --git a/acestep/device_map/errors.py b/acestep/device_map/errors.py new file mode 100644 index 000000000..12a671972 --- /dev/null +++ b/acestep/device_map/errors.py @@ -0,0 +1,5 @@ +"""Exceptions raised by device-map parsing and resolution.""" + + +class DeviceMapError(ValueError): + """Raised when a GPU mapping string cannot be parsed or validated.""" diff --git a/acestep/device_map/layout.py b/acestep/device_map/layout.py new file mode 100644 index 000000000..9ff7f5361 --- /dev/null +++ b/acestep/device_map/layout.py @@ -0,0 +1,100 @@ +"""VRAM-aware automatic multi-GPU layout.""" + +from __future__ import annotations + +from typing import List, Optional, Union + +from acestep.gpu_config import ( + DIT_INFERENCE_VRAM_PER_BATCH, + LM_VRAM, + MODEL_VRAM, + VRAM_SAFETY_MARGIN_GB, + get_lm_model_size, +) + +from acestep.device_map.types import ComponentDeviceMap, GpuInfo, LayoutError, LayoutRequest + + +def _dit_model_vram_key(dit_type: str) -> str: + if dit_type.startswith("dit_"): + return dit_type + return f"dit_{dit_type}" + + +def estimate_dit_peak_gb(dit_type: str, batch_size: int) -> float: + """Estimate DiT GPU memory including co-located VAE and text encoder.""" + model_key = _dit_model_vram_key(dit_type) + weights = MODEL_VRAM.get(model_key, MODEL_VRAM["dit_turbo"]) + per_batch = DIT_INFERENCE_VRAM_PER_BATCH.get( + dit_type, + DIT_INFERENCE_VRAM_PER_BATCH["turbo"], + ) + aux = ( + MODEL_VRAM["vae"] + + MODEL_VRAM["text_encoder"] + + MODEL_VRAM["cuda_context"] + + VRAM_SAFETY_MARGIN_GB + ) + return weights + (per_batch * max(1, batch_size)) + aux + + +def estimate_lm_total_gb(lm_model_path: Optional[str]) -> float: + """Estimate LM weights plus KV cache on the target GPU.""" + model_size = get_lm_model_size(lm_model_path or "acestep-5Hz-lm-1.7B") + lm_info = LM_VRAM.get(model_size, LM_VRAM["0.6B"]) + return lm_info["weights"] + lm_info["kv_cache_4k"] + 0.3 + + +def _sort_gpus_for_layout(gpus: List[GpuInfo]) -> List[GpuInfo]: + return sorted( + gpus, + key=lambda gpu: (-gpu.free_vram_gb, -gpu.total_vram_gb, gpu.logical_index), + ) + + +def compute_auto_device_map(request: LayoutRequest) -> Union[ComponentDeviceMap, LayoutError]: + """Place components across visible CUDA devices based on free VRAM.""" + gpus = _sort_gpus_for_layout(request.gpus) + if not gpus: + return LayoutError( + "No CUDA devices detected for auto layout", + suggestions=("Install CUDA drivers or pass an explicit --gpu-mapping",), + ) + + dit_need_gb = estimate_dit_peak_gb(request.dit_type, request.batch_size) + dit_gpu = next((gpu for gpu in gpus if gpu.free_vram_gb >= dit_need_gb), None) + if dit_gpu is None: + return LayoutError( + f"No GPU has {dit_need_gb:.1f}GB free for DiT ({request.dit_type})", + suggestions=( + "Reduce batch size", + "Use a smaller DiT checkpoint", + "Enable CPU offload", + "Pass an explicit gpu mapping", + ), + ) + + dit_device = f"cuda:{dit_gpu.logical_index}" + lm_device = None + if request.use_lm: + lm_need_gb = estimate_lm_total_gb(request.lm_model_path) + lm_candidates = [gpu for gpu in gpus if gpu.logical_index != dit_gpu.logical_index] + lm_candidates.extend(gpu for gpu in gpus if gpu.logical_index == dit_gpu.logical_index) + lm_gpu = next((gpu for gpu in lm_candidates if gpu.free_vram_gb >= lm_need_gb), None) + if lm_gpu is None: + return LayoutError( + f"No GPU has {lm_need_gb:.1f}GB free for LM ({request.lm_model_path or 'default'})", + suggestions=( + "Use a smaller LM model", + "Pass gpu_mapping with lm on a specific GPU", + "Disable LM initialization", + ), + ) + lm_device = f"cuda:{lm_gpu.logical_index}" + + return ComponentDeviceMap( + dit=dit_device, + vae=dit_device, + text_encoder=dit_device, + lm=lm_device, + ) diff --git a/acestep/device_map/parsing.py b/acestep/device_map/parsing.py new file mode 100644 index 000000000..b75d8f84e --- /dev/null +++ b/acestep/device_map/parsing.py @@ -0,0 +1,104 @@ +"""GPU mapping string parsing.""" + +from __future__ import annotations + +import os +from typing import Dict, Optional + +from acestep.device_map.constants import COMPONENT_KEYS, GPU_MAPPING_ENV, PAIR_PATTERN, SINGLE_PATTERN +from acestep.device_map.devices import device_type, is_cuda_device +from acestep.device_map.errors import DeviceMapError +from acestep.device_map.types import ComponentDeviceMap + + +def _format_device_for_backend(backend: str, index: int) -> str: + if backend == "cuda": + return f"cuda:{index}" + if backend in {"mps", "xpu", "cpu"}: + if index != 0: + raise DeviceMapError( + f"Component index {index} is invalid for backend '{backend}'" + ) + return backend if backend != "xpu" else "xpu:0" + raise DeviceMapError(f"Unsupported backend for component mapping: {backend!r}") + + +def _parse_mapping_pairs(mapping: str) -> Dict[str, int]: + pairs: Dict[str, int] = {} + for raw_part in mapping.split(","): + part = raw_part.strip() + if not part: + continue + match = PAIR_PATTERN.fullmatch(part) + if match is None: + raise DeviceMapError( + f"Invalid mapping segment {part!r}; expected 'component:index'" + ) + component, index_text = match.group(1), match.group(2) + if component not in COMPONENT_KEYS: + raise DeviceMapError( + f"Unknown component {component!r}; expected one of {COMPONENT_KEYS}" + ) + if component in pairs: + raise DeviceMapError(f"Duplicate component entry: {component!r}") + pairs[component] = int(index_text) + if not pairs: + raise DeviceMapError("GPU mapping must include at least one component") + return pairs + + +def parse_gpu_mapping( + mapping: Optional[str], + *, + default_device: str, +) -> Optional[ComponentDeviceMap]: + """ + Parse a GPU mapping string into a component device map. + + Returns None when mapping is unset or explicitly set to 'auto'. + """ + raw = (mapping or "").strip() + if not raw: + raw = os.environ.get(GPU_MAPPING_ENV, "").strip() + if not raw or raw.lower() == "auto": + return None + + backend = device_type(default_device) + if backend not in {"cuda", "mps", "xpu", "cpu"}: + backend = "cuda" if is_cuda_device(default_device) else device_type(default_device) + + single_match = SINGLE_PATTERN.fullmatch(raw) + if single_match is not None: + index = int(single_match.group(1)) + device = _format_device_for_backend(backend, index) + return ComponentDeviceMap.from_single_device(device) + + pairs = _parse_mapping_pairs(raw) + dit_index = pairs.get("dit") + if dit_index is None: + raise DeviceMapError("Explicit GPU mapping must include a 'dit' component") + + dit_device = _format_device_for_backend(backend, dit_index) + vae_device = _format_device_for_backend(backend, pairs.get("vae", dit_index)) + text_encoder_device = _format_device_for_backend( + backend, + pairs.get("text_encoder", dit_index), + ) + lm_device = None + if "lm" in pairs: + lm_device = _format_device_for_backend(backend, pairs["lm"]) + + return ComponentDeviceMap( + dit=dit_device, + vae=vae_device, + text_encoder=text_encoder_device, + lm=lm_device, + ) + + +def raw_gpu_mapping_value(gpu_mapping: Optional[str]) -> str: + """Return the effective raw mapping string from args or environment.""" + raw = (gpu_mapping or "").strip() + if not raw: + raw = os.environ.get(GPU_MAPPING_ENV, "").strip() + return raw diff --git a/acestep/device_map/resolve.py b/acestep/device_map/resolve.py new file mode 100644 index 000000000..be3011506 --- /dev/null +++ b/acestep/device_map/resolve.py @@ -0,0 +1,70 @@ +"""Resolve the effective component layout for service initialization.""" + +from __future__ import annotations + +import os +from typing import Optional + +from loguru import logger + +from acestep.gpu_config import get_dit_type_from_path + +from acestep.device_map.devices import is_cuda_device, normalize_component_device +from acestep.device_map.discovery import discover_gpus +from acestep.device_map.layout import compute_auto_device_map +from acestep.device_map.parsing import parse_gpu_mapping, raw_gpu_mapping_value +from acestep.device_map.types import ComponentDeviceMap, LayoutRequest + + +def resolve_component_device_map( + *, + requested_device: str, + gpu_mapping: Optional[str] = None, + config_path: Optional[str] = None, + lm_model_path: Optional[str] = None, + use_lm: bool = True, + batch_size: int = 1, +) -> ComponentDeviceMap: + """ + Resolve the effective component device map for service initialization. + + When no mapping is provided, all components share ``requested_device``. + When mapping is ``auto`` and multiple CUDA devices are visible, compute a + VRAM-aware layout; otherwise fall back to single-device placement. + """ + normalized_device = normalize_component_device(requested_device) + raw_mapping = raw_gpu_mapping_value(gpu_mapping) + + if raw_mapping.lower() == "auto" and is_cuda_device(normalized_device): + gpus = discover_gpus() + if len(gpus) >= 2: + dit_type = get_dit_type_from_path(config_path or "") + layout = compute_auto_device_map( + LayoutRequest( + gpus=gpus, + dit_type=dit_type, + lm_model_path=lm_model_path, + use_lm=use_lm, + batch_size=batch_size, + ) + ) + if isinstance(layout, ComponentDeviceMap): + logger.info("[device_map] Auto layout selected from {} GPU(s)", len(gpus)) + return layout + logger.warning( + "[device_map] Auto layout failed: {}. Suggestions: {}", + layout.message, + "; ".join(layout.suggestions), + ) + + parsed = parse_gpu_mapping(gpu_mapping, default_device=normalized_device) + if parsed is None: + return ComponentDeviceMap.from_single_device(normalized_device) + return parsed + + +def log_device_map(device_map: ComponentDeviceMap) -> None: + """Emit a startup log line describing the active component layout.""" + logger.info("[device_map] Active layout: {}", device_map.summary()) + if device_map.is_multi_device(): + logger.info("[device_map] Multi-device placement enabled with cross-GPU routing") diff --git a/acestep/device_map/status.py b/acestep/device_map/status.py new file mode 100644 index 000000000..841b385f9 --- /dev/null +++ b/acestep/device_map/status.py @@ -0,0 +1,81 @@ +"""CLI/API status serialization and deprecation helpers.""" + +from __future__ import annotations + +import os +from typing import Any, Dict, Optional + +from loguru import logger + +from acestep.device_map.constants import GPU_MAPPING_ENV, LM_DEVICE_ENV +from acestep.device_map.discovery import discover_gpus +from acestep.device_map.types import ComponentDeviceMap, GpuInfo + + +def gpu_info_to_dict(gpu: GpuInfo) -> Dict[str, object]: + """Serialize a :class:`GpuInfo` record for CLI and API responses.""" + payload: Dict[str, object] = { + "index": gpu.logical_index, + "name": gpu.name, + "total_vram_gb": gpu.total_vram_gb, + "free_vram_gb": gpu.free_vram_gb, + } + if gpu.compute_capability is not None: + payload["compute_capability"] = list(gpu.compute_capability) + return payload + + +def device_map_to_dict(device_map: ComponentDeviceMap) -> Dict[str, object]: + """Serialize an active component layout for status endpoints.""" + return { + "dit": device_map.dit, + "vae": device_map.vae, + "text_encoder": device_map.text_encoder, + "lm": device_map.lm, + "lm_tensor_parallel": device_map.lm_tensor_parallel, + "summary": device_map.summary(), + "multi_device": device_map.is_multi_device(), + } + + +def collect_gpu_runtime_status(handler: Any = None) -> Dict[str, object]: + """Collect GPU inventory and the active component layout for API status.""" + gpus = [gpu_info_to_dict(gpu) for gpu in discover_gpus()] + device_map = getattr(handler, "device_map", None) if handler is not None else None + mapping_env = os.environ.get(GPU_MAPPING_ENV, "").strip() or None + return { + "gpus": gpus, + "gpu_mapping": mapping_env, + "device_map": device_map_to_dict(device_map) if device_map is not None else None, + } + + +def log_lm_device_deprecation( + *, + explicit_lm_device: Optional[str] = None, + gpu_mapping_env: Optional[str] = None, + using_device_map_lm: bool = False, +) -> None: + """Warn when legacy ``ACESTEP_LM_DEVICE`` should be replaced by gpu mapping.""" + explicit = (explicit_lm_device or os.getenv(LM_DEVICE_ENV, "")).strip() + if not explicit: + return + + mapping = (gpu_mapping_env or os.environ.get(GPU_MAPPING_ENV, "")).strip() + if mapping and using_device_map_lm: + logger.warning( + "[device_map] ACESTEP_LM_DEVICE={} is ignored because {}={} " + "assigns the LM device. Use 'lm:N' in the mapping instead.", + explicit, + GPU_MAPPING_ENV, + mapping, + ) + return + + if not mapping: + logger.warning( + "[device_map] ACESTEP_LM_DEVICE is deprecated; use {} instead " + "(e.g. 'single:1' or 'dit:0,vae:0,text_encoder:0,lm:1'). Current value: {}", + GPU_MAPPING_ENV, + explicit, + ) diff --git a/acestep/device_map/types.py b/acestep/device_map/types.py new file mode 100644 index 000000000..0e8437dc4 --- /dev/null +++ b/acestep/device_map/types.py @@ -0,0 +1,93 @@ +"""Datatypes describing GPUs and per-component device layouts.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import List, Optional, Tuple + +from acestep.device_map.devices import device_index_label, normalize_component_device +from acestep.device_map.constants import COMPONENT_KEYS + + +@dataclass(frozen=True) +class GpuInfo: + """Describes one visible CUDA device using logical indices.""" + + logical_index: int + name: str + total_vram_gb: float + free_vram_gb: float + compute_capability: Optional[Tuple[int, int]] = None + + +@dataclass(frozen=True) +class ComponentDeviceMap: + """Per-component runtime device strings for inference.""" + + dit: str + vae: str + text_encoder: str + lm: Optional[str] = None + lm_tensor_parallel: int = 1 + + def device_for(self, component: str) -> str: + """Return the resolved device string for a component key.""" + key = component.strip().lower() + if key == "model": + key = "dit" + if key not in COMPONENT_KEYS: + raise KeyError(f"Unknown component: {component}") + value = getattr(self, key) + if value is None: + raise ValueError(f"Component '{component}' has no assigned device") + return value + + def is_multi_device(self) -> bool: + """Return True when components target more than one distinct device.""" + devices = {self.dit, self.vae, self.text_encoder} + if self.lm is not None: + devices.add(self.lm) + return len(devices) > 1 + + def summary(self) -> str: + """Return a compact human-readable layout description.""" + parts = [ + f"dit:{device_index_label(self.dit)}", + f"vae:{device_index_label(self.vae)}", + f"text_encoder:{device_index_label(self.text_encoder)}", + ] + if self.lm is not None: + parts.append(f"lm:{device_index_label(self.lm)}") + if self.lm_tensor_parallel > 1: + parts.append(f"lm_tp={self.lm_tensor_parallel}") + return ", ".join(parts) + + @classmethod + def from_single_device(cls, device: str) -> "ComponentDeviceMap": + """Place all inference components on one device.""" + normalized = normalize_component_device(device) + return cls( + dit=normalized, + vae=normalized, + text_encoder=normalized, + lm=normalized, + ) + + +@dataclass(frozen=True) +class LayoutRequest: + """Inputs used to compute an automatic multi-GPU layout.""" + + gpus: List[GpuInfo] + dit_type: str = "turbo" + lm_model_path: Optional[str] = None + use_lm: bool = True + batch_size: int = 1 + + +@dataclass(frozen=True) +class LayoutError: + """Returned when automatic layout cannot place all requested components.""" + + message: str + suggestions: Tuple[str, ...] = field(default_factory=tuple) diff --git a/acestep/test_device_map.py b/acestep/test_device_map.py index a2686aae0..8c55e2696 100644 --- a/acestep/test_device_map.py +++ b/acestep/test_device_map.py @@ -105,7 +105,7 @@ def test_resolve_component_device_map_auto_uses_multi_gpu_layout(self): GpuInfo(0, "GPU0", 24.0, 22.0), GpuInfo(1, "GPU1", 24.0, 23.0), ] - with patch("acestep.device_map.discover_gpus", return_value=gpus): + with patch("acestep.device_map.resolve.discover_gpus", return_value=gpus): device_map = resolve_component_device_map( requested_device="cuda:0", gpu_mapping="auto", @@ -117,7 +117,7 @@ def test_resolve_component_device_map_auto_uses_multi_gpu_layout(self): self.assertTrue(device_map.is_multi_device()) def test_resolve_component_device_map_auto_falls_back_on_single_gpu(self): - with patch("acestep.device_map.discover_gpus", return_value=[GpuInfo(0, "GPU0", 24.0, 22.0)]): + with patch("acestep.device_map.resolve.discover_gpus", return_value=[GpuInfo(0, "GPU0", 24.0, 22.0)]): device_map = resolve_component_device_map( requested_device="cuda:0", gpu_mapping="auto", From 7d2b30217c61d4bbec9c692423c6c5071c4b898d Mon Sep 17 00:00:00 2001 From: steve Date: Mon, 6 Jul 2026 15:07:43 -0400 Subject: [PATCH 04/18] fix(inference): complete cross-GPU diffusion routing and APG device index. Route generate_kwargs tensors to the DiT device, move latents during preprocess, and preserve CUDA device indices in APG project(). Co-authored-by: Cursor --- .../generation/handler/conditioning_embed.py | 6 ++ .../handler/init_service_memory_basic.py | 11 +++- .../generation/handler/init_service_setup.py | 41 ++++++++++++ .../generation/handler/init_service_test.py | 39 ++++++++++++ .../handler/service_generate_execute.py | 14 ++++- acestep/models/common/apg_guidance.py | 10 ++- acestep/models/common/apg_guidance_test.py | 63 +++++++++++++++++++ 7 files changed, 177 insertions(+), 7 deletions(-) create mode 100644 acestep/models/common/apg_guidance_test.py diff --git a/acestep/core/generation/handler/conditioning_embed.py b/acestep/core/generation/handler/conditioning_embed.py index 1c0c38292..184f7c1b4 100644 --- a/acestep/core/generation/handler/conditioning_embed.py +++ b/acestep/core/generation/handler/conditioning_embed.py @@ -135,6 +135,12 @@ def preprocess_batch(self, batch) -> Tuple: lyric_hidden_states = lyric_hidden_states.to(dit_device) text_attention_mask = text_attention_mask.to(dit_device) lyric_attention_mask = lyric_attention_mask.to(dit_device) + src_latents = src_latents.to(dit_device) + target_latents = target_latents.to(dit_device) + chunk_mask = chunk_mask.to(dit_device) + is_covers = is_covers.to(dit_device) if isinstance(is_covers, torch.Tensor) else is_covers + if precomputed_lm_hints_25hz is not None: + precomputed_lm_hints_25hz = precomputed_lm_hints_25hz.to(dit_device) if non_cover_text_hidden_states is not None: non_cover_text_hidden_states = non_cover_text_hidden_states.to(dit_device) if non_cover_text_attention_masks is not None: diff --git a/acestep/core/generation/handler/init_service_memory_basic.py b/acestep/core/generation/handler/init_service_memory_basic.py index b83fa52ea..442d5302a 100644 --- a/acestep/core/generation/handler/init_service_memory_basic.py +++ b/acestep/core/generation/handler/init_service_memory_basic.py @@ -152,10 +152,15 @@ def _has_quantized_params(self, module): return False def _ensure_silence_latent_on_device(self): - """Ensure ``silence_latent`` is on ``self.device``.""" + """Ensure ``silence_latent`` is on the active DiT device.""" if hasattr(self, "silence_latent") and self.silence_latent is not None: - if not self._is_on_target_device(self.silence_latent, self.device): - self.silence_latent = self.silence_latent.to(self.device).to(self.dtype) + target = ( + self._get_component_device("dit") + if getattr(self, "device_map", None) is not None + else self.device + ) + if not self._is_on_target_device(self.silence_latent, target): + self.silence_latent = self.silence_latent.to(target).to(self.dtype) @staticmethod def _get_rss_mb() -> float: diff --git a/acestep/core/generation/handler/init_service_setup.py b/acestep/core/generation/handler/init_service_setup.py index d7efc6963..a0fbc2f9d 100644 --- a/acestep/core/generation/handler/init_service_setup.py +++ b/acestep/core/generation/handler/init_service_setup.py @@ -138,6 +138,47 @@ def _route_service_payload_to_dit(self, payload: Dict[str, Any]) -> Dict[str, An routed[key] = self._to_component_device(routed[key], "dit") return routed + def _route_service_generate_kwargs_to_dit( + self, + generate_kwargs: Dict[str, Any], + payload: Dict[str, Any], + ) -> Dict[str, Any]: + """Move diffusion kwargs tensors onto the DiT device for multi-GPU runs.""" + device_map = getattr(self, "device_map", None) + if device_map is None or not device_map.is_multi_device(): + return generate_kwargs + + routed = dict(generate_kwargs) + payload_field_map = { + "text_hidden_states": "text_hidden_states", + "text_attention_mask": "text_attention_mask", + "lyric_hidden_states": "lyric_hidden_states", + "lyric_attention_mask": "lyric_attention_mask", + "refer_audio_acoustic_hidden_states_packed": "refer_audio_acoustic_hidden_states_packed", + "refer_audio_order_mask": "refer_audio_order_mask", + "src_latents": "src_latents", + "chunk_masks": "chunk_mask", + "is_covers": "is_covers", + "non_cover_text_hidden_states": "non_cover_text_hidden_states", + "non_cover_text_attention_mask": "non_cover_text_attention_masks", + "precomputed_lm_hints_25Hz": "precomputed_lm_hints_25Hz", + "repaint_mask": "repaint_mask", + "clean_src_latents": "target_latents", + } + for kwarg_key, payload_key in payload_field_map.items(): + if payload_key in payload: + routed[kwarg_key] = payload[payload_key] + + for key in ( + "silence_latent", + "timesteps", + "repaint_mask", + "clean_src_latents", + ): + if key in routed: + routed[key] = self._to_component_device(routed[key], "dit") + return routed + def _configure_initialize_runtime( self, *, diff --git a/acestep/core/generation/handler/init_service_test.py b/acestep/core/generation/handler/init_service_test.py index 69591fecd..cafa762cc 100644 --- a/acestep/core/generation/handler/init_service_test.py +++ b/acestep/core/generation/handler/init_service_test.py @@ -306,6 +306,45 @@ def test_route_service_payload_to_dit_moves_tensors(self): self.assertIs(routed["text_hidden_states"], tensor) self.assertEqual(routed["ignored"], "keep-me") + def test_route_service_generate_kwargs_to_dit_syncs_from_routed_payload(self): + """Diffusion kwargs should mirror routed payload tensors in multi-GPU mode.""" + host = _Host(project_root="K:/fake_root", device="cuda:3") + host.device_map = host._resolve_component_device_map( + resolved_device="cuda:0", + gpu_mapping="dit:3,vae:0,text_encoder:0,lm:1", + ) + src_on_dit = object() + payload = { + "text_hidden_states": object(), + "text_attention_mask": object(), + "lyric_hidden_states": object(), + "lyric_attention_mask": object(), + "refer_audio_acoustic_hidden_states_packed": object(), + "refer_audio_order_mask": object(), + "src_latents": src_on_dit, + "chunk_mask": object(), + "is_covers": object(), + "non_cover_text_hidden_states": None, + "non_cover_text_attention_masks": None, + "precomputed_lm_hints_25Hz": None, + "target_latents": object(), + } + generate_kwargs = { + "text_hidden_states": object(), + "src_latents": object(), + "chunk_masks": object(), + "silence_latent": object(), + "timesteps": object(), + } + with patch.object(host, "_to_component_device", side_effect=lambda value, component: f"moved:{component}") as move_mock: + routed_kwargs = host._route_service_generate_kwargs_to_dit(generate_kwargs, payload) + self.assertIs(routed_kwargs["src_latents"], src_on_dit) + self.assertEqual(routed_kwargs["chunk_masks"], payload["chunk_mask"]) + self.assertEqual(routed_kwargs["clean_src_latents"], "moved:dit") + self.assertEqual(routed_kwargs["silence_latent"], "moved:dit") + self.assertEqual(routed_kwargs["timesteps"], "moved:dit") + self.assertEqual(move_mock.call_count, 3) + def test_configure_initialize_runtime_redirects_compile_on_mps(self): """It converts MPS compile intent to MLX compile and disables quantization.""" host = _Host(project_root="K:/fake_root", device="mps") diff --git a/acestep/core/generation/handler/service_generate_execute.py b/acestep/core/generation/handler/service_generate_execute.py index 1cb8758e9..0500acdfc 100644 --- a/acestep/core/generation/handler/service_generate_execute.py +++ b/acestep/core/generation/handler/service_generate_execute.py @@ -132,7 +132,16 @@ def _build_service_generate_kwargs( "retake_variance": retake_variance, } if timesteps is not None: - kwargs["timesteps"] = torch.tensor(timesteps, dtype=torch.float32, device=self.device) + dit_device = ( + self._get_component_device("dit") + if getattr(self, "device_map", None) is not None + else self.device + ) + kwargs["timesteps"] = torch.tensor( + timesteps, + dtype=torch.float32, + device=dit_device, + ) return kwargs def _execute_service_generate_diffusion( @@ -156,6 +165,9 @@ def _execute_service_generate_diffusion( seed_param=seed_param, flow_edit_ctx=flow_edit_ctx, ) payload = self._route_service_payload_to_dit(payload) + generate_kwargs = self._route_service_generate_kwargs_to_dit(generate_kwargs, payload) + if hasattr(self, "silence_latent") and self.silence_latent is not None: + self.silence_latent = self._to_component_device(self.silence_latent, "dit") dit_backend = ( "MLX (native)" if (self.use_mlx_dit and self.mlx_decoder is not None) else f"PyTorch ({self.device})" ) diff --git a/acestep/models/common/apg_guidance.py b/acestep/models/common/apg_guidance.py index 3f114a281..6935ebaae 100644 --- a/acestep/models/common/apg_guidance.py +++ b/acestep/models/common/apg_guidance.py @@ -19,15 +19,19 @@ def project( dims=[-1], ): dtype = v0.dtype - device_type = v0.device.type - if device_type == "mps": + device = v0.device + if device.type == "mps": v0, v1 = v0.cpu(), v1.cpu() + device = v0.device v0, v1 = v0.double(), v1.double() v1 = torch.nn.functional.normalize(v1, dim=dims) v0_parallel = (v0 * v1).sum(dim=dims, keepdim=True) * v1 v0_orthogonal = v0 - v0_parallel - return v0_parallel.to(dtype).to(device_type), v0_orthogonal.to(dtype).to(device_type) + return ( + v0_parallel.to(device=device, dtype=dtype), + v0_orthogonal.to(device=device, dtype=dtype), + ) def apg_forward( diff --git a/acestep/models/common/apg_guidance_test.py b/acestep/models/common/apg_guidance_test.py new file mode 100644 index 000000000..22ae4a085 --- /dev/null +++ b/acestep/models/common/apg_guidance_test.py @@ -0,0 +1,63 @@ +"""Unit tests for APG guidance helpers.""" + +from __future__ import annotations + +import unittest + +import torch + +from acestep.models.common.apg_guidance import MomentumBuffer, apg_forward, project + + +class ApgGuidanceDeviceTests(unittest.TestCase): + """Verify guidance math preserves tensor device placement.""" + + def test_project_preserves_cpu_device(self): + v0 = torch.randn(2, 4, 8) + v1 = torch.randn(2, 4, 8) + parallel, orthogonal = project(v0, v1, dims=[1]) + self.assertEqual(parallel.device, v0.device) + self.assertEqual(orthogonal.device, v0.device) + + def test_apg_forward_preserves_cpu_device(self): + pred_cond = torch.randn(2, 4, 8) + pred_uncond = pred_cond + 0.25 + guided = apg_forward( + pred_cond=pred_cond, + pred_uncond=pred_uncond, + guidance_scale=7.0, + momentum_buffer=MomentumBuffer(), + dims=[1], + ) + self.assertEqual(guided.device, pred_cond.device) + + @unittest.skipUnless(torch.cuda.is_available(), "CUDA required") + def test_project_preserves_non_default_cuda_index(self): + if torch.cuda.device_count() < 2: + self.skipTest("Need at least 2 CUDA devices") + device = torch.device("cuda:1") + v0 = torch.randn(2, 4, 8, device=device) + v1 = torch.randn(2, 4, 8, device=device) + parallel, orthogonal = project(v0, v1, dims=[1]) + self.assertEqual(parallel.device, device) + self.assertEqual(orthogonal.device, device) + + @unittest.skipUnless(torch.cuda.is_available(), "CUDA required") + def test_apg_forward_preserves_non_default_cuda_index(self): + if torch.cuda.device_count() < 2: + self.skipTest("Need at least 2 CUDA devices") + device = torch.device("cuda:1") + pred_cond = torch.randn(2, 4, 8, device=device) + pred_uncond = pred_cond + 0.25 + guided = apg_forward( + pred_cond=pred_cond, + pred_uncond=pred_uncond, + guidance_scale=7.0, + momentum_buffer=MomentumBuffer(), + dims=[1], + ) + self.assertEqual(guided.device, device) + + +if __name__ == "__main__": + unittest.main() From 18c8d4a1db0fda5775dcbd625e2a701b8d5249e4 Mon Sep 17 00:00:00 2001 From: steve Date: Mon, 6 Jul 2026 17:09:23 -0400 Subject: [PATCH 05/18] fix(device_map): reserve DiT VRAM for co-located LM auto-layout fallback. Also narrow get_device_capability exception handling to RuntimeError. Co-authored-by: Cursor --- acestep/device_map/discovery.py | 2 +- acestep/device_map/layout.py | 16 +++++++++++++++- acestep/test_device_map.py | 13 +++++++++++++ 3 files changed, 29 insertions(+), 2 deletions(-) diff --git a/acestep/device_map/discovery.py b/acestep/device_map/discovery.py index 5a4d24430..84d3e15fa 100644 --- a/acestep/device_map/discovery.py +++ b/acestep/device_map/discovery.py @@ -20,7 +20,7 @@ def discover_gpus() -> List[GpuInfo]: free_bytes, total_bytes = torch.cuda.mem_get_info(index) try: capability = torch.cuda.get_device_capability(index) - except Exception: + except RuntimeError: capability = None gpus.append( GpuInfo( diff --git a/acestep/device_map/layout.py b/acestep/device_map/layout.py index 9ff7f5361..039086abb 100644 --- a/acestep/device_map/layout.py +++ b/acestep/device_map/layout.py @@ -52,6 +52,13 @@ def _sort_gpus_for_layout(gpus: List[GpuInfo]) -> List[GpuInfo]: ) +def _lm_free_vram_gb(gpu: GpuInfo, dit_gpu: GpuInfo, dit_need_gb: float) -> float: + """Return free VRAM on *gpu* available for LM after reserving the DiT stack.""" + if gpu.logical_index != dit_gpu.logical_index: + return gpu.free_vram_gb + return max(0.0, gpu.free_vram_gb - dit_need_gb) + + def compute_auto_device_map(request: LayoutRequest) -> Union[ComponentDeviceMap, LayoutError]: """Place components across visible CUDA devices based on free VRAM.""" gpus = _sort_gpus_for_layout(request.gpus) @@ -80,7 +87,14 @@ def compute_auto_device_map(request: LayoutRequest) -> Union[ComponentDeviceMap, lm_need_gb = estimate_lm_total_gb(request.lm_model_path) lm_candidates = [gpu for gpu in gpus if gpu.logical_index != dit_gpu.logical_index] lm_candidates.extend(gpu for gpu in gpus if gpu.logical_index == dit_gpu.logical_index) - lm_gpu = next((gpu for gpu in lm_candidates if gpu.free_vram_gb >= lm_need_gb), None) + lm_gpu = next( + ( + gpu + for gpu in lm_candidates + if _lm_free_vram_gb(gpu, dit_gpu, dit_need_gb) >= lm_need_gb + ), + None, + ) if lm_gpu is None: return LayoutError( f"No GPU has {lm_need_gb:.1f}GB free for LM ({request.lm_model_path or 'default'})", diff --git a/acestep/test_device_map.py b/acestep/test_device_map.py index 8c55e2696..068d457ac 100644 --- a/acestep/test_device_map.py +++ b/acestep/test_device_map.py @@ -153,6 +153,19 @@ def test_compute_auto_device_map_returns_error_when_dit_does_not_fit(self): self.assertIsInstance(layout, LayoutError) self.assertIn("No GPU has", layout.message) + def test_compute_auto_device_map_rejects_same_gpu_lm_without_headroom(self): + """Co-located LM fallback must reserve the DiT footprint on the shared GPU.""" + gpus = [GpuInfo(0, "GPU0", 24.0, 22.0)] + layout = compute_auto_device_map( + LayoutRequest( + gpus=gpus, + dit_type="xl_base", + lm_model_path="acestep-5Hz-lm-4B", + ) + ) + self.assertIsInstance(layout, LayoutError) + self.assertIn("No GPU has", layout.message) + def test_estimate_helpers_use_config_profiles(self): self.assertGreater(estimate_dit_peak_gb("xl_base", batch_size=2), estimate_dit_peak_gb("turbo", 1)) self.assertGreater( From c1d9d98a5d723ed33ac5ed414c05710e6fe6ca31 Mon Sep 17 00:00:00 2001 From: steve Date: Mon, 6 Jul 2026 17:14:50 -0400 Subject: [PATCH 06/18] fix: address CodeRabbit parsing dead code and APG MPS regression. Reject unsupported default_device backends explicitly in GPU mapping parsing, and preserve the original MPS device when project() uses CPU math fallback. Co-authored-by: Cursor --- acestep/device_map/parsing.py | 20 ++++++++++++++++---- acestep/models/common/apg_guidance.py | 9 ++++----- acestep/models/common/apg_guidance_test.py | 9 +++++++++ acestep/test_device_map.py | 4 ++++ 4 files changed, 33 insertions(+), 9 deletions(-) diff --git a/acestep/device_map/parsing.py b/acestep/device_map/parsing.py index b75d8f84e..74fde9200 100644 --- a/acestep/device_map/parsing.py +++ b/acestep/device_map/parsing.py @@ -6,7 +6,7 @@ from typing import Dict, Optional from acestep.device_map.constants import COMPONENT_KEYS, GPU_MAPPING_ENV, PAIR_PATTERN, SINGLE_PATTERN -from acestep.device_map.devices import device_type, is_cuda_device +from acestep.device_map.devices import device_type from acestep.device_map.errors import DeviceMapError from acestep.device_map.types import ComponentDeviceMap @@ -47,6 +47,20 @@ def _parse_mapping_pairs(mapping: str) -> Dict[str, int]: return pairs +_SUPPORTED_MAPPING_BACKENDS = frozenset({"cuda", "mps", "xpu", "cpu"}) + + +def _resolve_mapping_backend(default_device: str) -> str: + """Return the backend token used to interpret explicit GPU indices.""" + backend = device_type(default_device) + if backend in _SUPPORTED_MAPPING_BACKENDS: + return backend + raise DeviceMapError( + f"Unsupported default device {default_device!r} for GPU mapping; " + "expected cuda, mps, xpu, or cpu" + ) + + def parse_gpu_mapping( mapping: Optional[str], *, @@ -63,9 +77,7 @@ def parse_gpu_mapping( if not raw or raw.lower() == "auto": return None - backend = device_type(default_device) - if backend not in {"cuda", "mps", "xpu", "cpu"}: - backend = "cuda" if is_cuda_device(default_device) else device_type(default_device) + backend = _resolve_mapping_backend(default_device) single_match = SINGLE_PATTERN.fullmatch(raw) if single_match is not None: diff --git a/acestep/models/common/apg_guidance.py b/acestep/models/common/apg_guidance.py index 6935ebaae..0889420d2 100644 --- a/acestep/models/common/apg_guidance.py +++ b/acestep/models/common/apg_guidance.py @@ -19,18 +19,17 @@ def project( dims=[-1], ): dtype = v0.dtype - device = v0.device - if device.type == "mps": + target_device = v0.device + if target_device.type == "mps": v0, v1 = v0.cpu(), v1.cpu() - device = v0.device v0, v1 = v0.double(), v1.double() v1 = torch.nn.functional.normalize(v1, dim=dims) v0_parallel = (v0 * v1).sum(dim=dims, keepdim=True) * v1 v0_orthogonal = v0 - v0_parallel return ( - v0_parallel.to(device=device, dtype=dtype), - v0_orthogonal.to(device=device, dtype=dtype), + v0_parallel.to(device=target_device, dtype=dtype), + v0_orthogonal.to(device=target_device, dtype=dtype), ) diff --git a/acestep/models/common/apg_guidance_test.py b/acestep/models/common/apg_guidance_test.py index 22ae4a085..4c38d3bb8 100644 --- a/acestep/models/common/apg_guidance_test.py +++ b/acestep/models/common/apg_guidance_test.py @@ -58,6 +58,15 @@ def test_apg_forward_preserves_non_default_cuda_index(self): ) self.assertEqual(guided.device, device) + @unittest.skipUnless(hasattr(torch.backends, "mps") and torch.backends.mps.is_available(), "MPS required") + def test_project_returns_mps_tensors_after_cpu_math(self): + device = torch.device("mps") + v0 = torch.randn(2, 4, 8, device=device) + v1 = torch.randn(2, 4, 8, device=device) + parallel, orthogonal = project(v0, v1, dims=[1]) + self.assertEqual(parallel.device.type, "mps") + self.assertEqual(orthogonal.device.type, "mps") + if __name__ == "__main__": unittest.main() diff --git a/acestep/test_device_map.py b/acestep/test_device_map.py index 068d457ac..0dee4da9e 100644 --- a/acestep/test_device_map.py +++ b/acestep/test_device_map.py @@ -73,6 +73,10 @@ def test_parse_gpu_mapping_rejects_missing_dit(self): with self.assertRaises(DeviceMapError): parse_gpu_mapping("vae:0,lm:1", default_device="cuda:0") + def test_parse_gpu_mapping_rejects_unsupported_default_device(self): + with self.assertRaises(DeviceMapError): + parse_gpu_mapping("single:0", default_device="unknown-backend") + def test_parse_gpu_mapping_reads_env_when_argument_missing(self): with patch.dict(os.environ, {"ACESTEP_GPU_MAPPING": "single:3"}, clear=False): device_map = parse_gpu_mapping(None, default_device="cuda:0") From b6d40614d75288431343157911e20b4ee88cb220 Mon Sep 17 00:00:00 2001 From: steve Date: Mon, 6 Jul 2026 17:21:27 -0400 Subject: [PATCH 07/18] Fix silence_latent device check for multi-GPU CUDA indices. _is_on_target_device only compares backend types, so cuda:0 and cuda:3 look equivalent. Use exact device matching when placing silence_latent on the DiT device from device_map. Co-authored-by: Cursor --- .../handler/init_service_memory_basic.py | 14 ++++++++- .../generation/handler/init_service_test.py | 29 +++++++++++++++++++ 2 files changed, 42 insertions(+), 1 deletion(-) diff --git a/acestep/core/generation/handler/init_service_memory_basic.py b/acestep/core/generation/handler/init_service_memory_basic.py index 442d5302a..42a9881ca 100644 --- a/acestep/core/generation/handler/init_service_memory_basic.py +++ b/acestep/core/generation/handler/init_service_memory_basic.py @@ -111,6 +111,18 @@ def _is_on_target_device(self, tensor, target_device): return False return tensor.device.type == target_type + def _tensor_on_exact_device(self, tensor, target_device: str) -> bool: + """Return whether *tensor* is on the exact device string (including CUDA index).""" + if tensor is None: + return True + from acestep.device_map.devices import normalize_component_device + + try: + expected = torch.device(normalize_component_device(str(target_device))) + except Exception: + return False + return tensor.device == expected + @staticmethod def _get_affine_quantized_tensor_class(): """Return the AffineQuantizedTensor class from torchao, or None if unavailable.""" @@ -159,7 +171,7 @@ def _ensure_silence_latent_on_device(self): if getattr(self, "device_map", None) is not None else self.device ) - if not self._is_on_target_device(self.silence_latent, target): + if not self._tensor_on_exact_device(self.silence_latent, target): self.silence_latent = self.silence_latent.to(target).to(self.dtype) @staticmethod diff --git a/acestep/core/generation/handler/init_service_test.py b/acestep/core/generation/handler/init_service_test.py index cafa762cc..257dd0dc9 100644 --- a/acestep/core/generation/handler/init_service_test.py +++ b/acestep/core/generation/handler/init_service_test.py @@ -120,6 +120,35 @@ def test_is_on_target_device_malformed_target_logs_and_returns_false(self): self.assertFalse(host._is_on_target_device(t, ":0")) warning.assert_called_once() + def test_tensor_on_exact_device_distinguishes_cuda_indices(self): + """It compares full CUDA device strings, not just backend types.""" + host = _Host(project_root="K:/fake_root", device="cuda:0") + tensor = types.SimpleNamespace(device=torch.device("cuda:0")) + self.assertTrue(host._tensor_on_exact_device(tensor, "cuda:0")) + self.assertFalse(host._tensor_on_exact_device(tensor, "cuda:3")) + self.assertTrue(host._is_on_target_device(tensor, "cuda:3")) + + def test_ensure_silence_latent_on_device_moves_across_cuda_indices(self): + """``silence_latent`` must follow ``device_map.dit``, not just the CUDA backend.""" + host = _Host(project_root="K:/fake_root", device="cuda:3") + host.dtype = torch.float32 + host.device_map = host._resolve_component_device_map( + resolved_device="cuda:0", + gpu_mapping="dit:3,vae:0,text_encoder:0,lm:1", + ) + moved_to = [] + + class _Latent: + device = torch.device("cuda:0") + + def to(self, device_or_dtype, *args, **kwargs): + moved_to.append(device_or_dtype) + return self + + host.silence_latent = _Latent() + host._ensure_silence_latent_on_device() + self.assertEqual(moved_to[0], "cuda:3") + self.assertEqual(len(moved_to), 2) def test_get_auto_decode_chunk_size_uses_cuda_device_index(self): """It probes effective VRAM on the selected CUDA device index.""" From 0c21299fa9ed93516814b512d5f8062046938d8c Mon Sep 17 00:00:00 2001 From: steve Date: Mon, 6 Jul 2026 17:36:01 -0400 Subject: [PATCH 08/18] Hoist normalize_component_device import and narrow device errors. Module-level import avoids per-call overhead; catch DeviceMapError and RuntimeError instead of a broad Exception in _tensor_on_exact_device. Co-authored-by: Cursor --- .../core/generation/handler/init_service_memory_basic.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/acestep/core/generation/handler/init_service_memory_basic.py b/acestep/core/generation/handler/init_service_memory_basic.py index 42a9881ca..a516a7f3b 100644 --- a/acestep/core/generation/handler/init_service_memory_basic.py +++ b/acestep/core/generation/handler/init_service_memory_basic.py @@ -11,6 +11,9 @@ import torch from loguru import logger +from acestep.device_map.devices import normalize_component_device +from acestep.device_map.errors import DeviceMapError + # Cached libc handle for mallopt/malloc_trim calls (Linux only). _LIBC = None _MALLOPT_APPLIED = False @@ -115,11 +118,9 @@ def _tensor_on_exact_device(self, tensor, target_device: str) -> bool: """Return whether *tensor* is on the exact device string (including CUDA index).""" if tensor is None: return True - from acestep.device_map.devices import normalize_component_device - try: expected = torch.device(normalize_component_device(str(target_device))) - except Exception: + except (DeviceMapError, RuntimeError): return False return tensor.device == expected From f7ba68398141c0965823ec756f7e06b5b55483d4 Mon Sep 17 00:00:00 2001 From: steve Date: Mon, 6 Jul 2026 17:38:59 -0400 Subject: [PATCH 09/18] Catch specific exceptions in device string parsing helpers. Replace bare Exception handlers with DeviceMapError/RuntimeError/TypeError for exact device checks and RuntimeError/TypeError/ValueError for backend alias parsing in _is_on_target_device. Co-authored-by: Cursor --- acestep/core/generation/handler/init_service_memory_basic.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/acestep/core/generation/handler/init_service_memory_basic.py b/acestep/core/generation/handler/init_service_memory_basic.py index a516a7f3b..d0f5a2999 100644 --- a/acestep/core/generation/handler/init_service_memory_basic.py +++ b/acestep/core/generation/handler/init_service_memory_basic.py @@ -104,7 +104,7 @@ def _is_on_target_device(self, tensor, target_device): target_type = target_device.type else: target_type = torch.device(str(target_device)).type - except Exception: + except (RuntimeError, TypeError, ValueError): target_type = str(target_device).strip().lower().split(":", 1)[0] if not target_type: logger.warning( @@ -120,7 +120,7 @@ def _tensor_on_exact_device(self, tensor, target_device: str) -> bool: return True try: expected = torch.device(normalize_component_device(str(target_device))) - except (DeviceMapError, RuntimeError): + except (DeviceMapError, RuntimeError, TypeError): return False return tensor.device == expected From 8f1463f32d55d8c443cca1dca3bb4f9044533375 Mon Sep 17 00:00:00 2001 From: steve Date: Mon, 6 Jul 2026 18:22:26 -0400 Subject: [PATCH 10/18] Fix VAE VRAM preflight to use mapped component device index. vae_component_device was overwritten with None before the CUDA free-VRAM check, so multi-GPU layouts always queried cuda:0. Keep a separate vae_restore_device for the CPU offload restore path. Co-authored-by: Cursor --- .../handler/generate_music_decode.py | 16 +++--- .../handler/generate_music_decode_test.py | 50 ++++++++++++++++++- 2 files changed, 56 insertions(+), 10 deletions(-) diff --git a/acestep/core/generation/handler/generate_music_decode.py b/acestep/core/generation/handler/generate_music_decode.py index 3d58fcc39..93acf5c7b 100644 --- a/acestep/core/generation/handler/generate_music_decode.py +++ b/acestep/core/generation/handler/generate_music_decode.py @@ -130,11 +130,11 @@ def _decode_generate_music_pred_latents( with torch.inference_mode(): with self._load_model_context("vae"): pred_latents_cpu = pred_latents.detach().cpu() - vae_device = self._get_component_device("vae") + vae_component_device = self._get_component_device("vae") pred_latents_for_decode = ( pred_latents.transpose(1, 2) .contiguous() - .to(device=vae_device, dtype=self.vae.dtype) + .to(device=vae_component_device, dtype=self.vae.dtype) ) del pred_latents self._empty_cache() @@ -146,7 +146,7 @@ def _decode_generate_music_pred_latents( ) using_mlx_vae = self.use_mlx_vae and self.mlx_vae is not None vae_cpu = False - vae_device = None + vae_restore_device = None if not using_mlx_vae: vae_cpu = os.environ.get("ACESTEP_VAE_ON_CPU", "0").lower() in ("1", "true", "yes") if not vae_cpu: @@ -157,8 +157,8 @@ def _decode_generate_music_pred_latents( ) else: vae_cuda_index = ( - cuda_device_index(vae_device) - if is_cuda_device(vae_device) + cuda_device_index(vae_component_device) + if is_cuda_device(vae_component_device) else 0 ) effective_free = get_effective_free_vram_gb(vae_cuda_index) @@ -174,7 +174,7 @@ def _decode_generate_music_pred_latents( vae_cpu = True if vae_cpu: logger.info("[generate_music] Moving VAE to CPU for decode (ACESTEP_VAE_ON_CPU=1)...") - vae_device = next(self.vae.parameters()).device + vae_restore_device = next(self.vae.parameters()).device self.vae = self.vae.cpu() pred_latents_for_decode = pred_latents_for_decode.cpu() self._empty_cache() @@ -197,9 +197,9 @@ def _decode_generate_music_pred_latents( pred_wavs = decoder_output.sample del decoder_output finally: - if vae_cpu and vae_device is not None: + if vae_cpu and vae_restore_device is not None: logger.info("[generate_music] Restoring VAE to original device after CPU decode path...") - self.vae = self.vae.to(vae_device) + self.vae = self.vae.to(vae_restore_device) self._empty_cache() logger.debug( "[generate_music] After VAE decode: " diff --git a/acestep/core/generation/handler/generate_music_decode_test.py b/acestep/core/generation/handler/generate_music_decode_test.py index ece612ff1..57c280504 100644 --- a/acestep/core/generation/handler/generate_music_decode_test.py +++ b/acestep/core/generation/handler/generate_music_decode_test.py @@ -6,7 +6,7 @@ import unittest from contextlib import contextmanager from pathlib import Path -from unittest.mock import patch +from unittest.mock import MagicMock, patch import torch @@ -115,6 +115,11 @@ def _max_memory_allocated(self): """Return deterministic max-memory value for debug logging.""" return 0.0 + def _get_component_device(self, component: str) -> str: + """Return the host device when no component map is configured in tests.""" + _ = component + return self.device + def _mlx_vae_decode(self, latents): """Return deterministic decoded waveform for MLX decode branch.""" _ = latents @@ -291,7 +296,7 @@ def __init__(self): self.use_mlx_vae = False self.mlx_vae = None self.vae = _SuccessVae() - self.device = "cuda" + self.device = "cpu" host = _SuccessHost() pred_latents = torch.ones(1, 4, 3) @@ -310,6 +315,47 @@ def __init__(self): # The decoded waveform must be returned correctly. self.assertEqual(tuple(pred_wavs.shape), (1, 2, 8)) + def test_decode_pred_latents_queries_vram_on_mapped_vae_cuda_index(self): + """VRAM preflight should use the mapped VAE device, not cuda:0 by default.""" + + class _VramHost(_Host): + """Host that exposes a non-zero mapped VAE device for VRAM checks.""" + + def __init__(self): + super().__init__() + self.use_mlx_vae = False + self.mlx_vae = None + self.device = "cuda:0" + self.vae_component_device = "cuda:2" + + def _get_component_device(self, component: str) -> str: + if component == "vae": + return self.vae_component_device + return self.device + + host = _VramHost() + latent = MagicMock() + latent.detach.return_value = latent + latent.transpose.return_value = latent + latent.contiguous.return_value = latent + latent.to.return_value = latent + time_costs = {"total_time_cost": 1.0} + + with patch.object( + GENERATE_MUSIC_DECODE_MODULE, + "get_effective_free_vram_gb", + return_value=8.0, + ) as free_mock: + with patch.object(GENERATE_MUSIC_DECODE_MODULE.time, "time", side_effect=[10.0, 11.0]): + host._decode_generate_music_pred_latents( + pred_latents=latent, + progress=None, + use_tiled_decode=False, + time_costs=time_costs, + ) + + free_mock.assert_called_once_with(2) + if __name__ == "__main__": unittest.main() From 1d19b9b7731cb00c47f4b39e5736e84a2d7d3d21 Mon Sep 17 00:00:00 2001 From: steve Date: Mon, 6 Jul 2026 18:39:32 -0400 Subject: [PATCH 11/18] Raise DeviceMapError for malformed CUDA device indices. Wrap cuda_device_index int() parsing so strings like cuda:x raise the module's domain error instead of a raw ValueError. Co-authored-by: Cursor --- acestep/device_map/devices.py | 6 +++- acestep/test_device_map.py | 68 +++++++++++++++++++++++++++++++++++ 2 files changed, 73 insertions(+), 1 deletion(-) diff --git a/acestep/device_map/devices.py b/acestep/device_map/devices.py index fb19c76c7..430f19163 100644 --- a/acestep/device_map/devices.py +++ b/acestep/device_map/devices.py @@ -16,7 +16,11 @@ def cuda_device_index(device: str) -> int: if normalized == "cuda": return 0 if normalized.startswith("cuda:"): - return int(normalized.split(":", 1)[1]) + suffix = normalized.split(":", 1)[1] + try: + return int(suffix) + except ValueError as exc: + raise DeviceMapError(f"Not a CUDA device string: {device!r}") from exc raise DeviceMapError(f"Not a CUDA device string: {device!r}") diff --git a/acestep/test_device_map.py b/acestep/test_device_map.py index 0dee4da9e..cfe4a4f05 100644 --- a/acestep/test_device_map.py +++ b/acestep/test_device_map.py @@ -2,6 +2,7 @@ import os import unittest +from types import SimpleNamespace from unittest.mock import patch from acestep.device_map import ( @@ -10,12 +11,17 @@ GpuInfo, LayoutError, LayoutRequest, + collect_gpu_runtime_status, compute_auto_device_map, cuda_device_index, + device_map_to_dict, device_type, estimate_dit_peak_gb, estimate_lm_total_gb, + format_gpu_list_text, + gpu_info_to_dict, is_cuda_device, + log_lm_device_deprecation, normalize_component_device, parse_gpu_mapping, resolve_component_device_map, @@ -191,6 +197,68 @@ def test_cuda_device_index_helpers(self): self.assertEqual(cuda_device_index("cuda"), 0) self.assertEqual(cuda_device_index("cuda:2"), 2) + def test_cuda_device_index_raises_device_map_error_for_malformed_index(self): + with self.assertRaises(DeviceMapError): + cuda_device_index("cuda:x") + + +class GpuStatusHelpersTests(unittest.TestCase): + """Tests for CLI/API GPU listing and status helpers.""" + + def test_format_gpu_list_text_reports_no_devices(self): + self.assertEqual(format_gpu_list_text([]), "No CUDA devices detected.") + + def test_format_gpu_list_text_includes_device_rows(self): + text = format_gpu_list_text( + [GpuInfo(0, "RTX 3090", 24.0, 22.5, (8, 6))] + ) + self.assertIn("RTX 3090", text) + self.assertIn("24.00", text) + + def test_gpu_info_and_device_map_serialization(self): + gpu_payload = gpu_info_to_dict(GpuInfo(1, "GPU1", 24.0, 20.0, (8, 6))) + self.assertEqual(gpu_payload["index"], 1) + self.assertEqual(gpu_payload["compute_capability"], [8, 6]) + + device_map = ComponentDeviceMap( + dit="cuda:0", + vae="cuda:0", + text_encoder="cuda:0", + lm="cuda:1", + ) + map_payload = device_map_to_dict(device_map) + self.assertTrue(map_payload["multi_device"]) + self.assertIn("lm:1", map_payload["summary"]) + + def test_collect_gpu_runtime_status_includes_handler_layout(self): + handler = SimpleNamespace( + device_map=ComponentDeviceMap.from_single_device("cuda:2") + ) + with patch("acestep.device_map.status.discover_gpus", return_value=[]): + status = collect_gpu_runtime_status(handler=handler) + self.assertEqual(status["device_map"]["dit"], "cuda:2") + self.assertEqual(status["gpus"], []) + + def test_log_lm_device_deprecation_without_mapping(self): + with patch.dict(os.environ, {"ACESTEP_LM_DEVICE": "cuda:1"}, clear=False), patch( + "acestep.device_map.status.logger" + ) as mock_logger: + log_lm_device_deprecation() + mock_logger.warning.assert_called_once() + + def test_log_lm_device_deprecation_ignored_when_mapping_assigns_lm(self): + with patch.dict( + os.environ, + { + "ACESTEP_LM_DEVICE": "cuda:0", + "ACESTEP_GPU_MAPPING": "dit:0,vae:0,text_encoder:0,lm:1", + }, + clear=False, + ), patch("acestep.device_map.status.logger") as mock_logger: + log_lm_device_deprecation(using_device_map_lm=True) + mock_logger.warning.assert_called_once() + self.assertIn("ignored", mock_logger.warning.call_args[0][0]) + if __name__ == "__main__": unittest.main() From c3ef2dfd73331ac9c86f4d2f6c664d68e636a5f7 Mon Sep 17 00:00:00 2001 From: steve Date: Mon, 6 Jul 2026 19:13:53 -0400 Subject: [PATCH 12/18] Split generate_music_decode tests to satisfy 200 LOC cap. Move shared fixtures into generate_music_decode_test_support.py and split scenario tests into prepare, basic, CPU-offload, and VRAM modules. Co-authored-by: Cursor --- .../generate_music_decode_basic_test.py | 42 ++ .../generate_music_decode_cpu_offload_test.py | 100 +++++ .../generate_music_decode_prepare_test.py | 50 +++ .../handler/generate_music_decode_test.py | 362 ------------------ .../generate_music_decode_test_support.py | 111 ++++++ .../generate_music_decode_vram_test.py | 53 +++ 6 files changed, 356 insertions(+), 362 deletions(-) create mode 100644 acestep/core/generation/handler/generate_music_decode_basic_test.py create mode 100644 acestep/core/generation/handler/generate_music_decode_cpu_offload_test.py create mode 100644 acestep/core/generation/handler/generate_music_decode_prepare_test.py delete mode 100644 acestep/core/generation/handler/generate_music_decode_test.py create mode 100644 acestep/core/generation/handler/generate_music_decode_test_support.py create mode 100644 acestep/core/generation/handler/generate_music_decode_vram_test.py diff --git a/acestep/core/generation/handler/generate_music_decode_basic_test.py b/acestep/core/generation/handler/generate_music_decode_basic_test.py new file mode 100644 index 000000000..a3afb5f3d --- /dev/null +++ b/acestep/core/generation/handler/generate_music_decode_basic_test.py @@ -0,0 +1,42 @@ +"""Tests for the default ``generate_music`` latent decode path.""" + +import unittest +from unittest.mock import patch + +import torch + +from acestep.core.generation.handler.generate_music_decode_test_support import ( + GENERATE_MUSIC_DECODE_MODULE, + DecodeTestHost, +) + + +class GenerateMusicDecodeBasicTests(unittest.TestCase): + """Verify successful decode timing and output behavior.""" + + def test_decode_pred_latents_updates_decode_time_and_returns_cpu_latents(self): + host = DecodeTestHost() + pred_latents = torch.ones(1, 4, 3) + time_costs = {"total_time_cost": 1.0} + + def _progress(value, desc=None): + host.progress_calls.append((value, desc)) + + with patch.object(GENERATE_MUSIC_DECODE_MODULE.time, "time", side_effect=[10.0, 11.5]): + pred_wavs, pred_latents_cpu, updated_costs = host._decode_generate_music_pred_latents( + pred_latents=pred_latents, + progress=_progress, + use_tiled_decode=False, + time_costs=time_costs, + ) + + self.assertEqual(tuple(pred_wavs.shape), (1, 2, 8)) + self.assertEqual(pred_latents_cpu.device.type, "cpu") + self.assertAlmostEqual(updated_costs["vae_decode_time_cost"], 1.5, places=6) + self.assertAlmostEqual(updated_costs["total_time_cost"], 2.5, places=6) + self.assertAlmostEqual(updated_costs["offload_time_cost"], 0.25, places=6) + self.assertEqual(host.progress_calls[0][0], 0.8) + + +if __name__ == "__main__": + unittest.main() diff --git a/acestep/core/generation/handler/generate_music_decode_cpu_offload_test.py b/acestep/core/generation/handler/generate_music_decode_cpu_offload_test.py new file mode 100644 index 000000000..8df421f43 --- /dev/null +++ b/acestep/core/generation/handler/generate_music_decode_cpu_offload_test.py @@ -0,0 +1,100 @@ +"""Tests for CPU-offload and VAE restoration in ``generate_music`` decode.""" + +import unittest +from unittest.mock import patch + +import torch + +from acestep.core.generation.handler.generate_music_decode_test_support import ( + GENERATE_MUSIC_DECODE_MODULE, + DecodeTestHost, + FakeDecodeOutput, + FakeVae, +) + + +class GenerateMusicDecodeCpuOffloadTests(unittest.TestCase): + """Verify CPU-offload decode paths and device restoration.""" + + def test_decode_pred_latents_restores_vae_device_on_decode_error(self): + class FailingVae(FakeVae): + def __init__(self): + super().__init__() + self.cpu_calls = 0 + self.to_calls = [] + + def decode(self, latents: torch.Tensor): + _ = latents + raise RuntimeError("decode failed") + + def cpu(self): + self.cpu_calls += 1 + return self + + def to(self, *args, **kwargs): + self.to_calls.append((args, kwargs)) + return self + + class FailingHost(DecodeTestHost): + def __init__(self): + super().__init__() + self.use_mlx_vae = False + self.mlx_vae = None + self.vae = FailingVae() + self.empty_cache_calls = 0 + + def _empty_cache(self): + self.empty_cache_calls += 1 + + host = FailingHost() + with patch.dict(GENERATE_MUSIC_DECODE_MODULE.os.environ, {"ACESTEP_VAE_ON_CPU": "1"}, clear=False): + with self.assertRaisesRegex(RuntimeError, "decode failed"): + host._decode_generate_music_pred_latents( + pred_latents=torch.ones(1, 4, 3), + progress=None, + use_tiled_decode=False, + time_costs={"total_time_cost": 1.0}, + ) + + self.assertEqual(host.vae.cpu_calls, 1) + self.assertEqual(len(host.vae.to_calls), 1) + self.assertGreaterEqual(host.empty_cache_calls, 2) + + def test_decode_pred_latents_does_not_restore_latents_to_gpu_after_successful_cpu_decode(self): + class SuccessVae(FakeVae): + def __init__(self): + super().__init__() + self.vae_to_calls = [] + + def decode(self, latents: torch.Tensor): + return FakeDecodeOutput(torch.ones(latents.shape[0], 2, 8)) + + def cpu(self): + return self + + def to(self, *args, **kwargs): + self.vae_to_calls.append(args[0] if args else kwargs) + return self + + class SuccessHost(DecodeTestHost): + def __init__(self): + super().__init__() + self.use_mlx_vae = False + self.mlx_vae = None + self.vae = SuccessVae() + + host = SuccessHost() + with patch.dict(GENERATE_MUSIC_DECODE_MODULE.os.environ, {"ACESTEP_VAE_ON_CPU": "1"}, clear=False): + pred_wavs, _cpu_latents, _costs = host._decode_generate_music_pred_latents( + pred_latents=torch.ones(1, 4, 3), + progress=None, + use_tiled_decode=False, + time_costs={"total_time_cost": 1.0}, + ) + + self.assertEqual(len(host.vae.vae_to_calls), 1) + self.assertEqual(tuple(pred_wavs.shape), (1, 2, 8)) + + +if __name__ == "__main__": + unittest.main() diff --git a/acestep/core/generation/handler/generate_music_decode_prepare_test.py b/acestep/core/generation/handler/generate_music_decode_prepare_test.py new file mode 100644 index 000000000..5752631ba --- /dev/null +++ b/acestep/core/generation/handler/generate_music_decode_prepare_test.py @@ -0,0 +1,50 @@ +"""Tests for ``generate_music`` decode-state preparation.""" + +import unittest + +import torch + +from acestep.core.generation.handler.generate_music_decode_test_support import DecodeTestHost + + +class GenerateMusicDecodePrepareTests(unittest.TestCase): + """Verify decode-state preparation helper behavior.""" + + def test_prepare_decode_state_updates_progress_estimates(self): + host = DecodeTestHost() + outputs = { + "target_latents": torch.ones(1, 4, 3), + "time_costs": {"total_time_cost": 1.0, "diffusion_per_step_time_cost": 0.2}, + } + pred_latents, time_costs = host._prepare_generate_music_decode_state( + outputs=outputs, + infer_steps_for_progress=8, + actual_batch_size=1, + audio_duration=12.0, + latent_shift=0.0, + latent_rescale=1.0, + ) + self.assertEqual(tuple(pred_latents.shape), (1, 4, 3)) + self.assertEqual(time_costs["offload_time_cost"], 0.25) + self.assertEqual(host._last_diffusion_per_step_sec, 0.2) + self.assertEqual(host.estimate_calls[0]["infer_steps"], 8) + + def test_prepare_decode_state_raises_for_nan_latents(self): + host = DecodeTestHost() + outputs = { + "target_latents": torch.tensor([[[float("nan")]]]), + "time_costs": {"total_time_cost": 1.0}, + } + with self.assertRaises(RuntimeError): + host._prepare_generate_music_decode_state( + outputs=outputs, + infer_steps_for_progress=8, + actual_batch_size=1, + audio_duration=None, + latent_shift=0.0, + latent_rescale=1.0, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/acestep/core/generation/handler/generate_music_decode_test.py b/acestep/core/generation/handler/generate_music_decode_test.py deleted file mode 100644 index 57c280504..000000000 --- a/acestep/core/generation/handler/generate_music_decode_test.py +++ /dev/null @@ -1,362 +0,0 @@ -"""Tests for extracted ``generate_music`` decode helper mixin behavior.""" - -import importlib.util -import types -import sys -import unittest -from contextlib import contextmanager -from pathlib import Path -from unittest.mock import MagicMock, patch - -import torch - - -def _load_generate_music_decode_module(): - """Load ``generate_music_decode.py`` from disk and return its module object. - - Raises ``FileNotFoundError`` or ``ImportError`` when loading fails. - """ - repo_root = Path(__file__).resolve().parents[4] - if str(repo_root) not in sys.path: - sys.path.insert(0, str(repo_root)) - package_paths = { - "acestep": repo_root / "acestep", - "acestep.core": repo_root / "acestep" / "core", - "acestep.core.generation": repo_root / "acestep" / "core" / "generation", - "acestep.core.generation.handler": repo_root / "acestep" / "core" / "generation" / "handler", - } - for package_name, package_path in package_paths.items(): - if package_name in sys.modules: - continue - package_module = types.ModuleType(package_name) - package_module.__path__ = [str(package_path)] - sys.modules[package_name] = package_module - module_path = Path(__file__).with_name("generate_music_decode.py") - spec = importlib.util.spec_from_file_location( - "acestep.core.generation.handler.generate_music_decode", - module_path, - ) - module = importlib.util.module_from_spec(spec) - assert spec.loader is not None - spec.loader.exec_module(module) - return module - - -GENERATE_MUSIC_DECODE_MODULE = _load_generate_music_decode_module() -GenerateMusicDecodeMixin = GENERATE_MUSIC_DECODE_MODULE.GenerateMusicDecodeMixin - - -class _FakeDecodeOutput: - """Minimal VAE decode output container exposing ``sample`` attribute.""" - - def __init__(self, sample: torch.Tensor): - """Store decoded sample tensor for mixin decode flow.""" - self.sample = sample - - -class _FakeVae: - """Minimal VAE stand-in with dtype, decode, and parameter iteration hooks.""" - - def __init__(self): - """Initialize deterministic dtype/device state for decode tests.""" - self.dtype = torch.float32 - self._param = torch.nn.Parameter(torch.zeros(1)) - - def decode(self, latents: torch.Tensor): - """Return deterministic decoded waveform output.""" - return _FakeDecodeOutput(torch.ones(latents.shape[0], 2, 8)) - - def parameters(self): - """Yield one parameter so `.device` lookups remain valid.""" - yield self._param - - def cpu(self): - """Return self for test-only CPU transfer calls.""" - return self - - def to(self, *_args, **_kwargs): - """Return self for test-only device transfer calls.""" - return self - - -class _Host(GenerateMusicDecodeMixin): - """Minimal decode-mixin host exposing deterministic state for assertions.""" - - def __init__(self): - """Initialize deterministic runtime state for decode tests.""" - self.current_offload_cost = 0.25 - self.debug_stats = False - self._last_diffusion_per_step_sec = None - self.estimate_calls = [] - self.progress_calls = [] - self.device = "cpu" - self.use_mlx_vae = True - self.mlx_vae = object() - self.vae = _FakeVae() - - def _update_progress_estimate(self, **kwargs): - """Capture estimate updates for assertions.""" - self.estimate_calls.append(kwargs) - - @contextmanager - def _load_model_context(self, _model_name): - """Provide no-op model context manager for decode tests.""" - yield - - def _empty_cache(self): - """Provide no-op cache clear helper for decode tests.""" - return None - - def _memory_allocated(self): - """Return deterministic allocated-memory value for debug logging.""" - return 0.0 - - def _max_memory_allocated(self): - """Return deterministic max-memory value for debug logging.""" - return 0.0 - - def _get_component_device(self, component: str) -> str: - """Return the host device when no component map is configured in tests.""" - _ = component - return self.device - - def _mlx_vae_decode(self, latents): - """Return deterministic decoded waveform for MLX decode branch.""" - _ = latents - return torch.ones(1, 2, 8) - - def tiled_decode(self, latents): - """Return deterministic decoded waveform for tiled decode branch.""" - _ = latents - return torch.ones(1, 2, 8) - - -class GenerateMusicDecodeMixinTests(unittest.TestCase): - """Verify decode-state preparation and latent decode helper behavior.""" - - def test_prepare_decode_state_updates_progress_estimates(self): - """It updates timing fields and progress estimate metadata for valid latents.""" - host = _Host() - outputs = { - "target_latents": torch.ones(1, 4, 3), - "time_costs": {"total_time_cost": 1.0, "diffusion_per_step_time_cost": 0.2}, - } - pred_latents, time_costs = host._prepare_generate_music_decode_state( - outputs=outputs, - infer_steps_for_progress=8, - actual_batch_size=1, - audio_duration=12.0, - latent_shift=0.0, - latent_rescale=1.0, - ) - self.assertEqual(tuple(pred_latents.shape), (1, 4, 3)) - self.assertEqual(time_costs["offload_time_cost"], 0.25) - self.assertEqual(host._last_diffusion_per_step_sec, 0.2) - self.assertEqual(host.estimate_calls[0]["infer_steps"], 8) - - def test_prepare_decode_state_raises_for_nan_latents(self): - """It raises runtime error when diffusion latents contain NaN values.""" - host = _Host() - outputs = { - "target_latents": torch.tensor([[[float("nan")]]]), - "time_costs": {"total_time_cost": 1.0}, - } - with self.assertRaises(RuntimeError): - host._prepare_generate_music_decode_state( - outputs=outputs, - infer_steps_for_progress=8, - actual_batch_size=1, - audio_duration=None, - latent_shift=0.0, - latent_rescale=1.0, - ) - - def test_decode_pred_latents_updates_decode_time_and_returns_cpu_latents(self): - """It decodes latents and updates decode timing metrics in time_costs.""" - host = _Host() - pred_latents = torch.ones(1, 4, 3) - time_costs = {"total_time_cost": 1.0} - - def _progress(value, desc=None): - """Capture progress updates for assertions.""" - host.progress_calls.append((value, desc)) - - with patch.object(GENERATE_MUSIC_DECODE_MODULE.time, "time", side_effect=[10.0, 11.5]): - pred_wavs, pred_latents_cpu, updated_costs = host._decode_generate_music_pred_latents( - pred_latents=pred_latents, - progress=_progress, - use_tiled_decode=False, - time_costs=time_costs, - ) - - self.assertEqual(tuple(pred_wavs.shape), (1, 2, 8)) - self.assertEqual(pred_latents_cpu.device.type, "cpu") - self.assertAlmostEqual(updated_costs["vae_decode_time_cost"], 1.5, places=6) - self.assertAlmostEqual(updated_costs["total_time_cost"], 2.5, places=6) - self.assertAlmostEqual(updated_costs["offload_time_cost"], 0.25, places=6) - self.assertEqual(host.progress_calls[0][0], 0.8) - - def test_decode_pred_latents_restores_vae_device_on_decode_error(self): - """It restores VAE device in the CPU-offload path even when decode raises.""" - - class _FailingVae(_FakeVae): - """VAE double that raises during decode and records transfer calls.""" - - def __init__(self): - """Initialize transfer call trackers for restoration assertions.""" - super().__init__() - self.cpu_calls = 0 - self.to_calls = [] - - def decode(self, latents: torch.Tensor): - """Raise decode error to exercise restoration in finally branch.""" - _ = latents - raise RuntimeError("decode failed") - - def cpu(self): - """Record explicit CPU transfer and return self.""" - self.cpu_calls += 1 - return self - - def to(self, *args, **kwargs): - """Record restore transfer target and return self.""" - self.to_calls.append((args, kwargs)) - return self - - class _FailingHost(_Host): - """Host variant that forces non-MLX VAE decode and tracks cache clears.""" - - def __init__(self): - """Set non-MLX state so CPU offload path is exercised deterministically.""" - super().__init__() - self.use_mlx_vae = False - self.mlx_vae = None - self.vae = _FailingVae() - self.empty_cache_calls = 0 - - def _empty_cache(self): - """Count cache-clear calls to verify finally cleanup runs.""" - self.empty_cache_calls += 1 - - host = _FailingHost() - pred_latents = torch.ones(1, 4, 3) - time_costs = {"total_time_cost": 1.0} - - with patch.dict(GENERATE_MUSIC_DECODE_MODULE.os.environ, {"ACESTEP_VAE_ON_CPU": "1"}, clear=False): - with self.assertRaisesRegex(RuntimeError, "decode failed"): - host._decode_generate_music_pred_latents( - pred_latents=pred_latents, - progress=None, - use_tiled_decode=False, - time_costs=time_costs, - ) - - self.assertEqual(host.vae.cpu_calls, 1) - self.assertEqual(len(host.vae.to_calls), 1) - self.assertGreaterEqual(host.empty_cache_calls, 2) - - - def test_decode_pred_latents_does_not_restore_latents_to_gpu_after_successful_cpu_decode(self): - """It does not move pred_latents_for_decode back to GPU after a successful CPU decode. - - The removed ``pred_latents_for_decode = pred_latents_for_decode.to(vae_device)`` line - was causing a wasteful re-allocation of the already-decoded input tensor on the - GPU. After the fix, only the VAE itself is restored; the input latent is not. - """ - - class _SuccessVae(_FakeVae): - """VAE double that records transfer calls and succeeds on decode.""" - - def __init__(self): - """Initialize transfer call trackers.""" - super().__init__() - self.vae_to_calls = [] - self._device = "cuda" - - def decode(self, latents: torch.Tensor): - """Return a simple decoded output.""" - return _FakeDecodeOutput(torch.ones(latents.shape[0], 2, 8)) - - def cpu(self): - """Simulate VAE being moved to CPU.""" - self._device = "cpu" - return self - - def to(self, *args, **kwargs): - """Record VAE device-transfer destinations.""" - self.vae_to_calls.append(args[0] if args else kwargs) - return self - - class _SuccessHost(_Host): - """Host that forces non-MLX VAE so the CPU-decode path is exercised.""" - - def __init__(self): - """Configure non-MLX state and a tracking VAE.""" - super().__init__() - self.use_mlx_vae = False - self.mlx_vae = None - self.vae = _SuccessVae() - self.device = "cpu" - - host = _SuccessHost() - pred_latents = torch.ones(1, 4, 3) - time_costs = {"total_time_cost": 1.0} - - with patch.dict(GENERATE_MUSIC_DECODE_MODULE.os.environ, {"ACESTEP_VAE_ON_CPU": "1"}, clear=False): - pred_wavs, _cpu_latents, _costs = host._decode_generate_music_pred_latents( - pred_latents=pred_latents, - progress=None, - use_tiled_decode=False, - time_costs=time_costs, - ) - - # VAE itself must be restored to its original device. - self.assertEqual(len(host.vae.vae_to_calls), 1) - # The decoded waveform must be returned correctly. - self.assertEqual(tuple(pred_wavs.shape), (1, 2, 8)) - - def test_decode_pred_latents_queries_vram_on_mapped_vae_cuda_index(self): - """VRAM preflight should use the mapped VAE device, not cuda:0 by default.""" - - class _VramHost(_Host): - """Host that exposes a non-zero mapped VAE device for VRAM checks.""" - - def __init__(self): - super().__init__() - self.use_mlx_vae = False - self.mlx_vae = None - self.device = "cuda:0" - self.vae_component_device = "cuda:2" - - def _get_component_device(self, component: str) -> str: - if component == "vae": - return self.vae_component_device - return self.device - - host = _VramHost() - latent = MagicMock() - latent.detach.return_value = latent - latent.transpose.return_value = latent - latent.contiguous.return_value = latent - latent.to.return_value = latent - time_costs = {"total_time_cost": 1.0} - - with patch.object( - GENERATE_MUSIC_DECODE_MODULE, - "get_effective_free_vram_gb", - return_value=8.0, - ) as free_mock: - with patch.object(GENERATE_MUSIC_DECODE_MODULE.time, "time", side_effect=[10.0, 11.0]): - host._decode_generate_music_pred_latents( - pred_latents=latent, - progress=None, - use_tiled_decode=False, - time_costs=time_costs, - ) - - free_mock.assert_called_once_with(2) - - -if __name__ == "__main__": - unittest.main() - diff --git a/acestep/core/generation/handler/generate_music_decode_test_support.py b/acestep/core/generation/handler/generate_music_decode_test_support.py new file mode 100644 index 000000000..aa6f6c54a --- /dev/null +++ b/acestep/core/generation/handler/generate_music_decode_test_support.py @@ -0,0 +1,111 @@ +"""Shared fixtures for ``generate_music_decode`` mixin tests.""" + +import importlib.util +import sys +import types +from contextlib import contextmanager +from pathlib import Path + +import torch + + +def load_generate_music_decode_module(): + """Load ``generate_music_decode.py`` from disk and return its module object.""" + repo_root = Path(__file__).resolve().parents[4] + if str(repo_root) not in sys.path: + sys.path.insert(0, str(repo_root)) + package_paths = { + "acestep": repo_root / "acestep", + "acestep.core": repo_root / "acestep" / "core", + "acestep.core.generation": repo_root / "acestep" / "core" / "generation", + "acestep.core.generation.handler": repo_root / "acestep" / "core" / "generation" / "handler", + } + for package_name, package_path in package_paths.items(): + if package_name in sys.modules: + continue + package_module = types.ModuleType(package_name) + package_module.__path__ = [str(package_path)] + sys.modules[package_name] = package_module + module_path = Path(__file__).with_name("generate_music_decode.py") + spec = importlib.util.spec_from_file_location( + "acestep.core.generation.handler.generate_music_decode", + module_path, + ) + module = importlib.util.module_from_spec(spec) + assert spec.loader is not None + spec.loader.exec_module(module) + return module + + +GENERATE_MUSIC_DECODE_MODULE = load_generate_music_decode_module() +GenerateMusicDecodeMixin = GENERATE_MUSIC_DECODE_MODULE.GenerateMusicDecodeMixin + + +class FakeDecodeOutput: + """Minimal VAE decode output container exposing ``sample`` attribute.""" + + def __init__(self, sample: torch.Tensor): + self.sample = sample + + +class FakeVae: + """Minimal VAE stand-in with dtype, decode, and parameter iteration hooks.""" + + def __init__(self): + self.dtype = torch.float32 + self._param = torch.nn.Parameter(torch.zeros(1)) + + def decode(self, latents: torch.Tensor): + return FakeDecodeOutput(torch.ones(latents.shape[0], 2, 8)) + + def parameters(self): + yield self._param + + def cpu(self): + return self + + def to(self, *_args, **_kwargs): + return self + + +class DecodeTestHost(GenerateMusicDecodeMixin): + """Minimal decode-mixin host exposing deterministic state for assertions.""" + + def __init__(self): + self.current_offload_cost = 0.25 + self.debug_stats = False + self._last_diffusion_per_step_sec = None + self.estimate_calls = [] + self.progress_calls = [] + self.device = "cpu" + self.use_mlx_vae = True + self.mlx_vae = object() + self.vae = FakeVae() + + def _update_progress_estimate(self, **kwargs): + self.estimate_calls.append(kwargs) + + @contextmanager + def _load_model_context(self, _model_name): + yield + + def _empty_cache(self): + return None + + def _memory_allocated(self): + return 0.0 + + def _max_memory_allocated(self): + return 0.0 + + def _get_component_device(self, component: str) -> str: + _ = component + return self.device + + def _mlx_vae_decode(self, latents): + _ = latents + return torch.ones(1, 2, 8) + + def tiled_decode(self, latents): + _ = latents + return torch.ones(1, 2, 8) diff --git a/acestep/core/generation/handler/generate_music_decode_vram_test.py b/acestep/core/generation/handler/generate_music_decode_vram_test.py new file mode 100644 index 000000000..5721a42ce --- /dev/null +++ b/acestep/core/generation/handler/generate_music_decode_vram_test.py @@ -0,0 +1,53 @@ +"""Tests for multi-GPU VRAM preflight in ``generate_music`` decode.""" + +import unittest +from unittest.mock import MagicMock, patch + +from acestep.core.generation.handler.generate_music_decode_test_support import ( + GENERATE_MUSIC_DECODE_MODULE, + DecodeTestHost, +) + + +class GenerateMusicDecodeVramTests(unittest.TestCase): + """Verify decode VRAM checks honor mapped component devices.""" + + def test_decode_pred_latents_queries_vram_on_mapped_vae_cuda_index(self): + class VramHost(DecodeTestHost): + def __init__(self): + super().__init__() + self.use_mlx_vae = False + self.mlx_vae = None + self.device = "cuda:0" + self.vae_component_device = "cuda:2" + + def _get_component_device(self, component: str) -> str: + if component == "vae": + return self.vae_component_device + return self.device + + host = VramHost() + latent = MagicMock() + latent.detach.return_value = latent + latent.transpose.return_value = latent + latent.contiguous.return_value = latent + latent.to.return_value = latent + + with patch.object( + GENERATE_MUSIC_DECODE_MODULE, + "get_effective_free_vram_gb", + return_value=8.0, + ) as free_mock: + with patch.object(GENERATE_MUSIC_DECODE_MODULE.time, "time", side_effect=[10.0, 11.0]): + host._decode_generate_music_pred_latents( + pred_latents=latent, + progress=None, + use_tiled_decode=False, + time_costs={"total_time_cost": 1.0}, + ) + + free_mock.assert_called_once_with(2) + + +if __name__ == "__main__": + unittest.main() From 7aee60f52383d1d58593d6812f815ee8c85241a3 Mon Sep 17 00:00:00 2001 From: steve Date: Mon, 6 Jul 2026 19:23:26 -0400 Subject: [PATCH 13/18] Use normal imports in generate_music_decode test support. Drop the custom sys.modules package stubbing loader; import the mixin through the public package path like other handler tests. Co-authored-by: Cursor --- .../generate_music_decode_test_support.py | 38 +------------------ 1 file changed, 2 insertions(+), 36 deletions(-) diff --git a/acestep/core/generation/handler/generate_music_decode_test_support.py b/acestep/core/generation/handler/generate_music_decode_test_support.py index aa6f6c54a..222b85192 100644 --- a/acestep/core/generation/handler/generate_music_decode_test_support.py +++ b/acestep/core/generation/handler/generate_music_decode_test_support.py @@ -1,44 +1,10 @@ """Shared fixtures for ``generate_music_decode`` mixin tests.""" -import importlib.util -import sys -import types from contextlib import contextmanager -from pathlib import Path +import acestep.core.generation.handler.generate_music_decode as GENERATE_MUSIC_DECODE_MODULE import torch - - -def load_generate_music_decode_module(): - """Load ``generate_music_decode.py`` from disk and return its module object.""" - repo_root = Path(__file__).resolve().parents[4] - if str(repo_root) not in sys.path: - sys.path.insert(0, str(repo_root)) - package_paths = { - "acestep": repo_root / "acestep", - "acestep.core": repo_root / "acestep" / "core", - "acestep.core.generation": repo_root / "acestep" / "core" / "generation", - "acestep.core.generation.handler": repo_root / "acestep" / "core" / "generation" / "handler", - } - for package_name, package_path in package_paths.items(): - if package_name in sys.modules: - continue - package_module = types.ModuleType(package_name) - package_module.__path__ = [str(package_path)] - sys.modules[package_name] = package_module - module_path = Path(__file__).with_name("generate_music_decode.py") - spec = importlib.util.spec_from_file_location( - "acestep.core.generation.handler.generate_music_decode", - module_path, - ) - module = importlib.util.module_from_spec(spec) - assert spec.loader is not None - spec.loader.exec_module(module) - return module - - -GENERATE_MUSIC_DECODE_MODULE = load_generate_music_decode_module() -GenerateMusicDecodeMixin = GENERATE_MUSIC_DECODE_MODULE.GenerateMusicDecodeMixin +from acestep.core.generation.handler.generate_music_decode import GenerateMusicDecodeMixin class FakeDecodeOutput: From ec020947b3881270a4c791fe1dbb399f96ab720d Mon Sep 17 00:00:00 2001 From: steve Date: Sat, 11 Jul 2026 10:17:42 -0400 Subject: [PATCH 14/18] fix(inference): honor mapped cuda:N for LM init and Gradio wiring Resolve Gradio lm_device after DiT initialize_service so device_map.lm is applied, normalize/preserve indexed CUDA devices in llm_inference, and move text-encoder token ids onto the text-encoder component device. Co-authored-by: Cursor --- .../generation/handler/conditioning_embed.py | 10 ++- .../handler/conditioning_embed_test.py | 32 +++++++++- .../service_generate_flow_edit_source.py | 8 +-- acestep/llm_backend_compat.py | 4 +- acestep/llm_backend_compat_test.py | 13 ++++ acestep/llm_inference.py | 29 +++++---- acestep/llm_inference_cuda_index_test.py | 63 +++++++++++++++++++ .../gradio/events/generation/service_init.py | 34 +++++----- .../events/generation/service_init_test.py | 57 +++++++++++++++++ 9 files changed, 212 insertions(+), 38 deletions(-) create mode 100644 acestep/llm_inference_cuda_index_test.py diff --git a/acestep/core/generation/handler/conditioning_embed.py b/acestep/core/generation/handler/conditioning_embed.py index 184f7c1b4..6a38d5db9 100644 --- a/acestep/core/generation/handler/conditioning_embed.py +++ b/acestep/core/generation/handler/conditioning_embed.py @@ -75,12 +75,20 @@ def _ensure_latent_3d(z: torch.Tensor) -> torch.Tensor: return refer_audio_latents, refer_audio_order_mask def infer_text_embeddings(self, text_token_idss): - """Infer text-token embeddings via text encoder.""" + """Infer text-token embeddings via text encoder. + + Token ids are often built on the DiT device; move them onto the + text-encoder device so split maps like ``text_encoder:3`` work. + """ + encoder_device = self._get_component_device("text_encoder") + text_token_idss = text_token_idss.to(encoder_device) with torch.inference_mode(): return self.text_encoder(input_ids=text_token_idss, lyric_attention_mask=None).last_hidden_state def infer_lyric_embeddings(self, lyric_token_ids): """Infer lyric-token embeddings via text encoder embedding table.""" + encoder_device = self._get_component_device("text_encoder") + lyric_token_ids = lyric_token_ids.to(encoder_device) with torch.inference_mode(): return self.text_encoder.embed_tokens(lyric_token_ids) diff --git a/acestep/core/generation/handler/conditioning_embed_test.py b/acestep/core/generation/handler/conditioning_embed_test.py index 8cb85413a..3d7507aec 100644 --- a/acestep/core/generation/handler/conditioning_embed_test.py +++ b/acestep/core/generation/handler/conditioning_embed_test.py @@ -24,12 +24,21 @@ def embed_tokens(self, token_ids): class _Host(ConditioningEmbedMixin): """Minimal host implementing ConditioningEmbedMixin dependencies.""" - def __init__(self): - self.device = "cpu" + def __init__(self, text_encoder_device="cpu", dit_device="cpu"): + self.device = dit_device self.dtype = torch.float32 self.silence_latent = torch.zeros(1, 128, 6, dtype=torch.float32) self.text_encoder = _FakeTextEncoder() self.tiled_encode_calls = 0 + self._devices = { + "dit": dit_device, + "vae": dit_device, + "text_encoder": text_encoder_device, + "lm": dit_device, + } + + def _get_component_device(self, component): + return self._devices[component] def _ensure_silence_latent_on_device(self): return None @@ -111,6 +120,25 @@ def test_preprocess_batch_returns_expected_tuple_shape(self): self.assertEqual(result[0], ["k1", "k2"]) self.assertEqual(result[3].shape, (2, 128, 6)) + def test_infer_text_embeddings_moves_ids_to_text_encoder_device(self): + """Token ids on DiT device must be moved onto a remote text encoder.""" + if not torch.cuda.is_available() or torch.cuda.device_count() < 2: + self.skipTest("needs 2+ CUDA devices") + + seen = {} + + class _DeviceCheckingEncoder(_FakeTextEncoder): + def __call__(self, input_ids, lyric_attention_mask=None): + seen["device"] = str(input_ids.device) + return super().__call__(input_ids, lyric_attention_mask) + + host = _Host(text_encoder_device="cuda:1", dit_device="cuda:0") + host.text_encoder = _DeviceCheckingEncoder() + ids = torch.ones(1, 4, dtype=torch.long, device="cuda:0") + out = host.infer_text_embeddings(ids) + self.assertEqual(seen["device"], "cuda:1") + self.assertEqual(out.shape, (1, 4, 6)) + if __name__ == "__main__": unittest.main() diff --git a/acestep/core/generation/handler/service_generate_flow_edit_source.py b/acestep/core/generation/handler/service_generate_flow_edit_source.py index 638e10b1b..9f408af39 100644 --- a/acestep/core/generation/handler/service_generate_flow_edit_source.py +++ b/acestep/core/generation/handler/service_generate_flow_edit_source.py @@ -77,13 +77,9 @@ def embed_source( ) -> Tuple[torch.Tensor, torch.Tensor]: """Run text + lyric encoders on the source tokens. - Tokens come back from the tokenizer on CPU; the regular batch path - moves them to ``handler.device`` before encoding (see - ``preprocess_batch``), so we mirror that here. + Tokens may live on CPU or the DiT device; ``infer_*_embeddings`` + moves them onto the text-encoder component device. """ - device = handler.device - text_token_idss = text_token_idss.to(device=device) - lyric_token_idss = lyric_token_idss.to(device=device) with handler._load_model_context("text_encoder"): text_hs = handler.infer_text_embeddings(text_token_idss) lyric_hs = handler.infer_lyric_embeddings(lyric_token_idss) diff --git a/acestep/llm_backend_compat.py b/acestep/llm_backend_compat.py index 7bd6a2751..b2c4f9001 100644 --- a/acestep/llm_backend_compat.py +++ b/acestep/llm_backend_compat.py @@ -3,6 +3,8 @@ import importlib import sys +from acestep.device_map import is_cuda_device + def _has_working_triton_installation() -> bool: """Return whether the Triton modules required by nano-vllm import cleanly.""" @@ -25,7 +27,7 @@ def get_vllm_preflight_warning(*, device: str, platform: str | None = None) -> s A warning string when vLLM should fall back to PyTorch, otherwise ``None``. """ active_platform = sys.platform if platform is None else platform - if device != "cuda" or active_platform != "win32": + if not is_cuda_device(device) or active_platform != "win32": return None if _has_working_triton_installation(): return None diff --git a/acestep/llm_backend_compat_test.py b/acestep/llm_backend_compat_test.py index 4dab0887a..f1faf5b85 100644 --- a/acestep/llm_backend_compat_test.py +++ b/acestep/llm_backend_compat_test.py @@ -46,9 +46,22 @@ def test_get_vllm_preflight_warning_returns_none_outside_windows_cuda(self) -> N ): warning = get_vllm_preflight_warning(device="cpu", platform="win32") linux_warning = get_vllm_preflight_warning(device="cuda", platform="linux") + indexed_linux = get_vllm_preflight_warning(device="cuda:1", platform="linux") self.assertIsNone(warning) self.assertIsNone(linux_warning) + self.assertIsNone(indexed_linux) + + def test_get_vllm_preflight_warning_accepts_indexed_cuda_on_windows(self) -> None: + """Indexed CUDA devices must still trigger the Windows Triton preflight.""" + with patch( + "acestep.llm_backend_compat._has_working_triton_installation", + return_value=False, + ): + warning = get_vllm_preflight_warning(device="cuda:1", platform="win32") + + self.assertIsNotNone(warning) + self.assertIn("Windows", warning) @unittest.skipIf(LLMHandler is None, f"llm_inference import unavailable: {_IMPORT_ERROR}") diff --git a/acestep/llm_inference.py b/acestep/llm_inference.py index c8f73db56..7cdb12933 100644 --- a/acestep/llm_inference.py +++ b/acestep/llm_inference.py @@ -554,8 +554,6 @@ def initialize( else: logger.warning("[initialize] CUDA requested but unavailable. Falling back to CPU.") device = "cpu" - elif is_cuda_device(device): - device = normalize_component_device(device) elif device == "mps" and not (hasattr(torch.backends, "mps") and torch.backends.mps.is_available()): if torch.cuda.is_available(): logger.warning("[initialize] MPS requested but unavailable. Falling back to CUDA.") @@ -577,6 +575,10 @@ def initialize( logger.warning("[initialize] XPU requested but unavailable. Falling back to CPU.") device = "cpu" + # Normalize bare "cuda" → "cuda:0" and preserve mapped "cuda:N". + if is_cuda_device(device): + device = normalize_component_device(device) + self.device = device self.offload_to_cpu = offload_to_cpu @@ -659,9 +661,10 @@ def initialize( # RuntimeError: Offset increment outside graph capture encountered unexpectedly is_rocm = hasattr(torch.version, 'hip') and torch.version.hip is not None is_jetson = False - if device == "cuda" and torch.cuda.is_available(): + if is_cuda_device(device) and torch.cuda.is_available(): try: - dev_name = torch.cuda.get_device_name(0).lower() + dit_idx = cuda_device_index(device) + dev_name = torch.cuda.get_device_name(dit_idx).lower() is_jetson = any(k in dev_name for k in ("orin", "xavier", "tegra")) if is_jetson: logger.info(f"Jetson GPU detected ({dev_name}): disabling CUDA graph capture for nano-vllm") @@ -720,7 +723,7 @@ def initialize( status_msg = f"✅ 5Hz LM initialized (PyTorch fallback, MLX not available)\nModel: {full_lm_model_path}\nBackend: PyTorch" return status_msg, True - if backend == "vllm" and device != "cuda": + if backend == "vllm" and not is_cuda_device(device): logger.info( f"[initialize] vllm backend requires CUDA, using PyTorch backend for device={device}." ) @@ -738,19 +741,23 @@ def initialize( # Initialize based on user-selected backend if backend == "vllm": _warn_if_prerelease_python() - total_gb = get_gpu_memory_gb() if device == "cuda" else 0.0 + total_gb = 0.0 free_gb = 0.0 - if device == "cuda" and torch.cuda.is_available(): + if is_cuda_device(device) and torch.cuda.is_available(): try: + device_index = cuda_device_index(device) + total_gb = get_gpu_memory_gb() if hasattr(torch.cuda, "mem_get_info"): - free_bytes, _ = torch.cuda.mem_get_info() + free_bytes, _ = torch.cuda.mem_get_info(device_index) free_gb = free_bytes / (1024**3) else: - total_bytes = torch.cuda.get_device_properties(0).total_memory - free_gb = (total_bytes - torch.cuda.memory_reserved(0)) / (1024**3) + total_bytes = torch.cuda.get_device_properties(device_index).total_memory + free_gb = ( + total_bytes - torch.cuda.memory_reserved(device_index) + ) / (1024**3) except Exception: free_gb = 0.0 - if device == "cuda" and free_gb < VRAM_SAFE_FREE_GB: + if is_cuda_device(device) and free_gb < VRAM_SAFE_FREE_GB: logger.warning( f"vLLM disabled due to insufficient free VRAM (total={total_gb:.2f}GB, free={free_gb:.2f}GB, need>={VRAM_SAFE_FREE_GB}GB free) — falling back to PyTorch backend" ) diff --git a/acestep/llm_inference_cuda_index_test.py b/acestep/llm_inference_cuda_index_test.py new file mode 100644 index 000000000..d355614f6 --- /dev/null +++ b/acestep/llm_inference_cuda_index_test.py @@ -0,0 +1,63 @@ +"""Unit tests for indexed CUDA device resolution in LLMHandler.initialize.""" + +from __future__ import annotations + +import unittest +from unittest.mock import patch + +try: + from acestep.llm_inference import LLMHandler + from acestep.device_map import is_cuda_device + _IMPORT_ERROR = None +except ImportError as exc: # pragma: no cover + LLMHandler = None + is_cuda_device = None + _IMPORT_ERROR = exc + + +@unittest.skipIf(LLMHandler is None, f"llm_inference import unavailable: {_IMPORT_ERROR}") +class TestLmInitializeCudaIndex(unittest.TestCase): + """Verify auto/cuda:N normalization for LM multi-GPU placement.""" + + def _resolved_device(self, device: str) -> str: + """Run initialize far enough to resolve self.device, then abort cleanly.""" + handler = LLMHandler() + with patch("acestep.llm_inference.os.path.exists", return_value=False): + status, ok = handler.initialize( + checkpoint_dir="/tmp/checkpoints", + lm_model_path="acestep-5Hz-lm-1.7B", + backend="pt", + device=device, + offload_to_cpu=False, + dtype=None, + ) + self.assertFalse(ok) + self.assertIn("not found", status) + return handler.device + + @patch("acestep.llm_inference.torch.cuda.is_available", return_value=True) + def test_auto_normalizes_to_cuda0(self, _mock_cuda): + """auto on CUDA must become cuda:0, not bare cuda.""" + self.assertEqual(self._resolved_device("auto"), "cuda:0") + + @patch("acestep.llm_inference.torch.cuda.is_available", return_value=True) + def test_bare_cuda_normalizes_to_cuda0(self, _mock_cuda): + """Bare cuda must normalize to cuda:0.""" + self.assertEqual(self._resolved_device("cuda"), "cuda:0") + + @patch("acestep.llm_inference.torch.cuda.is_available", return_value=True) + def test_mapped_cuda_index_preserved(self, _mock_cuda): + """Explicit cuda:N from device_map must be preserved.""" + self.assertEqual(self._resolved_device("cuda:1"), "cuda:1") + + def test_vllm_gate_accepts_indexed_cuda(self): + """vLLM CUDA gate must treat cuda:N as CUDA (not force PT).""" + device = "cuda:1" + backend = "vllm" + if backend == "vllm" and not is_cuda_device(device): + backend = "pt" + self.assertEqual(backend, "vllm") + + +if __name__ == "__main__": + unittest.main() diff --git a/acestep/ui/gradio/events/generation/service_init.py b/acestep/ui/gradio/events/generation/service_init.py index eb7865421..ad0f88598 100644 --- a/acestep/ui/gradio/events/generation/service_init.py +++ b/acestep/ui/gradio/events/generation/service_init.py @@ -86,23 +86,8 @@ def init_service_wrapper( quantization = False quant_value = None - # Compute lm_device only when initializing the LLM to avoid overwriting a - # previously-resolved device (e.g. "cuda") with the raw UI value ("auto"). - # "auto" is resolved to the concrete device inside llm_handler.initialize(). - if init_llm: - if not gpu_config.available_lm_models: - logger.warning( - f"⚠️ GPU tier {gpu_config.tier} ({gpu_config.gpu_memory_gb:.1f}GB) does not support LM on GPU. " - "Falling back to CPU for LM initialization." - ) - lm_device = "cpu" - else: - device_map = getattr(dit_handler, "device_map", None) - if device_map is not None and device_map.lm is not None: - lm_device = device_map.lm - else: - lm_device = device - + # Tier / backend checks can run before DiT init; LM device must be resolved + # after initialize_service so device_map.lm from ACESTEP_GPU_MAPPING is set. if init_llm and lm_model_path and gpu_config.available_lm_models: if not is_lm_model_size_allowed(lm_model_path, gpu_config.available_lm_models): logger.warning( @@ -137,6 +122,21 @@ def init_service_wrapper( gpu_mapping=os.environ.get("ACESTEP_GPU_MAPPING"), ) + lm_device = device + if init_llm: + if not gpu_config.available_lm_models: + logger.warning( + f"⚠️ GPU tier {gpu_config.tier} ({gpu_config.gpu_memory_gb:.1f}GB) does not support LM on GPU. " + "Falling back to CPU for LM initialization." + ) + lm_device = "cpu" + else: + device_map = getattr(dit_handler, "device_map", None) + if device_map is not None and device_map.lm is not None: + lm_device = device_map.lm + else: + lm_device = device + if init_llm: checkpoint_dir = os.path.join(project_root, "checkpoints") diff --git a/acestep/ui/gradio/events/generation/service_init_test.py b/acestep/ui/gradio/events/generation/service_init_test.py index 427de34b0..33c7d5a3a 100644 --- a/acestep/ui/gradio/events/generation/service_init_test.py +++ b/acestep/ui/gradio/events/generation/service_init_test.py @@ -222,6 +222,7 @@ def test_init_llm_with_auto_device_calls_initialize(self, mock_gpu_config): dit_handler.initialize_service.return_value = ("ok", True) dit_handler.model = MagicMock() dit_handler.is_turbo_model.return_value = True + dit_handler.device_map = None llm_handler = MagicMock() llm_handler.llm_initialized = False @@ -251,6 +252,62 @@ def test_init_llm_with_auto_device_calls_initialize(self, mock_gpu_config): "initialize() must receive 'auto' so it can resolve to the best device", ) + @patch("acestep.ui.gradio.events.generation.service_init.get_global_gpu_config") + def test_init_llm_uses_device_map_lm_after_dit_init(self, mock_gpu_config): + """LM device must come from device_map after DiT initialize_service. + + Regression: resolving lm_device before initialize_service left device_map + empty on first Gradio init, so UI device='auto' collapsed the LM onto + bare cuda:0 even when ACESTEP_GPU_MAPPING set lm:1. + """ + module = self._import_module() + + mock_gpu_config.return_value = MagicMock( + available_lm_models=["acestep-5Hz-lm-1.7B"], + lm_backend_restriction=None, + tier="tier6", + gpu_memory_gb=24.0, + max_duration_with_lm=600, + max_duration_without_lm=600, + max_batch_size_with_lm=4, + max_batch_size_without_lm=8, + ) + + dit_handler = MagicMock() + dit_handler.device_map = None + + def _init_service(*_args, **_kwargs): + dit_handler.device_map = MagicMock(lm="cuda:1") + return ("ok", True) + + dit_handler.initialize_service.side_effect = _init_service + dit_handler.model = MagicMock() + dit_handler.is_turbo_model.return_value = True + + llm_handler = MagicMock() + llm_handler.llm_initialized = False + llm_handler.initialize.return_value = ("[OK] LLM initialized", True) + + module.init_service_wrapper( + dit_handler, + llm_handler, + "/some/project/checkpoints", + "acestep-v15-turbo", + "auto", + True, + "acestep-5Hz-lm-1.7B", + "pt", + use_flash_attention=False, + offload_to_cpu=False, + offload_dit_to_cpu=False, + compile_model=False, + quantization=False, + ) + + llm_handler.initialize.assert_called_once() + _, call_kwargs = llm_handler.initialize.call_args + self.assertEqual(call_kwargs.get("device"), "cuda:1") + @patch("acestep.ui.gradio.events.generation.service_init.get_global_gpu_config") def test_legacy_cuda_config_forces_pt_backend(self, mock_gpu_config): """Legacy CUDA restrictions should override a requested vllm backend.""" From a42ff621629d6e7911ddaa3333ceebc8883e7896 Mon Sep 17 00:00:00 2001 From: steve Date: Mon, 6 Jul 2026 14:24:36 -0400 Subject: [PATCH 15/18] Add multi-GPU CLI/API flags, status fields, and docs (PR3). Expose --gpu-mapping and --list-gpus on the Gradio and generation CLIs, surface GPU inventory and device_map in API health/model endpoints, log ACESTEP_LM_DEVICE deprecation, and document multi-GPU usage. Co-authored-by: Cursor --- acestep/acestep_v15_pipeline.py | 54 ++++++++- acestep/acestep_v15_pipeline_test.py | 112 +++++++++++++++++- acestep/api/http/model_init_service.py | 12 ++ acestep/api/http/model_service_routes.py | 5 + acestep/api/http/model_service_routes_test.py | 16 +++ acestep/api/startup_llm_init.py | 8 ++ acestep/test_device_map.py | 58 +++++++++ cli.py | 40 ++++++- docs/en/GPU_COMPATIBILITY.md | 2 + docs/en/MULTI_GPU.md | 103 ++++++++++++++++ 10 files changed, 403 insertions(+), 7 deletions(-) create mode 100644 docs/en/MULTI_GPU.md diff --git a/acestep/acestep_v15_pipeline.py b/acestep/acestep_v15_pipeline.py index 6b9450653..ea3b1422b 100644 --- a/acestep/acestep_v15_pipeline.py +++ b/acestep/acestep_v15_pipeline.py @@ -60,6 +60,11 @@ is_mps_platform, ) from .model_downloader import ensure_lm_model + from .device_map import ( + GPU_MAPPING_ENV, + format_gpu_list_text, + log_lm_device_deprecation, + ) except ImportError: # When executed as a script: `python acestep/acestep_v15_pipeline.py` project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) @@ -82,6 +87,11 @@ is_mps_platform, ) from acestep.model_downloader import ensure_lm_model + from acestep.device_map import ( + GPU_MAPPING_ENV, + format_gpu_list_text, + log_lm_device_deprecation, + ) def create_demo(init_params=None, language="en"): @@ -283,6 +293,22 @@ def main(): choices=["auto", "cuda", "mps", "xpu", "cpu"], help="Processing device (default: auto)", ) + parser.add_argument( + "--gpu-mapping", + dest="gpu_mapping", + type=str, + default=None, + metavar="MAPPING", + help=( + "Component GPU layout: 'auto', 'single:N', or explicit " + "'dit:0,vae:0,text_encoder:0,lm:1'. Also reads ACESTEP_GPU_MAPPING." + ), + ) + parser.add_argument( + "--list-gpus", + action="store_true", + help="List visible CUDA devices and exit", + ) parser.add_argument( "--init_llm", type=lambda x: x.lower() in ["true", "1", "yes"], @@ -390,6 +416,16 @@ def main(): args = parser.parse_args() + if args.list_gpus: + print(format_gpu_list_text()) + sys.exit(0) + + effective_gpu_mapping = args.gpu_mapping + if effective_gpu_mapping is None: + effective_gpu_mapping = os.environ.get(GPU_MAPPING_ENV) + elif effective_gpu_mapping: + os.environ[GPU_MAPPING_ENV] = effective_gpu_mapping + # Enable API requires init_service if args.enable_api: args.init_service = True @@ -511,6 +547,7 @@ def main(): offload_dit_to_cpu=args.offload_dit_to_cpu, quantization=args.quantization, prefer_source=prefer_source, + gpu_mapping=effective_gpu_mapping, ) if not enable_generate: @@ -566,14 +603,26 @@ def main(): file=sys.stderr, ) + lm_device = args.device + device_map = getattr(dit_handler, "device_map", None) + if device_map is not None and device_map.lm is not None: + lm_device = device_map.lm + log_lm_device_deprecation( + explicit_lm_device=os.environ.get("ACESTEP_LM_DEVICE"), + gpu_mapping_env=effective_gpu_mapping, + using_device_map_lm=bool( + device_map is not None and device_map.lm is not None + ), + ) + print( - f"Initializing 5Hz LM: {args.lm_model_path} on {args.device}..." + f"Initializing 5Hz LM: {args.lm_model_path} on {lm_device}..." ) lm_status, lm_success = llm_handler.initialize( checkpoint_dir=checkpoint_dir, lm_model_path=args.lm_model_path, backend=args.backend, - device=args.device, + device=lm_device, offload_to_cpu=args.offload_to_cpu, dtype=None, ) @@ -595,6 +644,7 @@ def main(): "checkpoint": args.checkpoint, "config_path": args.config_path, "device": args.device, + "gpu_mapping": effective_gpu_mapping, "init_llm": args.init_llm, "lm_model_path": args.lm_model_path, "backend": args.backend, diff --git a/acestep/acestep_v15_pipeline_test.py b/acestep/acestep_v15_pipeline_test.py index 48e9c5c3a..522d85a4f 100644 --- a/acestep/acestep_v15_pipeline_test.py +++ b/acestep/acestep_v15_pipeline_test.py @@ -90,11 +90,11 @@ def _create_demo(init_params=None, language="en"): ): acestep_v15_pipeline.main() - return llm_handler, captured + return llm_handler, dit_handler, captured def test_main_forces_pt_backend_for_explicit_vllm_argument(self) -> None: """Legacy CUDA startup should override an explicit CLI vLLM request.""" - llm_handler, captured = self._run_main( + llm_handler, _, captured = self._run_main( [ "acestep", "--init_service", @@ -115,7 +115,7 @@ def test_main_forces_pt_backend_for_explicit_vllm_argument(self) -> None: def test_main_forces_pt_backend_for_service_mode_backend_override(self) -> None: """Service mode should not re-enable vLLM on legacy CUDA hardware.""" - llm_handler, captured = self._run_main( + llm_handler, _, captured = self._run_main( ["acestep", "--service_mode", "true", "--init_llm", "true"], env={"SERVICE_MODE_BACKEND": "vllm"}, ) @@ -131,7 +131,7 @@ def test_main_forces_pt_backend_for_api_env_override(self) -> None: sys.modules, {"acestep.ui.gradio.api.api_routes": api_routes_module}, ), patch("time.sleep", side_effect=KeyboardInterrupt): - llm_handler, captured = self._run_main( + llm_handler, _, captured = self._run_main( [ "acestep", "--enable-api", @@ -149,5 +149,109 @@ def test_main_forces_pt_backend_for_api_env_override(self) -> None: self.assertEqual("pt", captured["init_params"]["backend"]) +class PipelineGpuMappingTests(unittest.TestCase): + """Verify multi-GPU CLI flags are wired into service initialization.""" + + def test_list_gpus_exits_after_printing_inventory(self) -> None: + with patch.object(sys, "argv", ["acestep", "--list-gpus"]), patch( + "acestep.acestep_v15_pipeline.format_gpu_list_text", + return_value="GPU TABLE", + ) as mock_format, patch( + "acestep.acestep_v15_pipeline.sys.exit", + side_effect=SystemExit(0), + ) as mock_exit: + with self.assertRaises(SystemExit): + acestep_v15_pipeline.main() + mock_format.assert_called_once() + mock_exit.assert_called_once_with(0) + + def test_gpu_mapping_passed_to_initialize_service(self) -> None: + gpu_config = SimpleNamespace( + gpu_memory_gb=24.0, + tier="tier6b", + max_duration_with_lm=480, + max_duration_without_lm=600, + max_batch_size_with_lm=8, + max_batch_size_without_lm=8, + init_lm_default=True, + available_lm_models=["acestep-5Hz-lm-0.6B"], + recommended_backend="vllm", + lm_backend_restriction=None, + offload_dit_to_cpu_default=False, + quantization_default=False, + ) + dit_handler = MagicMock() + dit_handler.get_available_acestep_v15_models.return_value = ["acestep-v15-turbo"] + dit_handler.is_flash_attention_available.return_value = False + dit_handler.initialize_service.return_value = ("ok", True) + dit_handler.device_map = SimpleNamespace(lm="cuda:1") + + llm_handler = MagicMock() + llm_handler.get_available_5hz_lm_models.return_value = ["acestep-5Hz-lm-0.6B"] + llm_handler.initialize.return_value = ("ok", True) + + demo = MagicMock() + demo.queue.return_value = demo + demo.launch.return_value = None + captured: dict[str, object] = {} + + def _create_demo(init_params=None, language="en"): + captured["init_params"] = init_params + return demo + + with patch.object( + sys, + "argv", + [ + "acestep", + "--init_service", + "true", + "--init_llm", + "true", + "--config_path", + "acestep-v15-turbo", + "--gpu-mapping", + "auto", + ], + ), patch.dict(os.environ, {}, clear=True), patch( + "acestep.acestep_v15_pipeline.get_gpu_config", + return_value=gpu_config, + ), patch( + "acestep.acestep_v15_pipeline.set_global_gpu_config" + ), patch( + "acestep.acestep_v15_pipeline.is_mps_platform", + return_value=False, + ), patch( + "acestep.acestep_v15_pipeline.get_i18n" + ), patch( + "acestep.acestep_v15_pipeline.available_languages_info", + return_value=[("en", "English", "English")], + ), patch( + "acestep.acestep_v15_pipeline.AceStepHandler", + return_value=dit_handler, + ), patch( + "acestep.acestep_v15_pipeline.LLMHandler", + return_value=llm_handler, + ), patch( + "acestep.acestep_v15_pipeline.create_demo", + side_effect=_create_demo, + ), patch( + "acestep.acestep_v15_pipeline.ensure_lm_model", + return_value=(True, "ok"), + ), patch( + "acestep.acestep_v15_pipeline.os.makedirs" + ), patch( + "acestep.acestep_v15_pipeline.log_lm_device_deprecation" + ): + acestep_v15_pipeline.main() + + self.assertEqual( + "auto", + dit_handler.initialize_service.call_args.kwargs["gpu_mapping"], + ) + self.assertEqual("cuda:1", llm_handler.initialize.call_args.kwargs["device"]) + self.assertEqual("auto", captured["init_params"]["gpu_mapping"]) + + if __name__ == "__main__": unittest.main() diff --git a/acestep/api/http/model_init_service.py b/acestep/api/http/model_init_service.py index 73e1bb8c3..a855efb95 100644 --- a/acestep/api/http/model_init_service.py +++ b/acestep/api/http/model_init_service.py @@ -117,6 +117,7 @@ def initialize_models_for_request( compile_model=compile_model, offload_to_cpu=offload_to_cpu, offload_dit_to_cpu=offload_dit_to_cpu, + gpu_mapping=os.getenv("ACESTEP_GPU_MAPPING"), ) if not ok: setattr(app_state, error_attr, status_msg) @@ -137,6 +138,17 @@ def initialize_models_for_request( lm_backend = resolve_lm_backend(os.getenv("ACESTEP_LM_BACKEND"), gpu_config) lm_device = os.getenv("ACESTEP_LM_DEVICE", device) + device_map = getattr(handler, "device_map", None) + using_device_map_lm = False + if device_map is not None and device_map.lm is not None: + lm_device = device_map.lm + using_device_map_lm = True + from acestep.device_map import log_lm_device_deprecation + + log_lm_device_deprecation( + explicit_lm_device=os.getenv("ACESTEP_LM_DEVICE"), + using_device_map_lm=using_device_map_lm, + ) lm_offload_env = os.getenv("ACESTEP_LM_OFFLOAD_TO_CPU") lm_offload = env_bool("ACESTEP_LM_OFFLOAD_TO_CPU", False) if lm_offload_env is not None else offload_to_cpu diff --git a/acestep/api/http/model_service_routes.py b/acestep/api/http/model_service_routes.py index c61a6a561..325ece3ad 100644 --- a/acestep/api/http/model_service_routes.py +++ b/acestep/api/http/model_service_routes.py @@ -12,6 +12,7 @@ from acestep.api.http.model_init_service import initialize_models_for_request from acestep.constants import TASK_TYPES_BASE, TASK_TYPES_TURBO +from acestep.device_map import collect_gpu_runtime_status class InitModelRequest(BaseModel): @@ -122,6 +123,7 @@ def _collect_model_inventory( "lm_models": lm_models, "loaded_lm_model": loaded_lm_model, "llm_initialized": llm_initialized, + **collect_gpu_runtime_status(getattr(app.state, "handler", None)), } @@ -157,6 +159,9 @@ async def health_check(): "llm_initialized": inventory["llm_initialized"], "loaded_model": inventory["default_model"], "loaded_lm_model": inventory["loaded_lm_model"], + "gpus": inventory["gpus"], + "gpu_mapping": inventory["gpu_mapping"], + "device_map": inventory["device_map"], } ) diff --git a/acestep/api/http/model_service_routes_test.py b/acestep/api/http/model_service_routes_test.py index f4c746a22..fda600689 100644 --- a/acestep/api/http/model_service_routes_test.py +++ b/acestep/api/http/model_service_routes_test.py @@ -103,6 +103,22 @@ def test_collect_model_inventory_merges_loaded_and_available_models(self): self.assertIn("acestep-v15-turbo", names) self.assertEqual("acestep-v15-base", inventory["default_model"]) self.assertTrue(inventory["llm_initialized"]) + self.assertIn("gpus", inventory) + self.assertIn("gpu_mapping", inventory) + self.assertIn("device_map", inventory) + + def test_health_route_includes_gpu_runtime_fields(self): + """Health endpoint should expose GPU inventory and mapping metadata.""" + + app = self._build_app() + endpoint = _get_endpoint(app, "/health", "GET") + with mock.patch("acestep.api.http.model_service_routes.os.path.isdir", return_value=False): + result = asyncio.run(endpoint()) + + self.assertEqual(200, result["code"]) + self.assertIn("gpus", result["data"]) + self.assertIn("gpu_mapping", result["data"]) + self.assertIn("device_map", result["data"]) def test_init_route_wraps_initializer_exception(self): """Init endpoint should convert initializer exceptions into wrapped code=500 payloads.""" diff --git a/acestep/api/startup_llm_init.py b/acestep/api/startup_llm_init.py index d0f0c1015..75be8134e 100644 --- a/acestep/api/startup_llm_init.py +++ b/acestep/api/startup_llm_init.py @@ -75,8 +75,16 @@ def initialize_llm_at_startup( lm_backend = resolve_lm_backend(os.getenv("ACESTEP_LM_BACKEND"), gpu_config) lm_device = os.getenv("ACESTEP_LM_DEVICE", device) device_map = getattr(dit_handler, "device_map", None) if dit_handler is not None else None + using_device_map_lm = False if device_map is not None and device_map.lm is not None: lm_device = device_map.lm + using_device_map_lm = True + from acestep.device_map import log_lm_device_deprecation + + log_lm_device_deprecation( + explicit_lm_device=os.getenv("ACESTEP_LM_DEVICE"), + using_device_map_lm=using_device_map_lm, + ) lm_offload_env = os.getenv("ACESTEP_LM_OFFLOAD_TO_CPU") lm_offload = env_bool("ACESTEP_LM_OFFLOAD_TO_CPU", False) if lm_offload_env is not None else offload_to_cpu diff --git a/acestep/test_device_map.py b/acestep/test_device_map.py index cfe4a4f05..48f74c439 100644 --- a/acestep/test_device_map.py +++ b/acestep/test_device_map.py @@ -260,5 +260,63 @@ def test_log_lm_device_deprecation_ignored_when_mapping_assigns_lm(self): self.assertIn("ignored", mock_logger.warning.call_args[0][0]) +class GpuStatusHelpersTests(unittest.TestCase): + """Tests for CLI/API GPU listing and status helpers.""" + + def test_format_gpu_list_text_reports_no_devices(self): + self.assertEqual(format_gpu_list_text([]), "No CUDA devices detected.") + + def test_format_gpu_list_text_includes_device_rows(self): + text = format_gpu_list_text( + [GpuInfo(0, "RTX 3090", 24.0, 22.5, (8, 6))] + ) + self.assertIn("RTX 3090", text) + self.assertIn("24.00", text) + + def test_gpu_info_and_device_map_serialization(self): + gpu_payload = gpu_info_to_dict(GpuInfo(1, "GPU1", 24.0, 20.0, (8, 6))) + self.assertEqual(gpu_payload["index"], 1) + self.assertEqual(gpu_payload["compute_capability"], [8, 6]) + + device_map = ComponentDeviceMap( + dit="cuda:0", + vae="cuda:0", + text_encoder="cuda:0", + lm="cuda:1", + ) + map_payload = device_map_to_dict(device_map) + self.assertTrue(map_payload["multi_device"]) + self.assertIn("lm:1", map_payload["summary"]) + + def test_collect_gpu_runtime_status_includes_handler_layout(self): + handler = SimpleNamespace( + device_map=ComponentDeviceMap.from_single_device("cuda:2") + ) + with patch("acestep.device_map.discover_gpus", return_value=[]): + status = collect_gpu_runtime_status(handler=handler) + self.assertEqual(status["device_map"]["dit"], "cuda:2") + self.assertEqual(status["gpus"], []) + + def test_log_lm_device_deprecation_without_mapping(self): + with patch.dict(os.environ, {"ACESTEP_LM_DEVICE": "cuda:1"}, clear=False), patch( + "acestep.device_map.logger" + ) as mock_logger: + log_lm_device_deprecation() + mock_logger.warning.assert_called_once() + + def test_log_lm_device_deprecation_ignored_when_mapping_assigns_lm(self): + with patch.dict( + os.environ, + { + "ACESTEP_LM_DEVICE": "cuda:0", + "ACESTEP_GPU_MAPPING": "dit:0,vae:0,text_encoder:0,lm:1", + }, + clear=False, + ), patch("acestep.device_map.logger") as mock_logger: + log_lm_device_deprecation(using_device_map_lm=True) + mock_logger.warning.assert_called_once() + self.assertIn("ignored", mock_logger.warning.call_args[0][0]) + + if __name__ == "__main__": unittest.main() diff --git a/cli.py b/cli.py index f6c7b0eb0..86a6e66c9 100644 --- a/cli.py +++ b/cli.py @@ -1051,8 +1051,30 @@ def main(): default="INFO", help="Logging level for internal modules (TRACE/DEBUG/INFO/WARNING/ERROR/CRITICAL).", ) + parser.add_argument( + "--gpu-mapping", + dest="gpu_mapping", + type=str, + default=None, + metavar="MAPPING", + help=( + "Component GPU layout: 'auto', 'single:N', or explicit " + "'dit:0,vae:0,text_encoder:0,lm:1'. Also reads ACESTEP_GPU_MAPPING." + ), + ) + parser.add_argument( + "--list-gpus", + action="store_true", + help="List visible CUDA devices and exit", + ) cli_args = parser.parse_args() + if cli_args.list_gpus: + from acestep.device_map import format_gpu_list_text + + print(format_gpu_list_text()) + sys.exit(0) + _configure_logging(level=cli_args.log_level) default_batch_size = 1 if not cli_args.config else config_defaults.batch_size @@ -1130,8 +1152,12 @@ def main(): "cfg_interval_end": params_defaults.cfg_interval_end, "lm_negative_prompt": params_defaults.lm_negative_prompt, "log_level": cli_args.log_level, + "gpu_mapping": cli_args.gpu_mapping or os.environ.get("ACESTEP_GPU_MAPPING"), } + if cli_args.gpu_mapping: + os.environ["ACESTEP_GPU_MAPPING"] = cli_args.gpu_mapping + args = argparse.Namespace(**defaults) args.config = None if cli_args.config: @@ -1416,6 +1442,7 @@ def main(): compile_model=compile_model, offload_to_cpu=args.offload_to_cpu, offload_dit_to_cpu=args.offload_dit_to_cpu, + gpu_mapping=getattr(args, "gpu_mapping", None), ) if requires_lm: @@ -1450,11 +1477,22 @@ def main(): parser.error(f"LM model '{lm_model_path}' not found locally and not in registry. Please provide a valid --lm_model_path.") print(f"Initializing LM handler with model: {args.lm_model_path}") + lm_device = device + device_map = getattr(dit_handler, "device_map", None) + if device_map is not None and device_map.lm is not None: + lm_device = device_map.lm + from acestep.device_map import log_lm_device_deprecation + + log_lm_device_deprecation( + explicit_lm_device=os.environ.get("ACESTEP_LM_DEVICE"), + gpu_mapping_env=getattr(args, "gpu_mapping", None), + using_device_map_lm=bool(device_map is not None and device_map.lm is not None), + ) llm_handler.initialize( checkpoint_dir=args.checkpoint_dir, lm_model_path=args.lm_model_path, backend=args.backend, - device=device, + device=lm_device, offload_to_cpu=args.offload_to_cpu, dtype=None, ) diff --git a/docs/en/GPU_COMPATIBILITY.md b/docs/en/GPU_COMPATIBILITY.md index c2de671c8..174f352fe 100644 --- a/docs/en/GPU_COMPATIBILITY.md +++ b/docs/en/GPU_COMPATIBILITY.md @@ -2,6 +2,8 @@ ACE-Step 1.5 automatically adapts to your GPU's available VRAM, adjusting generation limits, LM model availability, offloading strategies, and UI defaults accordingly. The system detects GPU memory at startup and configures optimal settings for your hardware. +> **Multiple GPUs:** To spread DiT and LM across cards (for example XL + 4B LM on a multi-GPU workstation), see the [Multi-GPU Inference Guide](MULTI_GPU.md). + ## GPU Tier Configuration | VRAM | Tier | XL (4B) DiT | LM Models | Recommended LM | Backend | Max Duration (LM / No LM) | Max Batch (LM / No LM) | Offload | Quantization | diff --git a/docs/en/MULTI_GPU.md b/docs/en/MULTI_GPU.md new file mode 100644 index 000000000..a2f4ff7c9 --- /dev/null +++ b/docs/en/MULTI_GPU.md @@ -0,0 +1,103 @@ +# Multi-GPU Inference Guide + +ACE-Step 1.5 can spread DiT, VAE, text encoder, and 5Hz LM components across multiple CUDA GPUs. This is useful when a single card cannot hold both an XL DiT checkpoint and a large LM (for example, 4× RTX 3090 with `acestep-v15-xl-sft` + `acestep-5Hz-lm-4B`). + +## Quick start + +```bash +# Automatic VRAM-aware layout (recommended on 2+ CUDA GPUs) +ACESTEP_GPU_MAPPING=auto ACESTEP_LM_MODEL_PATH=acestep-5Hz-lm-4B \ + uv run acestep --config-path acestep-v15-xl-sft --init-llm true + +# Explicit layout: DiT stack on GPU 0, LM on GPU 1 +ACESTEP_GPU_MAPPING=dit:0,vae:0,text_encoder:0,lm:1 \ + uv run acestep --config-path acestep-v15-xl-sft --init-llm true + +# List visible CUDA devices +uv run acestep --list-gpus +``` + +The same mapping can be passed on the CLI instead of the environment variable: + +```bash +uv run acestep --gpu-mapping auto --config-path acestep-v15-xl-sft --init-llm true +``` + +## Mapping formats + +| Format | Example | Behavior | +|--------|---------|----------| +| `auto` | `ACESTEP_GPU_MAPPING=auto` | VRAM-aware layout when 2+ CUDA GPUs are visible; otherwise single-GPU fallback | +| `single:N` | `single:1` | All components on `cuda:N` | +| Explicit | `dit:0,vae:0,text_encoder:0,lm:1` | Per-component placement (`dit` is required) | + +### Auto layout + +With `gpu_mapping=auto`, ACE-Step: + +1. Estimates DiT + VAE + text encoder peak VRAM from the selected checkpoint and batch size. +2. Places that stack on the GPU with the most free VRAM. +3. Places the LM on the next best GPU (preferring a different card when possible). + +If auto layout cannot fit all components, startup falls back to single-device placement and logs suggestions (smaller models, explicit mapping, CPU offload). + +## Environment variables + +| Variable | Description | +|----------|-------------| +| `ACESTEP_GPU_MAPPING` | Component layout (`auto`, `single:N`, or explicit map) | +| `ACESTEP_DEVICE` | Base device when no mapping is set (legacy single-GPU path) | +| `ACESTEP_LM_DEVICE` | **Deprecated** — use `ACESTEP_GPU_MAPPING` with `lm:N` instead | + +CLI flags `--gpu-mapping` and `--list-gpus` mirror the Gradio/API surface. When both CLI and env are set, the CLI value wins and is written to `ACESTEP_GPU_MAPPING`. + +## API status fields + +`/health`, `/v1/models`, and `/v1/model_inventory` include: + +- `gpus` — visible CUDA devices with free/total VRAM +- `gpu_mapping` — active `ACESTEP_GPU_MAPPING` value (if any) +- `device_map` — resolved component layout after initialization (`dit`, `vae`, `text_encoder`, `lm`, `summary`, `multi_device`) + +Example `device_map` payload: + +```json +{ + "dit": "cuda:0", + "vae": "cuda:0", + "text_encoder": "cuda:0", + "lm": "cuda:1", + "summary": "dit:0, vae:0, text_encoder:0, lm:1", + "multi_device": true +} +``` + +## Cross-GPU inference + +When components run on different GPUs, ACE-Step routes tensors at stage boundaries (conditioning → DiT, latents → VAE, audio codes, etc.). No extra configuration is required beyond the mapping. + +## Hardware examples + +### 4× RTX 3090 (24 GB each) + +Typical auto layout for XL SFT + 4B LM: + +``` +dit:0, vae:0, text_encoder:0, lm:1 +``` + +GPUs 2–3 remain available for future batch-serving work. + +### Single GPU + +Leave `ACESTEP_GPU_MAPPING` unset, or use `single:0`. Behavior matches pre–multi-GPU releases. + +## Troubleshooting + +| Symptom | Check | +|---------|-------| +| LM still on same GPU as DiT | Confirm `ACESTEP_GPU_MAPPING=auto` and 2+ visible GPUs (`--list-gpus`) | +| OOM during auto layout | Try explicit `lm:1`, smaller LM, or CPU offload | +| `ACESTEP_LM_DEVICE` ignored | Expected when `ACESTEP_GPU_MAPPING` assigns `lm`; use `lm:N` in the mapping | + +See also: [GPU Compatibility Guide](GPU_COMPATIBILITY.md) for per-tier VRAM limits and UI defaults. From 1eba94f98d2c8a6a7654782bddf1d296791ca931 Mon Sep 17 00:00:00 2001 From: steve Date: Mon, 6 Jul 2026 17:07:44 -0400 Subject: [PATCH 16/18] fix(device_map): address review feedback on layout and discovery. Reserve DiT VRAM before co-located LM fallback in auto-layout, narrow get_device_capability exception handling to RuntimeError, and fix MULTI_GPU.md fenced-code lint. Co-authored-by: Cursor --- docs/en/MULTI_GPU.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/en/MULTI_GPU.md b/docs/en/MULTI_GPU.md index a2f4ff7c9..403dc2116 100644 --- a/docs/en/MULTI_GPU.md +++ b/docs/en/MULTI_GPU.md @@ -82,7 +82,7 @@ When components run on different GPUs, ACE-Step routes tensors at stage boundari Typical auto layout for XL SFT + 4B LM: -``` +```text dit:0, vae:0, text_encoder:0, lm:1 ``` From 216d3b275cf1363a1322f7207c0036a556c0f089 Mon Sep 17 00:00:00 2001 From: steve Date: Sat, 11 Jul 2026 21:04:21 -0400 Subject: [PATCH 17/18] refactor: address CodeRabbit module-size and VRAM review notes Split Gradio pipeline CLI/startup into focused modules under the 200 LOC cap, extract device-resolution tests, add missing test helper docstrings, and read LM free/total VRAM from the mapped cuda:N device. Co-authored-by: Cursor --- acestep/acestep_v15_pipeline.py | 707 ++---------------- .../acestep_v15_pipeline_gpu_mapping_test.py | 150 ++++ acestep/acestep_v15_pipeline_test.py | 124 +-- .../handler/conditioning_embed_test.py | 10 + acestep/gradio_pipeline_banner.py | 52 ++ acestep/gradio_pipeline_cli.py | 140 ++++ acestep/gradio_pipeline_cli_service.py | 134 ++++ acestep/gradio_pipeline_launch.py | 69 ++ acestep/gradio_pipeline_mode_defaults.py | 52 ++ acestep/gradio_pipeline_startup.py | 181 +++++ acestep/llm_inference.py | 20 +- .../service_init_device_resolution_test.py | 200 +++++ .../events/generation/service_init_test.py | 272 +------ 13 files changed, 1097 insertions(+), 1014 deletions(-) create mode 100644 acestep/acestep_v15_pipeline_gpu_mapping_test.py create mode 100644 acestep/gradio_pipeline_banner.py create mode 100644 acestep/gradio_pipeline_cli.py create mode 100644 acestep/gradio_pipeline_cli_service.py create mode 100644 acestep/gradio_pipeline_launch.py create mode 100644 acestep/gradio_pipeline_mode_defaults.py create mode 100644 acestep/gradio_pipeline_startup.py create mode 100644 acestep/ui/gradio/events/generation/service_init_device_resolution_test.py diff --git a/acestep/acestep_v15_pipeline.py b/acestep/acestep_v15_pipeline.py index ea3b1422b..2aaceb833 100644 --- a/acestep/acestep_v15_pipeline.py +++ b/acestep/acestep_v15_pipeline.py @@ -8,7 +8,7 @@ # Load environment variables from .env file at most once per process to avoid # epoch-boundary stalls (e.g. on Windows when Gradio yields during training) -_env_loaded = False # module-level so we never reload .env in the same process +_env_loaded = False try: from dotenv import load_dotenv @@ -25,10 +25,8 @@ print(f"Loaded configuration from {_env_example_path} (fallback)") _env_loaded = True except ImportError: - # python-dotenv not installed, skip loading .env pass -# Clear proxy settings that may affect Gradio for proxy_var in [ "http_proxy", "https_proxy", @@ -38,80 +36,33 @@ ]: os.environ.pop(proxy_var, None) -# Force torchaudio to use ffmpeg backend (torchcodec not available on XPU/Windows) os.environ["TORCHAUDIO_USE_BACKEND"] = "ffmpeg" -try: - # When executed as a module: `python -m acestep.acestep_v15_pipeline` - from .cli_args import parse_quantization_arg - from .handler import AceStepHandler - from .llm_inference import LLMHandler - from .dataset_handler import DatasetHandler - from .ui.gradio import create_gradio_interface - from acestep.ui.gradio.i18n import get_i18n, available_languages_info - from .gpu_config import ( - get_gpu_config, - get_gpu_memory_gb, - resolve_lm_backend, - print_gpu_config_info, - set_global_gpu_config, - VRAM_16GB_MIN_GB, - VRAM_AUTO_OFFLOAD_THRESHOLD_GB, - is_mps_platform, - ) - from .model_downloader import ensure_lm_model - from .device_map import ( - GPU_MAPPING_ENV, - format_gpu_list_text, - log_lm_device_deprecation, - ) -except ImportError: - # When executed as a script: `python acestep/acestep_v15_pipeline.py` - project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) - if project_root not in sys.path: - sys.path.insert(0, project_root) - from acestep.cli_args import parse_quantization_arg - from acestep.handler import AceStepHandler - from acestep.llm_inference import LLMHandler - from acestep.dataset_handler import DatasetHandler - from acestep.ui.gradio import create_gradio_interface - from acestep.ui.gradio.i18n import get_i18n, available_languages_info - from acestep.gpu_config import ( - get_gpu_config, - get_gpu_memory_gb, - resolve_lm_backend, - print_gpu_config_info, - set_global_gpu_config, - VRAM_16GB_MIN_GB, - VRAM_AUTO_OFFLOAD_THRESHOLD_GB, - is_mps_platform, - ) - from acestep.model_downloader import ensure_lm_model - from acestep.device_map import ( - GPU_MAPPING_ENV, - format_gpu_list_text, - log_lm_device_deprecation, - ) +from acestep.dataset_handler import DatasetHandler +from acestep.gpu_config import ( + VRAM_AUTO_OFFLOAD_THRESHOLD_GB, + get_gpu_config, + is_mps_platform, + resolve_lm_backend, + set_global_gpu_config, +) +from acestep.gradio_pipeline_banner import print_gpu_banner +from acestep.gradio_pipeline_cli import ( + apply_gpu_mapping_args, + build_gradio_parser, + resolve_default_quantization, +) +from acestep.gradio_pipeline_launch import launch_gradio_demo +from acestep.gradio_pipeline_mode_defaults import apply_startup_mode_defaults +from acestep.gradio_pipeline_startup import initialize_from_cli +from acestep.handler import AceStepHandler +from acestep.llm_inference import LLMHandler +from acestep.ui.gradio import create_gradio_interface +from acestep.ui.gradio.i18n import get_i18n def create_demo(init_params=None, language="en"): - """ - Create Gradio demo interface - - Args: - init_params: Dictionary containing initialization parameters and state. - If None, service will not be pre-initialized. - Keys: 'pre_initialized' (bool), 'checkpoint', 'config_path', 'device', - 'init_llm', 'lm_model_path', 'backend', 'use_flash_attention', - 'offload_to_cpu', 'offload_dit_to_cpu', 'init_status', - 'dit_handler', 'llm_handler' (initialized handlers if pre-initialized), - 'language' (UI language code) - language: UI language code ('en', 'zh', 'ja', default: 'en') - - Returns: - Gradio Blocks instance - """ - # Use pre-initialized handlers if available, otherwise create new ones + """Create Gradio demo interface with optional pre-initialized handlers.""" if ( init_params and init_params.get("pre_initialized") @@ -120,13 +71,11 @@ def create_demo(init_params=None, language="en"): dit_handler = init_params["dit_handler"] llm_handler = init_params["llm_handler"] else: - dit_handler = AceStepHandler() # DiT handler - llm_handler = LLMHandler() # LM handler - - dataset_handler = DatasetHandler() # Dataset handler + dit_handler = AceStepHandler() + llm_handler = LLMHandler() - # Create Gradio interface with all handlers and initialization parameters - demo = create_gradio_interface( + dataset_handler = DatasetHandler() + return create_gradio_interface( dit_handler, llm_handler, dataset_handler, @@ -134,621 +83,87 @@ def create_demo(init_params=None, language="en"): language=language, ) - return demo - def _resolve_startup_lm_backend(requested_backend: str | None, gpu_config) -> str: """Resolve the startup LM backend against hardware compatibility restrictions.""" resolved_backend = resolve_lm_backend(requested_backend, gpu_config) normalized_backend = (requested_backend or "").strip().lower() - if normalized_backend and normalized_backend != resolved_backend: print( - f"Requested LM backend '{normalized_backend}' is not supported on this hardware. " - f"Using '{resolved_backend}' instead." + f"Requested LM backend '{normalized_backend}' is not supported on this " + f"hardware. Using '{resolved_backend}' instead." ) - return resolved_backend def main(): - """Main entry function""" - import argparse - - # Detect GPU memory and get configuration + """Main entry function for the Gradio demo.""" gpu_config = get_gpu_config() - set_global_gpu_config(gpu_config) # Set global config for use across modules + set_global_gpu_config(gpu_config) gpu_memory_gb = gpu_config.gpu_memory_gb - _is_mac = is_mps_platform() - # Enable auto-offload for GPUs below 20 GB. 16 GB GPUs cannot hold all - # models simultaneously (DiT ~4.7 + VAE ~0.3 + text_enc ~1.2 + LM ≥1.2 + - # activations) so they *must* offload. The old threshold of 16 GB caused - # 16 GB GPUs to never offload, leading to OOM. - # Mac (Apple Silicon) uses unified memory — offloading provides no benefit. + is_mac = is_mps_platform() auto_offload = ( - (not _is_mac) + (not is_mac) and gpu_memory_gb > 0 and gpu_memory_gb < VRAM_AUTO_OFFLOAD_THRESHOLD_GB ) - _default_backend = gpu_config.recommended_backend + default_backend = gpu_config.recommended_backend + print_gpu_banner(gpu_config, gpu_memory_gb, is_mac, auto_offload, default_backend) - # Print GPU configuration info - print(f"\n{'=' * 60}") - print("GPU Configuration Detected:") - print(f"{'=' * 60}") - print(f" GPU Memory: {gpu_memory_gb:.2f} GB") - print(f" Configuration Tier: {gpu_config.tier}") - print( - f" Max Duration (with LM): {gpu_config.max_duration_with_lm}s ({gpu_config.max_duration_with_lm // 60} min)" - ) - print( - f" Max Duration (without LM): {gpu_config.max_duration_without_lm}s ({gpu_config.max_duration_without_lm // 60} min)" - ) - print(f" Max Batch Size (with LM): {gpu_config.max_batch_size_with_lm}") - print(f" Max Batch Size (without LM): {gpu_config.max_batch_size_without_lm}") - print(f" Default LM Init: {gpu_config.init_lm_default}") - print(f" Available LM Models: {gpu_config.available_lm_models or 'None'}") - print(f"{'=' * 60}\n") - - if _is_mac: - print( - f"Apple Silicon (MPS) detected — unified memory {gpu_memory_gb:.1f}GB, no CPU offload needed, backend={_default_backend}" - ) - elif auto_offload: - print( - f"Auto-enabling CPU offload (GPU {gpu_memory_gb:.1f}GB < {VRAM_AUTO_OFFLOAD_THRESHOLD_GB}GB threshold)" - ) - elif gpu_memory_gb > 0: - print( - f"CPU offload disabled by default (GPU {gpu_memory_gb:.1f}GB >= {VRAM_AUTO_OFFLOAD_THRESHOLD_GB}GB threshold)" - ) - else: - print("No GPU detected, running on CPU") - - # Define local outputs directory project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) - output_dir = os.path.join(project_root, "gradio_outputs") - # Normalize path to use forward slashes for Gradio 6 compatibility on Windows - output_dir = output_dir.replace("\\", "/") + output_dir = os.path.join(project_root, "gradio_outputs").replace("\\", "/") os.makedirs(output_dir, exist_ok=True) print(f"Output directory: {output_dir}") - - # Initialize i18n with default language (en) get_i18n() - parser = argparse.ArgumentParser( - description="Gradio Demo for ACE-Step V1.5", - formatter_class=argparse.RawTextHelpFormatter, - ) - parser.add_argument( - "--port", type=int, default=7860, help="Port to run the gradio server on" - ) - parser.add_argument("--share", action="store_true", help="Create a public link") - parser.add_argument("--debug", action="store_true", help="Enable debug mode") - parser.add_argument( - "--server-name", - type=str, - default="127.0.0.1", - help="Server name (default: 127.0.0.1, use 0.0.0.0 for all interfaces)", - ) - - # language argument - available_languages = available_languages_info() - parser.add_argument( - "--language", - type=str, - default=os.environ.get("LANGUAGE", "en"), - choices=[language[0] for language in available_languages], - help="UI language:\n " - + "\n ".join( - ( - code - + f" ({native_name}" - + (f"/{name})" if name != native_name else ")") - for code, name, native_name in available_languages - ) + parser = build_gradio_parser( + auto_offload=auto_offload, + default_backend=default_backend, + default_offload_dit=( + gpu_config.offload_dit_to_cpu_default if not is_mac else False ), + default_quantization=resolve_default_quantization(gpu_config, is_mac=is_mac), ) - del available_languages - - parser.add_argument( - "--allowed-path", - action="append", - default=[], - help="Additional allowed file paths for Gradio (repeatable).", - ) - - # Service mode argument - parser.add_argument( - "--service_mode", - type=lambda x: x.lower() in ["true", "1", "yes"], - default=False, - help="Enable service mode (default: False). When enabled, uses preset models and restricts UI options.", - ) - - # Service initialization arguments - parser.add_argument( - "--init_service", - type=lambda x: x.lower() in ["true", "1", "yes"], - default=False, - help="Initialize service on startup (default: False)", - ) - parser.add_argument( - "--checkpoint", - type=str, - default=None, - help="Checkpoint file path (optional, for display purposes)", - ) - parser.add_argument( - "--config_path", - type=str, - default=None, - help="Main model path (e.g., 'acestep-v15-turbo')", - ) - parser.add_argument( - "--device", - type=str, - default="auto", - choices=["auto", "cuda", "mps", "xpu", "cpu"], - help="Processing device (default: auto)", - ) - parser.add_argument( - "--gpu-mapping", - dest="gpu_mapping", - type=str, - default=None, - metavar="MAPPING", - help=( - "Component GPU layout: 'auto', 'single:N', or explicit " - "'dit:0,vae:0,text_encoder:0,lm:1'. Also reads ACESTEP_GPU_MAPPING." - ), - ) - parser.add_argument( - "--list-gpus", - action="store_true", - help="List visible CUDA devices and exit", - ) - parser.add_argument( - "--init_llm", - type=lambda x: x.lower() in ["true", "1", "yes"], - default=None, - help="Initialize 5Hz LM (default: auto based on GPU memory)", - ) - parser.add_argument( - "--lm_model_path", - type=str, - default=None, - help="5Hz LM model path (e.g., 'acestep-5Hz-lm-0.6B')", - ) - parser.add_argument( - "--backend", - type=str, - default=_default_backend, - choices=["vllm", "pt", "mlx"], - help=f"5Hz LM backend (default: {_default_backend}, use 'mlx' for native Apple Silicon acceleration)", - ) - parser.add_argument( - "--use_flash_attention", - type=lambda x: x.lower() in ["true", "1", "yes"], - default=None, - help="Use flash attention (default: auto-detect)", - ) - parser.add_argument( - "--offload_to_cpu", - type=lambda x: x.lower() in ["true", "1", "yes"], - default=auto_offload, - help=f"Offload models to CPU (default: {'True' if auto_offload else 'False'}, auto-detected based on GPU VRAM)", - ) - _default_offload_dit = ( - gpu_config.offload_dit_to_cpu_default if not _is_mac else False - ) - parser.add_argument( - "--offload_dit_to_cpu", - type=lambda x: x.lower() in ["true", "1", "yes"], - default=_default_offload_dit, - help=f"Offload DiT to CPU after diffusion (default: {_default_offload_dit}, auto-detected based on GPU tier)", - ) - _default_quantization = None - if gpu_config.quantization_default and not _is_mac: - _default_quantization = "int8_weight_only" - try: - import torch - if torch.cuda.is_available(): - major, _ = torch.cuda.get_device_capability(0) - if major < 7: - _default_quantization = "w8a8_dynamic" - except Exception as exc: - logger.warning( - "[parse_args] CUDA capability probe failed while resolving quantization default: {}", - exc, - ) - parser.add_argument( - "--quantization", - type=parse_quantization_arg, - default=_default_quantization, - help=( - "DiT quantization method: int8_weight_only, fp8_weight_only, " - "w8a8_dynamic, or none " - f"(default: {_default_quantization}, auto-detected based on GPU tier)" - ), - ) - parser.add_argument( - "--download-source", - type=str, - default=None, - choices=["huggingface", "modelscope", "auto"], - help="Preferred model download source (default: auto-detect based on network)", - ) - parser.add_argument( - "--batch_size", - type=int, - default=None, - help="Default batch size for generation (1-8). Defaults to min(2, GPU_max) if not specified", - ) - - # API mode argument - parser.add_argument( - "--enable-api", - action="store_true", - help="Enable API endpoints (default: False)", - ) - - # Authentication arguments - parser.add_argument( - "--auth-username", - type=str, - default=None, - help="Username for Gradio authentication", - ) - parser.add_argument( - "--auth-password", - type=str, - default=None, - help="Password for Gradio authentication", - ) - parser.add_argument( - "--api-key", - type=str, - default=None, - help="API key for API endpoints authentication", - ) - args = parser.parse_args() + effective_gpu_mapping = apply_gpu_mapping_args(args) - if args.list_gpus: - print(format_gpu_list_text()) - sys.exit(0) - - effective_gpu_mapping = args.gpu_mapping - if effective_gpu_mapping is None: - effective_gpu_mapping = os.environ.get(GPU_MAPPING_ENV) - elif effective_gpu_mapping: - os.environ[GPU_MAPPING_ENV] = effective_gpu_mapping - - # Enable API requires init_service - if args.enable_api: - args.init_service = True - # Load config from .env if not specified - if args.config_path is None: - args.config_path = os.environ.get("ACESTEP_CONFIG_PATH") - if args.lm_model_path is None: - args.lm_model_path = os.environ.get("ACESTEP_LM_MODEL_PATH") - if os.environ.get("ACESTEP_LM_BACKEND"): - args.backend = os.environ.get("ACESTEP_LM_BACKEND") - - # Service mode defaults (can be configured via .env file) - if args.service_mode: - print("Service mode enabled - applying preset configurations...") - # Force init_service in service mode - args.init_service = True - # Default DiT model for service mode (from env or fallback) - if args.config_path is None: - args.config_path = os.environ.get( - "SERVICE_MODE_DIT_MODEL", "acestep-v15-turbo-fix-inst-shift-dynamic" - ) - # Default LM model for service mode (from env or fallback) - if args.lm_model_path is None: - args.lm_model_path = os.environ.get( - "SERVICE_MODE_LM_MODEL", "acestep-5Hz-lm-1.7B-v4-fix" - ) - # Backend for service mode (from env or fallback to vllm) - args.backend = os.environ.get("SERVICE_MODE_BACKEND", "vllm") - print(f" DiT model: {args.config_path}") - print(f" LM model: {args.lm_model_path}") - + apply_startup_mode_defaults(args, gpu_memory_gb) args.backend = _resolve_startup_lm_backend(args.backend, gpu_config) - if args.service_mode: print(f" Backend: {args.backend}") - # Auto-enable CPU offload for tier6 GPUs (16-24GB) when using the 4B LM model - # The 4B LM (~8GB) + DiT (~4.7GB) + VAE + text encoder exceeds 16-20GB with activations - if not args.offload_to_cpu and args.lm_model_path and "4B" in args.lm_model_path: - if 0 < gpu_memory_gb <= 24: - args.offload_to_cpu = True - print( - f"Auto-enabling CPU offload (4B LM model requires offloading on {gpu_memory_gb:.0f}GB GPU)" - ) - - # Safety: on 16GB GPUs, prevent selecting LM models that are too large. - # Even with offloading, a 4B LM (8 GB weights + KV cache) leaves almost no - # headroom for DiT activations on a 16 GB card. - if args.lm_model_path and 0 < gpu_memory_gb < VRAM_AUTO_OFFLOAD_THRESHOLD_GB: - if "4B" in args.lm_model_path: - # Downgrade to 1.7B if available - fallback = args.lm_model_path.replace("4B", "1.7B") - print( - f"WARNING: 4B LM model is too large for {gpu_memory_gb:.0f}GB GPU. " - f"Downgrading to 1.7B variant: {fallback}" - ) - args.lm_model_path = fallback - try: - init_params = None dit_handler = None llm_handler = None - - # If init_service is True, perform initialization before creating UI + init_params = None if args.init_service: - print("Initializing service from command line...") - - # Create handler instances for initialization - dit_handler = AceStepHandler() - llm_handler = LLMHandler() - - # Auto-select config_path if not provided - if args.config_path is None: - available_models = dit_handler.get_available_acestep_v15_models() - if available_models: - args.config_path = ( - "acestep-v15-turbo" - if "acestep-v15-turbo" in available_models - else available_models[0] - ) - print(f"Auto-selected config_path: {args.config_path}") - else: - print( - "Error: No available models found. Please specify --config_path", - file=sys.stderr, - ) - sys.exit(1) - - # Get project root (same logic as in handler) - current_file = os.path.abspath(__file__) - project_root = os.path.dirname(os.path.dirname(current_file)) - - # Determine flash attention setting - use_flash_attention = args.use_flash_attention - if use_flash_attention is None: - use_flash_attention = dit_handler.is_flash_attention_available( - args.device - ) - - # Determine download source preference - prefer_source = None - if args.download_source and args.download_source != "auto": - prefer_source = args.download_source - print(f"Using preferred download source: {prefer_source}") - - # Initialize DiT handler - print(f"Initializing DiT model: {args.config_path} on {args.device}...") - compile_model = os.environ.get( - "ACESTEP_COMPILE_MODEL", "" - ).strip().lower() in {"1", "true", "yes", "y", "on"} - - init_status, enable_generate = dit_handler.initialize_service( + init_params, dit_handler, llm_handler = initialize_from_cli( + args, + gpu_config=gpu_config, + effective_gpu_mapping=effective_gpu_mapping, project_root=project_root, - config_path=args.config_path, - device=args.device, - use_flash_attention=use_flash_attention, - compile_model=compile_model, - offload_to_cpu=args.offload_to_cpu, - offload_dit_to_cpu=args.offload_dit_to_cpu, - quantization=args.quantization, - prefer_source=prefer_source, - gpu_mapping=effective_gpu_mapping, + output_dir=output_dir, ) - - if not enable_generate: - print(f"Error initializing DiT model: {init_status}", file=sys.stderr) - sys.exit(1) - - print(f"DiT model initialized successfully") - - # Initialize LM handler if requested - # Auto-determine init_llm based on GPU config if not explicitly set - if args.init_llm is None: - args.init_llm = gpu_config.init_lm_default - print( - f"Auto-setting init_llm to {args.init_llm} based on GPU configuration" - ) - - lm_status = "" - if args.init_llm: - if args.lm_model_path is None: - # Try to get default LM model - available_lm_models = llm_handler.get_available_5hz_lm_models() - if available_lm_models: - args.lm_model_path = available_lm_models[0] - print(f"Using default LM model: {args.lm_model_path}") - else: - print( - "Warning: No LM models available, skipping LM initialization", - file=sys.stderr, - ) - args.init_llm = False - - if args.init_llm and args.lm_model_path: - checkpoint_dir = os.path.join(project_root, "checkpoints") - - # Ensure LM model is downloaded before initialization - prefer_source = None - if args.download_source and args.download_source != "auto": - prefer_source = args.download_source - try: - dl_ok, dl_msg = ensure_lm_model( - model_name=args.lm_model_path, - checkpoints_dir=checkpoint_dir, - prefer_source=prefer_source, - ) - if not dl_ok: - print( - f"Warning: LM model download failed: {dl_msg}", - file=sys.stderr, - ) - except Exception as e: - print( - f"Warning: Failed to download LM model: {e}", - file=sys.stderr, - ) - - lm_device = args.device - device_map = getattr(dit_handler, "device_map", None) - if device_map is not None and device_map.lm is not None: - lm_device = device_map.lm - log_lm_device_deprecation( - explicit_lm_device=os.environ.get("ACESTEP_LM_DEVICE"), - gpu_mapping_env=effective_gpu_mapping, - using_device_map_lm=bool( - device_map is not None and device_map.lm is not None - ), - ) - - print( - f"Initializing 5Hz LM: {args.lm_model_path} on {lm_device}..." - ) - lm_status, lm_success = llm_handler.initialize( - checkpoint_dir=checkpoint_dir, - lm_model_path=args.lm_model_path, - backend=args.backend, - device=lm_device, - offload_to_cpu=args.offload_to_cpu, - dtype=None, - ) - - if lm_success: - print(f"5Hz LM initialized successfully") - init_status += f"\n{lm_status}" - else: - print( - f"Warning: 5Hz LM initialization failed: {lm_status}", - file=sys.stderr, - ) - init_status += f"\n{lm_status}" - - # Prepare initialization parameters for UI - init_params = { - "pre_initialized": True, - "service_mode": args.service_mode, - "checkpoint": args.checkpoint, - "config_path": args.config_path, - "device": args.device, - "gpu_mapping": effective_gpu_mapping, - "init_llm": args.init_llm, - "lm_model_path": args.lm_model_path, - "backend": args.backend, - "use_flash_attention": use_flash_attention, - "offload_to_cpu": args.offload_to_cpu, - "offload_dit_to_cpu": args.offload_dit_to_cpu, - "quantization": args.quantization, - "init_status": init_status, - "enable_generate": enable_generate, - "dit_handler": dit_handler, - "llm_handler": llm_handler, - "language": args.language, - "gpu_config": gpu_config, # Pass GPU config to UI - "output_dir": output_dir, # Pass output dir to UI - "default_batch_size": args.batch_size, # Pass user-specified default batch size - } - - print("Service initialization completed successfully!") - - # Create and launch demo - print(f"Creating Gradio interface with language: {args.language}...") - - # If not using init_service, still pass gpu_config to init_params if init_params is None: init_params = { "gpu_config": gpu_config, "language": args.language, - "output_dir": output_dir, # Pass output dir to UI - "default_batch_size": args.batch_size, # Pass user-specified default batch size + "output_dir": output_dir, + "default_batch_size": args.batch_size, } + print(f"Creating Gradio interface with language: {args.language}...") demo = create_demo(init_params=init_params, language=args.language) - - # Enable queue for multi-user support - # This ensures proper request queuing and prevents concurrent generation conflicts - print("Enabling queue for multi-user support...") - demo.queue( - max_size=20, # Maximum queue size (adjust based on your needs) - status_update_rate="auto", # Update rate for queue status - default_concurrency_limit=1, # Prevents VRAM saturation + launch_gradio_demo( + demo, + args, + output_dir=output_dir, + dit_handler=dit_handler, + llm_handler=llm_handler, ) - - print(f"Launching server on {args.server_name}:{args.port}...") - - # Setup authentication if provided - auth = None - if args.auth_username and args.auth_password: - auth = (args.auth_username, args.auth_password) - print("Authentication enabled") - - allowed_paths = [output_dir] - for p in args.allowed_path: - if p and p not in allowed_paths: - allowed_paths.append(p) - - # Enable API endpoints if requested - if args.enable_api: - print("Enabling API endpoints...") - from acestep.ui.gradio.api.api_routes import setup_api_routes - - # Launch Gradio first with prevent_thread_lock=True - demo.launch( - server_name=args.server_name, - server_port=args.port, - share=args.share, - debug=args.debug, - show_error=True, - prevent_thread_lock=True, # Don't block, so we can add routes - inbrowser=False, - auth=auth, - allowed_paths=allowed_paths, # include output_dir + user-provided - ) - - # Now add API routes to Gradio's FastAPI app (app is available after launch) - setup_api_routes(demo, dit_handler, llm_handler, api_key=args.api_key) - - if args.api_key: - print("API authentication enabled") - print( - "API endpoints enabled: /health, /v1/models, /release_task, /query_result, /create_random_sample, /format_lyrics" - ) - - # Keep the main thread alive - try: - while True: - import time - - time.sleep(1) - except KeyboardInterrupt: - print("\nShutting down...") - else: - demo.launch( - server_name=args.server_name, - server_port=args.port, - share=args.share, - debug=args.debug, - show_error=True, - prevent_thread_lock=False, - inbrowser=False, - auth=auth, - allowed_paths=allowed_paths, # include output_dir + user-provided - ) - except Exception as e: - print(f"Error launching Gradio: {e}", file=sys.stderr) + except Exception as exc: + print(f"Error launching Gradio: {exc}", file=sys.stderr) import traceback traceback.print_exc() diff --git a/acestep/acestep_v15_pipeline_gpu_mapping_test.py b/acestep/acestep_v15_pipeline_gpu_mapping_test.py new file mode 100644 index 000000000..2dd07c3d8 --- /dev/null +++ b/acestep/acestep_v15_pipeline_gpu_mapping_test.py @@ -0,0 +1,150 @@ +"""Unit tests for multi-GPU Gradio CLI flags in the pipeline.""" + +from __future__ import annotations + +import os +import sys +import unittest +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +from acestep import acestep_v15_pipeline + + +class PipelineGpuMappingTests(unittest.TestCase): + """Verify multi-GPU CLI flags are wired into service initialization.""" + + def test_list_gpus_exits_after_printing_inventory(self) -> None: + """``--list-gpus`` should print inventory and exit without launching UI.""" + with patch.object(sys, "argv", ["acestep", "--list-gpus"]), patch( + "acestep.gradio_pipeline_cli.format_gpu_list_text", + return_value="GPU TABLE", + ) as mock_format, patch( + "acestep.gradio_pipeline_cli.sys.exit", + side_effect=SystemExit(0), + ) as mock_exit, patch( + "acestep.acestep_v15_pipeline.get_gpu_config", + return_value=SimpleNamespace( + gpu_memory_gb=24.0, + tier="tier6b", + max_duration_with_lm=480, + max_duration_without_lm=600, + max_batch_size_with_lm=8, + max_batch_size_without_lm=8, + init_lm_default=True, + available_lm_models=["acestep-5Hz-lm-0.6B"], + recommended_backend="vllm", + lm_backend_restriction=None, + offload_dit_to_cpu_default=False, + quantization_default=False, + ), + ), patch( + "acestep.acestep_v15_pipeline.set_global_gpu_config" + ), patch( + "acestep.acestep_v15_pipeline.is_mps_platform", + return_value=False, + ), patch( + "acestep.acestep_v15_pipeline.get_i18n" + ), patch( + "acestep.gradio_pipeline_cli.available_languages_info", + return_value=[("en", "English", "English")], + ), patch( + "acestep.acestep_v15_pipeline.os.makedirs" + ): + with self.assertRaises(SystemExit): + acestep_v15_pipeline.main() + mock_format.assert_called_once() + mock_exit.assert_called_once_with(0) + + def test_gpu_mapping_passed_to_initialize_service(self) -> None: + """``--gpu-mapping`` must reach DiT init and drive LM device selection.""" + gpu_config = SimpleNamespace( + gpu_memory_gb=24.0, + tier="tier6b", + max_duration_with_lm=480, + max_duration_without_lm=600, + max_batch_size_with_lm=8, + max_batch_size_without_lm=8, + init_lm_default=True, + available_lm_models=["acestep-5Hz-lm-0.6B"], + recommended_backend="vllm", + lm_backend_restriction=None, + offload_dit_to_cpu_default=False, + quantization_default=False, + ) + dit_handler = MagicMock() + dit_handler.get_available_acestep_v15_models.return_value = ["acestep-v15-turbo"] + dit_handler.is_flash_attention_available.return_value = False + dit_handler.initialize_service.return_value = ("ok", True) + dit_handler.device_map = SimpleNamespace(lm="cuda:1") + + llm_handler = MagicMock() + llm_handler.get_available_5hz_lm_models.return_value = ["acestep-5Hz-lm-0.6B"] + llm_handler.initialize.return_value = ("ok", True) + + demo = MagicMock() + demo.queue.return_value = demo + demo.launch.return_value = None + captured: dict[str, object] = {} + + def _create_demo(init_params=None, language="en"): + """Capture init_params while returning a stub Gradio demo.""" + captured["init_params"] = init_params + return demo + + with patch.object( + sys, + "argv", + [ + "acestep", + "--init_service", + "true", + "--init_llm", + "true", + "--config_path", + "acestep-v15-turbo", + "--gpu-mapping", + "auto", + ], + ), patch.dict(os.environ, {}, clear=True), patch( + "acestep.acestep_v15_pipeline.get_gpu_config", + return_value=gpu_config, + ), patch( + "acestep.acestep_v15_pipeline.set_global_gpu_config" + ), patch( + "acestep.acestep_v15_pipeline.is_mps_platform", + return_value=False, + ), patch( + "acestep.acestep_v15_pipeline.get_i18n" + ), patch( + "acestep.gradio_pipeline_cli.available_languages_info", + return_value=[("en", "English", "English")], + ), patch( + "acestep.gradio_pipeline_startup.AceStepHandler", + return_value=dit_handler, + ), patch( + "acestep.gradio_pipeline_startup.LLMHandler", + return_value=llm_handler, + ), patch( + "acestep.acestep_v15_pipeline.create_demo", + side_effect=_create_demo, + ), patch( + "acestep.gradio_pipeline_startup.ensure_lm_model", + return_value=(True, "ok"), + ), patch( + "acestep.acestep_v15_pipeline.os.makedirs" + ), patch( + "acestep.gradio_pipeline_startup.log_lm_device_deprecation" + ): + acestep_v15_pipeline.main() + + self.assertEqual( + "auto", + dit_handler.initialize_service.call_args.kwargs["gpu_mapping"], + ) + self.assertEqual("cuda:1", llm_handler.initialize.call_args.kwargs["device"]) + self.assertEqual("auto", captured["init_params"]["gpu_mapping"]) + + +if __name__ == "__main__": + unittest.main() diff --git a/acestep/acestep_v15_pipeline_test.py b/acestep/acestep_v15_pipeline_test.py index 522d85a4f..28eaf4cc3 100644 --- a/acestep/acestep_v15_pipeline_test.py +++ b/acestep/acestep_v15_pipeline_test.py @@ -37,7 +37,7 @@ def _run_main( argv: list[str], *, env: dict[str, str] | None = None, - ) -> tuple[MagicMock, dict[str, object]]: + ) -> tuple[MagicMock, MagicMock, dict[str, object]]: """Run ``main`` with heavy dependencies stubbed and capture startup state.""" gpu_config = self._legacy_gpu_config() dit_handler = MagicMock() @@ -52,15 +52,17 @@ def _run_main( demo = MagicMock() demo.queue.return_value = demo demo.launch.return_value = None - captured: dict[str, object] = {} def _create_demo(init_params=None, language="en"): + """Capture init_params while returning a stub Gradio demo.""" captured["init_params"] = init_params captured["language"] = language return demo - with patch.object(sys, "argv", argv), patch.dict(os.environ, env or {}, clear=True), patch( + with patch.object(sys, "argv", argv), patch.dict( + os.environ, env or {}, clear=True + ), patch( "acestep.acestep_v15_pipeline.get_gpu_config", return_value=gpu_config, ), patch( @@ -71,19 +73,19 @@ def _create_demo(init_params=None, language="en"): ), patch( "acestep.acestep_v15_pipeline.get_i18n" ), patch( - "acestep.acestep_v15_pipeline.available_languages_info", + "acestep.gradio_pipeline_cli.available_languages_info", return_value=[("en", "English", "English")], ), patch( - "acestep.acestep_v15_pipeline.AceStepHandler", + "acestep.gradio_pipeline_startup.AceStepHandler", return_value=dit_handler, ), patch( - "acestep.acestep_v15_pipeline.LLMHandler", + "acestep.gradio_pipeline_startup.LLMHandler", return_value=llm_handler, ), patch( "acestep.acestep_v15_pipeline.create_demo", side_effect=_create_demo, ), patch( - "acestep.acestep_v15_pipeline.ensure_lm_model", + "acestep.gradio_pipeline_startup.ensure_lm_model", return_value=(True, "ok"), ), patch( "acestep.acestep_v15_pipeline.os.makedirs" @@ -109,7 +111,6 @@ def test_main_forces_pt_backend_for_explicit_vllm_argument(self) -> None: "vllm", ] ) - self.assertEqual("pt", llm_handler.initialize.call_args.kwargs["backend"]) self.assertEqual("pt", captured["init_params"]["backend"]) @@ -119,14 +120,12 @@ def test_main_forces_pt_backend_for_service_mode_backend_override(self) -> None: ["acestep", "--service_mode", "true", "--init_llm", "true"], env={"SERVICE_MODE_BACKEND": "vllm"}, ) - self.assertEqual("pt", llm_handler.initialize.call_args.kwargs["backend"]) self.assertEqual("pt", captured["init_params"]["backend"]) def test_main_forces_pt_backend_for_api_env_override(self) -> None: """API-mode env overrides should still resolve to the safe startup backend.""" api_routes_module = types.SimpleNamespace(setup_api_routes=MagicMock()) - with patch.dict( sys.modules, {"acestep.ui.gradio.api.api_routes": api_routes_module}, @@ -144,114 +143,9 @@ def test_main_forces_pt_backend_for_api_env_override(self) -> None: ], env={"ACESTEP_LM_BACKEND": "vllm"}, ) - self.assertEqual("pt", llm_handler.initialize.call_args.kwargs["backend"]) self.assertEqual("pt", captured["init_params"]["backend"]) -class PipelineGpuMappingTests(unittest.TestCase): - """Verify multi-GPU CLI flags are wired into service initialization.""" - - def test_list_gpus_exits_after_printing_inventory(self) -> None: - with patch.object(sys, "argv", ["acestep", "--list-gpus"]), patch( - "acestep.acestep_v15_pipeline.format_gpu_list_text", - return_value="GPU TABLE", - ) as mock_format, patch( - "acestep.acestep_v15_pipeline.sys.exit", - side_effect=SystemExit(0), - ) as mock_exit: - with self.assertRaises(SystemExit): - acestep_v15_pipeline.main() - mock_format.assert_called_once() - mock_exit.assert_called_once_with(0) - - def test_gpu_mapping_passed_to_initialize_service(self) -> None: - gpu_config = SimpleNamespace( - gpu_memory_gb=24.0, - tier="tier6b", - max_duration_with_lm=480, - max_duration_without_lm=600, - max_batch_size_with_lm=8, - max_batch_size_without_lm=8, - init_lm_default=True, - available_lm_models=["acestep-5Hz-lm-0.6B"], - recommended_backend="vllm", - lm_backend_restriction=None, - offload_dit_to_cpu_default=False, - quantization_default=False, - ) - dit_handler = MagicMock() - dit_handler.get_available_acestep_v15_models.return_value = ["acestep-v15-turbo"] - dit_handler.is_flash_attention_available.return_value = False - dit_handler.initialize_service.return_value = ("ok", True) - dit_handler.device_map = SimpleNamespace(lm="cuda:1") - - llm_handler = MagicMock() - llm_handler.get_available_5hz_lm_models.return_value = ["acestep-5Hz-lm-0.6B"] - llm_handler.initialize.return_value = ("ok", True) - - demo = MagicMock() - demo.queue.return_value = demo - demo.launch.return_value = None - captured: dict[str, object] = {} - - def _create_demo(init_params=None, language="en"): - captured["init_params"] = init_params - return demo - - with patch.object( - sys, - "argv", - [ - "acestep", - "--init_service", - "true", - "--init_llm", - "true", - "--config_path", - "acestep-v15-turbo", - "--gpu-mapping", - "auto", - ], - ), patch.dict(os.environ, {}, clear=True), patch( - "acestep.acestep_v15_pipeline.get_gpu_config", - return_value=gpu_config, - ), patch( - "acestep.acestep_v15_pipeline.set_global_gpu_config" - ), patch( - "acestep.acestep_v15_pipeline.is_mps_platform", - return_value=False, - ), patch( - "acestep.acestep_v15_pipeline.get_i18n" - ), patch( - "acestep.acestep_v15_pipeline.available_languages_info", - return_value=[("en", "English", "English")], - ), patch( - "acestep.acestep_v15_pipeline.AceStepHandler", - return_value=dit_handler, - ), patch( - "acestep.acestep_v15_pipeline.LLMHandler", - return_value=llm_handler, - ), patch( - "acestep.acestep_v15_pipeline.create_demo", - side_effect=_create_demo, - ), patch( - "acestep.acestep_v15_pipeline.ensure_lm_model", - return_value=(True, "ok"), - ), patch( - "acestep.acestep_v15_pipeline.os.makedirs" - ), patch( - "acestep.acestep_v15_pipeline.log_lm_device_deprecation" - ): - acestep_v15_pipeline.main() - - self.assertEqual( - "auto", - dit_handler.initialize_service.call_args.kwargs["gpu_mapping"], - ) - self.assertEqual("cuda:1", llm_handler.initialize.call_args.kwargs["device"]) - self.assertEqual("auto", captured["init_params"]["gpu_mapping"]) - - if __name__ == "__main__": unittest.main() diff --git a/acestep/core/generation/handler/conditioning_embed_test.py b/acestep/core/generation/handler/conditioning_embed_test.py index 3d7507aec..fc6c31515 100644 --- a/acestep/core/generation/handler/conditioning_embed_test.py +++ b/acestep/core/generation/handler/conditioning_embed_test.py @@ -12,11 +12,13 @@ class _FakeTextEncoder: """Minimal text encoder stub for preprocess tests.""" def __call__(self, input_ids, lyric_attention_mask=None): + """Return a zero embedding matching the input sequence length.""" del lyric_attention_mask b, t = input_ids.shape return type("O", (), {"last_hidden_state": torch.zeros(b, t, 6, dtype=torch.float32)}) def embed_tokens(self, token_ids): + """Return zero token embeddings with a fixed hidden size.""" b, t = token_ids.shape return torch.zeros(b, t, 6, dtype=torch.float32) @@ -25,6 +27,7 @@ class _Host(ConditioningEmbedMixin): """Minimal host implementing ConditioningEmbedMixin dependencies.""" def __init__(self, text_encoder_device="cpu", dit_device="cpu"): + """Build a host with optional per-component device placement.""" self.device = dit_device self.dtype = torch.float32 self.silence_latent = torch.zeros(1, 128, 6, dtype=torch.float32) @@ -38,16 +41,20 @@ def __init__(self, text_encoder_device="cpu", dit_device="cpu"): } def _get_component_device(self, component): + """Return the stub device string for a named component.""" return self._devices[component] def _ensure_silence_latent_on_device(self): + """No-op silence-latent placement for unit tests.""" return None @contextmanager def _load_model_context(self, _name): + """Yield without loading real model weights.""" yield def tiled_encode(self, audio, offload_latent_to_cpu=True): + """Record encode calls and return a zero latent of matching length.""" del offload_latent_to_cpu self.tiled_encode_calls += 1 t = max(1, audio.shape[-1] // 1920) @@ -128,7 +135,10 @@ def test_infer_text_embeddings_moves_ids_to_text_encoder_device(self): seen = {} class _DeviceCheckingEncoder(_FakeTextEncoder): + """Text encoder stub that records the device of incoming token ids.""" + def __call__(self, input_ids, lyric_attention_mask=None): + """Record ``input_ids.device`` then delegate to the fake encoder.""" seen["device"] = str(input_ids.device) return super().__call__(input_ids, lyric_attention_mask) diff --git a/acestep/gradio_pipeline_banner.py b/acestep/gradio_pipeline_banner.py new file mode 100644 index 000000000..cc515bc4a --- /dev/null +++ b/acestep/gradio_pipeline_banner.py @@ -0,0 +1,52 @@ +"""GPU configuration banner printed at Gradio demo startup.""" + +from __future__ import annotations + +from typing import Any + +from acestep.gpu_config import VRAM_AUTO_OFFLOAD_THRESHOLD_GB + + +def print_gpu_banner( + gpu_config: Any, + gpu_memory_gb: float, + is_mac: bool, + auto_offload: bool, + default_backend: str, +) -> None: + """Print the detected GPU configuration summary used at Gradio startup.""" + print(f"\n{'=' * 60}") + print("GPU Configuration Detected:") + print(f"{'=' * 60}") + print(f" GPU Memory: {gpu_memory_gb:.2f} GB") + print(f" Configuration Tier: {gpu_config.tier}") + print( + f" Max Duration (with LM): {gpu_config.max_duration_with_lm}s " + f"({gpu_config.max_duration_with_lm // 60} min)" + ) + print( + f" Max Duration (without LM): {gpu_config.max_duration_without_lm}s " + f"({gpu_config.max_duration_without_lm // 60} min)" + ) + print(f" Max Batch Size (with LM): {gpu_config.max_batch_size_with_lm}") + print(f" Max Batch Size (without LM): {gpu_config.max_batch_size_without_lm}") + print(f" Default LM Init: {gpu_config.init_lm_default}") + print(f" Available LM Models: {gpu_config.available_lm_models or 'None'}") + print(f"{'=' * 60}\n") + if is_mac: + print( + f"Apple Silicon (MPS) detected — unified memory {gpu_memory_gb:.1f}GB, " + f"no CPU offload needed, backend={default_backend}" + ) + elif auto_offload: + print( + f"Auto-enabling CPU offload (GPU {gpu_memory_gb:.1f}GB < " + f"{VRAM_AUTO_OFFLOAD_THRESHOLD_GB}GB threshold)" + ) + elif gpu_memory_gb > 0: + print( + f"CPU offload disabled by default (GPU {gpu_memory_gb:.1f}GB >= " + f"{VRAM_AUTO_OFFLOAD_THRESHOLD_GB}GB threshold)" + ) + else: + print("No GPU detected, running on CPU") diff --git a/acestep/gradio_pipeline_cli.py b/acestep/gradio_pipeline_cli.py new file mode 100644 index 000000000..43c010c43 --- /dev/null +++ b/acestep/gradio_pipeline_cli.py @@ -0,0 +1,140 @@ +"""Gradio demo CLI argument construction and GPU-mapping env wiring.""" + +from __future__ import annotations + +import argparse +import os +import sys +from typing import Any + +from loguru import logger + +from acestep.device_map import GPU_MAPPING_ENV, format_gpu_list_text +from acestep.gradio_pipeline_cli_service import add_service_init_args +from acestep.ui.gradio.i18n import available_languages_info + + +def build_gradio_parser( + *, + auto_offload: bool, + default_backend: str, + default_offload_dit: bool, + default_quantization: str | None, +) -> argparse.ArgumentParser: + """Build the Gradio demo ArgumentParser with service and GPU-mapping flags.""" + parser = argparse.ArgumentParser( + description="Gradio Demo for ACE-Step V1.5", + formatter_class=argparse.RawTextHelpFormatter, + ) + _add_server_args(parser) + add_service_init_args( + parser, + auto_offload=auto_offload, + default_backend=default_backend, + default_offload_dit=default_offload_dit, + default_quantization=default_quantization, + ) + _add_auth_and_api_args(parser) + return parser + + +def apply_gpu_mapping_args(args: argparse.Namespace) -> str | None: + """Apply ``--list-gpus`` / ``--gpu-mapping`` and return the effective mapping. + + Side effects: + * Prints GPU inventory and exits when ``--list-gpus`` is set. + * Writes ``ACESTEP_GPU_MAPPING`` when an explicit mapping is provided. + """ + if args.list_gpus: + print(format_gpu_list_text()) + sys.exit(0) + + effective = args.gpu_mapping + if effective is None: + return os.environ.get(GPU_MAPPING_ENV) + if effective: + os.environ[GPU_MAPPING_ENV] = effective + return effective + + +def resolve_default_quantization(gpu_config: Any, *, is_mac: bool) -> str | None: + """Choose the CLI default quantization method from GPU tier and capability.""" + if not gpu_config.quantization_default or is_mac: + return None + default = "int8_weight_only" + try: + import torch + + if torch.cuda.is_available(): + major, _ = torch.cuda.get_device_capability(0) + if major < 7: + default = "w8a8_dynamic" + except Exception as exc: + logger.warning( + "[parse_args] CUDA capability probe failed while resolving " + "quantization default: {}", + exc, + ) + return default + + +def _add_server_args(parser: argparse.ArgumentParser) -> None: + """Register server, language, and path-related Gradio flags.""" + parser.add_argument( + "--port", type=int, default=7860, help="Port to run the gradio server on" + ) + parser.add_argument("--share", action="store_true", help="Create a public link") + parser.add_argument("--debug", action="store_true", help="Enable debug mode") + parser.add_argument( + "--server-name", + type=str, + default="127.0.0.1", + help="Server name (default: 127.0.0.1, use 0.0.0.0 for all interfaces)", + ) + languages = available_languages_info() + parser.add_argument( + "--language", + type=str, + default=os.environ.get("LANGUAGE", "en"), + choices=[language[0] for language in languages], + help="UI language:\n " + + "\n ".join( + code + + f" ({native_name}" + + (f"/{name})" if name != native_name else ")") + for code, name, native_name in languages + ), + ) + parser.add_argument( + "--allowed-path", + action="append", + default=[], + help="Additional allowed file paths for Gradio (repeatable).", + ) + + +def _add_auth_and_api_args(parser: argparse.ArgumentParser) -> None: + """Register API enablement and Gradio authentication flags.""" + parser.add_argument( + "--enable-api", + action="store_true", + help="Enable API endpoints (default: False)", + ) + parser.add_argument( + "--auth-username", + type=str, + default=None, + help="Username for Gradio authentication", + ) + parser.add_argument( + "--auth-password", + type=str, + default=None, + help="Password for Gradio authentication", + ) + parser.add_argument( + "--api-key", + type=str, + default=None, + help="API key for API endpoints authentication", + ) diff --git a/acestep/gradio_pipeline_cli_service.py b/acestep/gradio_pipeline_cli_service.py new file mode 100644 index 000000000..85e2c6b6b --- /dev/null +++ b/acestep/gradio_pipeline_cli_service.py @@ -0,0 +1,134 @@ +"""Service-initialization flags for the Gradio demo CLI.""" + +from __future__ import annotations + +import argparse + +from acestep.cli_args import parse_quantization_arg + + +def add_service_init_args( + parser: argparse.ArgumentParser, + *, + auto_offload: bool, + default_backend: str, + default_offload_dit: bool, + default_quantization: str | None, +) -> None: + """Register service init, model, offload, and multi-GPU mapping flags.""" + parser.add_argument( + "--service_mode", + type=lambda x: x.lower() in ["true", "1", "yes"], + default=False, + help="Enable service mode (default: False). When enabled, uses preset models and restricts UI options.", + ) + parser.add_argument( + "--init_service", + type=lambda x: x.lower() in ["true", "1", "yes"], + default=False, + help="Initialize service on startup (default: False)", + ) + parser.add_argument( + "--checkpoint", + type=str, + default=None, + help="Checkpoint file path (optional, for display purposes)", + ) + parser.add_argument( + "--config_path", + type=str, + default=None, + help="Main model path (e.g., 'acestep-v15-turbo')", + ) + parser.add_argument( + "--device", + type=str, + default="auto", + choices=["auto", "cuda", "mps", "xpu", "cpu"], + help="Processing device (default: auto)", + ) + parser.add_argument( + "--gpu-mapping", + dest="gpu_mapping", + type=str, + default=None, + metavar="MAPPING", + help=( + "Component GPU layout: 'auto', 'single:N', or explicit " + "'dit:0,vae:0,text_encoder:0,lm:1'. Also reads ACESTEP_GPU_MAPPING." + ), + ) + parser.add_argument( + "--list-gpus", + action="store_true", + help="List visible CUDA devices and exit", + ) + parser.add_argument( + "--init_llm", + type=lambda x: x.lower() in ["true", "1", "yes"], + default=None, + help="Initialize 5Hz LM (default: auto based on GPU memory)", + ) + parser.add_argument( + "--lm_model_path", + type=str, + default=None, + help="5Hz LM model path (e.g., 'acestep-5Hz-lm-0.6B')", + ) + parser.add_argument( + "--backend", + type=str, + default=default_backend, + choices=["vllm", "pt", "mlx"], + help=( + f"5Hz LM backend (default: {default_backend}, " + "use 'mlx' for native Apple Silicon acceleration)" + ), + ) + parser.add_argument( + "--use_flash_attention", + type=lambda x: x.lower() in ["true", "1", "yes"], + default=None, + help="Use flash attention (default: auto-detect)", + ) + parser.add_argument( + "--offload_to_cpu", + type=lambda x: x.lower() in ["true", "1", "yes"], + default=auto_offload, + help=( + f"Offload models to CPU (default: {'True' if auto_offload else 'False'}, " + "auto-detected based on GPU VRAM)" + ), + ) + parser.add_argument( + "--offload_dit_to_cpu", + type=lambda x: x.lower() in ["true", "1", "yes"], + default=default_offload_dit, + help=( + f"Offload DiT to CPU after diffusion (default: {default_offload_dit}, " + "auto-detected based on GPU tier)" + ), + ) + parser.add_argument( + "--quantization", + type=parse_quantization_arg, + default=default_quantization, + help=( + "DiT quantization method: int8_weight_only, fp8_weight_only, " + "w8a8_dynamic, or none " + f"(default: {default_quantization}, auto-detected based on GPU tier)" + ), + ) + parser.add_argument( + "--download-source", + type=str, + default=None, + choices=["huggingface", "modelscope", "auto"], + help="Preferred model download source (default: auto-detect based on network)", + ) + parser.add_argument( + "--batch_size", + type=int, + default=None, + help="Default batch size for generation (1-8). Defaults to min(2, GPU_max) if not specified", + ) diff --git a/acestep/gradio_pipeline_launch.py b/acestep/gradio_pipeline_launch.py new file mode 100644 index 000000000..105ebf255 --- /dev/null +++ b/acestep/gradio_pipeline_launch.py @@ -0,0 +1,69 @@ +"""Gradio demo launch helper (queue, auth, optional API routes).""" + +from __future__ import annotations + +import argparse +from typing import Any, Optional + +from acestep.handler import AceStepHandler +from acestep.llm_inference import LLMHandler + + +def launch_gradio_demo( + demo: Any, + args: argparse.Namespace, + *, + output_dir: str, + dit_handler: Optional[AceStepHandler], + llm_handler: Optional[LLMHandler], +) -> None: + """Queue and launch the Gradio demo, optionally attaching API routes.""" + print("Enabling queue for multi-user support...") + demo.queue( + max_size=20, + status_update_rate="auto", + default_concurrency_limit=1, + ) + print(f"Launching server on {args.server_name}:{args.port}...") + auth = None + if args.auth_username and args.auth_password: + auth = (args.auth_username, args.auth_password) + print("Authentication enabled") + + allowed_paths = [output_dir] + for path in args.allowed_path: + if path and path not in allowed_paths: + allowed_paths.append(path) + + launch_kwargs = { + "server_name": args.server_name, + "server_port": args.port, + "share": args.share, + "debug": args.debug, + "show_error": True, + "inbrowser": False, + "auth": auth, + "allowed_paths": allowed_paths, + } + if not args.enable_api: + demo.launch(prevent_thread_lock=False, **launch_kwargs) + return + + print("Enabling API endpoints...") + from acestep.ui.gradio.api.api_routes import setup_api_routes + + demo.launch(prevent_thread_lock=True, **launch_kwargs) + setup_api_routes(demo, dit_handler, llm_handler, api_key=args.api_key) + if args.api_key: + print("API authentication enabled") + print( + "API endpoints enabled: /health, /v1/models, /release_task, " + "/query_result, /create_random_sample, /format_lyrics" + ) + try: + import time + + while True: + time.sleep(1) + except KeyboardInterrupt: + print("\nShutting down...") diff --git a/acestep/gradio_pipeline_mode_defaults.py b/acestep/gradio_pipeline_mode_defaults.py new file mode 100644 index 000000000..8b2d249a2 --- /dev/null +++ b/acestep/gradio_pipeline_mode_defaults.py @@ -0,0 +1,52 @@ +"""API/service-mode defaults and 4B LM safety rules for Gradio startup.""" + +from __future__ import annotations + +import argparse +import os + +from acestep.gpu_config import VRAM_AUTO_OFFLOAD_THRESHOLD_GB + + +def apply_startup_mode_defaults(args: argparse.Namespace, gpu_memory_gb: float) -> None: + """Apply API/service-mode env defaults and 4B LM offload safety rules.""" + if args.enable_api: + args.init_service = True + if args.config_path is None: + args.config_path = os.environ.get("ACESTEP_CONFIG_PATH") + if args.lm_model_path is None: + args.lm_model_path = os.environ.get("ACESTEP_LM_MODEL_PATH") + if os.environ.get("ACESTEP_LM_BACKEND"): + args.backend = os.environ.get("ACESTEP_LM_BACKEND") + + if args.service_mode: + print("Service mode enabled - applying preset configurations...") + args.init_service = True + if args.config_path is None: + args.config_path = os.environ.get( + "SERVICE_MODE_DIT_MODEL", "acestep-v15-turbo-fix-inst-shift-dynamic" + ) + if args.lm_model_path is None: + args.lm_model_path = os.environ.get( + "SERVICE_MODE_LM_MODEL", "acestep-5Hz-lm-1.7B-v4-fix" + ) + args.backend = os.environ.get("SERVICE_MODE_BACKEND", "vllm") + print(f" DiT model: {args.config_path}") + print(f" LM model: {args.lm_model_path}") + + if not args.offload_to_cpu and args.lm_model_path and "4B" in args.lm_model_path: + if 0 < gpu_memory_gb <= 24: + args.offload_to_cpu = True + print( + f"Auto-enabling CPU offload (4B LM model requires offloading " + f"on {gpu_memory_gb:.0f}GB GPU)" + ) + + if args.lm_model_path and 0 < gpu_memory_gb < VRAM_AUTO_OFFLOAD_THRESHOLD_GB: + if "4B" in args.lm_model_path: + fallback = args.lm_model_path.replace("4B", "1.7B") + print( + f"WARNING: 4B LM model is too large for {gpu_memory_gb:.0f}GB GPU. " + f"Downgrading to 1.7B variant: {fallback}" + ) + args.lm_model_path = fallback diff --git a/acestep/gradio_pipeline_startup.py b/acestep/gradio_pipeline_startup.py new file mode 100644 index 000000000..e27adefee --- /dev/null +++ b/acestep/gradio_pipeline_startup.py @@ -0,0 +1,181 @@ +"""Gradio demo startup: CLI DiT/LM model initialization.""" + +from __future__ import annotations + +import argparse +import os +import sys +from typing import Any, Optional + +from acestep.device_map import log_lm_device_deprecation +from acestep.handler import AceStepHandler +from acestep.llm_inference import LLMHandler +from acestep.model_downloader import ensure_lm_model + + +def initialize_from_cli( + args: argparse.Namespace, + *, + gpu_config: Any, + effective_gpu_mapping: Optional[str], + project_root: str, + output_dir: str, +) -> tuple[dict[str, Any], AceStepHandler, LLMHandler]: + """Initialize DiT/LM handlers from CLI args and return UI init_params.""" + print("Initializing service from command line...") + dit_handler = AceStepHandler() + llm_handler = LLMHandler() + + if args.config_path is None: + available_models = dit_handler.get_available_acestep_v15_models() + if not available_models: + print( + "Error: No available models found. Please specify --config_path", + file=sys.stderr, + ) + sys.exit(1) + args.config_path = ( + "acestep-v15-turbo" + if "acestep-v15-turbo" in available_models + else available_models[0] + ) + print(f"Auto-selected config_path: {args.config_path}") + + use_flash_attention = args.use_flash_attention + if use_flash_attention is None: + use_flash_attention = dit_handler.is_flash_attention_available(args.device) + + prefer_source = None + if args.download_source and args.download_source != "auto": + prefer_source = args.download_source + print(f"Using preferred download source: {prefer_source}") + + print(f"Initializing DiT model: {args.config_path} on {args.device}...") + compile_model = os.environ.get("ACESTEP_COMPILE_MODEL", "").strip().lower() in { + "1", + "true", + "yes", + "y", + "on", + } + init_status, enable_generate = dit_handler.initialize_service( + project_root=project_root, + config_path=args.config_path, + device=args.device, + use_flash_attention=use_flash_attention, + compile_model=compile_model, + offload_to_cpu=args.offload_to_cpu, + offload_dit_to_cpu=args.offload_dit_to_cpu, + quantization=args.quantization, + prefer_source=prefer_source, + gpu_mapping=effective_gpu_mapping, + ) + if not enable_generate: + print(f"Error initializing DiT model: {init_status}", file=sys.stderr) + sys.exit(1) + print("DiT model initialized successfully") + + if args.init_llm is None: + args.init_llm = gpu_config.init_lm_default + print(f"Auto-setting init_llm to {args.init_llm} based on GPU configuration") + + if args.init_llm: + init_status = _initialize_lm_from_cli( + args, + dit_handler=dit_handler, + llm_handler=llm_handler, + project_root=project_root, + effective_gpu_mapping=effective_gpu_mapping, + init_status=init_status, + prefer_source=prefer_source, + ) + + init_params = { + "pre_initialized": True, + "service_mode": args.service_mode, + "checkpoint": args.checkpoint, + "config_path": args.config_path, + "device": args.device, + "gpu_mapping": effective_gpu_mapping, + "init_llm": args.init_llm, + "lm_model_path": args.lm_model_path, + "backend": args.backend, + "use_flash_attention": use_flash_attention, + "offload_to_cpu": args.offload_to_cpu, + "offload_dit_to_cpu": args.offload_dit_to_cpu, + "quantization": args.quantization, + "init_status": init_status, + "enable_generate": enable_generate, + "dit_handler": dit_handler, + "llm_handler": llm_handler, + "language": args.language, + "gpu_config": gpu_config, + "output_dir": output_dir, + "default_batch_size": args.batch_size, + } + print("Service initialization completed successfully!") + return init_params, dit_handler, llm_handler + + +def _initialize_lm_from_cli( + args: argparse.Namespace, + *, + dit_handler: AceStepHandler, + llm_handler: LLMHandler, + project_root: str, + effective_gpu_mapping: Optional[str], + init_status: str, + prefer_source: Optional[str], +) -> str: + """Download/initialize the 5Hz LM using the mapped device when available.""" + if args.lm_model_path is None: + available_lm_models = llm_handler.get_available_5hz_lm_models() + if available_lm_models: + args.lm_model_path = available_lm_models[0] + print(f"Using default LM model: {args.lm_model_path}") + else: + print( + "Warning: No LM models available, skipping LM initialization", + file=sys.stderr, + ) + args.init_llm = False + return init_status + + if not (args.init_llm and args.lm_model_path): + return init_status + + checkpoint_dir = os.path.join(project_root, "checkpoints") + try: + dl_ok, dl_msg = ensure_lm_model( + model_name=args.lm_model_path, + checkpoints_dir=checkpoint_dir, + prefer_source=prefer_source, + ) + if not dl_ok: + print(f"Warning: LM model download failed: {dl_msg}", file=sys.stderr) + except Exception as exc: + print(f"Warning: Failed to download LM model: {exc}", file=sys.stderr) + + lm_device = args.device + device_map = getattr(dit_handler, "device_map", None) + if device_map is not None and device_map.lm is not None: + lm_device = device_map.lm + log_lm_device_deprecation( + explicit_lm_device=os.environ.get("ACESTEP_LM_DEVICE"), + gpu_mapping_env=effective_gpu_mapping, + using_device_map_lm=bool(device_map is not None and device_map.lm is not None), + ) + print(f"Initializing 5Hz LM: {args.lm_model_path} on {lm_device}...") + lm_status, lm_success = llm_handler.initialize( + checkpoint_dir=checkpoint_dir, + lm_model_path=args.lm_model_path, + backend=args.backend, + device=lm_device, + offload_to_cpu=args.offload_to_cpu, + dtype=None, + ) + if lm_success: + print("5Hz LM initialized successfully") + else: + print(f"Warning: 5Hz LM initialization failed: {lm_status}", file=sys.stderr) + return f"{init_status}\n{lm_status}" diff --git a/acestep/llm_inference.py b/acestep/llm_inference.py index 7cdb12933..c24f46f56 100644 --- a/acestep/llm_inference.py +++ b/acestep/llm_inference.py @@ -25,7 +25,12 @@ from acestep.llm_backend_compat import get_vllm_preflight_warning from acestep.constrained_logits_processor import MetadataConstrainedLogitsProcessor from acestep.constants import DEFAULT_LM_INSTRUCTION, DEFAULT_LM_UNDERSTAND_INSTRUCTION, DEFAULT_LM_INSPIRED_INSTRUCTION, DEFAULT_LM_REWRITE_INSTRUCTION, DURATION_MIN, DURATION_MAX -from acestep.gpu_config import get_lm_gpu_memory_ratio, get_gpu_memory_gb, get_lm_model_size, get_global_gpu_config +from acestep.gpu_config import ( + DEBUG_MAX_CUDA_VRAM_ENV, + get_lm_gpu_memory_ratio, + get_lm_model_size, + get_global_gpu_config, +) from acestep.device_map import ( cuda_device_index, is_cuda_device, @@ -746,14 +751,21 @@ def initialize( if is_cuda_device(device) and torch.cuda.is_available(): try: device_index = cuda_device_index(device) - total_gb = get_gpu_memory_gb() + debug_vram = os.environ.get(DEBUG_MAX_CUDA_VRAM_ENV) + if debug_vram is not None: + total_gb = float(debug_vram) + else: + props = torch.cuda.get_device_properties(device_index) + total_gb = props.total_memory / (1024**3) if hasattr(torch.cuda, "mem_get_info"): free_bytes, _ = torch.cuda.mem_get_info(device_index) free_gb = free_bytes / (1024**3) else: - total_bytes = torch.cuda.get_device_properties(device_index).total_memory free_gb = ( - total_bytes - torch.cuda.memory_reserved(device_index) + torch.cuda.get_device_properties( + device_index + ).total_memory + - torch.cuda.memory_reserved(device_index) ) / (1024**3) except Exception: free_gb = 0.0 diff --git a/acestep/ui/gradio/events/generation/service_init_device_resolution_test.py b/acestep/ui/gradio/events/generation/service_init_device_resolution_test.py new file mode 100644 index 000000000..068550b1f --- /dev/null +++ b/acestep/ui/gradio/events/generation/service_init_device_resolution_test.py @@ -0,0 +1,200 @@ +"""Device-resolution tests for service_init.init_service_wrapper.""" + +from __future__ import annotations + +import unittest +from unittest.mock import MagicMock, patch + +def _stub_gpu_config(**overrides): + """Build a MagicMock GPU config with common Gradio defaults.""" + values = dict( + available_lm_models=["acestep-5Hz-lm-1.7B"], + lm_backend_restriction=None, + tier="tier6", + gpu_memory_gb=24.0, + max_duration_with_lm=600, + max_duration_without_lm=600, + max_batch_size_with_lm=4, + max_batch_size_without_lm=8, + ) + values.update(overrides) + return MagicMock(**values) +class InitServiceWrapperDeviceResolutionTests(unittest.TestCase): + """Verify auto-device handling and device_map LM placement after DiT init.""" + + def _import_module(self): + """Import service_init lazily to avoid heavy transitive imports.""" + from acestep.ui.gradio.events.generation import service_init + + return service_init + + @patch("acestep.ui.gradio.events.generation.service_init.get_global_gpu_config") + def test_reinit_without_llm_preserves_resolved_device(self, mock_gpu_config): + """init_llm=False must not overwrite a previously resolved llm_handler.device.""" + module = self._import_module() + mock_gpu_config.return_value = _stub_gpu_config() + + dit_handler = MagicMock() + dit_handler.initialize_service.return_value = ("ok", True) + dit_handler.model = MagicMock() + dit_handler.is_turbo_model.return_value = True + + llm_handler = MagicMock() + llm_handler.llm_initialized = True + llm_handler.device = "cuda" + + module.init_service_wrapper( + dit_handler, + llm_handler, + "/some/project/checkpoints", + "acestep-v15-turbo", + "auto", + False, + None, + "vllm", + use_flash_attention=False, + offload_to_cpu=False, + offload_dit_to_cpu=False, + compile_model=False, + quantization=False, + ) + + llm_handler.initialize.assert_not_called() + self.assertEqual( + llm_handler.device, + "cuda", + "llm_handler.device must remain 'cuda' when init_llm=False, " + f"got '{llm_handler.device}' instead", + ) + + @patch("acestep.ui.gradio.events.generation.service_init.get_global_gpu_config") + def test_init_llm_with_auto_device_calls_initialize(self, mock_gpu_config): + """init_llm=True with device='auto' must pass 'auto' into initialize().""" + module = self._import_module() + mock_gpu_config.return_value = _stub_gpu_config() + + dit_handler = MagicMock() + dit_handler.initialize_service.return_value = ("ok", True) + dit_handler.model = MagicMock() + dit_handler.is_turbo_model.return_value = True + dit_handler.device_map = None + + llm_handler = MagicMock() + llm_handler.llm_initialized = False + llm_handler.initialize.return_value = ("[OK] LLM initialized", True) + + module.init_service_wrapper( + dit_handler, + llm_handler, + "/some/project/checkpoints", + "acestep-v15-turbo", + "auto", + True, + "acestep-5Hz-lm-1.7B", + "pt", + use_flash_attention=False, + offload_to_cpu=False, + offload_dit_to_cpu=False, + compile_model=False, + quantization=False, + ) + + llm_handler.initialize.assert_called_once() + _, call_kwargs = llm_handler.initialize.call_args + self.assertEqual( + call_kwargs.get("device"), + "auto", + "initialize() must receive 'auto' so it can resolve to the best device", + ) + + @patch("acestep.ui.gradio.events.generation.service_init.get_global_gpu_config") + def test_init_llm_uses_device_map_lm_after_dit_init(self, mock_gpu_config): + """LM device must come from device_map after DiT initialize_service. + + Regression: resolving lm_device before initialize_service left device_map + empty on first Gradio init, so UI device='auto' collapsed the LM onto + bare cuda:0 even when ACESTEP_GPU_MAPPING set lm:1. + """ + module = self._import_module() + mock_gpu_config.return_value = _stub_gpu_config() + + dit_handler = MagicMock() + dit_handler.device_map = None + + def _init_service(*_args, **_kwargs): + """Populate device_map only after DiT initialize_service runs.""" + dit_handler.device_map = MagicMock(lm="cuda:1") + return ("ok", True) + + dit_handler.initialize_service.side_effect = _init_service + dit_handler.model = MagicMock() + dit_handler.is_turbo_model.return_value = True + + llm_handler = MagicMock() + llm_handler.llm_initialized = False + llm_handler.initialize.return_value = ("[OK] LLM initialized", True) + + module.init_service_wrapper( + dit_handler, + llm_handler, + "/some/project/checkpoints", + "acestep-v15-turbo", + "auto", + True, + "acestep-5Hz-lm-1.7B", + "pt", + use_flash_attention=False, + offload_to_cpu=False, + offload_dit_to_cpu=False, + compile_model=False, + quantization=False, + ) + + llm_handler.initialize.assert_called_once() + _, call_kwargs = llm_handler.initialize.call_args + self.assertEqual(call_kwargs.get("device"), "cuda:1") + + @patch("acestep.ui.gradio.events.generation.service_init.get_global_gpu_config") + def test_legacy_cuda_config_forces_pt_backend(self, mock_gpu_config): + """Legacy CUDA restrictions should override a requested vllm backend.""" + module = self._import_module() + mock_gpu_config.return_value = _stub_gpu_config( + available_lm_models=["acestep-5Hz-lm-0.6B"], + lm_backend_restriction="pt_only", + recommended_backend="pt", + tier="tier5", + gpu_memory_gb=12.0, + max_duration_with_lm=480, + max_batch_size_without_lm=4, + ) + + dit_handler = MagicMock() + dit_handler.initialize_service.return_value = ("ok", True) + dit_handler.model = MagicMock() + dit_handler.is_turbo_model.return_value = True + + llm_handler = MagicMock() + llm_handler.llm_initialized = False + llm_handler.initialize.return_value = ("ok", True) + + module.init_service_wrapper( + dit_handler, + llm_handler, + "/some/project/checkpoints", + "acestep-v15-turbo", + "cuda", + True, + "acestep-5Hz-lm-0.6B", + "vllm", + use_flash_attention=False, + offload_to_cpu=False, + offload_dit_to_cpu=False, + compile_model=False, + quantization=False, + ) + + _, call_kwargs = llm_handler.initialize.call_args + self.assertEqual("pt", call_kwargs.get("backend")) + +if __name__ == "__main__": + unittest.main() diff --git a/acestep/ui/gradio/events/generation/service_init_test.py b/acestep/ui/gradio/events/generation/service_init_test.py index 33c7d5a3a..098853524 100644 --- a/acestep/ui/gradio/events/generation/service_init_test.py +++ b/acestep/ui/gradio/events/generation/service_init_test.py @@ -13,6 +13,7 @@ class InitServiceWrapperPathTests(unittest.TestCase): def _import_module(self): """Import service_init lazily to avoid heavy transitive imports.""" from acestep.ui.gradio.events.generation import service_init + return service_init @patch("acestep.ui.gradio.events.generation.service_init.get_global_gpu_config") @@ -26,7 +27,6 @@ def test_passes_project_root_not_checkpoint_dir(self, mock_gpu_config): """ module = self._import_module() - # Stub GPU config mock_gpu_config.return_value = MagicMock( available_lm_models=["acestep-5Hz-lm-1.7B"], lm_backend_restriction=None, @@ -46,7 +46,6 @@ def test_passes_project_root_not_checkpoint_dir(self, mock_gpu_config): llm_handler = MagicMock() llm_handler.llm_initialized = False - # Simulate the checkpoint dropdown value: full path to checkpoints dir checkpoint_value = "/some/project/checkpoints" module.init_service_wrapper( @@ -55,23 +54,19 @@ def test_passes_project_root_not_checkpoint_dir(self, mock_gpu_config): checkpoint_value, "acestep-v15-turbo", "cpu", - False, # init_llm - None, # lm_model_path - "vllm", # backend - False, # use_flash_attention - False, # offload_to_cpu - False, # offload_dit_to_cpu - False, # compile_model - False, # quantization + False, + None, + "vllm", + False, + False, + False, + False, + False, ) - # The first positional arg to initialize_service must be the project root, - # NOT the checkpoints directory. call_args = dit_handler.initialize_service.call_args actual_project_root = call_args[0][0] - # It should be computed from __file__, not from the checkpoint dropdown. - # Critically, it must NOT end with "checkpoints". self.assertFalse( actual_project_root.rstrip("/").endswith("checkpoints"), f"project_root must not be the checkpoints dir, got: {actual_project_root}", @@ -104,22 +99,26 @@ def test_project_root_is_consistent_with_checkpoint_dir(self, mock_gpu_config): module.init_service_wrapper( dit_handler, llm_handler, - "/any/path/checkpoints", # checkpoint dropdown value (unused now) + "/any/path/checkpoints", "acestep-v15-turbo", "cpu", - False, None, "vllm", False, False, False, False, False, + False, + None, + "vllm", + False, + False, + False, + False, + False, ) call_args = dit_handler.initialize_service.call_args actual_project_root = call_args[0][0] - - # The project_root + "checkpoints" should form a valid checkpoints path expected_checkpoints = os.path.join(actual_project_root, "checkpoints") self.assertTrue( os.path.isabs(expected_checkpoints) or actual_project_root, "project_root should be a meaningful path", ) - # It should NOT contain double "checkpoints" self.assertNotIn( "checkpoints/checkpoints", expected_checkpoints, @@ -127,248 +126,22 @@ def test_project_root_is_consistent_with_checkpoint_dir(self, mock_gpu_config): ) -class InitServiceWrapperDeviceResolutionTests(unittest.TestCase): - """Verify that 'auto' device is not written back to llm_handler when init_llm=False. - - Regression test for: Auto-labelling broke after recent update if auto is - chosen for device. When the user re-initialises the service without the - 'Init LLM' checkbox ticked, the previously-resolved device (e.g. 'cuda') - must not be overwritten with the raw UI value 'auto'. - """ - - def _import_module(self): - """Import service_init lazily to avoid heavy transitive imports.""" - from acestep.ui.gradio.events.generation import service_init - return service_init - - @patch("acestep.ui.gradio.events.generation.service_init.get_global_gpu_config") - def test_reinit_without_llm_preserves_resolved_device(self, mock_gpu_config): - """Calling init_service_wrapper with init_llm=False must not overwrite llm_handler.device. - - Scenario: LLM was previously initialized (llm_initialized=True, device='cuda'). - User re-initialises the service (e.g. to change checkpoint) with init_llm=False - and device='auto'. The LLM handler's resolved device must remain 'cuda'. - """ - module = self._import_module() - - mock_gpu_config.return_value = MagicMock( - available_lm_models=["acestep-5Hz-lm-1.7B"], - lm_backend_restriction=None, - tier="tier6", - gpu_memory_gb=24.0, - max_duration_with_lm=600, - max_duration_without_lm=600, - max_batch_size_with_lm=4, - max_batch_size_without_lm=8, - ) - - dit_handler = MagicMock() - dit_handler.initialize_service.return_value = ("ok", True) - dit_handler.model = MagicMock() - dit_handler.is_turbo_model.return_value = True - - # Simulate LLM previously initialised with resolved device="cuda" - llm_handler = MagicMock() - llm_handler.llm_initialized = True - llm_handler.device = "cuda" # previously resolved from "auto" -> "cuda" - - module.init_service_wrapper( - dit_handler, - llm_handler, - "/some/project/checkpoints", - "acestep-v15-turbo", - "auto", # raw UI value -- must NOT overwrite the resolved "cuda" - False, # init_llm=False: do not re-initialize LLM - None, # lm_model_path - "vllm", # backend - use_flash_attention=False, - offload_to_cpu=False, - offload_dit_to_cpu=False, - compile_model=False, - quantization=False, - ) - - # llm_handler.initialize must NOT have been called (init_llm=False) - llm_handler.initialize.assert_not_called() - # The previously-resolved device must be preserved - self.assertEqual( - llm_handler.device, - "cuda", - "llm_handler.device must remain 'cuda' when init_llm=False, " - f"got '{llm_handler.device}' instead", - ) - - @patch("acestep.ui.gradio.events.generation.service_init.get_global_gpu_config") - def test_init_llm_with_auto_device_calls_initialize(self, mock_gpu_config): - """When init_llm=True and device='auto', initialize() must be called with 'auto' device. - - The 'auto' -> concrete device resolution happens inside initialize(), so we - must pass 'auto' through correctly. - """ - module = self._import_module() - - mock_gpu_config.return_value = MagicMock( - available_lm_models=["acestep-5Hz-lm-1.7B"], - lm_backend_restriction=None, - tier="tier6", - gpu_memory_gb=24.0, - max_duration_with_lm=600, - max_duration_without_lm=600, - max_batch_size_with_lm=4, - max_batch_size_without_lm=8, - ) - - dit_handler = MagicMock() - dit_handler.initialize_service.return_value = ("ok", True) - dit_handler.model = MagicMock() - dit_handler.is_turbo_model.return_value = True - dit_handler.device_map = None - - llm_handler = MagicMock() - llm_handler.llm_initialized = False - llm_handler.initialize.return_value = ("[OK] LLM initialized", True) - - module.init_service_wrapper( - dit_handler, - llm_handler, - "/some/project/checkpoints", - "acestep-v15-turbo", - "auto", # raw UI value - True, # init_llm=True - "acestep-5Hz-lm-1.7B", - "pt", - use_flash_attention=False, - offload_to_cpu=False, - offload_dit_to_cpu=False, - compile_model=False, - quantization=False, - ) - - llm_handler.initialize.assert_called_once() - _, call_kwargs = llm_handler.initialize.call_args - self.assertEqual( - call_kwargs.get("device"), - "auto", - "initialize() must receive 'auto' so it can resolve to the best device", - ) - - @patch("acestep.ui.gradio.events.generation.service_init.get_global_gpu_config") - def test_init_llm_uses_device_map_lm_after_dit_init(self, mock_gpu_config): - """LM device must come from device_map after DiT initialize_service. - - Regression: resolving lm_device before initialize_service left device_map - empty on first Gradio init, so UI device='auto' collapsed the LM onto - bare cuda:0 even when ACESTEP_GPU_MAPPING set lm:1. - """ - module = self._import_module() - - mock_gpu_config.return_value = MagicMock( - available_lm_models=["acestep-5Hz-lm-1.7B"], - lm_backend_restriction=None, - tier="tier6", - gpu_memory_gb=24.0, - max_duration_with_lm=600, - max_duration_without_lm=600, - max_batch_size_with_lm=4, - max_batch_size_without_lm=8, - ) - - dit_handler = MagicMock() - dit_handler.device_map = None - - def _init_service(*_args, **_kwargs): - dit_handler.device_map = MagicMock(lm="cuda:1") - return ("ok", True) - - dit_handler.initialize_service.side_effect = _init_service - dit_handler.model = MagicMock() - dit_handler.is_turbo_model.return_value = True - - llm_handler = MagicMock() - llm_handler.llm_initialized = False - llm_handler.initialize.return_value = ("[OK] LLM initialized", True) - - module.init_service_wrapper( - dit_handler, - llm_handler, - "/some/project/checkpoints", - "acestep-v15-turbo", - "auto", - True, - "acestep-5Hz-lm-1.7B", - "pt", - use_flash_attention=False, - offload_to_cpu=False, - offload_dit_to_cpu=False, - compile_model=False, - quantization=False, - ) - - llm_handler.initialize.assert_called_once() - _, call_kwargs = llm_handler.initialize.call_args - self.assertEqual(call_kwargs.get("device"), "cuda:1") - - @patch("acestep.ui.gradio.events.generation.service_init.get_global_gpu_config") - def test_legacy_cuda_config_forces_pt_backend(self, mock_gpu_config): - """Legacy CUDA restrictions should override a requested vllm backend.""" - module = self._import_module() - - mock_gpu_config.return_value = MagicMock( - available_lm_models=["acestep-5Hz-lm-0.6B"], - lm_backend_restriction="pt_only", - recommended_backend="pt", - tier="tier5", - gpu_memory_gb=12.0, - max_duration_with_lm=480, - max_duration_without_lm=600, - max_batch_size_with_lm=4, - max_batch_size_without_lm=4, - ) - - dit_handler = MagicMock() - dit_handler.initialize_service.return_value = ("ok", True) - dit_handler.model = MagicMock() - dit_handler.is_turbo_model.return_value = True - - llm_handler = MagicMock() - llm_handler.llm_initialized = False - llm_handler.initialize.return_value = ("ok", True) - - module.init_service_wrapper( - dit_handler, - llm_handler, - "/some/project/checkpoints", - "acestep-v15-turbo", - "cuda", - True, - "acestep-5Hz-lm-0.6B", - "vllm", - use_flash_attention=False, - offload_to_cpu=False, - offload_dit_to_cpu=False, - compile_model=False, - quantization=False, - ) - - _, call_kwargs = llm_handler.initialize.call_args - self.assertEqual("pt", call_kwargs.get("backend")) - - - class QuantizationSelectionTests(unittest.TestCase): """Verify pre-Ampere quantization mode selection.""" def _import_module(self): """Import service_init lazily to avoid heavy transitive imports.""" from acestep.ui.gradio.events.generation import service_init + return service_init def test_select_quantization_value_uses_dynamic_mode_for_pre_ampere_cuda(self): """It selects ``w8a8_dynamic`` for pre-Ampere CUDA devices.""" module = self._import_module() - with patch("torch.cuda.is_available", return_value=True), \ - patch("torch.cuda.get_device_capability", return_value=(6, 1)): + with patch("torch.cuda.is_available", return_value=True), patch( + "torch.cuda.get_device_capability", return_value=(6, 1) + ): self.assertEqual( module._select_quantization_value( quantization_enabled=True, @@ -407,5 +180,6 @@ def fake_import(name, globals_=None, locals_=None, fromlist=(), level=0): if removed_torch_nn_module is not None: sys.modules["torch.nn"] = removed_torch_nn_module + if __name__ == "__main__": unittest.main() From 7085b679aa6a9f7f4c1cda716c00d629e88797dd Mon Sep 17 00:00:00 2001 From: steve Date: Sat, 11 Jul 2026 21:23:54 -0400 Subject: [PATCH 18/18] fix(cli): match LM size token as -4B, not substring 4B Avoid false positives on names like 14B when auto-offloading or downgrading large LMs on limited VRAM at Gradio startup. Co-authored-by: Cursor --- acestep/gradio_pipeline_mode_defaults.py | 6 +- acestep/gradio_pipeline_mode_defaults_test.py | 64 +++++++++++++++++++ 2 files changed, 67 insertions(+), 3 deletions(-) create mode 100644 acestep/gradio_pipeline_mode_defaults_test.py diff --git a/acestep/gradio_pipeline_mode_defaults.py b/acestep/gradio_pipeline_mode_defaults.py index 8b2d249a2..71dbbe11f 100644 --- a/acestep/gradio_pipeline_mode_defaults.py +++ b/acestep/gradio_pipeline_mode_defaults.py @@ -34,7 +34,7 @@ def apply_startup_mode_defaults(args: argparse.Namespace, gpu_memory_gb: float) print(f" DiT model: {args.config_path}") print(f" LM model: {args.lm_model_path}") - if not args.offload_to_cpu and args.lm_model_path and "4B" in args.lm_model_path: + if not args.offload_to_cpu and args.lm_model_path and "-4B" in args.lm_model_path: if 0 < gpu_memory_gb <= 24: args.offload_to_cpu = True print( @@ -43,8 +43,8 @@ def apply_startup_mode_defaults(args: argparse.Namespace, gpu_memory_gb: float) ) if args.lm_model_path and 0 < gpu_memory_gb < VRAM_AUTO_OFFLOAD_THRESHOLD_GB: - if "4B" in args.lm_model_path: - fallback = args.lm_model_path.replace("4B", "1.7B") + if "-4B" in args.lm_model_path: + fallback = args.lm_model_path.replace("-4B", "-1.7B") print( f"WARNING: 4B LM model is too large for {gpu_memory_gb:.0f}GB GPU. " f"Downgrading to 1.7B variant: {fallback}" diff --git a/acestep/gradio_pipeline_mode_defaults_test.py b/acestep/gradio_pipeline_mode_defaults_test.py new file mode 100644 index 000000000..ed372c7de --- /dev/null +++ b/acestep/gradio_pipeline_mode_defaults_test.py @@ -0,0 +1,64 @@ +"""Unit tests for Gradio startup LM size safety defaults.""" + +from __future__ import annotations + +import unittest +from argparse import Namespace + +from acestep.gradio_pipeline_mode_defaults import apply_startup_mode_defaults + + +class StartupModeDefaultsLmSizeTests(unittest.TestCase): + """Verify 4B LM detection does not false-match larger size tokens.""" + + def test_downgrade_only_matches_hyphen_4b_token(self) -> None: + """``-4B`` downgrades; a hypothetical ``-14B`` name must stay untouched.""" + four_b = Namespace( + enable_api=False, + service_mode=False, + offload_to_cpu=False, + lm_model_path="acestep-5Hz-lm-4B", + config_path=None, + backend="pt", + ) + apply_startup_mode_defaults(four_b, gpu_memory_gb=16.0) + self.assertEqual(four_b.lm_model_path, "acestep-5Hz-lm-1.7B") + + fourteen_b = Namespace( + enable_api=False, + service_mode=False, + offload_to_cpu=False, + lm_model_path="acestep-5Hz-lm-14B", + config_path=None, + backend="pt", + ) + apply_startup_mode_defaults(fourteen_b, gpu_memory_gb=16.0) + self.assertEqual(fourteen_b.lm_model_path, "acestep-5Hz-lm-14B") + + def test_offload_only_matches_hyphen_4b_token(self) -> None: + """CPU offload auto-enable must key off ``-4B``, not substring ``4B``.""" + four_b = Namespace( + enable_api=False, + service_mode=False, + offload_to_cpu=False, + lm_model_path="acestep-5Hz-lm-4B", + config_path=None, + backend="pt", + ) + apply_startup_mode_defaults(four_b, gpu_memory_gb=24.0) + self.assertTrue(four_b.offload_to_cpu) + + fourteen_b = Namespace( + enable_api=False, + service_mode=False, + offload_to_cpu=False, + lm_model_path="acestep-5Hz-lm-14B", + config_path=None, + backend="pt", + ) + apply_startup_mode_defaults(fourteen_b, gpu_memory_gb=24.0) + self.assertFalse(fourteen_b.offload_to_cpu) + + +if __name__ == "__main__": + unittest.main()