Spaces:
Running
Running
| """Initialization-adjacent utility mixin for AceStepHandler.""" | |
| import os | |
| import time | |
| from contextlib import contextmanager | |
| from typing import List, Optional | |
| import torch | |
| from loguru import logger | |
| class InitServiceMixin: | |
| def _device_type(self) -> str: | |
| """Normalize the host device value to a backend type string.""" | |
| if isinstance(self.device, str): | |
| return self.device.split(":", 1)[0] | |
| return self.device.type | |
| def get_available_checkpoints(self) -> List[str]: | |
| """Return available checkpoint directory paths under the project root. | |
| Uses ``self._get_project_root()`` to resolve the checkpoints directory and | |
| returns a single-item list when present, otherwise an empty list. | |
| """ | |
| # Get project root (handler.py is in acestep/, so go up two levels to project root) | |
| project_root = self._get_project_root() | |
| # default checkpoints | |
| checkpoint_dir = os.path.join(project_root, "checkpoints") | |
| if os.path.exists(checkpoint_dir): | |
| return [checkpoint_dir] | |
| else: | |
| return [] | |
| def get_available_acestep_v15_models(self) -> List[str]: | |
| """Scan and return all model directory names starting with 'acestep-v15-'""" | |
| # Get project root | |
| project_root = self._get_project_root() | |
| checkpoint_dir = os.path.join(project_root, "checkpoints") | |
| models = [] | |
| if os.path.exists(checkpoint_dir): | |
| # Scan all directories starting with 'acestep-v15-' in checkpoints folder | |
| for item in os.listdir(checkpoint_dir): | |
| item_path = os.path.join(checkpoint_dir, item) | |
| if os.path.isdir(item_path) and item.startswith("acestep-v15-"): | |
| models.append(item) | |
| # Sort by name | |
| models.sort() | |
| return models | |
| 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] | |
| if target_device == "auto": | |
| if not torch.cuda.is_available(): | |
| return False | |
| else: | |
| if target_device != "cuda" or not torch.cuda.is_available(): | |
| return False | |
| # FlashAttention requires Ampere (compute capability >= 8.0) or newer | |
| try: | |
| major, _ = torch.cuda.get_device_capability() | |
| if major < 8: | |
| logger.info( | |
| f"[is_flash_attention_available] GPU compute capability {major}.x < 8.0 " | |
| f"(pre-Ampere) — FlashAttention not supported, will use SDPA instead." | |
| ) | |
| return False | |
| except Exception: | |
| return False | |
| try: | |
| import flash_attn | |
| return True | |
| except ImportError: | |
| return False | |
| def is_turbo_model(self) -> bool: | |
| """Check if the currently loaded model is a turbo model""" | |
| if self.config is None: | |
| return False | |
| return getattr(self.config, "is_turbo", False) | |
| def _empty_cache(self): | |
| """Clear accelerator memory cache (CUDA, XPU, or MPS).""" | |
| device_type = self._device_type() | |
| if device_type == "cuda" and torch.cuda.is_available(): | |
| torch.cuda.empty_cache() | |
| elif device_type == "xpu" and hasattr(torch, "xpu") and torch.xpu.is_available(): | |
| torch.xpu.empty_cache() | |
| elif device_type == "mps" and hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): | |
| torch.mps.empty_cache() | |
| def _synchronize(self): | |
| """Synchronize accelerator operations (CUDA, XPU, or MPS).""" | |
| device_type = self._device_type() | |
| if device_type == "cuda" and torch.cuda.is_available(): | |
| torch.cuda.synchronize() | |
| elif device_type == "xpu" and hasattr(torch, "xpu") and torch.xpu.is_available(): | |
| torch.xpu.synchronize() | |
| elif device_type == "mps" and hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): | |
| torch.mps.synchronize() | |
| def _memory_allocated(self): | |
| """Get current accelerator memory usage in bytes, or 0 for unsupported backends.""" | |
| device_type = self._device_type() | |
| if device_type == "cuda" and torch.cuda.is_available(): | |
| return torch.cuda.memory_allocated() | |
| # MPS and XPU don't expose per-tensor memory tracking | |
| return 0 | |
| def _max_memory_allocated(self): | |
| """Get peak accelerator memory usage in bytes, or 0 for unsupported backends.""" | |
| device_type = self._device_type() | |
| if device_type == "cuda" and torch.cuda.is_available(): | |
| return torch.cuda.max_memory_allocated() | |
| return 0 | |
| def _is_on_target_device(self, tensor, target_device): | |
| """Check if tensor is on the target device (handles cuda vs cuda:0 comparison).""" | |
| if tensor is None: | |
| return True | |
| try: | |
| if isinstance(target_device, torch.device): | |
| target_type = target_device.type | |
| else: | |
| target_type = torch.device(str(target_device)).type | |
| except Exception: | |
| # Keep fallback conservative: derive backend token instead of assuming CUDA. | |
| target_type = str(target_device).strip().lower().split(":", 1)[0] | |
| if not target_type: | |
| logger.warning( | |
| "[_is_on_target_device] Malformed target device value: {!r}", | |
| target_device, | |
| ) | |
| return False | |
| return tensor.device.type == target_type | |
| def _get_affine_quantized_tensor_class(): | |
| """Return the AffineQuantizedTensor class from torchao, or None if unavailable. | |
| Supports both old (torchao.quantization.affine_quantized) and new | |
| (torchao.dtypes.affine_quantized_tensor) import paths across torchao versions. | |
| """ | |
| try: | |
| from torchao.dtypes.affine_quantized_tensor import AffineQuantizedTensor | |
| return AffineQuantizedTensor | |
| except ImportError: | |
| pass | |
| try: | |
| from torchao.quantization.affine_quantized import AffineQuantizedTensor | |
| return AffineQuantizedTensor | |
| except ImportError: | |
| pass | |
| return None | |
| def _is_quantized_tensor(self, t): | |
| """True if t is a torchao AffineQuantizedTensor (calling .to() on it can raise NotImplementedError).""" | |
| if t is None: | |
| return False | |
| cls = self._get_affine_quantized_tensor_class() | |
| if cls is None: | |
| return False | |
| return isinstance(t, cls) | |
| def _has_quantized_params(self, module): | |
| """True if module (or any submodule) has at least one AffineQuantizedTensor parameter.""" | |
| cls = self._get_affine_quantized_tensor_class() | |
| if cls is None: | |
| return False | |
| for _, param in module.named_parameters(): | |
| if param is not None and isinstance(param, cls): | |
| return True | |
| return False | |
| def _ensure_silence_latent_on_device(self): | |
| """Ensure silence_latent is on the correct device (self.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) | |
| def _move_module_recursive(self, module, target_device, dtype=None, visited=None): | |
| """ | |
| Recursively move a module and all its submodules to the target device. | |
| This handles modules that may not be properly registered. | |
| """ | |
| if visited is None: | |
| visited = set() | |
| module_id = id(module) | |
| if module_id in visited: | |
| return | |
| visited.add(module_id) | |
| # Move the module itself | |
| module.to(target_device) | |
| if dtype is not None: | |
| module.to(dtype) | |
| # Move all direct parameters | |
| for param_name, param in module._parameters.items(): | |
| if param is not None and not self._is_on_target_device(param, target_device): | |
| if self._is_quantized_tensor(param): | |
| moved_param = self._move_quantized_param(param, target_device) | |
| else: | |
| moved_param = torch.nn.Parameter( | |
| param.data.to(target_device), requires_grad=param.requires_grad | |
| ) | |
| if dtype is not None and moved_param.is_floating_point(): | |
| moved_param = torch.nn.Parameter( | |
| moved_param.data.to(dtype), requires_grad=param.requires_grad | |
| ) | |
| module._parameters[param_name] = moved_param | |
| # Move all direct buffers | |
| for buf_name, buf in module._buffers.items(): | |
| if buf is not None and not self._is_on_target_device(buf, target_device): | |
| module._buffers[buf_name] = buf.to(target_device) | |
| # Recursively process all submodules (registered and unregistered) | |
| for name, child in module._modules.items(): | |
| if child is not None: | |
| self._move_module_recursive(child, target_device, dtype, visited) | |
| # Also check for any nn.Module attributes that might not be in _modules | |
| for attr_name in dir(module): | |
| if attr_name.startswith('_'): | |
| continue | |
| try: | |
| attr = getattr(module, attr_name, None) | |
| if isinstance(attr, torch.nn.Module) and id(attr) not in visited: | |
| self._move_module_recursive(attr, target_device, dtype, visited) | |
| except Exception: | |
| pass | |
| def _move_quantized_param(self, param, target_device): | |
| """Move an AffineQuantizedTensor to target_device using _apply_fn_to_data. | |
| This is the safe fallback for older torch versions where model.to(device) raises | |
| NotImplementedError on AffineQuantizedTensor (because aten._has_compatible_shallow_copy_type | |
| is not implemented). _apply_fn_to_data recursively applies a function to all inner | |
| tensors (int_data, scale, zero_point, etc.) without going through Module._apply. | |
| """ | |
| if hasattr(param, '_apply_fn_to_data'): | |
| return torch.nn.Parameter( | |
| param._apply_fn_to_data(lambda x: x.to(target_device)), | |
| requires_grad=param.requires_grad, | |
| ) | |
| # Last resort: try direct .to() (may raise), but preserve Parameter registration. | |
| moved = param.to(target_device) | |
| return torch.nn.Parameter(moved, requires_grad=param.requires_grad) | |
| def _recursive_to_device(self, model, device, dtype=None): | |
| """ | |
| Recursively move all parameters and buffers of a model to the specified device. | |
| This is more thorough than model.to() for some custom HuggingFace models. | |
| Handles torchao AffineQuantizedTensor parameters that may raise NotImplementedError | |
| on model.to(device) in older torch versions (where Module._apply calls | |
| _has_compatible_shallow_copy_type, which is not implemented for AffineQuantizedTensor). | |
| In that case, falls back to moving quantized parameters individually via _apply_fn_to_data. | |
| """ | |
| target_device = torch.device(device) if isinstance(device, str) else device | |
| # Method 1: Standard .to() call — works on newer torch where _apply uses swap_tensors | |
| try: | |
| model.to(target_device) | |
| if dtype is not None: | |
| model.to(dtype) | |
| except NotImplementedError: | |
| # Older torch: Module._apply calls _has_compatible_shallow_copy_type which is | |
| # not implemented for AffineQuantizedTensor. Move parameters manually. | |
| logger.info( | |
| "[_recursive_to_device] model.to() raised NotImplementedError " | |
| "(AffineQuantizedTensor on older torch). Moving parameters individually." | |
| ) | |
| for module in model.modules(): | |
| # Move non-quantized parameters and buffers directly | |
| for param_name, param in module._parameters.items(): | |
| if param is None: | |
| continue | |
| if self._is_on_target_device(param, target_device): | |
| continue | |
| if self._is_quantized_tensor(param): | |
| module._parameters[param_name] = self._move_quantized_param(param, target_device) | |
| else: | |
| module._parameters[param_name] = torch.nn.Parameter( | |
| param.data.to(target_device), requires_grad=param.requires_grad | |
| ) | |
| if dtype is not None: | |
| module._parameters[param_name] = torch.nn.Parameter( | |
| module._parameters[param_name].data.to(dtype), | |
| requires_grad=param.requires_grad, | |
| ) | |
| for buf_name, buf in module._buffers.items(): | |
| if buf is not None and not self._is_on_target_device(buf, target_device): | |
| module._buffers[buf_name] = buf.to(target_device) | |
| # Method 2: Use our thorough recursive moving for any missed modules | |
| # (skip if model.to() failed — we already moved everything above) | |
| try: | |
| self._move_module_recursive(model, target_device, dtype) | |
| except NotImplementedError: | |
| pass # Already handled above | |
| # Method 3: Force move via state_dict if there are still parameters on wrong device | |
| wrong_device_params = [] | |
| for name, param in model.named_parameters(): | |
| if not self._is_on_target_device(param, device): | |
| wrong_device_params.append(name) | |
| if wrong_device_params and device != "cpu": | |
| logger.warning(f"[_recursive_to_device] {len(wrong_device_params)} parameters on wrong device after initial move, retrying individually") | |
| for module in model.modules(): | |
| for param_name, param in module._parameters.items(): | |
| if param is None or self._is_on_target_device(param, target_device): | |
| continue | |
| if self._is_quantized_tensor(param): | |
| module._parameters[param_name] = self._move_quantized_param(param, target_device) | |
| else: | |
| module._parameters[param_name] = torch.nn.Parameter( | |
| param.data.to(target_device), requires_grad=param.requires_grad | |
| ) | |
| if dtype is not None and module._parameters[param_name].is_floating_point(): | |
| module._parameters[param_name] = torch.nn.Parameter( | |
| module._parameters[param_name].data.to(dtype), | |
| requires_grad=param.requires_grad, | |
| ) | |
| # Synchronize accelerator to ensure all transfers are complete | |
| if device != "cpu": | |
| self._synchronize() | |
| # Final verification | |
| if device != "cpu": | |
| still_wrong = [] | |
| for name, param in model.named_parameters(): | |
| if not self._is_on_target_device(param, device): | |
| still_wrong.append(f"{name} on {param.device}") | |
| if still_wrong: | |
| logger.error(f"[_recursive_to_device] CRITICAL: {len(still_wrong)} parameters still on wrong device: {still_wrong[:10]}") | |
| def _load_model_context(self, model_name: str): | |
| """ | |
| Context manager to load a model to GPU and offload it back to CPU after use. | |
| Args: | |
| model_name: Name of the model to load ("text_encoder", "vae", "model") | |
| """ | |
| if not self.offload_to_cpu: | |
| yield | |
| return | |
| # If model is DiT ("model") and offload_dit_to_cpu is False, do not offload | |
| if model_name == "model" and not self.offload_dit_to_cpu: | |
| # Ensure it's on device if not already (should be handled by init, but safe to check) | |
| model = getattr(self, model_name, None) | |
| if model is not None: | |
| # Check if model is on CPU, if so move to device (one-time move if it was somehow on CPU) | |
| # We check the first parameter's device | |
| 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) | |
| if hasattr(self, "silence_latent"): | |
| self.silence_latent = self.silence_latent.to(self.device).to(self.dtype) | |
| except StopIteration: | |
| pass | |
| yield | |
| return | |
| model = getattr(self, model_name, None) | |
| if model is None: | |
| yield | |
| return | |
| # Load to GPU | |
| logger.info(f"[_load_model_context] Loading {model_name} to {self.device}") | |
| start_time = time.time() | |
| if model_name == "vae": | |
| vae_dtype = self._get_vae_dtype() | |
| self._recursive_to_device(model, self.device, vae_dtype) | |
| else: | |
| self._recursive_to_device(model, self.device, self.dtype) | |
| if model_name == "model" and hasattr(self, "silence_latent"): | |
| self.silence_latent = self.silence_latent.to(self.device).to(self.dtype) | |
| load_time = time.time() - start_time | |
| self.current_offload_cost += load_time | |
| logger.info(f"[_load_model_context] Loaded {model_name} to {self.device} in {load_time:.4f}s") | |
| try: | |
| yield | |
| finally: | |
| # Offload to CPU | |
| logger.info(f"[_load_model_context] Offloading {model_name} to CPU") | |
| start_time = time.time() | |
| if model_name == "vae": | |
| self._recursive_to_device(model, "cpu", self._get_vae_dtype("cpu")) | |
| else: | |
| self._recursive_to_device(model, "cpu") | |
| # NOTE: Do NOT offload silence_latent to CPU here! | |
| # silence_latent is used in many places outside of model context, | |
| # so it should stay on GPU to avoid device mismatch errors. | |
| self._empty_cache() | |
| offload_time = time.time() - start_time | |
| self.current_offload_cost += offload_time | |
| logger.info(f"[_load_model_context] Offloaded {model_name} to CPU in {offload_time:.4f}s") | |