kenyonxu
7fb8d794e8
download_image() wrote directly to dest_path without ensuring the parent directory exists. This caused FileNotFoundError (caught by OSError handler) when the per-hash image subdirectory under cache/civitai/ had not been created yet, silently failing all image downloads. Add os.makedirs(..., exist_ok=True) before open() to fix this.
107 lines
3.0 KiB
Python
107 lines
3.0 KiB
Python
import asyncio
|
|
import hashlib
|
|
import os
|
|
|
|
import aiohttp
|
|
|
|
from .config import get_proxy_url, get_timeout
|
|
|
|
# Module-level in-memory hash cache: filepath:mtime -> hash
|
|
_hash_cache: dict[str, str] = {}
|
|
|
|
# Shared aiohttp session
|
|
_session: aiohttp.ClientSession | None = None
|
|
|
|
CIVITAI_API_URL = "https://civitai.com/api/v1/model-versions/by-hash"
|
|
|
|
|
|
def get_session() -> aiohttp.ClientSession:
|
|
"""Get or create a shared aiohttp ClientSession.
|
|
|
|
Proxy is read from config.json or environment variables.
|
|
"""
|
|
global _session
|
|
if _session is None or _session.closed:
|
|
_session = aiohttp.ClientSession(
|
|
timeout=aiohttp.ClientTimeout(total=get_timeout()),
|
|
trust_env=True,
|
|
)
|
|
return _session
|
|
|
|
|
|
def _hash_file(path: str) -> str:
|
|
"""Shared SHA256 hashing logic (synchronous)."""
|
|
h = hashlib.sha256()
|
|
with open(path, "rb") as f:
|
|
for chunk in iter(lambda: f.read(8192), b""):
|
|
h.update(chunk)
|
|
return h.hexdigest()
|
|
|
|
|
|
def compute_file_hash_sync(path: str) -> str:
|
|
"""Compute SHA256 synchronously. Used by node execution path (runs in worker thread)."""
|
|
mtime = os.path.getmtime(path)
|
|
cache_key = f"{path}:{mtime}"
|
|
if cache_key in _hash_cache:
|
|
return _hash_cache[cache_key]
|
|
file_hash = _hash_file(path)
|
|
_hash_cache[cache_key] = file_hash
|
|
return file_hash
|
|
|
|
|
|
async def compute_file_hash(path: str) -> str:
|
|
"""Compute SHA256 asynchronously. Uses in-memory cache."""
|
|
mtime = os.path.getmtime(path)
|
|
cache_key = f"{path}:{mtime}"
|
|
if cache_key in _hash_cache:
|
|
return _hash_cache[cache_key]
|
|
|
|
file_hash = await asyncio.to_thread(_hash_file, path)
|
|
_hash_cache[cache_key] = file_hash
|
|
return file_hash
|
|
|
|
|
|
async def fetch_model_by_hash(file_hash: str) -> dict:
|
|
"""Fetch model metadata from Civitai API by SHA256 hash.
|
|
|
|
Returns:
|
|
dict: The parsed JSON response from Civitai.
|
|
|
|
Raises:
|
|
aiohttp.ClientResponseError: On HTTP error (404, 429, etc.)
|
|
aiohttp.ClientError: On network/connection errors
|
|
"""
|
|
session = get_session()
|
|
url = f"{CIVITAI_API_URL}/{file_hash}"
|
|
proxy = get_proxy_url()
|
|
kwargs = {}
|
|
if proxy:
|
|
kwargs["proxy"] = proxy
|
|
async with session.get(url, **kwargs) as response:
|
|
response.raise_for_status()
|
|
data = await response.json()
|
|
return data
|
|
|
|
|
|
async def download_image(url: str, dest_path: str) -> bool:
|
|
"""Download an image from URL to dest_path.
|
|
|
|
Returns True on success, False on failure.
|
|
"""
|
|
session = get_session()
|
|
proxy = get_proxy_url()
|
|
kwargs = {}
|
|
if proxy:
|
|
kwargs["proxy"] = proxy
|
|
try:
|
|
async with session.get(url, **kwargs) as response:
|
|
if response.status != 200:
|
|
return False
|
|
os.makedirs(os.path.dirname(dest_path), exist_ok=True)
|
|
with open(dest_path, "wb") as f:
|
|
async for chunk in response.content.iter_chunked(8192):
|
|
f.write(chunk)
|
|
return True
|
|
except (aiohttp.ClientError, OSError):
|
|
return False
|