alphard/ComfyUI-Lora-Manager
1
0
Fork 0
mirror of https://github.com/willmiao/ComfyUI-Lora-Manager.git synced 2026-03-25 15:15:44 -03:00
Code Issues Packages Projects Releases Wiki Activity

feat: Optimize session management and enhance download functionality with resumable support

Browse Source
This commit is contained in:
Will Miao
2025-04-13 21:51:21 +08:00
parent 76fc9e5a3d
commit 40381d4b11
1 changed files with 35 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

41
py/services/civitai_client.py
View File

@@ -34,14 +34,29 @@ class CivitaiClient:
'User-Agent': 'ComfyUI-LoRA-Manager/1.0' 'User-Agent': 'ComfyUI-LoRA-Manager/1.0'
} }
self._session = None self._session = None
# Set default buffer size to 1MB for higher throughput
self.chunk_size = 1024 * 1024
@property @property
async def session(self) -> aiohttp.ClientSession: async def session(self) -> aiohttp.ClientSession:
"""Lazy initialize the session""" """Lazy initialize the session"""
if self._session is None: if self._session is None:
connector = aiohttp.TCPConnector(ssl=True) # Optimize TCP connection parameters
trust_env = True # 允许使用系统环境变量中的代理设置 connector = aiohttp.TCPConnector(
self._session = aiohttp.ClientSession(connector=connector, trust_env=trust_env) ssl=True,
limit=10, # Increase parallel connections
ttl_dns_cache=300, # DNS cache time
force_close=False, # Keep connections for reuse
enable_cleanup_closed=True
)
trust_env = True # Allow using system environment proxy settings
# Configure timeout parameters
timeout = aiohttp.ClientTimeout(total=None, connect=60, sock_read=60)
self._session = aiohttp.ClientSession(
connector=connector,
trust_env=trust_env,
timeout=timeout
)
return self._session return self._session
def _parse_content_disposition(self, header: str) -> str: def _parse_content_disposition(self, header: str) -> str:
@@ -91,6 +106,10 @@ class CivitaiClient:
session = await self.session session = await self.session
try: try:
headers = self._get_request_headers() headers = self._get_request_headers()
# Add Range header to allow resumable downloads
headers['Accept-Encoding'] = 'identity' # Disable compression for better chunked downloads
async with session.get(url, headers=headers, allow_redirects=True) as response: async with session.get(url, headers=headers, allow_redirects=True) as response:
if response.status != 200: if response.status != 200:
# Handle 401 unauthorized responses # Handle 401 unauthorized responses
@@ -118,16 +137,23 @@ class CivitaiClient:
# Get total file size for progress calculation # Get total file size for progress calculation
total_size = int(response.headers.get('content-length', 0)) total_size = int(response.headers.get('content-length', 0))
current_size = 0 current_size = 0
last_progress_report_time = datetime.now()
# Stream download to file with progress updates # Stream download to file with progress updates using larger buffer
with open(save_path, 'wb') as f: with open(save_path, 'wb') as f:
async for chunk in response.content.iter_chunked(8192): async for chunk in response.content.iter_chunked(self.chunk_size):
if chunk: if chunk:
f.write(chunk) f.write(chunk)
current_size += len(chunk) current_size += len(chunk)
if progress_callback and total_size:
# Limit progress update frequency to reduce overhead
now = datetime.now()
time_diff = (now - last_progress_report_time).total_seconds()
if progress_callback and total_size and time_diff >= 0.5:
progress = (current_size / total_size) * 100 progress = (current_size / total_size) * 100
await progress_callback(progress) await progress_callback(progress)
last_progress_report_time = now
# Ensure 100% progress is reported # Ensure 100% progress is reported
if progress_callback: if progress_callback:
@@ -135,6 +161,9 @@ class CivitaiClient:
return True, save_path return True, save_path
except aiohttp.ClientError as e:
logger.error(f"Network error during download: {e}")
return False, f"Network error: {str(e)}"
except Exception as e: except Exception as e:
logger.error(f"Download error: {e}") logger.error(f"Download error: {e}")
return False, str(e) return False, str(e)