301 lines
12 KiB
Python
301 lines
12 KiB
Python
"""
|
|
LoRA file download and cache management.
|
|
"""
|
|
import os
|
|
import json
|
|
import urllib.request
|
|
import ssl
|
|
from typing import Optional, Dict, Tuple
|
|
|
|
import folder_paths
|
|
from .civitai_api import get_civitai_api
|
|
|
|
|
|
class LoRACacheManager:
|
|
"""Manages downloading and caching of LoRA files."""
|
|
|
|
def __init__(self, cache_dir: str = None):
|
|
if cache_dir is None:
|
|
# Use ComfyUI's standard loras folder
|
|
lora_paths = folder_paths.get_folder_paths("loras")
|
|
if lora_paths:
|
|
cache_dir = lora_paths[0] # Use the first/primary loras folder
|
|
else:
|
|
# Fallback to default location
|
|
cache_dir = os.path.join(folder_paths.models_dir, "loras")
|
|
self.cache_dir = cache_dir
|
|
# Store manifest in plugin directory to track which files we downloaded
|
|
self.manifest_file = os.path.join(os.path.dirname(__file__), "cache", "manifest.json")
|
|
os.makedirs(cache_dir, exist_ok=True)
|
|
os.makedirs(os.path.dirname(self.manifest_file), exist_ok=True)
|
|
|
|
self._manifest = self._load_manifest()
|
|
self._download_progress: Dict[int, float] = {}
|
|
|
|
def _load_manifest(self) -> Dict:
|
|
"""Load the manifest of downloaded LoRAs."""
|
|
try:
|
|
if os.path.exists(self.manifest_file):
|
|
with open(self.manifest_file, 'r', encoding='utf-8') as f:
|
|
return json.load(f)
|
|
except Exception as e:
|
|
print(f"[SEngine] Error loading manifest: {e}")
|
|
return {"files": {}}
|
|
|
|
def _save_manifest(self):
|
|
"""Save the manifest to disk."""
|
|
try:
|
|
with open(self.manifest_file, 'w', encoding='utf-8') as f:
|
|
json.dump(self._manifest, f, indent=2)
|
|
except Exception as e:
|
|
print(f"[SEngine] Error saving manifest: {e}")
|
|
|
|
def _get_full_path(self, file_info: Dict) -> str:
|
|
"""Construct full path from manifest entry."""
|
|
# Support both old format (local_path) and new format (file_name)
|
|
if "local_path" in file_info:
|
|
# Old format - check if file exists at stored path
|
|
if os.path.exists(file_info["local_path"]):
|
|
return file_info["local_path"]
|
|
# New format or fallback - construct from cache_dir + filename
|
|
filename = file_info.get("file_name", "")
|
|
if filename:
|
|
return os.path.join(self.cache_dir, filename)
|
|
return ""
|
|
|
|
def is_downloaded(self, version_id: int) -> bool:
|
|
"""Check if a LoRA version is already downloaded."""
|
|
str_id = str(version_id)
|
|
if str_id not in self._manifest.get("files", {}):
|
|
return False
|
|
file_info = self._manifest["files"][str_id]
|
|
local_path = self._get_full_path(file_info)
|
|
return local_path and os.path.exists(local_path)
|
|
|
|
def get_local_path(self, version_id: int) -> Optional[str]:
|
|
"""Get the local path for a downloaded LoRA, or None if not downloaded."""
|
|
str_id = str(version_id)
|
|
if str_id in self._manifest.get("files", {}):
|
|
local_path = self._get_full_path(self._manifest["files"][str_id])
|
|
if local_path and os.path.exists(local_path):
|
|
return local_path
|
|
return None
|
|
|
|
def get_download_progress(self, version_id: int) -> float:
|
|
"""Get download progress for a LoRA (0-1, or -1 if not downloading)."""
|
|
return self._download_progress.get(version_id, -1)
|
|
|
|
def download_lora_sync(
|
|
self,
|
|
version_id: int,
|
|
file_name: str,
|
|
api_key: str = "",
|
|
progress_callback=None,
|
|
download_url: str = ""
|
|
) -> Tuple[bool, str]:
|
|
"""
|
|
Download a LoRA file synchronously.
|
|
|
|
Args:
|
|
version_id: The Civitai model version ID
|
|
file_name: The original filename
|
|
api_key: Civitai API key for authenticated downloads
|
|
progress_callback: Optional callback(version_id, progress)
|
|
download_url: Direct download URL (preferred)
|
|
|
|
Returns:
|
|
Tuple of (success, local_path or error_message)
|
|
"""
|
|
# Check if already downloaded
|
|
existing_path = self.get_local_path(version_id)
|
|
if existing_path:
|
|
return (True, existing_path)
|
|
|
|
# Use provided URL or construct one
|
|
if not download_url:
|
|
api = get_civitai_api(api_key)
|
|
download_url = api.get_download_url(version_id)
|
|
|
|
# Determine local filename
|
|
safe_filename = f"{version_id}_{file_name}"
|
|
local_path = os.path.join(self.cache_dir, safe_filename)
|
|
|
|
self._download_progress[version_id] = 0.0
|
|
|
|
try:
|
|
# Add token to URL for Civitai
|
|
if api_key:
|
|
separator = "&" if "?" in download_url else "?"
|
|
download_url = f"{download_url}{separator}token={api_key}"
|
|
|
|
print(f"[SEngine] Downloading from: {download_url[:100]}...")
|
|
|
|
# Create request with headers
|
|
headers = {
|
|
"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36",
|
|
"Content-Type": "application/json",
|
|
}
|
|
|
|
request = urllib.request.Request(download_url, headers=headers)
|
|
|
|
# Create SSL context
|
|
ssl_context = ssl.create_default_context()
|
|
|
|
# Download
|
|
with urllib.request.urlopen(request, context=ssl_context) as response:
|
|
total_size = int(response.headers.get('content-length', 0))
|
|
content_type = response.headers.get('content-type', '')
|
|
downloaded = 0
|
|
|
|
print(f"[SEngine] Response content-type: {content_type}")
|
|
print(f"[SEngine] Expected size: {total_size / (1024*1024):.1f} MB")
|
|
|
|
# Check if we got an HTML page instead of a file
|
|
if 'text/html' in content_type.lower():
|
|
return (False, "Download URL returned HTML page instead of file. Check API key or URL.")
|
|
|
|
with open(local_path, 'wb') as f:
|
|
while True:
|
|
chunk = response.read(65536) # 64KB chunks
|
|
if not chunk:
|
|
break
|
|
f.write(chunk)
|
|
downloaded += len(chunk)
|
|
|
|
if total_size > 0:
|
|
progress = downloaded / total_size
|
|
self._download_progress[version_id] = progress
|
|
if progress_callback:
|
|
progress_callback(version_id, progress)
|
|
|
|
# Print progress periodically
|
|
if total_size > 0 and downloaded % (1024 * 1024) < 65536:
|
|
mb_done = downloaded / (1024 * 1024)
|
|
mb_total = total_size / (1024 * 1024)
|
|
print(f"[SEngine] Progress: {mb_done:.1f}/{mb_total:.1f} MB")
|
|
|
|
# Verify download
|
|
if os.path.exists(local_path):
|
|
actual_size = os.path.getsize(local_path)
|
|
|
|
if actual_size == 0:
|
|
os.remove(local_path)
|
|
return (False, "Download produced empty file")
|
|
|
|
# Verify size matches expected if we got content-length
|
|
if total_size > 0 and actual_size != total_size:
|
|
print(f"[SEngine] Size mismatch: expected {total_size}, got {actual_size}")
|
|
os.remove(local_path)
|
|
return (False, f"Download incomplete: {actual_size}/{total_size} bytes")
|
|
|
|
print(f"[SEngine] Download verification passed ({actual_size} bytes)")
|
|
|
|
# Update manifest - store only filename, not full path
|
|
if "files" not in self._manifest:
|
|
self._manifest["files"] = {}
|
|
|
|
self._manifest["files"][str(version_id)] = {
|
|
"file_name": safe_filename, # Just the filename, not full path
|
|
"original_name": file_name,
|
|
"version_id": version_id,
|
|
}
|
|
self._save_manifest()
|
|
|
|
if version_id in self._download_progress:
|
|
del self._download_progress[version_id]
|
|
|
|
print(f"[SEngine] Download complete: {local_path}")
|
|
return (True, local_path)
|
|
else:
|
|
return (False, "Download produced no file")
|
|
|
|
except urllib.error.HTTPError as e:
|
|
error_msg = f"HTTP Error {e.code}: {e.reason}"
|
|
print(f"[SEngine] {error_msg}")
|
|
if os.path.exists(local_path):
|
|
try:
|
|
os.remove(local_path)
|
|
except:
|
|
pass
|
|
if version_id in self._download_progress:
|
|
del self._download_progress[version_id]
|
|
return (False, error_msg)
|
|
|
|
except Exception as e:
|
|
error_msg = str(e)
|
|
print(f"[SEngine] Download error: {error_msg}")
|
|
if os.path.exists(local_path):
|
|
try:
|
|
os.remove(local_path)
|
|
except:
|
|
pass
|
|
if version_id in self._download_progress:
|
|
del self._download_progress[version_id]
|
|
return (False, error_msg)
|
|
|
|
# Async version for server routes
|
|
async def download_lora(
|
|
self,
|
|
version_id: int,
|
|
file_name: str,
|
|
api_key: str = "",
|
|
progress_callback=None,
|
|
download_url: str = ""
|
|
) -> Tuple[bool, str]:
|
|
"""Async wrapper that calls sync download in executor."""
|
|
import asyncio
|
|
import functools
|
|
loop = asyncio.get_event_loop()
|
|
func = functools.partial(
|
|
self.download_lora_sync,
|
|
version_id,
|
|
file_name,
|
|
api_key,
|
|
progress_callback,
|
|
download_url
|
|
)
|
|
return await loop.run_in_executor(None, func)
|
|
|
|
def clear_cache(self):
|
|
"""Clear all cached LoRA files."""
|
|
for str_id, info in list(self._manifest.get("files", {}).items()):
|
|
local_path = self._get_full_path(info)
|
|
if local_path and os.path.exists(local_path):
|
|
try:
|
|
os.remove(local_path)
|
|
except Exception as e:
|
|
print(f"[SEngine] Error removing {local_path}: {e}")
|
|
|
|
self._manifest = {"files": {}}
|
|
self._save_manifest()
|
|
|
|
def get_cache_size(self) -> int:
|
|
"""Get total size of cached files in bytes."""
|
|
total = 0
|
|
for str_id, info in self._manifest.get("files", {}).items():
|
|
local_path = self._get_full_path(info)
|
|
if local_path and os.path.exists(local_path):
|
|
total += os.path.getsize(local_path)
|
|
return total
|
|
|
|
def get_cached_count(self) -> int:
|
|
"""Get number of cached LoRA files."""
|
|
count = 0
|
|
for str_id, info in self._manifest.get("files", {}).items():
|
|
local_path = self._get_full_path(info)
|
|
if local_path and os.path.exists(local_path):
|
|
count += 1
|
|
return count
|
|
|
|
|
|
# Global instance
|
|
_cache_manager: Optional[LoRACacheManager] = None
|
|
|
|
|
|
def get_cache_manager() -> LoRACacheManager:
|
|
"""Get or create the global LoRACacheManager instance."""
|
|
global _cache_manager
|
|
if _cache_manager is None:
|
|
_cache_manager = LoRACacheManager()
|
|
return _cache_manager
|