mirror of
https://github.com/jamiepine/voicebox.git
synced 2026-09-19 14:50:38 -07:00
chore(backend): repair test suite and bring ruff to green
The suite hadn't run green since the routes refactor: - test_profile_duplicate_names.py imported the pre-refactor module layout and broke collection; now imports backend.services.profiles - tests/conftest.py puts the repo root and backend dir on sys.path so files collect standalone instead of depending on run order - test_cors.py tested a hand-copied mirror of the origin list that had drifted from app.py (missing http://tauri.localhost); it now builds the app via the real create_app() factory - test_progress.py simulated a 1KB download, below the tracker's 1MB reporting threshold; simulation raised to 5MB - slow/timeout markers registered in pyproject Ruff: ~900 violations auto-fixed (typing modernization, import sorting, unused imports, whitespace). The remaining rules are baselined in pyproject.toml with per-rule counts to burn down, plus per-file carve-outs for deliberate env-before-import ordering. ruff check is now clean; suite is 134 passed, 2 skipped.
This commit is contained in:
+16
-16
@@ -2,10 +2,10 @@
|
||||
Audio processing utilities.
|
||||
"""
|
||||
|
||||
|
||||
import librosa
|
||||
import numpy as np
|
||||
import soundfile as sf
|
||||
import librosa
|
||||
from typing import Tuple, Optional
|
||||
|
||||
|
||||
def normalize_audio(
|
||||
@@ -15,32 +15,32 @@ def normalize_audio(
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
Normalize audio to target loudness with peak limiting.
|
||||
|
||||
|
||||
Args:
|
||||
audio: Input audio array
|
||||
target_db: Target RMS level in dB
|
||||
peak_limit: Peak limit (0.0-1.0)
|
||||
|
||||
|
||||
Returns:
|
||||
Normalized audio array
|
||||
"""
|
||||
# Convert to float32
|
||||
audio = audio.astype(np.float32)
|
||||
|
||||
|
||||
# Calculate current RMS
|
||||
rms = np.sqrt(np.mean(audio**2))
|
||||
|
||||
|
||||
# Calculate target RMS
|
||||
target_rms = 10**(target_db / 20)
|
||||
|
||||
|
||||
# Apply gain
|
||||
if rms > 0:
|
||||
gain = target_rms / rms
|
||||
audio = audio * gain
|
||||
|
||||
|
||||
# Peak limiting
|
||||
audio = np.clip(audio, -peak_limit, peak_limit)
|
||||
|
||||
|
||||
return audio
|
||||
|
||||
|
||||
@@ -48,15 +48,15 @@ def load_audio(
|
||||
path: str,
|
||||
sample_rate: int = 24000,
|
||||
mono: bool = True,
|
||||
) -> Tuple[np.ndarray, int]:
|
||||
) -> tuple[np.ndarray, int]:
|
||||
"""
|
||||
Load audio file with normalization.
|
||||
|
||||
|
||||
Args:
|
||||
path: Path to audio file
|
||||
sample_rate: Target sample rate
|
||||
mono: Convert to mono
|
||||
|
||||
|
||||
Returns:
|
||||
Tuple of (audio_array, sample_rate)
|
||||
"""
|
||||
@@ -84,8 +84,8 @@ def save_audio(
|
||||
Raises:
|
||||
OSError: If file cannot be written
|
||||
"""
|
||||
from pathlib import Path
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
temp_path = f"{path}.tmp"
|
||||
try:
|
||||
@@ -264,7 +264,7 @@ def validate_reference_audio(
|
||||
min_duration: float = 2.0,
|
||||
max_duration: float = 30.0,
|
||||
min_rms: float = 0.01,
|
||||
) -> Tuple[bool, Optional[str]]:
|
||||
) -> tuple[bool, str | None]:
|
||||
"""
|
||||
Validate reference audio for voice cloning.
|
||||
|
||||
@@ -288,7 +288,7 @@ def validate_and_load_reference_audio(
|
||||
min_duration: float = 2.0,
|
||||
max_duration: float = 30.0,
|
||||
min_rms: float = 0.01,
|
||||
) -> Tuple[bool, Optional[str], Optional[np.ndarray], Optional[int]]:
|
||||
) -> tuple[bool, str | None, np.ndarray | None, int | None]:
|
||||
"""
|
||||
Validate and load reference audio in a single pass.
|
||||
|
||||
@@ -315,4 +315,4 @@ def validate_and_load_reference_audio(
|
||||
|
||||
return True, None, audio, sr
|
||||
except Exception as e:
|
||||
return False, f"Error validating audio: {str(e)}", None, None
|
||||
return False, f"Error validating audio: {e!s}", None, None
|
||||
|
||||
@@ -4,9 +4,10 @@ Voice prompt caching utilities.
|
||||
|
||||
import hashlib
|
||||
import logging
|
||||
import torch
|
||||
from pathlib import Path
|
||||
from typing import Optional, Union, Dict, Any
|
||||
from typing import Any, Union
|
||||
|
||||
import torch
|
||||
|
||||
from .. import config
|
||||
|
||||
@@ -19,7 +20,7 @@ def _get_cache_dir() -> Path:
|
||||
|
||||
|
||||
# In-memory cache - can store dict (voice prompt) or tensor (legacy)
|
||||
_memory_cache: dict[str, Union[torch.Tensor, Dict[str, Any]]] = {}
|
||||
_memory_cache: dict[str, Union[torch.Tensor, dict[str, Any]]] = {}
|
||||
|
||||
|
||||
def get_cache_key(audio_path: str, reference_text: str) -> str:
|
||||
@@ -46,7 +47,7 @@ def get_cache_key(audio_path: str, reference_text: str) -> str:
|
||||
|
||||
def get_cached_voice_prompt(
|
||||
cache_key: str,
|
||||
) -> Optional[Union[torch.Tensor, Dict[str, Any]]]:
|
||||
) -> Union[torch.Tensor, dict[str, Any]] | None:
|
||||
"""
|
||||
Get cached voice prompt if available.
|
||||
|
||||
@@ -76,7 +77,7 @@ def get_cached_voice_prompt(
|
||||
|
||||
def cache_voice_prompt(
|
||||
cache_key: str,
|
||||
voice_prompt: Union[torch.Tensor, Dict[str, Any]],
|
||||
voice_prompt: Union[torch.Tensor, dict[str, Any]],
|
||||
) -> None:
|
||||
"""
|
||||
Cache voice prompt to memory and disk.
|
||||
|
||||
@@ -4,7 +4,6 @@ from __future__ import annotations
|
||||
|
||||
import sys
|
||||
|
||||
|
||||
MAC_PUSH_TO_TALK = ["MetaRight", "AltGr"]
|
||||
MAC_TOGGLE_TO_TALK = ["MetaRight", "AltGr", "Space"]
|
||||
NON_MAC_PUSH_TO_TALK = ["ControlRight", "ShiftRight"]
|
||||
|
||||
@@ -11,7 +11,6 @@ overhead.
|
||||
|
||||
import logging
|
||||
import re
|
||||
from typing import List, Tuple
|
||||
|
||||
import numpy as np
|
||||
|
||||
@@ -58,7 +57,7 @@ _ABBREVIATIONS = frozenset(
|
||||
_PARA_TAG_RE = re.compile(r"\[[^\]]*\]")
|
||||
|
||||
|
||||
def split_text_into_chunks(text: str, max_chars: int = DEFAULT_MAX_CHUNK_CHARS) -> List[str]:
|
||||
def split_text_into_chunks(text: str, max_chars: int = DEFAULT_MAX_CHUNK_CHARS) -> list[str]:
|
||||
"""Split *text* at natural boundaries into chunks of at most *max_chars*.
|
||||
|
||||
Priority: sentence-end (``.!?`` not preceded by an abbreviation and not
|
||||
@@ -73,7 +72,7 @@ def split_text_into_chunks(text: str, max_chars: int = DEFAULT_MAX_CHUNK_CHARS)
|
||||
if len(text) <= max_chars:
|
||||
return [text]
|
||||
|
||||
chunks: List[str] = []
|
||||
chunks: list[str] = []
|
||||
remaining = text
|
||||
|
||||
while remaining:
|
||||
@@ -170,7 +169,7 @@ def _safe_hard_cut(segment: str, max_chars: int) -> int:
|
||||
|
||||
|
||||
def concatenate_audio_chunks(
|
||||
chunks: List[np.ndarray],
|
||||
chunks: list[np.ndarray],
|
||||
sample_rate: int,
|
||||
crossfade_ms: int = 50,
|
||||
) -> np.ndarray:
|
||||
@@ -211,7 +210,7 @@ async def generate_chunked(
|
||||
max_chunk_chars: int = DEFAULT_MAX_CHUNK_CHARS,
|
||||
crossfade_ms: int = 50,
|
||||
trim_fn=None,
|
||||
) -> Tuple[np.ndarray, int]:
|
||||
) -> tuple[np.ndarray, int]:
|
||||
"""Generate audio with automatic chunking for long text.
|
||||
|
||||
For text shorter than *max_chunk_chars* this is a thin wrapper around
|
||||
@@ -266,7 +265,7 @@ async def generate_chunked(
|
||||
len(chunks),
|
||||
max_chunk_chars,
|
||||
)
|
||||
audio_chunks: List[np.ndarray] = []
|
||||
audio_chunks: list[np.ndarray] = []
|
||||
sample_rate: int | None = None
|
||||
|
||||
for i, chunk_text in enumerate(chunks):
|
||||
|
||||
@@ -21,7 +21,6 @@ import types
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
# ── Snake activation (from dac/nn/layers.py) ────────────────────────
|
||||
|
||||
# NOTE: The original DAC code uses @torch.jit.script here for a 1.4x
|
||||
|
||||
+12
-13
@@ -19,24 +19,23 @@ Supported effect types:
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import numpy as np
|
||||
from typing import Any, Dict, List, Optional
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
from pedalboard import (
|
||||
Pedalboard,
|
||||
Chorus,
|
||||
Reverb,
|
||||
Compressor,
|
||||
Delay,
|
||||
Gain,
|
||||
HighpassFilter,
|
||||
LowpassFilter,
|
||||
Delay,
|
||||
Pedalboard,
|
||||
PitchShift,
|
||||
Reverb,
|
||||
)
|
||||
|
||||
|
||||
# Each param definition: (default, min, max, description)
|
||||
EFFECT_REGISTRY: Dict[str, Dict[str, Any]] = {
|
||||
EFFECT_REGISTRY: dict[str, dict[str, Any]] = {
|
||||
"chorus": {
|
||||
"cls": Chorus,
|
||||
"label": "Chorus / Flanger",
|
||||
@@ -147,7 +146,7 @@ EFFECT_REGISTRY: Dict[str, Dict[str, Any]] = {
|
||||
}
|
||||
|
||||
|
||||
BUILTIN_PRESETS: Dict[str, Dict[str, Any]] = {
|
||||
BUILTIN_PRESETS: dict[str, dict[str, Any]] = {
|
||||
"robotic": {
|
||||
"name": "Robotic",
|
||||
"sort_order": 0,
|
||||
@@ -255,7 +254,7 @@ BUILTIN_PRESETS: Dict[str, Dict[str, Any]] = {
|
||||
}
|
||||
|
||||
|
||||
def get_available_effects() -> List[Dict[str, Any]]:
|
||||
def get_available_effects() -> list[dict[str, Any]]:
|
||||
"""Return the list of available effect types with their parameter definitions.
|
||||
|
||||
Used by the frontend to build the effects chain editor UI.
|
||||
@@ -273,12 +272,12 @@ def get_available_effects() -> List[Dict[str, Any]]:
|
||||
return result
|
||||
|
||||
|
||||
def get_builtin_presets() -> Dict[str, Dict[str, Any]]:
|
||||
def get_builtin_presets() -> dict[str, dict[str, Any]]:
|
||||
"""Return all built-in effect presets."""
|
||||
return BUILTIN_PRESETS
|
||||
|
||||
|
||||
def validate_effects_chain(effects_chain: List[Dict[str, Any]]) -> Optional[str]:
|
||||
def validate_effects_chain(effects_chain: list[dict[str, Any]]) -> str | None:
|
||||
"""Validate an effects chain configuration.
|
||||
|
||||
Returns None if valid, or an error message string.
|
||||
@@ -315,7 +314,7 @@ def validate_effects_chain(effects_chain: List[Dict[str, Any]]) -> Optional[str]
|
||||
return None
|
||||
|
||||
|
||||
def build_pedalboard(effects_chain: List[Dict[str, Any]]) -> Pedalboard:
|
||||
def build_pedalboard(effects_chain: list[dict[str, Any]]) -> Pedalboard:
|
||||
"""Build a Pedalboard instance from an effects chain config.
|
||||
|
||||
Skips effects where ``enabled`` is ``False``.
|
||||
@@ -342,7 +341,7 @@ def build_pedalboard(effects_chain: List[Dict[str, Any]]) -> Pedalboard:
|
||||
def apply_effects(
|
||||
audio: np.ndarray,
|
||||
sample_rate: int,
|
||||
effects_chain: List[Dict[str, Any]],
|
||||
effects_chain: list[dict[str, Any]],
|
||||
) -> np.ndarray:
|
||||
"""Apply an effects chain to audio data.
|
||||
|
||||
|
||||
@@ -9,7 +9,7 @@ import os
|
||||
import threading
|
||||
from contextlib import contextmanager
|
||||
from pathlib import Path
|
||||
from typing import Optional, Union
|
||||
from typing import Union
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -25,9 +25,9 @@ logger = logging.getLogger(__name__)
|
||||
|
||||
_offline_lock = threading.RLock()
|
||||
_offline_refcount = 0
|
||||
_saved_env: Optional[str] = None
|
||||
_saved_hf_const: Optional[bool] = None
|
||||
_saved_transformers_const: Optional[bool] = None
|
||||
_saved_env: str | None = None
|
||||
_saved_hf_const: bool | None = None
|
||||
_saved_transformers_const: bool | None = None
|
||||
|
||||
|
||||
@contextmanager
|
||||
@@ -61,8 +61,8 @@ def force_offline_if_cached(is_cached: bool, model_label: str = ""):
|
||||
# bumping the refcount — a persistent offline leak that outlives
|
||||
# the process and is miserable to debug.
|
||||
prev_env = os.environ.get("HF_HUB_OFFLINE")
|
||||
prev_hf: Optional[bool] = None
|
||||
prev_tf: Optional[bool] = None
|
||||
prev_hf: bool | None = None
|
||||
prev_tf: bool | None = None
|
||||
try:
|
||||
try:
|
||||
import huggingface_hub.constants as hf_const
|
||||
@@ -206,8 +206,8 @@ def patch_huggingface_hub_offline():
|
||||
repo_id: str,
|
||||
filename: str,
|
||||
cache_dir: Union[str, Path, None] = None,
|
||||
revision: Optional[str] = None,
|
||||
repo_type: Optional[str] = None,
|
||||
revision: str | None = None,
|
||||
repo_type: str | None = None,
|
||||
):
|
||||
result = original_try_load(
|
||||
repo_id=repo_id,
|
||||
|
||||
@@ -2,11 +2,11 @@
|
||||
HuggingFace Hub download progress tracking.
|
||||
"""
|
||||
|
||||
from typing import Optional, Callable
|
||||
from contextlib import contextmanager
|
||||
import logging
|
||||
import threading
|
||||
import sys
|
||||
import threading
|
||||
from collections.abc import Callable
|
||||
from contextlib import contextmanager
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -14,7 +14,7 @@ logger = logging.getLogger(__name__)
|
||||
class HFProgressTracker:
|
||||
"""Tracks HuggingFace Hub download progress by intercepting tqdm."""
|
||||
|
||||
def __init__(self, progress_callback: Optional[Callable] = None, filter_non_downloads: bool = False):
|
||||
def __init__(self, progress_callback: Callable | None = None, filter_non_downloads: bool = False):
|
||||
self.progress_callback = progress_callback
|
||||
self.filter_non_downloads = filter_non_downloads # Only filter if True
|
||||
self._original_tqdm_class = None
|
||||
|
||||
+20
-20
@@ -1,7 +1,7 @@
|
||||
"""Image processing utilities for avatar uploads."""
|
||||
|
||||
from pathlib import Path
|
||||
from typing import Optional, Tuple
|
||||
|
||||
from PIL import Image
|
||||
|
||||
# JPEG can be reported as 'JPEG' or 'MPO' (for multi-picture format from some cameras)
|
||||
@@ -10,46 +10,46 @@ MAX_SIZE = 512
|
||||
MAX_FILE_SIZE = 5 * 1024 * 1024 # 5MB
|
||||
|
||||
|
||||
def validate_image(file_path: str) -> Tuple[bool, Optional[str]]:
|
||||
def validate_image(file_path: str) -> tuple[bool, str | None]:
|
||||
"""
|
||||
Validate image format and file size.
|
||||
|
||||
|
||||
Args:
|
||||
file_path: Path to image file
|
||||
|
||||
|
||||
Returns:
|
||||
Tuple of (is_valid, error_message)
|
||||
"""
|
||||
path = Path(file_path)
|
||||
|
||||
|
||||
# Check file size
|
||||
if path.stat().st_size > MAX_FILE_SIZE:
|
||||
return False, f"File size exceeds maximum of {MAX_FILE_SIZE // (1024 * 1024)}MB"
|
||||
|
||||
|
||||
try:
|
||||
with Image.open(file_path) as img:
|
||||
# Verify the image can be loaded
|
||||
img.load()
|
||||
|
||||
|
||||
# Check format (normalize JPEG variants)
|
||||
img_format = img.format
|
||||
if img_format in ('MPO', 'JPG'):
|
||||
img_format = 'JPEG'
|
||||
|
||||
|
||||
if img_format not in {'PNG', 'JPEG', 'WEBP'}:
|
||||
return False, f"Invalid format '{img_format}'. Allowed formats: PNG, JPEG, WEBP"
|
||||
|
||||
|
||||
return True, None
|
||||
except Exception as e:
|
||||
return False, f"Invalid image file: {str(e)}"
|
||||
return False, f"Invalid image file: {e!s}"
|
||||
|
||||
|
||||
def process_avatar(input_path: str, output_path: str, max_size: int = MAX_SIZE) -> None:
|
||||
"""
|
||||
Process avatar image: resize and optimize.
|
||||
|
||||
|
||||
Resizes image to fit within max_size x max_size while maintaining aspect ratio.
|
||||
|
||||
|
||||
Args:
|
||||
input_path: Path to input image
|
||||
output_path: Path to save processed image
|
||||
@@ -59,7 +59,7 @@ def process_avatar(input_path: str, output_path: str, max_size: int = MAX_SIZE)
|
||||
# Handle EXIF orientation for JPEG images
|
||||
try:
|
||||
from PIL import ExifTags
|
||||
for orientation in ExifTags.TAGS.keys():
|
||||
for orientation in ExifTags.TAGS:
|
||||
if ExifTags.TAGS[orientation] == 'Orientation':
|
||||
break
|
||||
exif = img._getexif()
|
||||
@@ -74,7 +74,7 @@ def process_avatar(input_path: str, output_path: str, max_size: int = MAX_SIZE)
|
||||
except (AttributeError, KeyError, IndexError, TypeError):
|
||||
# No EXIF data or orientation tag
|
||||
pass
|
||||
|
||||
|
||||
# Convert to RGB if necessary (handles RGBA, P, CMYK, etc.)
|
||||
if img.mode not in ('RGB', 'L'):
|
||||
if img.mode == 'RGBA':
|
||||
@@ -90,25 +90,25 @@ def process_avatar(input_path: str, output_path: str, max_size: int = MAX_SIZE)
|
||||
img = img.convert('RGB')
|
||||
else:
|
||||
img = img.convert('RGB')
|
||||
|
||||
|
||||
# Calculate new size maintaining aspect ratio
|
||||
img.thumbnail((max_size, max_size), Image.Resampling.LANCZOS)
|
||||
|
||||
|
||||
# Determine output format from extension
|
||||
output_ext = Path(output_path).suffix.lower()
|
||||
|
||||
|
||||
format_map = {
|
||||
'.png': 'PNG',
|
||||
'.jpeg': 'JPEG',
|
||||
'.jpg': 'JPEG',
|
||||
'.webp': 'WEBP'
|
||||
}
|
||||
|
||||
|
||||
output_format = format_map.get(output_ext, 'PNG')
|
||||
|
||||
|
||||
# Save with optimization
|
||||
save_kwargs = {'optimize': True}
|
||||
if output_format == 'JPEG':
|
||||
save_kwargs['quality'] = 90
|
||||
|
||||
|
||||
img.save(output_path, format=output_format, **save_kwargs)
|
||||
|
||||
+39
-41
@@ -2,8 +2,6 @@
|
||||
Progress tracking for model downloads using Server-Sent Events.
|
||||
"""
|
||||
|
||||
from typing import Optional, Callable, Dict, List
|
||||
from fastapi.responses import StreamingResponse
|
||||
import asyncio
|
||||
import json
|
||||
import threading
|
||||
@@ -12,34 +10,34 @@ from datetime import datetime
|
||||
|
||||
class ProgressManager:
|
||||
"""Manages download progress for multiple models.
|
||||
|
||||
|
||||
Thread-safe: can be called from background threads (e.g., via asyncio.to_thread).
|
||||
"""
|
||||
|
||||
|
||||
# Throttle settings to prevent overwhelming SSE clients
|
||||
THROTTLE_INTERVAL_SECONDS = 0.5 # Minimum time between updates
|
||||
THROTTLE_PROGRESS_DELTA = 1.0 # Minimum progress change (%) to force update
|
||||
|
||||
|
||||
def __init__(self):
|
||||
self._progress: Dict[str, Dict] = {}
|
||||
self._listeners: Dict[str, list] = {}
|
||||
self._progress: dict[str, dict] = {}
|
||||
self._listeners: dict[str, list] = {}
|
||||
self._lock = threading.Lock() # Thread-safe lock for progress dict
|
||||
self._main_loop: Optional[asyncio.AbstractEventLoop] = None
|
||||
self._last_notify_time: Dict[str, float] = {} # Last notification time per model
|
||||
self._last_notify_progress: Dict[str, float] = {} # Last notified progress per model
|
||||
|
||||
self._main_loop: asyncio.AbstractEventLoop | None = None
|
||||
self._last_notify_time: dict[str, float] = {} # Last notification time per model
|
||||
self._last_notify_progress: dict[str, float] = {} # Last notified progress per model
|
||||
|
||||
def _set_main_loop(self, loop: asyncio.AbstractEventLoop):
|
||||
"""Set the main event loop for thread-safe operations."""
|
||||
self._main_loop = loop
|
||||
|
||||
def _notify_listeners_threadsafe(self, model_name: str, progress_data: Dict):
|
||||
|
||||
def _notify_listeners_threadsafe(self, model_name: str, progress_data: dict):
|
||||
"""Notify listeners in a thread-safe manner."""
|
||||
import logging
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
if model_name not in self._listeners:
|
||||
return
|
||||
|
||||
|
||||
for queue in self._listeners[model_name]:
|
||||
try:
|
||||
# Check if we're in the main event loop thread
|
||||
@@ -66,14 +64,14 @@ class ProgressManager:
|
||||
model_name: str,
|
||||
current: int,
|
||||
total: int,
|
||||
filename: Optional[str] = None,
|
||||
filename: str | None = None,
|
||||
status: str = "downloading",
|
||||
):
|
||||
"""
|
||||
Update progress for a model download.
|
||||
|
||||
Thread-safe: can be called from background threads.
|
||||
|
||||
|
||||
Progress updates are throttled to prevent overwhelming SSE clients.
|
||||
Updates are sent at most every THROTTLE_INTERVAL_SECONDS, or when
|
||||
progress changes by at least THROTTLE_PROGRESS_DELTA percent.
|
||||
@@ -116,20 +114,20 @@ class ProgressManager:
|
||||
current_time = time.time()
|
||||
last_time = self._last_notify_time.get(model_name, 0)
|
||||
last_progress = self._last_notify_progress.get(model_name, -100)
|
||||
|
||||
|
||||
time_delta = current_time - last_time
|
||||
progress_delta = abs(progress_pct - last_progress)
|
||||
|
||||
|
||||
# Always notify for complete/error status, or if throttle conditions are met
|
||||
should_notify = (
|
||||
status in ("complete", "error") or
|
||||
time_delta >= self.THROTTLE_INTERVAL_SECONDS or
|
||||
progress_delta >= self.THROTTLE_PROGRESS_DELTA
|
||||
)
|
||||
|
||||
|
||||
if not should_notify:
|
||||
return # Skip this update (throttled)
|
||||
|
||||
|
||||
# Update throttle tracking
|
||||
self._last_notify_time[model_name] = current_time
|
||||
self._last_notify_progress[model_name] = progress_pct
|
||||
@@ -142,14 +140,14 @@ class ProgressManager:
|
||||
self._notify_listeners_threadsafe(model_name, progress_data)
|
||||
else:
|
||||
logger.debug(f"No listeners for {model_name}, progress update stored: {progress_pct:.1f}%")
|
||||
|
||||
def get_progress(self, model_name: str) -> Optional[Dict]:
|
||||
|
||||
def get_progress(self, model_name: str) -> dict | None:
|
||||
"""Get current progress for a model. Thread-safe."""
|
||||
with self._lock:
|
||||
progress = self._progress.get(model_name)
|
||||
return progress.copy() if progress else None
|
||||
|
||||
def get_all_active(self) -> List[Dict]:
|
||||
|
||||
def get_all_active(self) -> list[dict]:
|
||||
"""Get all active downloads (status is 'downloading' or 'extracting'). Thread-safe."""
|
||||
active = []
|
||||
with self._lock:
|
||||
@@ -158,25 +156,25 @@ class ProgressManager:
|
||||
if status in ("downloading", "extracting"):
|
||||
active.append(progress.copy())
|
||||
return active
|
||||
|
||||
def create_progress_callback(self, model_name: str, filename: Optional[str] = None):
|
||||
|
||||
def create_progress_callback(self, model_name: str, filename: str | None = None):
|
||||
"""
|
||||
Create a progress callback function for HuggingFace downloads.
|
||||
|
||||
|
||||
Args:
|
||||
model_name: Name of the model
|
||||
filename: Optional filename filter
|
||||
|
||||
|
||||
Returns:
|
||||
Callback function
|
||||
"""
|
||||
def callback(progress: Dict):
|
||||
def callback(progress: dict):
|
||||
"""HuggingFace Hub progress callback."""
|
||||
if "total" in progress and "current" in progress:
|
||||
current = progress.get("current", 0)
|
||||
total = progress.get("total", 0)
|
||||
file_name = progress.get("filename", filename)
|
||||
|
||||
|
||||
self.update_progress(
|
||||
model_name=model_name,
|
||||
current=current,
|
||||
@@ -184,9 +182,9 @@ class ProgressManager:
|
||||
filename=file_name,
|
||||
status="downloading",
|
||||
)
|
||||
|
||||
|
||||
return callback
|
||||
|
||||
|
||||
async def subscribe(self, model_name: str):
|
||||
"""
|
||||
Subscribe to progress updates for a model.
|
||||
@@ -195,7 +193,7 @@ class ProgressManager:
|
||||
"""
|
||||
import logging
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# Store the main event loop for thread-safe operations
|
||||
try:
|
||||
self._main_loop = asyncio.get_running_loop()
|
||||
@@ -217,7 +215,7 @@ class ProgressManager:
|
||||
initial_progress = self._progress.get(model_name)
|
||||
if initial_progress:
|
||||
initial_progress = initial_progress.copy()
|
||||
|
||||
|
||||
if initial_progress:
|
||||
status = initial_progress.get('status')
|
||||
# Only send initial progress if download is actually in progress
|
||||
@@ -242,7 +240,7 @@ class ProgressManager:
|
||||
if progress.get("status") in ("complete", "error"):
|
||||
logger.info(f"Download {progress.get('status')} for {model_name}, closing SSE connection")
|
||||
break
|
||||
except asyncio.TimeoutError:
|
||||
except TimeoutError:
|
||||
# Send heartbeat
|
||||
yield ": heartbeat\n\n"
|
||||
continue
|
||||
@@ -255,7 +253,7 @@ class ProgressManager:
|
||||
if not self._listeners[model_name]:
|
||||
del self._listeners[model_name]
|
||||
logger.info(f"SSE client unsubscribed from {model_name}, remaining listeners: {len(self._listeners.get(model_name, []))}")
|
||||
|
||||
|
||||
def mark_complete(self, model_name: str):
|
||||
"""Mark a model download as complete. Thread-safe."""
|
||||
import logging
|
||||
@@ -269,11 +267,11 @@ class ProgressManager:
|
||||
else:
|
||||
logger.warning(f"Cannot mark {model_name} as complete: not found in progress")
|
||||
return
|
||||
|
||||
|
||||
logger.info(f"Marked {model_name} as complete")
|
||||
# Notify listeners (thread-safe)
|
||||
self._notify_listeners_threadsafe(model_name, progress_data)
|
||||
|
||||
|
||||
def mark_error(self, model_name: str, error: str):
|
||||
"""Mark a model download as failed. Thread-safe."""
|
||||
import logging
|
||||
@@ -297,14 +295,14 @@ class ProgressManager:
|
||||
"timestamp": datetime.now().isoformat(),
|
||||
}
|
||||
self._progress[model_name] = progress_data
|
||||
|
||||
|
||||
logger.error(f"Marked {model_name} as error: {error}")
|
||||
# Notify listeners (thread-safe)
|
||||
self._notify_listeners_threadsafe(model_name, progress_data)
|
||||
|
||||
|
||||
# Global progress manager instance
|
||||
_progress_manager: Optional[ProgressManager] = None
|
||||
_progress_manager: ProgressManager | None = None
|
||||
|
||||
|
||||
def get_progress_manager() -> ProgressManager:
|
||||
|
||||
+17
-18
@@ -2,9 +2,8 @@
|
||||
Task tracking for active downloads and generations.
|
||||
"""
|
||||
|
||||
from typing import Optional, Dict, List
|
||||
from datetime import datetime
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -13,7 +12,7 @@ class DownloadTask:
|
||||
model_name: str
|
||||
status: str = "downloading" # downloading, extracting, complete, error
|
||||
started_at: datetime = field(default_factory=datetime.utcnow)
|
||||
error: Optional[str] = None
|
||||
error: str | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -27,29 +26,29 @@ class GenerationTask:
|
||||
|
||||
class TaskManager:
|
||||
"""Manages active downloads and generations."""
|
||||
|
||||
|
||||
def __init__(self):
|
||||
self._active_downloads: Dict[str, DownloadTask] = {}
|
||||
self._active_generations: Dict[str, GenerationTask] = {}
|
||||
|
||||
self._active_downloads: dict[str, DownloadTask] = {}
|
||||
self._active_generations: dict[str, GenerationTask] = {}
|
||||
|
||||
def start_download(self, model_name: str) -> None:
|
||||
"""Mark a download as started."""
|
||||
self._active_downloads[model_name] = DownloadTask(
|
||||
model_name=model_name,
|
||||
status="downloading",
|
||||
)
|
||||
|
||||
|
||||
def complete_download(self, model_name: str) -> None:
|
||||
"""Mark a download as complete."""
|
||||
if model_name in self._active_downloads:
|
||||
del self._active_downloads[model_name]
|
||||
|
||||
|
||||
def error_download(self, model_name: str, error: str) -> None:
|
||||
"""Mark a download as failed."""
|
||||
if model_name in self._active_downloads:
|
||||
self._active_downloads[model_name].status = "error"
|
||||
self._active_downloads[model_name].error = error
|
||||
|
||||
|
||||
def start_generation(self, task_id: str, profile_id: str, text: str) -> None:
|
||||
"""Mark a generation as started."""
|
||||
text_preview = text[:50] + "..." if len(text) > 50 else text
|
||||
@@ -58,20 +57,20 @@ class TaskManager:
|
||||
profile_id=profile_id,
|
||||
text_preview=text_preview,
|
||||
)
|
||||
|
||||
|
||||
def complete_generation(self, task_id: str) -> None:
|
||||
"""Mark a generation as complete."""
|
||||
if task_id in self._active_generations:
|
||||
del self._active_generations[task_id]
|
||||
|
||||
def get_active_downloads(self) -> List[DownloadTask]:
|
||||
|
||||
def get_active_downloads(self) -> list[DownloadTask]:
|
||||
"""Get all active downloads."""
|
||||
return list(self._active_downloads.values())
|
||||
|
||||
def get_active_generations(self) -> List[GenerationTask]:
|
||||
|
||||
def get_active_generations(self) -> list[GenerationTask]:
|
||||
"""Get all active generations."""
|
||||
return list(self._active_generations.values())
|
||||
|
||||
|
||||
def cancel_download(self, model_name: str) -> bool:
|
||||
"""Cancel/dismiss a download task (removes it from active list)."""
|
||||
return self._active_downloads.pop(model_name, None) is not None
|
||||
@@ -84,14 +83,14 @@ class TaskManager:
|
||||
def is_download_active(self, model_name: str) -> bool:
|
||||
"""Check if a download is active."""
|
||||
return model_name in self._active_downloads
|
||||
|
||||
|
||||
def is_generation_active(self, task_id: str) -> bool:
|
||||
"""Check if a generation is active."""
|
||||
return task_id in self._active_generations
|
||||
|
||||
|
||||
# Global task manager instance
|
||||
_task_manager: Optional[TaskManager] = None
|
||||
_task_manager: TaskManager | None = None
|
||||
|
||||
|
||||
def get_task_manager() -> TaskManager:
|
||||
|
||||
Reference in New Issue
Block a user