From c4c967172bae4ed5a16b859c1bb875464dcf9103 Mon Sep 17 00:00:00 2001 From: silveroxides Date: Wed, 8 Apr 2026 15:04:24 +0200 Subject: [PATCH 1/7] Add Model preview/example image/video downloading, model hashing, preview/example workflow extraction from images and videos followed by compression of images to save storage, saving of hash database and much more --- __init__.py | 7 + nodes/downloader_nodes.py | 83 ++++++ nodes/downloader_utils.py | 567 ++++++++++++++++++++++++++++++++++++++ standalone_downloader.py | 561 +++++++++++++++++++++++++++++++++++++ 4 files changed, 1218 insertions(+) create mode 100644 nodes/downloader_nodes.py create mode 100644 nodes/downloader_utils.py create mode 100644 standalone_downloader.py diff --git a/__init__.py b/__init__.py index e29d660..a8a2706 100644 --- a/__init__.py +++ b/__init__.py @@ -29,6 +29,10 @@ from .nodes.lora_resize import ( LoRAResizeFrobenius, LoRAResizeCumulative, LoRAMergeToModel ) +from .nodes.downloader_nodes import ( + CheckpointDownloader, LoRADownloader, EmbeddingDownloader, + VAEDownloader, ControlNetDownloader, ManualPathDownloader +) class ModelUtilsExtension(ComfyExtension): @@ -60,6 +64,9 @@ class ModelUtilsExtension(ComfyExtension): LoRAMultiMerge, LoRAMultiMergeDARE, # LoRA Merge To Model LoRAMergeToModel, + # Downloaders + CheckpointDownloader, LoRADownloader, EmbeddingDownloader, + VAEDownloader, ControlNetDownloader, ManualPathDownloader, ] diff --git a/nodes/downloader_nodes.py b/nodes/downloader_nodes.py new file mode 100644 index 0000000..25a36d6 --- /dev/null +++ b/nodes/downloader_nodes.py @@ -0,0 +1,83 @@ +import folder_paths +from comfy_api.latest import io +from .downloader_utils import scan_and_process + +class BaseDownloaderNode(io.ComfyNode): + CATEGORY = "ModelUtils/Downloader" + OUTPUT_NODE = True + + @classmethod + def get_common_inputs(cls): + return { + "recursive": io.BOOLEAN(default=True), + "nsfw_level": io.COMBO(["None", "Soft", "Mature", "X", "XXX", "All"], default="All"), + "max_examples": io.INT(default=0, min=0, max=100), + "threads": io.INT(default=4, min=1, max=32), + "api_key": io.STRING(default="") + } + + @classmethod + def INPUT_TYPES(cls): + return {"required": cls.get_common_inputs()} + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("status",) + FUNCTION = "process" + + def process(self, recursive, nsfw_level, max_examples, threads, api_key, **kwargs): + scan_dirs = self.get_scan_dirs(**kwargs) + if not scan_dirs: + return ("No directories to scan.",) + + # Call the scan_and_process function + scan_and_process( + scan_dirs=scan_dirs, + recursive=recursive, + nsfw_level=nsfw_level, + max_examples=max_examples, + api_key=api_key if api_key else None, + threads=threads + ) + + return (f"Completed scanning and downloading for {len(scan_dirs)} directories.",) + + def get_scan_dirs(self, **kwargs) -> list[str]: + return [] + + +class CheckpointDownloader(BaseDownloaderNode): + def get_scan_dirs(self, **kwargs): + return folder_paths.get_folder_paths("checkpoints") + + +class LoRADownloader(BaseDownloaderNode): + def get_scan_dirs(self, **kwargs): + return folder_paths.get_folder_paths("loras") + + +class EmbeddingDownloader(BaseDownloaderNode): + def get_scan_dirs(self, **kwargs): + return folder_paths.get_folder_paths("embeddings") + + +class VAEDownloader(BaseDownloaderNode): + def get_scan_dirs(self, **kwargs): + return folder_paths.get_folder_paths("vae") + + +class ControlNetDownloader(BaseDownloaderNode): + def get_scan_dirs(self, **kwargs): + return folder_paths.get_folder_paths("controlnet") + + +class ManualPathDownloader(BaseDownloaderNode): + @classmethod + def INPUT_TYPES(cls): + inputs = cls.get_common_inputs() + inputs["scan_dir"] = io.STRING(default="") + return {"required": inputs} + + def get_scan_dirs(self, scan_dir, **kwargs): + if scan_dir.strip(): + return [scan_dir.strip()] + return [] diff --git a/nodes/downloader_utils.py b/nodes/downloader_utils.py new file mode 100644 index 0000000..3ae40dd --- /dev/null +++ b/nodes/downloader_utils.py @@ -0,0 +1,567 @@ +import os +import sys +import hashlib +import json +import sqlite3 +import requests +import glob +from pathlib import Path +import threading +import concurrent.futures +import codecs +import subprocess +import shutil + +try: + from PIL import Image + HAS_PIL = True +except ImportError: + HAS_PIL = False + print("Warning: Pillow not found. Image workflow extraction disabled.") + +try: + from mutagen.mp4 import MP4 + HAS_MUTAGEN = True +except ImportError: + HAS_MUTAGEN = False + print("Warning: mutagen not found. MP4 tag extraction will rely entirely on ffmpeg fallback.") + +# Add the reference module to the Python path +# Assuming downloader_utils.py is inside nodes/, the root is one level up +root_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +sys.path.append(os.path.join(root_dir, 'unifiedefficientloader-reference')) +try: + from unifiedefficientloader import UnifiedSafetensorsLoader + HAS_UNIFIED_LOADER = True +except ImportError: + HAS_UNIFIED_LOADER = False + print("Warning: unifiedefficientloader not found. Safetensors metadata extraction disabled.") + +# --- Database Path Logic --- +try: + import folder_paths + base_user_dir = folder_paths.get_public_user_directory("default") + if not base_user_dir: + base_user_dir = root_dir +except (ImportError, AttributeError): + base_user_dir = root_dir + +PLUGIN_DATA_DIR = os.path.join(base_user_dir, "ComfyUI-ModelUtils") +os.makedirs(PLUGIN_DATA_DIR, exist_ok=True) +DEFAULT_DB_PATH = os.path.join(PLUGIN_DATA_DIR, "cache.db") + + +# --- Workflow Extraction Logic --- + +def unescape_and_parse_nested_json(json_string): + if not isinstance(json_string, str): return None + stripped = json_string.strip() + if not (stripped.startswith('{') or stripped.startswith('[')): return None + try: return json.loads(stripped) + except json.JSONDecodeError: pass + try: + first_parse = json.loads(stripped) + if isinstance(first_parse, str): return json.loads(first_parse) + except json.JSONDecodeError: pass + try: + decoded = codecs.decode(stripped, "unicode_escape") + return json.loads(decoded) + except (json.JSONDecodeError, UnicodeDecodeError, ValueError): pass + try: + unescaped = stripped.replace('\\\\"', '\\"').replace('\\"', '"').replace("\\\\", "\\") + unescaped = unescaped.replace("\\n", "\n").replace("\\r", "\r").replace("\\t", "\t") + return json.loads(unescaped) + except json.JSONDecodeError: pass + return None + +def get_workflow_blob(metadata): + for key in ['workflow', 'prompt']: + val = metadata.get(key) + if val: + parsed = unescape_and_parse_nested_json(val) + if parsed: return parsed + for val in metadata.values(): + parsed = unescape_and_parse_nested_json(val) + if isinstance(parsed, dict) and (any(k.isdigit() for k in parsed.keys()) or 'nodes' in parsed): + return parsed + return None + +def process_image_workflow_and_compress(filepath_str): + if not HAS_PIL: return + p = Path(filepath_str) + if p.suffix.lower() not in ['.png', '.jpg', '.jpeg']: return + + workflow = None + actual_format = None + try: + with Image.open(p) as img: + actual_format = img.format + # CivitAI sometimes serves PNGs with .jpeg extension in URL + if actual_format == 'PNG': + workflow = get_workflow_blob(img.info) + if workflow: + # Save workflow + json_path = p.with_suffix('.json') + if not json_path.exists(): + with open(json_path, 'w', encoding='utf-8') as f: + json.dump(workflow, f, indent=2, ensure_ascii=False) + print(f"Extracted workflow to {json_path}") + + # Compress image to WebP + webp_path = p.with_suffix('.webp') + if not webp_path.exists(): + img.save(webp_path, 'WEBP', quality=85, method=4) + print(f"Compressed {p.name} (was actually {actual_format}) to {webp_path.name}") + + # Delete original if we created a webp replacement + if workflow and p.with_suffix('.webp').exists() and actual_format == 'PNG': + os.remove(p) + print(f"Deleted original uncompressed image: {p.name}") + except Exception as e: + print(f"Failed to process image {p.name}: {e}") + +# --- Video Workflow Logic --- +def extract_metadata_ffmpeg(mp4_path): + if not shutil.which('ffmpeg'): return None + try: + cmd = ['ffmpeg', '-i', str(mp4_path), '-f', 'ffmetadata', '-'] + result = subprocess.run(cmd, capture_output=True, text=True, encoding='utf-8', errors='replace') + if result.returncode != 0: return None + metadata, current_key, current_value = {}, None, [] + for line in result.stdout.splitlines(): + if line.startswith(';FFMETADATA'): continue + if '=' in line: + if current_key: metadata[current_key] = '\n'.join(current_value) + parts = line.split('=', 1) + key_lower = parts[0].lower() + if key_lower in ['workflow', 'prompt', 'comment']: + current_key, current_value = key_lower, [parts[1]] + else: current_key = None + elif current_key: + current_value.append(line) + if current_key: metadata[current_key] = '\n'.join(current_value) + return {k: v.replace(r'\=', '=').replace(r'\;', ';').replace(r'\#', '#').replace(r'\\', '\\') for k, v in metadata.items()} + except Exception as e: + print(f"ffmpeg extraction error for {mp4_path}: {e}") + return None + +def get_tag_value(mp4, keys): + for key in keys: + if key in mp4.tags: return mp4.tags[key][0] if mp4.tags[key] else None + for key in keys: + search_suffix = key.split(':')[-1].lower() + for tag_key, val in mp4.tags.items(): + if tag_key.lower().endswith(f":{search_suffix}"): return val[0] if val else None + return None + +def process_video_workflow(filepath_str): + p = Path(filepath_str) + if p.suffix.lower() != '.mp4': return + + json_path = p.with_suffix('.json') + if json_path.exists(): return + + workflow_data = None + mp4 = None + if HAS_MUTAGEN: + try: mp4 = MP4(p) + except: pass + + if mp4: + for tag_keys in [['----:QuickTime:Workflow', 'workflow'], ['----:QuickTime:Prompt', 'prompt'], ['©cmt']]: + val = get_tag_value(mp4, tag_keys) + if val: + if isinstance(val, bytes): val = val.decode('utf-8', errors='ignore') + try: + parsed = json.loads(val) + if isinstance(parsed, dict) and ('prompt' in parsed or any(k.isdigit() for k in parsed.keys())): + if 'prompt' in parsed and isinstance(parsed['prompt'], str): + try: workflow_data = json.loads(parsed['prompt']) + except: workflow_data = parsed + else: + workflow_data = parsed + break + except json.JSONDecodeError: pass + + if not workflow_data: + ffmpeg_data = extract_metadata_ffmpeg(p) + if ffmpeg_data: + for key in ['workflow', 'prompt', 'comment']: + if key in ffmpeg_data: + try: + parsed = json.loads(ffmpeg_data[key]) + if isinstance(parsed, dict): + if 'prompt' in parsed and isinstance(parsed['prompt'], str): + try: workflow_data = json.loads(parsed['prompt']) + except: workflow_data = parsed + else: + workflow_data = parsed + break + except json.JSONDecodeError: pass + + if workflow_data: + try: + with open(json_path, 'w', encoding='utf-8') as f: + json.dump(workflow_data, f, indent=2, ensure_ascii=False) + print(f"Extracted workflow to {json_path.name}") + except Exception as e: + print(f"Error saving workflow for {p.name}: {e}") + +# --- Constants for Civitai NSFW levels --- +NSFW_LEVELS = { + "None": 1, + "Soft": 2, + "Mature": 4, + "X": 8, + "XXX": 16, + "All": 32 # Anything +} + +class HashCache: + def __init__(self, db_path=None): + self.db_path = db_path if db_path else DEFAULT_DB_PATH + self.is_sqlite = self.db_path.endswith('.db') or self.db_path.endswith('.sqlite') + self.lock = threading.Lock() + if self.is_sqlite: + self._init_sqlite() + else: + self.cache = self._load_json() + + def _init_sqlite(self): + with sqlite3.connect(self.db_path) as conn: + cursor = conn.cursor() + cursor.execute(''' + CREATE TABLE IF NOT EXISTS file_hashes ( + filepath TEXT PRIMARY KEY, + mtime REAL, + size INTEGER, + hash TEXT + ) + ''') + conn.commit() + + def _load_json(self): + if os.path.exists(self.db_path): + with open(self.db_path, "r", encoding="utf-8") as f: + return json.load(f) + return {} + + def _save_json(self): + with open(self.db_path, "w", encoding="utf-8") as f: + json.dump(self.cache, f, indent=4) + + def _get_sqlite_hash(self, filepath, mtime, size): + with sqlite3.connect(self.db_path) as conn: + cursor = conn.cursor() + cursor.execute("SELECT mtime, size, hash FROM file_hashes WHERE filepath=?", (filepath,)) + row = cursor.fetchone() + if row: + db_mtime, db_size, db_hash = row + if db_mtime == mtime and db_size == size: + return db_hash + return None + + def _set_sqlite_hash(self, filepath, mtime, size, file_hash): + with sqlite3.connect(self.db_path, timeout=10.0) as conn: + cursor = conn.cursor() + cursor.execute(''' + INSERT OR REPLACE INTO file_hashes (filepath, mtime, size, hash) + VALUES (?, ?, ?, ?) + ''', (filepath, mtime, size, file_hash)) + conn.commit() + + def get_hash(self, filepath): + if not os.path.exists(filepath): + return None + mtime = os.path.getmtime(filepath) + size = os.path.getsize(filepath) + key = filepath + + with self.lock: + if self.is_sqlite: + cached_hash = self._get_sqlite_hash(key, mtime, size) + if cached_hash: return cached_hash + else: + if key in self.cache: + entry = self.cache[key] + if entry["mtime"] == mtime and entry["size"] == size: + return entry["hash"] + + # Calculate new hash (outside the lock to allow parallel hashing) + print(f"Calculating SHA256 for: {filepath}") + sha256_hash = self.calculate_sha256(filepath) + + with self.lock: + if self.is_sqlite: + self._set_sqlite_hash(key, mtime, size, sha256_hash) + else: + self.cache[key] = { + "mtime": mtime, + "size": size, + "hash": sha256_hash + } + self._save_json() + + return sha256_hash + + @staticmethod + def calculate_sha256(filename): + sha256_hash = hashlib.sha256() + with open(filename, "rb") as f: + # Read and update hash string value in blocks of 4K + for byte_block in iter(lambda: f.read(4096), b""): + sha256_hash.update(byte_block) + return sha256_hash.hexdigest() + +class CivitaiAPI: + def __init__(self, api_key=None): + self.api_key = api_key + self.base_url = "https://civitai.com/api/v1" + + def get_headers(self): + headers = {} + if self.api_key: + headers["Authorization"] = f"Bearer {self.api_key}" + return headers + + def get_model_info_by_hash(self, file_hash): + url = f"{self.base_url}/model-versions/by-hash/{file_hash}" + response = requests.get(url, headers=self.get_headers()) + if response.status_code == 200: + return response.json() + elif response.status_code == 404: + return None + else: + response.raise_for_status() + +class Downloader: + @staticmethod + def download_file(url, target_path, headers=None): + if headers is None: + headers = {} + + temp_path = target_path + ".downloading" + + downloaded_size = 0 + if os.path.exists(temp_path): + downloaded_size = os.path.getsize(temp_path) + + if downloaded_size > 0: + headers["Range"] = f"bytes={downloaded_size}-" + + response = requests.get(url, headers=headers, stream=True) + + if response.status_code == 416: # Range not satisfiable, file probably changed or completed + os.remove(temp_path) + downloaded_size = 0 + if "Range" in headers: + del headers["Range"] + response = requests.get(url, headers=headers, stream=True) + + if response.status_code not in [200, 206]: + print(f"Failed to download from {url}, status code: {response.status_code}") + return False + + mode = "ab" if response.status_code == 206 else "wb" + + total_size = int(response.headers.get('content-length', 0)) + downloaded_size + + print(f"Downloading: {target_path} ({downloaded_size}/{total_size} bytes)") + + with open(temp_path, mode) as f: + for chunk in response.iter_content(chunk_size=8192): + if chunk: + f.write(chunk) + + os.rename(temp_path, target_path) + return True + +def get_potential_preview_files(base_path): + exts = ["png", "jpg", "jpeg", "webp", "mp4"] + files = [] + for ext in exts: + files.append(f"{base_path}.preview.{ext}") + files.append(f"{base_path}.{ext}") + return files + +def preview_exists(base_path): + previews = get_potential_preview_files(base_path) + for p in previews: + if os.path.exists(p): + return True + return False + +def next_example_image_path(base_path, ext): + i = 0 + while len(glob.glob(f"{base_path}.example.{i}.*")) > 0: + i += 1 + return f"{base_path}.example.{i}.{ext}" + +def filter_images_by_nsfw(images, threshold_level): + valid_images = [] + for img in images: + rating = img.get("nsfwLevel", 1) + if rating <= threshold_level: + valid_images.append(img) + return valid_images + +def extract_safetensors_metadata(filepath): + if not HAS_UNIFIED_LOADER: + return None + try: + # Use low_memory=True to only read the header and metadata without loading tensors + loader = UnifiedSafetensorsLoader(filepath, low_memory=True) + metadata = loader.metadata() + loader.close() + return metadata if metadata else None + except Exception as e: + print(f"Failed to read safetensors metadata from {filepath} using unifiedefficientloader: {e}") + return None + +def process_file(filepath, cache, api, nsfw_threshold, max_examples): + filepath_str = str(filepath) + base_path = os.path.splitext(filepath_str)[0] + is_safetensors = filepath_str.lower().endswith(".safetensors") + + # Check if all targets already exist + info_path = f"{base_path}.civitai.info" + st_metadata_path = f"{base_path}.metadata.json" + + has_info = os.path.exists(info_path) + has_preview = preview_exists(base_path) + has_st_metadata = not is_safetensors or os.path.exists(st_metadata_path) + + # Process existing images/videos for workflow extraction & compression BEFORE we possibly skip the file + for p in get_potential_preview_files(base_path): + if os.path.exists(p): + ext_lower = p.lower() + if ext_lower.endswith(('.png', '.jpg', '.jpeg')): process_image_workflow_and_compress(p) + elif ext_lower.endswith('.mp4'): process_video_workflow(p) + + for ex in glob.glob(f"{base_path}.example.*"): + ext_lower = ex.lower() + if ext_lower.endswith(('.png', '.jpg', '.jpeg')): process_image_workflow_and_compress(ex) + elif ext_lower.endswith('.mp4'): process_video_workflow(ex) + + # If everything exists and we don't want examples or already have enough examples, we might skip API logic + if has_info and has_preview and has_st_metadata and max_examples == 0: + print(f"Skipping fully populated model: {filepath_str}") + return + + print(f"Processing model: {filepath_str}") + + # Safetensors built-in metadata extraction + if is_safetensors and not has_st_metadata: + metadata = extract_safetensors_metadata(filepath_str) + if metadata: + with open(st_metadata_path, "w", encoding="utf-8") as f: + json.dump(metadata, f, indent=4) + print(f"Saved safetensors metadata to {st_metadata_path}") + + file_hash = cache.get_hash(filepath_str) + if not file_hash: + return + + try: + model_info = api.get_model_info_by_hash(file_hash) + except Exception as e: + print(f"API Error for {filepath_str}: {e}") + return + + if not model_info: + print(f"Model not found on Civitai: {filepath_str}") + return + + # 1. Save Info + if not has_info: + with open(info_path, "w", encoding="utf-8") as f: + json.dump(model_info, f, indent=4) + print(f"Saved metadata to {info_path}") + + images = model_info.get("images", []) + valid_images = filter_images_by_nsfw(images, nsfw_threshold) + + if not valid_images: + print(f"No valid images found matching NSFW criteria for {filepath_str}") + return + + # 2. Save Preview + if not has_preview and len(valid_images) > 0: + preview_img = valid_images[0] + url = preview_img.get("url") + ext = "jpeg" # Default fallback + if ".png" in url.lower(): ext = "png" + elif ".mp4" in url.lower(): ext = "mp4" + elif ".jpg" in url.lower() or ".jpeg" in url.lower(): ext = "jpeg" + + preview_path = f"{base_path}.preview.{ext}" + if Downloader.download_file(url, preview_path): + if ext in ("png", "jpg", "jpeg"): process_image_workflow_and_compress(preview_path) + elif ext == "mp4": process_video_workflow(preview_path) + + # 3. Save Examples + if max_examples > 0: + existing_examples_count = len(glob.glob(f"{base_path}.example.*")) + # We divide by 2 roughly if there are json files, let's just count unique example indices + existing_indices = set() + for f in glob.glob(f"{base_path}.example.*"): + parts = f.split('.example.') + if len(parts) > 1: + idx = parts[1].split('.')[0] + if idx.isdigit(): existing_indices.add(idx) + + needed_examples = max_examples - len(existing_indices) + + if needed_examples > 0: + example_images = valid_images[1:max_examples+1] + # Skip images we probably already downloaded by skipping the first N + example_images_to_dl = example_images[len(existing_indices):] + + for img in example_images_to_dl: + url = img.get("url") + ext = "jpeg" + if ".png" in url.lower(): ext = "png" + elif ".mp4" in url.lower(): ext = "mp4" + elif ".jpg" in url.lower() or ".jpeg" in url.lower(): ext = "jpeg" + + example_path = next_example_image_path(base_path, ext) + if Downloader.download_file(url, example_path): + if ext in ("png", "jpg", "jpeg"): process_image_workflow_and_compress(example_path) + elif ext == "mp4": process_video_workflow(example_path) + +def scan_and_process(scan_dirs, recursive=True, db_path=None, nsfw_level="All", max_examples=0, api_key=None, threads=4): + """ + Scans directories for models and fetches their metadata/previews. + Adapted for direct Python API calls rather than argparse CLI. + """ + db_path = db_path if db_path else DEFAULT_DB_PATH + cache = HashCache(db_path) + api = CivitaiAPI(api_key) + + model_exts = [".safetensors", ".ckpt", ".pt"] + + files_to_process = [] + if isinstance(scan_dirs, str): + scan_dirs = [scan_dirs] + + for scan_dir in scan_dirs: + path = Path(scan_dir) + if recursive: + for ext in model_exts: + files_to_process.extend(path.rglob(f"*{ext}")) + else: + for ext in model_exts: + files_to_process.extend(path.glob(f"*{ext}")) + + nsfw_threshold = NSFW_LEVELS.get(nsfw_level, NSFW_LEVELS["All"]) + + print(f"Found {len(files_to_process)} models to process.") + + if threads > 1: + print(f"Starting {threads} threads...") + with concurrent.futures.ThreadPoolExecutor(max_workers=threads) as executor: + futures = [executor.submit(process_file, fp, cache, api, nsfw_threshold, max_examples) for fp in files_to_process] + concurrent.futures.wait(futures) + else: + for filepath in files_to_process: + process_file(filepath, cache, api, nsfw_threshold, max_examples) diff --git a/standalone_downloader.py b/standalone_downloader.py new file mode 100644 index 0000000..c06d965 --- /dev/null +++ b/standalone_downloader.py @@ -0,0 +1,561 @@ +import os +import sys +import argparse +import hashlib +import json +import sqlite3 +import requests +import glob +from pathlib import Path +import threading +import concurrent.futures +import codecs +import subprocess +import shutil + +try: + from PIL import Image + HAS_PIL = True +except ImportError: + HAS_PIL = False + print("Warning: Pillow not found. Image workflow extraction disabled.") + +try: + from mutagen.mp4 import MP4 + HAS_MUTAGEN = True +except ImportError: + HAS_MUTAGEN = False + print("Warning: mutagen not found. MP4 tag extraction will rely entirely on ffmpeg fallback.") + +# Add the reference module to the Python path +sys.path.append(os.path.join(os.path.dirname(__file__), 'unifiedefficientloader-reference')) +try: + from unifiedefficientloader import UnifiedSafetensorsLoader + HAS_UNIFIED_LOADER = True +except ImportError: + HAS_UNIFIED_LOADER = False + print("Warning: unifiedefficientloader not found. Safetensors metadata extraction disabled.") + +# --- Workflow Extraction Logic --- + +def unescape_and_parse_nested_json(json_string): + if not isinstance(json_string, str): return None + stripped = json_string.strip() + if not (stripped.startswith('{') or stripped.startswith('[')): return None + try: return json.loads(stripped) + except json.JSONDecodeError: pass + try: + first_parse = json.loads(stripped) + if isinstance(first_parse, str): return json.loads(first_parse) + except json.JSONDecodeError: pass + try: + decoded = codecs.decode(stripped, "unicode_escape") + return json.loads(decoded) + except (json.JSONDecodeError, UnicodeDecodeError, ValueError): pass + try: + unescaped = stripped.replace('\\\\"', '\\"').replace('\\"', '"').replace("\\\\", "\\") + unescaped = unescaped.replace("\\n", "\n").replace("\\r", "\r").replace("\\t", "\t") + return json.loads(unescaped) + except json.JSONDecodeError: pass + return None + +def get_workflow_blob(metadata): + for key in ['workflow', 'prompt']: + val = metadata.get(key) + if val: + parsed = unescape_and_parse_nested_json(val) + if parsed: return parsed + for val in metadata.values(): + parsed = unescape_and_parse_nested_json(val) + if isinstance(parsed, dict) and (any(k.isdigit() for k in parsed.keys()) or 'nodes' in parsed): + return parsed + return None + +def process_image_workflow_and_compress(filepath_str): + if not HAS_PIL: return + p = Path(filepath_str) + if p.suffix.lower() not in ['.png', '.jpg', '.jpeg']: return + + workflow = None + actual_format = None + try: + with Image.open(p) as img: + actual_format = img.format + # CivitAI sometimes serves PNGs with .jpeg extension in URL + if actual_format == 'PNG': + workflow = get_workflow_blob(img.info) + if workflow: + # Save workflow + json_path = p.with_suffix('.json') + if not json_path.exists(): + with open(json_path, 'w', encoding='utf-8') as f: + json.dump(workflow, f, indent=2, ensure_ascii=False) + print(f"Extracted workflow to {json_path}") + + # Compress image to WebP + webp_path = p.with_suffix('.webp') + if not webp_path.exists(): + img.save(webp_path, 'WEBP', quality=85, method=4) + print(f"Compressed {p.name} (was actually {actual_format}) to {webp_path.name}") + + # Delete original if we created a webp replacement + if workflow and p.with_suffix('.webp').exists() and actual_format == 'PNG': + os.remove(p) + print(f"Deleted original uncompressed image: {p.name}") + except Exception as e: + print(f"Failed to process image {p.name}: {e}") + +# --- Video Workflow Logic --- +def extract_metadata_ffmpeg(mp4_path): + if not shutil.which('ffmpeg'): return None + try: + cmd = ['ffmpeg', '-i', str(mp4_path), '-f', 'ffmetadata', '-'] + result = subprocess.run(cmd, capture_output=True, text=True, encoding='utf-8', errors='replace') + if result.returncode != 0: return None + metadata, current_key, current_value = {}, None, [] + for line in result.stdout.splitlines(): + if line.startswith(';FFMETADATA'): continue + if '=' in line: + if current_key: metadata[current_key] = '\n'.join(current_value) + parts = line.split('=', 1) + key_lower = parts[0].lower() + if key_lower in ['workflow', 'prompt', 'comment']: + current_key, current_value = key_lower, [parts[1]] + else: current_key = None + elif current_key: + current_value.append(line) + if current_key: metadata[current_key] = '\n'.join(current_value) + return {k: v.replace(r'\=', '=').replace(r'\;', ';').replace(r'\#', '#').replace(r'\\', '\\') for k, v in metadata.items()} + except Exception as e: + print(f"ffmpeg extraction error for {mp4_path}: {e}") + return None + +def get_tag_value(mp4, keys): + for key in keys: + if key in mp4.tags: return mp4.tags[key][0] if mp4.tags[key] else None + for key in keys: + search_suffix = key.split(':')[-1].lower() + for tag_key, val in mp4.tags.items(): + if tag_key.lower().endswith(f":{search_suffix}"): return val[0] if val else None + return None + +def process_video_workflow(filepath_str): + p = Path(filepath_str) + if p.suffix.lower() != '.mp4': return + + json_path = p.with_suffix('.json') + if json_path.exists(): return + + workflow_data = None + mp4 = None + if HAS_MUTAGEN: + try: mp4 = MP4(p) + except: pass + + if mp4: + for tag_keys in [['----:QuickTime:Workflow', 'workflow'], ['----:QuickTime:Prompt', 'prompt'], ['©cmt']]: + val = get_tag_value(mp4, tag_keys) + if val: + if isinstance(val, bytes): val = val.decode('utf-8', errors='ignore') + try: + parsed = json.loads(val) + if isinstance(parsed, dict) and ('prompt' in parsed or any(k.isdigit() for k in parsed.keys())): + if 'prompt' in parsed and isinstance(parsed['prompt'], str): + try: workflow_data = json.loads(parsed['prompt']) + except: workflow_data = parsed + else: + workflow_data = parsed + break + except json.JSONDecodeError: pass + + if not workflow_data: + ffmpeg_data = extract_metadata_ffmpeg(p) + if ffmpeg_data: + for key in ['workflow', 'prompt', 'comment']: + if key in ffmpeg_data: + try: + parsed = json.loads(ffmpeg_data[key]) + if isinstance(parsed, dict): + if 'prompt' in parsed and isinstance(parsed['prompt'], str): + try: workflow_data = json.loads(parsed['prompt']) + except: workflow_data = parsed + else: + workflow_data = parsed + break + except json.JSONDecodeError: pass + + if workflow_data: + try: + with open(json_path, 'w', encoding='utf-8') as f: + json.dump(workflow_data, f, indent=2, ensure_ascii=False) + print(f"Extracted workflow to {json_path.name}") + except Exception as e: + print(f"Error saving workflow for {p.name}: {e}") + +# --- Constants for Civitai NSFW levels --- +NSFW_LEVELS = { + "None": 1, + "Soft": 2, + "Mature": 4, + "X": 8, + "XXX": 16, + "All": 32 # Anything +} + +class HashCache: + def __init__(self, db_path): + self.db_path = db_path + self.is_sqlite = db_path.endswith('.db') or db_path.endswith('.sqlite') + self.lock = threading.Lock() + if self.is_sqlite: + self._init_sqlite() + else: + self.cache = self._load_json() + + def _init_sqlite(self): + with sqlite3.connect(self.db_path) as conn: + cursor = conn.cursor() + cursor.execute(''' + CREATE TABLE IF NOT EXISTS file_hashes ( + filepath TEXT PRIMARY KEY, + mtime REAL, + size INTEGER, + hash TEXT + ) + ''') + conn.commit() + + def _load_json(self): + if os.path.exists(self.db_path): + with open(self.db_path, "r", encoding="utf-8") as f: + return json.load(f) + return {} + + def _save_json(self): + with open(self.db_path, "w", encoding="utf-8") as f: + json.dump(self.cache, f, indent=4) + + def _get_sqlite_hash(self, filepath, mtime, size): + with sqlite3.connect(self.db_path) as conn: + cursor = conn.cursor() + cursor.execute("SELECT mtime, size, hash FROM file_hashes WHERE filepath=?", (filepath,)) + row = cursor.fetchone() + if row: + db_mtime, db_size, db_hash = row + if db_mtime == mtime and db_size == size: + return db_hash + return None + + def _set_sqlite_hash(self, filepath, mtime, size, file_hash): + with sqlite3.connect(self.db_path, timeout=10.0) as conn: + cursor = conn.cursor() + cursor.execute(''' + INSERT OR REPLACE INTO file_hashes (filepath, mtime, size, hash) + VALUES (?, ?, ?, ?) + ''', (filepath, mtime, size, file_hash)) + conn.commit() + + def get_hash(self, filepath): + if not os.path.exists(filepath): + return None + mtime = os.path.getmtime(filepath) + size = os.path.getsize(filepath) + key = filepath + + with self.lock: + if self.is_sqlite: + cached_hash = self._get_sqlite_hash(key, mtime, size) + if cached_hash: return cached_hash + else: + if key in self.cache: + entry = self.cache[key] + if entry["mtime"] == mtime and entry["size"] == size: + return entry["hash"] + + # Calculate new hash (outside the lock to allow parallel hashing) + print(f"Calculating SHA256 for: {filepath}") + sha256_hash = self.calculate_sha256(filepath) + + with self.lock: + if self.is_sqlite: + self._set_sqlite_hash(key, mtime, size, sha256_hash) + else: + self.cache[key] = { + "mtime": mtime, + "size": size, + "hash": sha256_hash + } + self._save_json() + + return sha256_hash + + @staticmethod + def calculate_sha256(filename): + sha256_hash = hashlib.sha256() + with open(filename, "rb") as f: + # Read and update hash string value in blocks of 4K + for byte_block in iter(lambda: f.read(4096), b""): + sha256_hash.update(byte_block) + return sha256_hash.hexdigest() + +class CivitaiAPI: + def __init__(self, api_key=None): + self.api_key = api_key + self.base_url = "https://civitai.com/api/v1" + + def get_headers(self): + headers = {} + if self.api_key: + headers["Authorization"] = f"Bearer {self.api_key}" + return headers + + def get_model_info_by_hash(self, file_hash): + url = f"{self.base_url}/model-versions/by-hash/{file_hash}" + response = requests.get(url, headers=self.get_headers()) + if response.status_code == 200: + return response.json() + elif response.status_code == 404: + return None + else: + response.raise_for_status() + +class Downloader: + @staticmethod + def download_file(url, target_path, headers=None): + if headers is None: + headers = {} + + temp_path = target_path + ".downloading" + + downloaded_size = 0 + if os.path.exists(temp_path): + downloaded_size = os.path.getsize(temp_path) + + if downloaded_size > 0: + headers["Range"] = f"bytes={downloaded_size}-" + + response = requests.get(url, headers=headers, stream=True) + + if response.status_code == 416: # Range not satisfiable, file probably changed or completed + os.remove(temp_path) + downloaded_size = 0 + if "Range" in headers: + del headers["Range"] + response = requests.get(url, headers=headers, stream=True) + + if response.status_code not in [200, 206]: + print(f"Failed to download from {url}, status code: {response.status_code}") + return False + + mode = "ab" if response.status_code == 206 else "wb" + + total_size = int(response.headers.get('content-length', 0)) + downloaded_size + + print(f"Downloading: {target_path} ({downloaded_size}/{total_size} bytes)") + + with open(temp_path, mode) as f: + for chunk in response.iter_content(chunk_size=8192): + if chunk: + f.write(chunk) + + os.rename(temp_path, target_path) + return True + +def get_potential_preview_files(base_path): + exts = ["png", "jpg", "jpeg", "webp", "mp4"] + files = [] + for ext in exts: + files.append(f"{base_path}.preview.{ext}") + files.append(f"{base_path}.{ext}") + return files + +def preview_exists(base_path): + previews = get_potential_preview_files(base_path) + for p in previews: + if os.path.exists(p): + return True + return False + +def next_example_image_path(base_path, ext): + i = 0 + while len(glob.glob(f"{base_path}.example.{i}.*")) > 0: + i += 1 + return f"{base_path}.example.{i}.{ext}" + +def filter_images_by_nsfw(images, threshold_level): + valid_images = [] + for img in images: + rating = img.get("nsfwLevel", 1) + if rating <= threshold_level: + valid_images.append(img) + return valid_images + +def extract_safetensors_metadata(filepath): + if not HAS_UNIFIED_LOADER: + return None + try: + # Use low_memory=True to only read the header and metadata without loading tensors + loader = UnifiedSafetensorsLoader(filepath, low_memory=True) + metadata = loader.metadata() + loader.close() + return metadata if metadata else None + except Exception as e: + print(f"Failed to read safetensors metadata from {filepath} using unifiedefficientloader: {e}") + return None + +def process_file(filepath, args, cache, api, nsfw_threshold): + filepath_str = str(filepath) + base_path = os.path.splitext(filepath_str)[0] + is_safetensors = filepath_str.lower().endswith(".safetensors") + + # Check if all targets already exist + info_path = f"{base_path}.civitai.info" + st_metadata_path = f"{base_path}.metadata.json" + + has_info = os.path.exists(info_path) + has_preview = preview_exists(base_path) + has_st_metadata = not is_safetensors or os.path.exists(st_metadata_path) + + # Process existing images/videos for workflow extraction & compression BEFORE we possibly skip the file + for p in get_potential_preview_files(base_path): + if os.path.exists(p): + ext_lower = p.lower() + if ext_lower.endswith(('.png', '.jpg', '.jpeg')): process_image_workflow_and_compress(p) + elif ext_lower.endswith('.mp4'): process_video_workflow(p) + + for ex in glob.glob(f"{base_path}.example.*"): + ext_lower = ex.lower() + if ext_lower.endswith(('.png', '.jpg', '.jpeg')): process_image_workflow_and_compress(ex) + elif ext_lower.endswith('.mp4'): process_video_workflow(ex) + + # If everything exists and we don't want examples or already have enough examples, we might skip API logic + if has_info and has_preview and has_st_metadata and args.max_examples == 0: + print(f"Skipping fully populated model: {filepath_str}") + return + + print(f"Processing model: {filepath_str}") + + # Safetensors built-in metadata extraction + if is_safetensors and not has_st_metadata: + metadata = extract_safetensors_metadata(filepath_str) + if metadata: + with open(st_metadata_path, "w", encoding="utf-8") as f: + json.dump(metadata, f, indent=4) + print(f"Saved safetensors metadata to {st_metadata_path}") + + file_hash = cache.get_hash(filepath_str) + if not file_hash: + return + + try: + model_info = api.get_model_info_by_hash(file_hash) + except Exception as e: + print(f"API Error for {filepath_str}: {e}") + return + + if not model_info: + print(f"Model not found on Civitai: {filepath_str}") + return + + # 1. Save Info + if not has_info: + with open(info_path, "w", encoding="utf-8") as f: + json.dump(model_info, f, indent=4) + print(f"Saved metadata to {info_path}") + + images = model_info.get("images", []) + valid_images = filter_images_by_nsfw(images, nsfw_threshold) + + if not valid_images: + print(f"No valid images found matching NSFW criteria for {filepath_str}") + return + + # 2. Save Preview + if not has_preview and len(valid_images) > 0: + preview_img = valid_images[0] + url = preview_img.get("url") + ext = "jpeg" # Default fallback + if ".png" in url.lower(): ext = "png" + elif ".mp4" in url.lower(): ext = "mp4" + elif ".jpg" in url.lower() or ".jpeg" in url.lower(): ext = "jpeg" + + preview_path = f"{base_path}.preview.{ext}" + if Downloader.download_file(url, preview_path): + if ext in ("png", "jpg", "jpeg"): process_image_workflow_and_compress(preview_path) + elif ext == "mp4": process_video_workflow(preview_path) + + # 3. Save Examples + if args.max_examples > 0: + existing_examples_count = len(glob.glob(f"{base_path}.example.*")) + # We divide by 2 roughly if there are json files, let's just count unique example indices + existing_indices = set() + for f in glob.glob(f"{base_path}.example.*"): + parts = f.split('.example.') + if len(parts) > 1: + idx = parts[1].split('.')[0] + if idx.isdigit(): existing_indices.add(idx) + + needed_examples = args.max_examples - len(existing_indices) + + if needed_examples > 0: + example_images = valid_images[1:args.max_examples+1] + # Skip images we probably already downloaded by skipping the first N + example_images_to_dl = example_images[len(existing_indices):] + + for img in example_images_to_dl: + url = img.get("url") + ext = "jpeg" + if ".png" in url.lower(): ext = "png" + elif ".mp4" in url.lower(): ext = "mp4" + elif ".jpg" in url.lower() or ".jpeg" in url.lower(): ext = "jpeg" + + example_path = next_example_image_path(base_path, ext) + if Downloader.download_file(url, example_path): + if ext in ("png", "jpg", "jpeg"): process_image_workflow_and_compress(example_path) + elif ext == "mp4": process_video_workflow(example_path) + +def scan_and_process(args): + cache = HashCache(args.db_path) + api = CivitaiAPI(args.api_key) + + model_exts = [".safetensors", ".ckpt", ".pt"] + + files_to_process = [] + for scan_dir in args.scan_dirs: + path = Path(scan_dir) + if args.recursive: + for ext in model_exts: + files_to_process.extend(path.rglob(f"*{ext}")) + else: + for ext in model_exts: + files_to_process.extend(path.glob(f"*{ext}")) + + nsfw_threshold = NSFW_LEVELS.get(args.nsfw_level, 1) + + print(f"Found {len(files_to_process)} models to process.") + + if args.threads > 1: + print(f"Starting {args.threads} threads...") + with concurrent.futures.ThreadPoolExecutor(max_workers=args.threads) as executor: + futures = [executor.submit(process_file, fp, args, cache, api, nsfw_threshold) for fp in files_to_process] + concurrent.futures.wait(futures) + else: + for filepath in files_to_process: + process_file(filepath, args, cache, api, nsfw_threshold) + +def main(): + parser = argparse.ArgumentParser(description="Standalone Civitai Model Downloader") + parser.add_argument("--scan-dirs", nargs="+", required=True, help="Directories to scan for models") + parser.add_argument("--recursive", action="store_true", help="Recursively scan directories") + parser.add_argument("--db-path", default="model_hash_cache.json", help="Path to hash cache JSON") + parser.add_argument("--nsfw-level", choices=list(NSFW_LEVELS.keys()), default="Soft", help="Max NSFW level for images") + parser.add_argument("--max-examples", type=int, default=0, help="Max number of example images to download") + parser.add_argument("--max-size-preview", action="store_true", help="Download max resolution previews (Not fully implemented, fallback to default URL)") + parser.add_argument("--api-key", default=None, help="Civitai API Key") + parser.add_argument("--threads", type=int, default=1, help="Number of concurrent threads for hashing and downloading") + + args = parser.parse_args() + scan_and_process(args) + +if __name__ == "__main__": + main() \ No newline at end of file From 4076c14cb54bdc57a29a16e709e3101313f0598e Mon Sep 17 00:00:00 2001 From: silveroxides Date: Wed, 8 Apr 2026 15:10:06 +0200 Subject: [PATCH 2/7] Add requirements.txt entries --- nodes/downloader_nodes.py | 89 +++++++++++++++++++++++++-------------- requirements.txt | 5 ++- 2 files changed, 61 insertions(+), 33 deletions(-) diff --git a/nodes/downloader_nodes.py b/nodes/downloader_nodes.py index 25a36d6..bc13b70 100644 --- a/nodes/downloader_nodes.py +++ b/nodes/downloader_nodes.py @@ -9,25 +9,29 @@ class BaseDownloaderNode(io.ComfyNode): @classmethod def get_common_inputs(cls): return { - "recursive": io.BOOLEAN(default=True), - "nsfw_level": io.COMBO(["None", "Soft", "Mature", "X", "XXX", "All"], default="All"), - "max_examples": io.INT(default=0, min=0, max=100), - "threads": io.INT(default=4, min=1, max=32), - "api_key": io.STRING(default="") + "recursive": io.Boolean.Input("recursive", default=True), + "nsfw_level": io.Combo.Input("nsfw_level", options=["None", "Soft", "Mature", "X", "XXX", "All"], default="All"), + "max_examples": io.Int.Input("max_examples", default=0, min=0, max=100), + "threads": io.Int.Input("threads", default=4, min=1, max=32), + "api_key": io.String.Input("api_key", default="") } @classmethod - def INPUT_TYPES(cls): - return {"required": cls.get_common_inputs()} + def define_schema(cls): + return io.Schema( + node_id=cls.__name__, + display_name=cls.__name__.replace("Downloader", " Downloader"), + category=cls.CATEGORY, + inputs=list(cls.get_common_inputs().values()), + outputs=[io.String.Output(display_name="status")], + is_output_node=True, + ) - RETURN_TYPES = ("STRING",) - RETURN_NAMES = ("status",) - FUNCTION = "process" - - def process(self, recursive, nsfw_level, max_examples, threads, api_key, **kwargs): - scan_dirs = self.get_scan_dirs(**kwargs) + @classmethod + def execute(cls, recursive, nsfw_level, max_examples, threads, api_key, **kwargs): + scan_dirs = cls.get_scan_dirs(**kwargs) if not scan_dirs: - return ("No directories to scan.",) + return io.NodeOutput("No directories to scan.") # Call the scan_and_process function scan_and_process( @@ -39,45 +43,66 @@ class BaseDownloaderNode(io.ComfyNode): threads=threads ) - return (f"Completed scanning and downloading for {len(scan_dirs)} directories.",) + return io.NodeOutput(f"Completed scanning and downloading for {len(scan_dirs)} directories.") - def get_scan_dirs(self, **kwargs) -> list[str]: + @classmethod + def get_scan_dirs(cls, **kwargs) -> list[str]: return [] - class CheckpointDownloader(BaseDownloaderNode): - def get_scan_dirs(self, **kwargs): + @classmethod + def get_scan_dirs(cls, **kwargs): return folder_paths.get_folder_paths("checkpoints") - class LoRADownloader(BaseDownloaderNode): - def get_scan_dirs(self, **kwargs): + @classmethod + def get_scan_dirs(cls, **kwargs): return folder_paths.get_folder_paths("loras") - class EmbeddingDownloader(BaseDownloaderNode): - def get_scan_dirs(self, **kwargs): + @classmethod + def get_scan_dirs(cls, **kwargs): return folder_paths.get_folder_paths("embeddings") - class VAEDownloader(BaseDownloaderNode): - def get_scan_dirs(self, **kwargs): + @classmethod + def get_scan_dirs(cls, **kwargs): return folder_paths.get_folder_paths("vae") - class ControlNetDownloader(BaseDownloaderNode): - def get_scan_dirs(self, **kwargs): + @classmethod + def get_scan_dirs(cls, **kwargs): return folder_paths.get_folder_paths("controlnet") - class ManualPathDownloader(BaseDownloaderNode): @classmethod - def INPUT_TYPES(cls): - inputs = cls.get_common_inputs() - inputs["scan_dir"] = io.STRING(default="") - return {"required": inputs} + def define_schema(cls): + inputs = list(cls.get_common_inputs().values()) + inputs.append(io.String.Input("scan_dir", default="")) + return io.Schema( + node_id=cls.__name__, + display_name="Manual Path Downloader", + category=cls.CATEGORY, + inputs=inputs, + outputs=[io.String.Output(display_name="status")], + is_output_node=True, + ) - def get_scan_dirs(self, scan_dir, **kwargs): + @classmethod + def execute(cls, recursive, nsfw_level, max_examples, threads, api_key, scan_dir, **kwargs): + return super().execute( + recursive=recursive, + nsfw_level=nsfw_level, + max_examples=max_examples, + threads=threads, + api_key=api_key, + scan_dir=scan_dir, + **kwargs + ) + + @classmethod + def get_scan_dirs(cls, **kwargs): + scan_dir = kwargs.get("scan_dir", "") if scan_dir.strip(): return [scan_dir.strip()] return [] diff --git a/requirements.txt b/requirements.txt index d2ce485..a3a3c01 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1 +1,4 @@ -# No external dependencies required +requests +Pillow +mutagen +unifiedefficientloader From 378c3274063c7e5f76ace71fb42134dae3a29c27 Mon Sep 17 00:00:00 2001 From: silveroxides Date: Wed, 8 Apr 2026 15:27:12 +0200 Subject: [PATCH 3/7] Add subdirectory selector --- __init__.py | 4 +- nodes/downloader_nodes.py | 231 +++++++++++++++++++++++++++++++++++++- nodes/downloader_utils.py | 9 +- 3 files changed, 235 insertions(+), 9 deletions(-) diff --git a/__init__.py b/__init__.py index a8a2706..f78e99f 100644 --- a/__init__.py +++ b/__init__.py @@ -30,7 +30,7 @@ from .nodes.lora_resize import ( LoRAMergeToModel ) from .nodes.downloader_nodes import ( - CheckpointDownloader, LoRADownloader, EmbeddingDownloader, + CheckpointDownloader, DiffusionModelDownloader, LoRADownloader, EmbeddingDownloader, VAEDownloader, ControlNetDownloader, ManualPathDownloader ) @@ -65,7 +65,7 @@ class ModelUtilsExtension(ComfyExtension): # LoRA Merge To Model LoRAMergeToModel, # Downloaders - CheckpointDownloader, LoRADownloader, EmbeddingDownloader, + CheckpointDownloader, DiffusionModelDownloader, LoRADownloader, EmbeddingDownloader, VAEDownloader, ControlNetDownloader, ManualPathDownloader, ] diff --git a/nodes/downloader_nodes.py b/nodes/downloader_nodes.py index bc13b70..6099656 100644 --- a/nodes/downloader_nodes.py +++ b/nodes/downloader_nodes.py @@ -1,7 +1,31 @@ +import os import folder_paths from comfy_api.latest import io from .downloader_utils import scan_and_process +def get_subdirectories(folder_name: str) -> list[str]: + """Helper to get a list of subdirectories for a given model folder name.""" + paths = folder_paths.get_folder_paths(folder_name) + if not paths: + return ["/"] + + subdirs = set() + for base_path in paths: + if not os.path.exists(base_path): + continue + for root, dirs, _ in os.walk(base_path): + # Ignore hidden directories directly during walk + dirs[:] = [d for d in dirs if not d.startswith('.')] + + rel_path = os.path.relpath(root, base_path) + if rel_path == ".": + subdirs.add("/") + else: + # Store with forward slashes for consistency + subdirs.add("/" + rel_path.replace("\\", "/")) + + return sorted(list(subdirs)) + class BaseDownloaderNode(io.ComfyNode): CATEGORY = "ModelUtils/Downloader" OUTPUT_NODE = True @@ -50,29 +74,226 @@ class BaseDownloaderNode(io.ComfyNode): return [] class CheckpointDownloader(BaseDownloaderNode): + @classmethod + def define_schema(cls): + inputs = list(cls.get_common_inputs().values()) + inputs.insert(0, io.Combo.Input("subdirectory", options=get_subdirectories("checkpoints"), default="/")) + return io.Schema( + node_id=cls.__name__, + display_name="Checkpoint Downloader", + category=cls.CATEGORY, + inputs=inputs, + outputs=[io.String.Output(display_name="status")], + is_output_node=True, + ) + + @classmethod + def execute(cls, subdirectory, recursive, nsfw_level, max_examples, threads, api_key, **kwargs): + return super().execute( + recursive=recursive, nsfw_level=nsfw_level, max_examples=max_examples, + threads=threads, api_key=api_key, subdirectory=subdirectory, **kwargs + ) + @classmethod def get_scan_dirs(cls, **kwargs): - return folder_paths.get_folder_paths("checkpoints") + base_paths = folder_paths.get_folder_paths("checkpoints") + if not base_paths: + return [] + sub = kwargs.get("subdirectory", "/") + if sub == "/": + return base_paths + + target_paths = [] + for bp in base_paths: + target_path = os.path.join(bp, sub.lstrip("/")) + if os.path.exists(target_path): + target_paths.append(target_path) + return target_paths + +class DiffusionModelDownloader(BaseDownloaderNode): + @classmethod + def define_schema(cls): + inputs = list(cls.get_common_inputs().values()) + inputs.insert(0, io.Combo.Input("subdirectory", options=get_subdirectories("diffusion_models"), default="/")) + return io.Schema( + node_id=cls.__name__, + display_name="Diffusion Model Downloader", + category=cls.CATEGORY, + inputs=inputs, + outputs=[io.String.Output(display_name="status")], + is_output_node=True, + ) + + @classmethod + def execute(cls, subdirectory, recursive, nsfw_level, max_examples, threads, api_key, **kwargs): + return super().execute( + recursive=recursive, nsfw_level=nsfw_level, max_examples=max_examples, + threads=threads, api_key=api_key, subdirectory=subdirectory, **kwargs + ) + + @classmethod + def get_scan_dirs(cls, **kwargs): + base_paths = folder_paths.get_folder_paths("diffusion_models") + if not base_paths: + return [] + sub = kwargs.get("subdirectory", "/") + if sub == "/": + return base_paths + + target_paths = [] + for bp in base_paths: + target_path = os.path.join(bp, sub.lstrip("/")) + if os.path.exists(target_path): + target_paths.append(target_path) + return target_paths class LoRADownloader(BaseDownloaderNode): + @classmethod + def define_schema(cls): + inputs = list(cls.get_common_inputs().values()) + inputs.insert(0, io.Combo.Input("subdirectory", options=get_subdirectories("loras"), default="/")) + return io.Schema( + node_id=cls.__name__, + display_name="LoRA Downloader", + category=cls.CATEGORY, + inputs=inputs, + outputs=[io.String.Output(display_name="status")], + is_output_node=True, + ) + + @classmethod + def execute(cls, subdirectory, recursive, nsfw_level, max_examples, threads, api_key, **kwargs): + return super().execute( + recursive=recursive, nsfw_level=nsfw_level, max_examples=max_examples, + threads=threads, api_key=api_key, subdirectory=subdirectory, **kwargs + ) + @classmethod def get_scan_dirs(cls, **kwargs): - return folder_paths.get_folder_paths("loras") + base_paths = folder_paths.get_folder_paths("loras") + if not base_paths: + return [] + sub = kwargs.get("subdirectory", "/") + if sub == "/": + return base_paths + + target_paths = [] + for bp in base_paths: + target_path = os.path.join(bp, sub.lstrip("/")) + if os.path.exists(target_path): + target_paths.append(target_path) + return target_paths class EmbeddingDownloader(BaseDownloaderNode): + @classmethod + def define_schema(cls): + inputs = list(cls.get_common_inputs().values()) + inputs.insert(0, io.Combo.Input("subdirectory", options=get_subdirectories("embeddings"), default="/")) + return io.Schema( + node_id=cls.__name__, + display_name="Embedding Downloader", + category=cls.CATEGORY, + inputs=inputs, + outputs=[io.String.Output(display_name="status")], + is_output_node=True, + ) + + @classmethod + def execute(cls, subdirectory, recursive, nsfw_level, max_examples, threads, api_key, **kwargs): + return super().execute( + recursive=recursive, nsfw_level=nsfw_level, max_examples=max_examples, + threads=threads, api_key=api_key, subdirectory=subdirectory, **kwargs + ) + @classmethod def get_scan_dirs(cls, **kwargs): - return folder_paths.get_folder_paths("embeddings") + base_paths = folder_paths.get_folder_paths("embeddings") + if not base_paths: + return [] + sub = kwargs.get("subdirectory", "/") + if sub == "/": + return base_paths + + target_paths = [] + for bp in base_paths: + target_path = os.path.join(bp, sub.lstrip("/")) + if os.path.exists(target_path): + target_paths.append(target_path) + return target_paths class VAEDownloader(BaseDownloaderNode): + @classmethod + def define_schema(cls): + inputs = list(cls.get_common_inputs().values()) + inputs.insert(0, io.Combo.Input("subdirectory", options=get_subdirectories("vae"), default="/")) + return io.Schema( + node_id=cls.__name__, + display_name="VAE Downloader", + category=cls.CATEGORY, + inputs=inputs, + outputs=[io.String.Output(display_name="status")], + is_output_node=True, + ) + + @classmethod + def execute(cls, subdirectory, recursive, nsfw_level, max_examples, threads, api_key, **kwargs): + return super().execute( + recursive=recursive, nsfw_level=nsfw_level, max_examples=max_examples, + threads=threads, api_key=api_key, subdirectory=subdirectory, **kwargs + ) + @classmethod def get_scan_dirs(cls, **kwargs): - return folder_paths.get_folder_paths("vae") + base_paths = folder_paths.get_folder_paths("vae") + if not base_paths: + return [] + sub = kwargs.get("subdirectory", "/") + if sub == "/": + return base_paths + + target_paths = [] + for bp in base_paths: + target_path = os.path.join(bp, sub.lstrip("/")) + if os.path.exists(target_path): + target_paths.append(target_path) + return target_paths class ControlNetDownloader(BaseDownloaderNode): + @classmethod + def define_schema(cls): + inputs = list(cls.get_common_inputs().values()) + inputs.insert(0, io.Combo.Input("subdirectory", options=get_subdirectories("controlnet"), default="/")) + return io.Schema( + node_id=cls.__name__, + display_name="ControlNet Downloader", + category=cls.CATEGORY, + inputs=inputs, + outputs=[io.String.Output(display_name="status")], + is_output_node=True, + ) + + @classmethod + def execute(cls, subdirectory, recursive, nsfw_level, max_examples, threads, api_key, **kwargs): + return super().execute( + recursive=recursive, nsfw_level=nsfw_level, max_examples=max_examples, + threads=threads, api_key=api_key, subdirectory=subdirectory, **kwargs + ) + @classmethod def get_scan_dirs(cls, **kwargs): - return folder_paths.get_folder_paths("controlnet") + base_paths = folder_paths.get_folder_paths("controlnet") + if not base_paths: + return [] + sub = kwargs.get("subdirectory", "/") + if sub == "/": + return base_paths + + target_paths = [] + for bp in base_paths: + target_path = os.path.join(bp, sub.lstrip("/")) + if os.path.exists(target_path): + target_paths.append(target_path) + return target_paths class ManualPathDownloader(BaseDownloaderNode): @classmethod diff --git a/nodes/downloader_utils.py b/nodes/downloader_utils.py index 3ae40dd..7694525 100644 --- a/nodes/downloader_utils.py +++ b/nodes/downloader_utils.py @@ -548,10 +548,15 @@ def scan_and_process(scan_dirs, recursive=True, db_path=None, nsfw_level="All", path = Path(scan_dir) if recursive: for ext in model_exts: - files_to_process.extend(path.rglob(f"*{ext}")) + # rglob includes all directories, so we must filter out hidden ones + for f in path.rglob(f"*{ext}"): + if not any(part.startswith('.') for part in f.parts if part != '.' and part != '..'): + files_to_process.append(f) else: for ext in model_exts: - files_to_process.extend(path.glob(f"*{ext}")) + for f in path.glob(f"*{ext}"): + if not any(part.startswith('.') for part in f.parts if part != '.' and part != '..'): + files_to_process.append(f) nsfw_threshold = NSFW_LEVELS.get(nsfw_level, NSFW_LEVELS["All"]) From 2ac0d46935eacca85f6763cd90fe662820bf8b3f Mon Sep 17 00:00:00 2001 From: silveroxides Date: Wed, 8 Apr 2026 15:43:05 +0200 Subject: [PATCH 4/7] Add progress bar --- nodes/downloader_utils.py | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/nodes/downloader_utils.py b/nodes/downloader_utils.py index 7694525..1793a54 100644 --- a/nodes/downloader_utils.py +++ b/nodes/downloader_utils.py @@ -562,11 +562,26 @@ def scan_and_process(scan_dirs, recursive=True, db_path=None, nsfw_level="All", print(f"Found {len(files_to_process)} models to process.") + # Import and use ComfyUI's ProgressBar + try: + import comfy.utils + pbar = comfy.utils.ProgressBar(len(files_to_process)) + except ImportError: + pbar = None + if threads > 1: print(f"Starting {threads} threads...") with concurrent.futures.ThreadPoolExecutor(max_workers=threads) as executor: futures = [executor.submit(process_file, fp, cache, api, nsfw_threshold, max_examples) for fp in files_to_process] - concurrent.futures.wait(futures) + for future in concurrent.futures.as_completed(futures): + try: + future.result() + except Exception as e: + print(f"Error in thread: {e}") + if pbar: + pbar.update(1) else: for filepath in files_to_process: process_file(filepath, cache, api, nsfw_threshold, max_examples) + if pbar: + pbar.update(1) From 39380f533d3b7bcc18f2e37dfd7dd9e22f6959ff Mon Sep 17 00:00:00 2001 From: silveroxides Date: Wed, 8 Apr 2026 16:28:45 +0200 Subject: [PATCH 5/7] Remove subprocess, use av. More info in status output from Downloader nodes --- nodes/downloader_nodes.py | 37 ++- nodes/downloader_utils.py | 263 +++++++++++------- nodes/merger_utils.py | 4 +- requirements.txt | 1 + standalone_downloader.py | 561 -------------------------------------- 5 files changed, 189 insertions(+), 677 deletions(-) delete mode 100644 standalone_downloader.py diff --git a/nodes/downloader_nodes.py b/nodes/downloader_nodes.py index 6099656..439fc3f 100644 --- a/nodes/downloader_nodes.py +++ b/nodes/downloader_nodes.py @@ -8,7 +8,7 @@ def get_subdirectories(folder_name: str) -> list[str]: paths = folder_paths.get_folder_paths(folder_name) if not paths: return ["/"] - + subdirs = set() for base_path in paths: if not os.path.exists(base_path): @@ -16,14 +16,14 @@ def get_subdirectories(folder_name: str) -> list[str]: for root, dirs, _ in os.walk(base_path): # Ignore hidden directories directly during walk dirs[:] = [d for d in dirs if not d.startswith('.')] - + rel_path = os.path.relpath(root, base_path) if rel_path == ".": subdirs.add("/") else: # Store with forward slashes for consistency subdirs.add("/" + rel_path.replace("\\", "/")) - + return sorted(list(subdirs)) class BaseDownloaderNode(io.ComfyNode): @@ -56,9 +56,9 @@ class BaseDownloaderNode(io.ComfyNode): scan_dirs = cls.get_scan_dirs(**kwargs) if not scan_dirs: return io.NodeOutput("No directories to scan.") - + # Call the scan_and_process function - scan_and_process( + stats = scan_and_process( scan_dirs=scan_dirs, recursive=recursive, nsfw_level=nsfw_level, @@ -66,8 +66,19 @@ class BaseDownloaderNode(io.ComfyNode): api_key=api_key if api_key else None, threads=threads ) - - return io.NodeOutput(f"Completed scanning and downloading for {len(scan_dirs)} directories.") + + mb_saved = stats.get("space_saved", 0) / (1024 * 1024) + + output_msg = ( + f"Scan Complete for {len(scan_dirs)} directory/directories.\n" + f"Total Models Processed: {stats.get('total_processed', 0)}\n" + f"New Models Hashed: {stats.get('new_hashes', 0)}\n" + f"Images Processed: {stats.get('images_processed', 0)}\n" + f"Workflows Extracted: {stats.get('workflows_extracted', 0)}\n" + f"Space Saved via Compression: {mb_saved:.2f} MB" + ) + + return io.NodeOutput(output_msg) @classmethod def get_scan_dirs(cls, **kwargs) -> list[str]: @@ -102,7 +113,7 @@ class CheckpointDownloader(BaseDownloaderNode): sub = kwargs.get("subdirectory", "/") if sub == "/": return base_paths - + target_paths = [] for bp in base_paths: target_path = os.path.join(bp, sub.lstrip("/")) @@ -139,7 +150,7 @@ class DiffusionModelDownloader(BaseDownloaderNode): sub = kwargs.get("subdirectory", "/") if sub == "/": return base_paths - + target_paths = [] for bp in base_paths: target_path = os.path.join(bp, sub.lstrip("/")) @@ -176,7 +187,7 @@ class LoRADownloader(BaseDownloaderNode): sub = kwargs.get("subdirectory", "/") if sub == "/": return base_paths - + target_paths = [] for bp in base_paths: target_path = os.path.join(bp, sub.lstrip("/")) @@ -213,7 +224,7 @@ class EmbeddingDownloader(BaseDownloaderNode): sub = kwargs.get("subdirectory", "/") if sub == "/": return base_paths - + target_paths = [] for bp in base_paths: target_path = os.path.join(bp, sub.lstrip("/")) @@ -250,7 +261,7 @@ class VAEDownloader(BaseDownloaderNode): sub = kwargs.get("subdirectory", "/") if sub == "/": return base_paths - + target_paths = [] for bp in base_paths: target_path = os.path.join(bp, sub.lstrip("/")) @@ -287,7 +298,7 @@ class ControlNetDownloader(BaseDownloaderNode): sub = kwargs.get("subdirectory", "/") if sub == "/": return base_paths - + target_paths = [] for bp in base_paths: target_path = os.path.join(bp, sub.lstrip("/")) diff --git a/nodes/downloader_utils.py b/nodes/downloader_utils.py index 1793a54..8917afa 100644 --- a/nodes/downloader_utils.py +++ b/nodes/downloader_utils.py @@ -9,7 +9,6 @@ from pathlib import Path import threading import concurrent.futures import codecs -import subprocess import shutil try: @@ -24,7 +23,14 @@ try: HAS_MUTAGEN = True except ImportError: HAS_MUTAGEN = False - print("Warning: mutagen not found. MP4 tag extraction will rely entirely on ffmpeg fallback.") + print("Warning: mutagen not found. MP4 tag extraction will rely entirely on av fallback.") + +try: + import av + HAS_AV = True +except ImportError: + HAS_AV = False + print("Warning: av (PyAV) not found. MP4 metadata extraction fallback disabled.") # Add the reference module to the Python path # Assuming downloader_utils.py is inside nodes/, the root is one level up @@ -87,12 +93,13 @@ def get_workflow_blob(metadata): return None def process_image_workflow_and_compress(filepath_str): - if not HAS_PIL: return + if not HAS_PIL: return {"images_processed": 0, "workflows_extracted": 0, "space_saved": 0} p = Path(filepath_str) - if p.suffix.lower() not in ['.png', '.jpg', '.jpeg']: return - + if p.suffix.lower() not in ['.png', '.jpg', '.jpeg']: return {"images_processed": 0, "workflows_extracted": 0, "space_saved": 0} + workflow = None actual_format = None + stats = {"images_processed": 1, "workflows_extracted": 0, "space_saved": 0} try: with Image.open(p) as img: actual_format = img.format @@ -106,45 +113,26 @@ def process_image_workflow_and_compress(filepath_str): with open(json_path, 'w', encoding='utf-8') as f: json.dump(workflow, f, indent=2, ensure_ascii=False) print(f"Extracted workflow to {json_path}") - + stats["workflows_extracted"] = 1 + # Compress image to WebP webp_path = p.with_suffix('.webp') if not webp_path.exists(): + orig_size = os.path.getsize(p) img.save(webp_path, 'WEBP', quality=85, method=4) + new_size = os.path.getsize(webp_path) + stats["space_saved"] = max(0, orig_size - new_size) print(f"Compressed {p.name} (was actually {actual_format}) to {webp_path.name}") - + # Delete original if we created a webp replacement if workflow and p.with_suffix('.webp').exists() and actual_format == 'PNG': os.remove(p) print(f"Deleted original uncompressed image: {p.name}") except Exception as e: print(f"Failed to process image {p.name}: {e}") + return stats # --- Video Workflow Logic --- -def extract_metadata_ffmpeg(mp4_path): - if not shutil.which('ffmpeg'): return None - try: - cmd = ['ffmpeg', '-i', str(mp4_path), '-f', 'ffmetadata', '-'] - result = subprocess.run(cmd, capture_output=True, text=True, encoding='utf-8', errors='replace') - if result.returncode != 0: return None - metadata, current_key, current_value = {}, None, [] - for line in result.stdout.splitlines(): - if line.startswith(';FFMETADATA'): continue - if '=' in line: - if current_key: metadata[current_key] = '\n'.join(current_value) - parts = line.split('=', 1) - key_lower = parts[0].lower() - if key_lower in ['workflow', 'prompt', 'comment']: - current_key, current_value = key_lower, [parts[1]] - else: current_key = None - elif current_key: - current_value.append(line) - if current_key: metadata[current_key] = '\n'.join(current_value) - return {k: v.replace(r'\=', '=').replace(r'\;', ';').replace(r'\#', '#').replace(r'\\', '\\') for k, v in metadata.items()} - except Exception as e: - print(f"ffmpeg extraction error for {mp4_path}: {e}") - return None - def get_tag_value(mp4, keys): for key in keys: if key in mp4.tags: return mp4.tags[key][0] if mp4.tags[key] else None @@ -154,19 +142,29 @@ def get_tag_value(mp4, keys): if tag_key.lower().endswith(f":{search_suffix}"): return val[0] if val else None return None +def extract_metadata_av(mp4_path): + if not HAS_AV: return None + try: + with av.open(str(mp4_path)) as container: + return {k.lower(): v for k, v in container.metadata.items()} + except Exception as e: + print(f"av extraction error for {mp4_path}: {e}") + return None + def process_video_workflow(filepath_str): p = Path(filepath_str) - if p.suffix.lower() != '.mp4': return - + if p.suffix.lower() != '.mp4': return {"images_processed": 0, "workflows_extracted": 0, "space_saved": 0} + json_path = p.with_suffix('.json') - if json_path.exists(): return - + if json_path.exists(): return {"images_processed": 0, "workflows_extracted": 0, "space_saved": 0} + workflow_data = None mp4 = None + stats = {"images_processed": 0, "workflows_extracted": 0, "space_saved": 0} if HAS_MUTAGEN: try: mp4 = MP4(p) except: pass - + if mp4: for tag_keys in [['----:QuickTime:Workflow', 'workflow'], ['----:QuickTime:Prompt', 'prompt'], ['©cmt']]: val = get_tag_value(mp4, tag_keys) @@ -174,22 +172,28 @@ def process_video_workflow(filepath_str): if isinstance(val, bytes): val = val.decode('utf-8', errors='ignore') try: parsed = json.loads(val) - if isinstance(parsed, dict) and ('prompt' in parsed or any(k.isdigit() for k in parsed.keys())): + if isinstance(parsed, dict) and ('prompt' in parsed or any(k.isdigit() for k in parsed.keys()) or 'nodes' in parsed): if 'prompt' in parsed and isinstance(parsed['prompt'], str): try: workflow_data = json.loads(parsed['prompt']) except: workflow_data = parsed else: workflow_data = parsed break - except json.JSONDecodeError: pass - + except json.JSONDecodeError: + # Try more aggressive parsing if it's a nested string + parsed = unescape_and_parse_nested_json(val) + if parsed: + workflow_data = parsed + break + if not workflow_data: - ffmpeg_data = extract_metadata_ffmpeg(p) - if ffmpeg_data: - for key in ['workflow', 'prompt', 'comment']: - if key in ffmpeg_data: + av_data = extract_metadata_av(p) + if av_data: + for key in ['workflow', 'prompt', 'comment', 'description']: + if key in av_data: + val = av_data[key] try: - parsed = json.loads(ffmpeg_data[key]) + parsed = json.loads(val) if isinstance(parsed, dict): if 'prompt' in parsed and isinstance(parsed['prompt'], str): try: workflow_data = json.loads(parsed['prompt']) @@ -197,15 +201,21 @@ def process_video_workflow(filepath_str): else: workflow_data = parsed break - except json.JSONDecodeError: pass - + except json.JSONDecodeError: + parsed = unescape_and_parse_nested_json(val) + if parsed: + workflow_data = parsed + break + if workflow_data: try: with open(json_path, 'w', encoding='utf-8') as f: json.dump(workflow_data, f, indent=2, ensure_ascii=False) print(f"Extracted workflow to {json_path.name}") + stats["workflows_extracted"] = 1 except Exception as e: print(f"Error saving workflow for {p.name}: {e}") + return stats # --- Constants for Civitai NSFW levels --- NSFW_LEVELS = { @@ -276,7 +286,7 @@ class HashCache: mtime = os.path.getmtime(filepath) size = os.path.getsize(filepath) key = filepath - + with self.lock: if self.is_sqlite: cached_hash = self._get_sqlite_hash(key, mtime, size) @@ -286,11 +296,11 @@ class HashCache: entry = self.cache[key] if entry["mtime"] == mtime and entry["size"] == size: return entry["hash"] - + # Calculate new hash (outside the lock to allow parallel hashing) print(f"Calculating SHA256 for: {filepath}") sha256_hash = self.calculate_sha256(filepath) - + with self.lock: if self.is_sqlite: self._set_sqlite_hash(key, mtime, size, sha256_hash) @@ -301,7 +311,7 @@ class HashCache: "hash": sha256_hash } self._save_json() - + return sha256_hash @staticmethod @@ -339,40 +349,40 @@ class Downloader: def download_file(url, target_path, headers=None): if headers is None: headers = {} - + temp_path = target_path + ".downloading" - + downloaded_size = 0 if os.path.exists(temp_path): downloaded_size = os.path.getsize(temp_path) - + if downloaded_size > 0: headers["Range"] = f"bytes={downloaded_size}-" - + response = requests.get(url, headers=headers, stream=True) - + if response.status_code == 416: # Range not satisfiable, file probably changed or completed os.remove(temp_path) downloaded_size = 0 if "Range" in headers: del headers["Range"] response = requests.get(url, headers=headers, stream=True) - + if response.status_code not in [200, 206]: print(f"Failed to download from {url}, status code: {response.status_code}") return False mode = "ab" if response.status_code == 206 else "wb" - + total_size = int(response.headers.get('content-length', 0)) + downloaded_size - + print(f"Downloading: {target_path} ({downloaded_size}/{total_size} bytes)") - + with open(temp_path, mode) as f: for chunk in response.iter_content(chunk_size=8192): if chunk: f.write(chunk) - + os.rename(temp_path, target_path) return True @@ -422,34 +432,44 @@ def process_file(filepath, cache, api, nsfw_threshold, max_examples): filepath_str = str(filepath) base_path = os.path.splitext(filepath_str)[0] is_safetensors = filepath_str.lower().endswith(".safetensors") - + + stats = {"new_hashes": 0, "images_processed": 0, "workflows_extracted": 0, "space_saved": 0} + # Check if all targets already exist info_path = f"{base_path}.civitai.info" st_metadata_path = f"{base_path}.metadata.json" - + has_info = os.path.exists(info_path) has_preview = preview_exists(base_path) has_st_metadata = not is_safetensors or os.path.exists(st_metadata_path) - + # Process existing images/videos for workflow extraction & compression BEFORE we possibly skip the file for p in get_potential_preview_files(base_path): if os.path.exists(p): ext_lower = p.lower() - if ext_lower.endswith(('.png', '.jpg', '.jpeg')): process_image_workflow_and_compress(p) - elif ext_lower.endswith('.mp4'): process_video_workflow(p) - + if ext_lower.endswith(('.png', '.jpg', '.jpeg')): + res = process_image_workflow_and_compress(p) + for k in stats: stats[k] += res.get(k, 0) + elif ext_lower.endswith('.mp4'): + res = process_video_workflow(p) + for k in stats: stats[k] += res.get(k, 0) + for ex in glob.glob(f"{base_path}.example.*"): ext_lower = ex.lower() - if ext_lower.endswith(('.png', '.jpg', '.jpeg')): process_image_workflow_and_compress(ex) - elif ext_lower.endswith('.mp4'): process_video_workflow(ex) - + if ext_lower.endswith(('.png', '.jpg', '.jpeg')): + res = process_image_workflow_and_compress(ex) + for k in stats: stats[k] += res.get(k, 0) + elif ext_lower.endswith('.mp4'): + res = process_video_workflow(ex) + for k in stats: stats[k] += res.get(k, 0) + # If everything exists and we don't want examples or already have enough examples, we might skip API logic if has_info and has_preview and has_st_metadata and max_examples == 0: print(f"Skipping fully populated model: {filepath_str}") - return - + return stats + print(f"Processing model: {filepath_str}") - + # Safetensors built-in metadata extraction if is_safetensors and not has_st_metadata: metadata = extract_safetensors_metadata(filepath_str) @@ -457,34 +477,52 @@ def process_file(filepath, cache, api, nsfw_threshold, max_examples): with open(st_metadata_path, "w", encoding="utf-8") as f: json.dump(metadata, f, indent=4) print(f"Saved safetensors metadata to {st_metadata_path}") - + + mtime = os.path.getmtime(filepath_str) + size = os.path.getsize(filepath_str) + + # We want to know if it's a NEW hash or cached + is_cached = False + with cache.lock: + if cache.is_sqlite: + cached_hash = cache._get_sqlite_hash(filepath_str, mtime, size) + if cached_hash: is_cached = True + else: + if filepath_str in cache.cache: + entry = cache.cache[filepath_str] + if entry["mtime"] == mtime and entry["size"] == size: + is_cached = True + file_hash = cache.get_hash(filepath_str) if not file_hash: - return - + return stats + + if not is_cached: + stats["new_hashes"] += 1 + try: model_info = api.get_model_info_by_hash(file_hash) except Exception as e: print(f"API Error for {filepath_str}: {e}") - return - + return stats + if not model_info: print(f"Model not found on Civitai: {filepath_str}") - return - + return stats + # 1. Save Info if not has_info: with open(info_path, "w", encoding="utf-8") as f: json.dump(model_info, f, indent=4) print(f"Saved metadata to {info_path}") - + images = model_info.get("images", []) valid_images = filter_images_by_nsfw(images, nsfw_threshold) - + if not valid_images: print(f"No valid images found matching NSFW criteria for {filepath_str}") - return - + return stats + # 2. Save Preview if not has_preview and len(valid_images) > 0: preview_img = valid_images[0] @@ -493,12 +531,16 @@ def process_file(filepath, cache, api, nsfw_threshold, max_examples): if ".png" in url.lower(): ext = "png" elif ".mp4" in url.lower(): ext = "mp4" elif ".jpg" in url.lower() or ".jpeg" in url.lower(): ext = "jpeg" - + preview_path = f"{base_path}.preview.{ext}" if Downloader.download_file(url, preview_path): - if ext in ("png", "jpg", "jpeg"): process_image_workflow_and_compress(preview_path) - elif ext == "mp4": process_video_workflow(preview_path) - + if ext in ("png", "jpg", "jpeg"): + res = process_image_workflow_and_compress(preview_path) + for k in stats: stats[k] += res.get(k, 0) + elif ext == "mp4": + res = process_video_workflow(preview_path) + for k in stats: stats[k] += res.get(k, 0) + # 3. Save Examples if max_examples > 0: existing_examples_count = len(glob.glob(f"{base_path}.example.*")) @@ -509,25 +551,30 @@ def process_file(filepath, cache, api, nsfw_threshold, max_examples): if len(parts) > 1: idx = parts[1].split('.')[0] if idx.isdigit(): existing_indices.add(idx) - + needed_examples = max_examples - len(existing_indices) - + if needed_examples > 0: example_images = valid_images[1:max_examples+1] # Skip images we probably already downloaded by skipping the first N example_images_to_dl = example_images[len(existing_indices):] - + for img in example_images_to_dl: url = img.get("url") ext = "jpeg" if ".png" in url.lower(): ext = "png" elif ".mp4" in url.lower(): ext = "mp4" elif ".jpg" in url.lower() or ".jpeg" in url.lower(): ext = "jpeg" - + example_path = next_example_image_path(base_path, ext) if Downloader.download_file(url, example_path): - if ext in ("png", "jpg", "jpeg"): process_image_workflow_and_compress(example_path) - elif ext == "mp4": process_video_workflow(example_path) + if ext in ("png", "jpg", "jpeg"): + res = process_image_workflow_and_compress(example_path) + for k in stats: stats[k] += res.get(k, 0) + elif ext == "mp4": + res = process_video_workflow(example_path) + for k in stats: stats[k] += res.get(k, 0) + return stats def scan_and_process(scan_dirs, recursive=True, db_path=None, nsfw_level="All", max_examples=0, api_key=None, threads=4): """ @@ -537,13 +584,13 @@ def scan_and_process(scan_dirs, recursive=True, db_path=None, nsfw_level="All", db_path = db_path if db_path else DEFAULT_DB_PATH cache = HashCache(db_path) api = CivitaiAPI(api_key) - + model_exts = [".safetensors", ".ckpt", ".pt"] - + files_to_process = [] if isinstance(scan_dirs, str): scan_dirs = [scan_dirs] - + for scan_dir in scan_dirs: path = Path(scan_dir) if recursive: @@ -557,31 +604,45 @@ def scan_and_process(scan_dirs, recursive=True, db_path=None, nsfw_level="All", for f in path.glob(f"*{ext}"): if not any(part.startswith('.') for part in f.parts if part != '.' and part != '..'): files_to_process.append(f) - + nsfw_threshold = NSFW_LEVELS.get(nsfw_level, NSFW_LEVELS["All"]) print(f"Found {len(files_to_process)} models to process.") - + + global_stats = { + "total_processed": len(files_to_process), + "new_hashes": 0, + "images_processed": 0, + "workflows_extracted": 0, + "space_saved": 0 + } + # Import and use ComfyUI's ProgressBar try: import comfy.utils pbar = comfy.utils.ProgressBar(len(files_to_process)) except ImportError: pbar = None - + if threads > 1: print(f"Starting {threads} threads...") with concurrent.futures.ThreadPoolExecutor(max_workers=threads) as executor: futures = [executor.submit(process_file, fp, cache, api, nsfw_threshold, max_examples) for fp in files_to_process] for future in concurrent.futures.as_completed(futures): try: - future.result() + res = future.result() + if res: + for k in global_stats: global_stats[k] += res.get(k, 0) except Exception as e: print(f"Error in thread: {e}") if pbar: pbar.update(1) else: for filepath in files_to_process: - process_file(filepath, cache, api, nsfw_threshold, max_examples) + res = process_file(filepath, cache, api, nsfw_threshold, max_examples) + if res: + for k in global_stats: global_stats[k] += res.get(k, 0) if pbar: pbar.update(1) + + return global_stats diff --git a/nodes/merger_utils.py b/nodes/merger_utils.py index d68e200..dbe56a9 100644 --- a/nodes/merger_utils.py +++ b/nodes/merger_utils.py @@ -65,8 +65,8 @@ def transfer_to_gpu_pinned( class MemoryEfficientSafeOpen: """Memory-efficient safetensors file reader with streaming and pinned memory support. - - Operates strictly in streaming mode (low memory). Tensors are read byte-by-byte from + + Operates strictly in streaming mode (low memory). Tensors are read byte-by-byte from their exact offsets without using safe_open or memory mapping, ensuring tensors do not remain persistently in OS page caches or RAM. """ diff --git a/requirements.txt b/requirements.txt index a3a3c01..1faba79 100644 --- a/requirements.txt +++ b/requirements.txt @@ -2,3 +2,4 @@ requests Pillow mutagen unifiedefficientloader +av diff --git a/standalone_downloader.py b/standalone_downloader.py deleted file mode 100644 index c06d965..0000000 --- a/standalone_downloader.py +++ /dev/null @@ -1,561 +0,0 @@ -import os -import sys -import argparse -import hashlib -import json -import sqlite3 -import requests -import glob -from pathlib import Path -import threading -import concurrent.futures -import codecs -import subprocess -import shutil - -try: - from PIL import Image - HAS_PIL = True -except ImportError: - HAS_PIL = False - print("Warning: Pillow not found. Image workflow extraction disabled.") - -try: - from mutagen.mp4 import MP4 - HAS_MUTAGEN = True -except ImportError: - HAS_MUTAGEN = False - print("Warning: mutagen not found. MP4 tag extraction will rely entirely on ffmpeg fallback.") - -# Add the reference module to the Python path -sys.path.append(os.path.join(os.path.dirname(__file__), 'unifiedefficientloader-reference')) -try: - from unifiedefficientloader import UnifiedSafetensorsLoader - HAS_UNIFIED_LOADER = True -except ImportError: - HAS_UNIFIED_LOADER = False - print("Warning: unifiedefficientloader not found. Safetensors metadata extraction disabled.") - -# --- Workflow Extraction Logic --- - -def unescape_and_parse_nested_json(json_string): - if not isinstance(json_string, str): return None - stripped = json_string.strip() - if not (stripped.startswith('{') or stripped.startswith('[')): return None - try: return json.loads(stripped) - except json.JSONDecodeError: pass - try: - first_parse = json.loads(stripped) - if isinstance(first_parse, str): return json.loads(first_parse) - except json.JSONDecodeError: pass - try: - decoded = codecs.decode(stripped, "unicode_escape") - return json.loads(decoded) - except (json.JSONDecodeError, UnicodeDecodeError, ValueError): pass - try: - unescaped = stripped.replace('\\\\"', '\\"').replace('\\"', '"').replace("\\\\", "\\") - unescaped = unescaped.replace("\\n", "\n").replace("\\r", "\r").replace("\\t", "\t") - return json.loads(unescaped) - except json.JSONDecodeError: pass - return None - -def get_workflow_blob(metadata): - for key in ['workflow', 'prompt']: - val = metadata.get(key) - if val: - parsed = unescape_and_parse_nested_json(val) - if parsed: return parsed - for val in metadata.values(): - parsed = unescape_and_parse_nested_json(val) - if isinstance(parsed, dict) and (any(k.isdigit() for k in parsed.keys()) or 'nodes' in parsed): - return parsed - return None - -def process_image_workflow_and_compress(filepath_str): - if not HAS_PIL: return - p = Path(filepath_str) - if p.suffix.lower() not in ['.png', '.jpg', '.jpeg']: return - - workflow = None - actual_format = None - try: - with Image.open(p) as img: - actual_format = img.format - # CivitAI sometimes serves PNGs with .jpeg extension in URL - if actual_format == 'PNG': - workflow = get_workflow_blob(img.info) - if workflow: - # Save workflow - json_path = p.with_suffix('.json') - if not json_path.exists(): - with open(json_path, 'w', encoding='utf-8') as f: - json.dump(workflow, f, indent=2, ensure_ascii=False) - print(f"Extracted workflow to {json_path}") - - # Compress image to WebP - webp_path = p.with_suffix('.webp') - if not webp_path.exists(): - img.save(webp_path, 'WEBP', quality=85, method=4) - print(f"Compressed {p.name} (was actually {actual_format}) to {webp_path.name}") - - # Delete original if we created a webp replacement - if workflow and p.with_suffix('.webp').exists() and actual_format == 'PNG': - os.remove(p) - print(f"Deleted original uncompressed image: {p.name}") - except Exception as e: - print(f"Failed to process image {p.name}: {e}") - -# --- Video Workflow Logic --- -def extract_metadata_ffmpeg(mp4_path): - if not shutil.which('ffmpeg'): return None - try: - cmd = ['ffmpeg', '-i', str(mp4_path), '-f', 'ffmetadata', '-'] - result = subprocess.run(cmd, capture_output=True, text=True, encoding='utf-8', errors='replace') - if result.returncode != 0: return None - metadata, current_key, current_value = {}, None, [] - for line in result.stdout.splitlines(): - if line.startswith(';FFMETADATA'): continue - if '=' in line: - if current_key: metadata[current_key] = '\n'.join(current_value) - parts = line.split('=', 1) - key_lower = parts[0].lower() - if key_lower in ['workflow', 'prompt', 'comment']: - current_key, current_value = key_lower, [parts[1]] - else: current_key = None - elif current_key: - current_value.append(line) - if current_key: metadata[current_key] = '\n'.join(current_value) - return {k: v.replace(r'\=', '=').replace(r'\;', ';').replace(r'\#', '#').replace(r'\\', '\\') for k, v in metadata.items()} - except Exception as e: - print(f"ffmpeg extraction error for {mp4_path}: {e}") - return None - -def get_tag_value(mp4, keys): - for key in keys: - if key in mp4.tags: return mp4.tags[key][0] if mp4.tags[key] else None - for key in keys: - search_suffix = key.split(':')[-1].lower() - for tag_key, val in mp4.tags.items(): - if tag_key.lower().endswith(f":{search_suffix}"): return val[0] if val else None - return None - -def process_video_workflow(filepath_str): - p = Path(filepath_str) - if p.suffix.lower() != '.mp4': return - - json_path = p.with_suffix('.json') - if json_path.exists(): return - - workflow_data = None - mp4 = None - if HAS_MUTAGEN: - try: mp4 = MP4(p) - except: pass - - if mp4: - for tag_keys in [['----:QuickTime:Workflow', 'workflow'], ['----:QuickTime:Prompt', 'prompt'], ['©cmt']]: - val = get_tag_value(mp4, tag_keys) - if val: - if isinstance(val, bytes): val = val.decode('utf-8', errors='ignore') - try: - parsed = json.loads(val) - if isinstance(parsed, dict) and ('prompt' in parsed or any(k.isdigit() for k in parsed.keys())): - if 'prompt' in parsed and isinstance(parsed['prompt'], str): - try: workflow_data = json.loads(parsed['prompt']) - except: workflow_data = parsed - else: - workflow_data = parsed - break - except json.JSONDecodeError: pass - - if not workflow_data: - ffmpeg_data = extract_metadata_ffmpeg(p) - if ffmpeg_data: - for key in ['workflow', 'prompt', 'comment']: - if key in ffmpeg_data: - try: - parsed = json.loads(ffmpeg_data[key]) - if isinstance(parsed, dict): - if 'prompt' in parsed and isinstance(parsed['prompt'], str): - try: workflow_data = json.loads(parsed['prompt']) - except: workflow_data = parsed - else: - workflow_data = parsed - break - except json.JSONDecodeError: pass - - if workflow_data: - try: - with open(json_path, 'w', encoding='utf-8') as f: - json.dump(workflow_data, f, indent=2, ensure_ascii=False) - print(f"Extracted workflow to {json_path.name}") - except Exception as e: - print(f"Error saving workflow for {p.name}: {e}") - -# --- Constants for Civitai NSFW levels --- -NSFW_LEVELS = { - "None": 1, - "Soft": 2, - "Mature": 4, - "X": 8, - "XXX": 16, - "All": 32 # Anything -} - -class HashCache: - def __init__(self, db_path): - self.db_path = db_path - self.is_sqlite = db_path.endswith('.db') or db_path.endswith('.sqlite') - self.lock = threading.Lock() - if self.is_sqlite: - self._init_sqlite() - else: - self.cache = self._load_json() - - def _init_sqlite(self): - with sqlite3.connect(self.db_path) as conn: - cursor = conn.cursor() - cursor.execute(''' - CREATE TABLE IF NOT EXISTS file_hashes ( - filepath TEXT PRIMARY KEY, - mtime REAL, - size INTEGER, - hash TEXT - ) - ''') - conn.commit() - - def _load_json(self): - if os.path.exists(self.db_path): - with open(self.db_path, "r", encoding="utf-8") as f: - return json.load(f) - return {} - - def _save_json(self): - with open(self.db_path, "w", encoding="utf-8") as f: - json.dump(self.cache, f, indent=4) - - def _get_sqlite_hash(self, filepath, mtime, size): - with sqlite3.connect(self.db_path) as conn: - cursor = conn.cursor() - cursor.execute("SELECT mtime, size, hash FROM file_hashes WHERE filepath=?", (filepath,)) - row = cursor.fetchone() - if row: - db_mtime, db_size, db_hash = row - if db_mtime == mtime and db_size == size: - return db_hash - return None - - def _set_sqlite_hash(self, filepath, mtime, size, file_hash): - with sqlite3.connect(self.db_path, timeout=10.0) as conn: - cursor = conn.cursor() - cursor.execute(''' - INSERT OR REPLACE INTO file_hashes (filepath, mtime, size, hash) - VALUES (?, ?, ?, ?) - ''', (filepath, mtime, size, file_hash)) - conn.commit() - - def get_hash(self, filepath): - if not os.path.exists(filepath): - return None - mtime = os.path.getmtime(filepath) - size = os.path.getsize(filepath) - key = filepath - - with self.lock: - if self.is_sqlite: - cached_hash = self._get_sqlite_hash(key, mtime, size) - if cached_hash: return cached_hash - else: - if key in self.cache: - entry = self.cache[key] - if entry["mtime"] == mtime and entry["size"] == size: - return entry["hash"] - - # Calculate new hash (outside the lock to allow parallel hashing) - print(f"Calculating SHA256 for: {filepath}") - sha256_hash = self.calculate_sha256(filepath) - - with self.lock: - if self.is_sqlite: - self._set_sqlite_hash(key, mtime, size, sha256_hash) - else: - self.cache[key] = { - "mtime": mtime, - "size": size, - "hash": sha256_hash - } - self._save_json() - - return sha256_hash - - @staticmethod - def calculate_sha256(filename): - sha256_hash = hashlib.sha256() - with open(filename, "rb") as f: - # Read and update hash string value in blocks of 4K - for byte_block in iter(lambda: f.read(4096), b""): - sha256_hash.update(byte_block) - return sha256_hash.hexdigest() - -class CivitaiAPI: - def __init__(self, api_key=None): - self.api_key = api_key - self.base_url = "https://civitai.com/api/v1" - - def get_headers(self): - headers = {} - if self.api_key: - headers["Authorization"] = f"Bearer {self.api_key}" - return headers - - def get_model_info_by_hash(self, file_hash): - url = f"{self.base_url}/model-versions/by-hash/{file_hash}" - response = requests.get(url, headers=self.get_headers()) - if response.status_code == 200: - return response.json() - elif response.status_code == 404: - return None - else: - response.raise_for_status() - -class Downloader: - @staticmethod - def download_file(url, target_path, headers=None): - if headers is None: - headers = {} - - temp_path = target_path + ".downloading" - - downloaded_size = 0 - if os.path.exists(temp_path): - downloaded_size = os.path.getsize(temp_path) - - if downloaded_size > 0: - headers["Range"] = f"bytes={downloaded_size}-" - - response = requests.get(url, headers=headers, stream=True) - - if response.status_code == 416: # Range not satisfiable, file probably changed or completed - os.remove(temp_path) - downloaded_size = 0 - if "Range" in headers: - del headers["Range"] - response = requests.get(url, headers=headers, stream=True) - - if response.status_code not in [200, 206]: - print(f"Failed to download from {url}, status code: {response.status_code}") - return False - - mode = "ab" if response.status_code == 206 else "wb" - - total_size = int(response.headers.get('content-length', 0)) + downloaded_size - - print(f"Downloading: {target_path} ({downloaded_size}/{total_size} bytes)") - - with open(temp_path, mode) as f: - for chunk in response.iter_content(chunk_size=8192): - if chunk: - f.write(chunk) - - os.rename(temp_path, target_path) - return True - -def get_potential_preview_files(base_path): - exts = ["png", "jpg", "jpeg", "webp", "mp4"] - files = [] - for ext in exts: - files.append(f"{base_path}.preview.{ext}") - files.append(f"{base_path}.{ext}") - return files - -def preview_exists(base_path): - previews = get_potential_preview_files(base_path) - for p in previews: - if os.path.exists(p): - return True - return False - -def next_example_image_path(base_path, ext): - i = 0 - while len(glob.glob(f"{base_path}.example.{i}.*")) > 0: - i += 1 - return f"{base_path}.example.{i}.{ext}" - -def filter_images_by_nsfw(images, threshold_level): - valid_images = [] - for img in images: - rating = img.get("nsfwLevel", 1) - if rating <= threshold_level: - valid_images.append(img) - return valid_images - -def extract_safetensors_metadata(filepath): - if not HAS_UNIFIED_LOADER: - return None - try: - # Use low_memory=True to only read the header and metadata without loading tensors - loader = UnifiedSafetensorsLoader(filepath, low_memory=True) - metadata = loader.metadata() - loader.close() - return metadata if metadata else None - except Exception as e: - print(f"Failed to read safetensors metadata from {filepath} using unifiedefficientloader: {e}") - return None - -def process_file(filepath, args, cache, api, nsfw_threshold): - filepath_str = str(filepath) - base_path = os.path.splitext(filepath_str)[0] - is_safetensors = filepath_str.lower().endswith(".safetensors") - - # Check if all targets already exist - info_path = f"{base_path}.civitai.info" - st_metadata_path = f"{base_path}.metadata.json" - - has_info = os.path.exists(info_path) - has_preview = preview_exists(base_path) - has_st_metadata = not is_safetensors or os.path.exists(st_metadata_path) - - # Process existing images/videos for workflow extraction & compression BEFORE we possibly skip the file - for p in get_potential_preview_files(base_path): - if os.path.exists(p): - ext_lower = p.lower() - if ext_lower.endswith(('.png', '.jpg', '.jpeg')): process_image_workflow_and_compress(p) - elif ext_lower.endswith('.mp4'): process_video_workflow(p) - - for ex in glob.glob(f"{base_path}.example.*"): - ext_lower = ex.lower() - if ext_lower.endswith(('.png', '.jpg', '.jpeg')): process_image_workflow_and_compress(ex) - elif ext_lower.endswith('.mp4'): process_video_workflow(ex) - - # If everything exists and we don't want examples or already have enough examples, we might skip API logic - if has_info and has_preview and has_st_metadata and args.max_examples == 0: - print(f"Skipping fully populated model: {filepath_str}") - return - - print(f"Processing model: {filepath_str}") - - # Safetensors built-in metadata extraction - if is_safetensors and not has_st_metadata: - metadata = extract_safetensors_metadata(filepath_str) - if metadata: - with open(st_metadata_path, "w", encoding="utf-8") as f: - json.dump(metadata, f, indent=4) - print(f"Saved safetensors metadata to {st_metadata_path}") - - file_hash = cache.get_hash(filepath_str) - if not file_hash: - return - - try: - model_info = api.get_model_info_by_hash(file_hash) - except Exception as e: - print(f"API Error for {filepath_str}: {e}") - return - - if not model_info: - print(f"Model not found on Civitai: {filepath_str}") - return - - # 1. Save Info - if not has_info: - with open(info_path, "w", encoding="utf-8") as f: - json.dump(model_info, f, indent=4) - print(f"Saved metadata to {info_path}") - - images = model_info.get("images", []) - valid_images = filter_images_by_nsfw(images, nsfw_threshold) - - if not valid_images: - print(f"No valid images found matching NSFW criteria for {filepath_str}") - return - - # 2. Save Preview - if not has_preview and len(valid_images) > 0: - preview_img = valid_images[0] - url = preview_img.get("url") - ext = "jpeg" # Default fallback - if ".png" in url.lower(): ext = "png" - elif ".mp4" in url.lower(): ext = "mp4" - elif ".jpg" in url.lower() or ".jpeg" in url.lower(): ext = "jpeg" - - preview_path = f"{base_path}.preview.{ext}" - if Downloader.download_file(url, preview_path): - if ext in ("png", "jpg", "jpeg"): process_image_workflow_and_compress(preview_path) - elif ext == "mp4": process_video_workflow(preview_path) - - # 3. Save Examples - if args.max_examples > 0: - existing_examples_count = len(glob.glob(f"{base_path}.example.*")) - # We divide by 2 roughly if there are json files, let's just count unique example indices - existing_indices = set() - for f in glob.glob(f"{base_path}.example.*"): - parts = f.split('.example.') - if len(parts) > 1: - idx = parts[1].split('.')[0] - if idx.isdigit(): existing_indices.add(idx) - - needed_examples = args.max_examples - len(existing_indices) - - if needed_examples > 0: - example_images = valid_images[1:args.max_examples+1] - # Skip images we probably already downloaded by skipping the first N - example_images_to_dl = example_images[len(existing_indices):] - - for img in example_images_to_dl: - url = img.get("url") - ext = "jpeg" - if ".png" in url.lower(): ext = "png" - elif ".mp4" in url.lower(): ext = "mp4" - elif ".jpg" in url.lower() or ".jpeg" in url.lower(): ext = "jpeg" - - example_path = next_example_image_path(base_path, ext) - if Downloader.download_file(url, example_path): - if ext in ("png", "jpg", "jpeg"): process_image_workflow_and_compress(example_path) - elif ext == "mp4": process_video_workflow(example_path) - -def scan_and_process(args): - cache = HashCache(args.db_path) - api = CivitaiAPI(args.api_key) - - model_exts = [".safetensors", ".ckpt", ".pt"] - - files_to_process = [] - for scan_dir in args.scan_dirs: - path = Path(scan_dir) - if args.recursive: - for ext in model_exts: - files_to_process.extend(path.rglob(f"*{ext}")) - else: - for ext in model_exts: - files_to_process.extend(path.glob(f"*{ext}")) - - nsfw_threshold = NSFW_LEVELS.get(args.nsfw_level, 1) - - print(f"Found {len(files_to_process)} models to process.") - - if args.threads > 1: - print(f"Starting {args.threads} threads...") - with concurrent.futures.ThreadPoolExecutor(max_workers=args.threads) as executor: - futures = [executor.submit(process_file, fp, args, cache, api, nsfw_threshold) for fp in files_to_process] - concurrent.futures.wait(futures) - else: - for filepath in files_to_process: - process_file(filepath, args, cache, api, nsfw_threshold) - -def main(): - parser = argparse.ArgumentParser(description="Standalone Civitai Model Downloader") - parser.add_argument("--scan-dirs", nargs="+", required=True, help="Directories to scan for models") - parser.add_argument("--recursive", action="store_true", help="Recursively scan directories") - parser.add_argument("--db-path", default="model_hash_cache.json", help="Path to hash cache JSON") - parser.add_argument("--nsfw-level", choices=list(NSFW_LEVELS.keys()), default="Soft", help="Max NSFW level for images") - parser.add_argument("--max-examples", type=int, default=0, help="Max number of example images to download") - parser.add_argument("--max-size-preview", action="store_true", help="Download max resolution previews (Not fully implemented, fallback to default URL)") - parser.add_argument("--api-key", default=None, help="Civitai API Key") - parser.add_argument("--threads", type=int, default=1, help="Number of concurrent threads for hashing and downloading") - - args = parser.parse_args() - scan_and_process(args) - -if __name__ == "__main__": - main() \ No newline at end of file From dc43040562709549d97775530f2d31a4eb142db8 Mon Sep 17 00:00:00 2001 From: silveroxides Date: Wed, 8 Apr 2026 16:46:35 +0200 Subject: [PATCH 6/7] Rename downloader nodes and add nodes for displaying previews, metadata of model and workflows extracted from images/videos --- __init__.py | 15 +++- nodes/downloader_nodes.py | 27 +++--- nodes/downloader_utils.py | 54 +++++++++++ nodes/model_info_nodes.py | 184 ++++++++++++++++++++++++++++++++++++++ 4 files changed, 263 insertions(+), 17 deletions(-) create mode 100644 nodes/model_info_nodes.py diff --git a/__init__.py b/__init__.py index f78e99f..8896830 100644 --- a/__init__.py +++ b/__init__.py @@ -30,8 +30,12 @@ from .nodes.lora_resize import ( LoRAMergeToModel ) from .nodes.downloader_nodes import ( - CheckpointDownloader, DiffusionModelDownloader, LoRADownloader, EmbeddingDownloader, - VAEDownloader, ControlNetDownloader, ManualPathDownloader + CheckpointInfoMetaDownloader, DiffusionModelInfoMetaDownloader, LoRAInfoMetaDownloader, EmbeddingInfoMetaDownloader, + VAEInfoMetaDownloader, ControlNetInfoMetaDownloader, ManualPathInfoMetaDownloader +) +from .nodes.model_info_nodes import ( + CheckpointInfoLoader, LoRAInfoLoader, EmbeddingInfoLoader, + VAEInfoLoader, ControlNetInfoLoader, DiffusionModelInfoLoader ) @@ -65,8 +69,11 @@ class ModelUtilsExtension(ComfyExtension): # LoRA Merge To Model LoRAMergeToModel, # Downloaders - CheckpointDownloader, DiffusionModelDownloader, LoRADownloader, EmbeddingDownloader, - VAEDownloader, ControlNetDownloader, ManualPathDownloader, + CheckpointInfoMetaDownloader, DiffusionModelInfoMetaDownloader, LoRAInfoMetaDownloader, EmbeddingInfoMetaDownloader, + VAEInfoMetaDownloader, ControlNetInfoMetaDownloader, ManualPathInfoMetaDownloader, + # Info Loaders + CheckpointInfoLoader, LoRAInfoLoader, EmbeddingInfoLoader, + VAEInfoLoader, ControlNetInfoLoader, DiffusionModelInfoLoader, ] diff --git a/nodes/downloader_nodes.py b/nodes/downloader_nodes.py index 439fc3f..f6f15ef 100644 --- a/nodes/downloader_nodes.py +++ b/nodes/downloader_nodes.py @@ -26,18 +26,18 @@ def get_subdirectories(folder_name: str) -> list[str]: return sorted(list(subdirs)) -class BaseDownloaderNode(io.ComfyNode): +class BaseInfoMetaDownloaderNode(io.ComfyNode): CATEGORY = "ModelUtils/Downloader" OUTPUT_NODE = True @classmethod def get_common_inputs(cls): return { - "recursive": io.Boolean.Input("recursive", default=True), - "nsfw_level": io.Combo.Input("nsfw_level", options=["None", "Soft", "Mature", "X", "XXX", "All"], default="All"), - "max_examples": io.Int.Input("max_examples", default=0, min=0, max=100), - "threads": io.Int.Input("threads", default=4, min=1, max=32), - "api_key": io.String.Input("api_key", default="") + "recursive": io.Boolean.Input("recursive", default=True, tooltip="Whether to scan subdirectories recursively."), + "nsfw_level": io.Combo.Input("nsfw_level", options=["None", "Soft", "Mature", "X", "XXX", "All"], default="All", tooltip="Filter models based on NSFW level in metadata."), + "max_examples": io.Int.Input("max_examples", default=0, min=0, max=100, tooltip="Maximum number of examples to process."), + "threads": io.Int.Input("threads", default=4, min=1, max=32, tooltip="Number of threads to use for processing."), + "api_key": io.String.Input("api_key", default="", tooltip="API key for accessing the model repository.") } @classmethod @@ -49,6 +49,7 @@ class BaseDownloaderNode(io.ComfyNode): inputs=list(cls.get_common_inputs().values()), outputs=[io.String.Output(display_name="status")], is_output_node=True, + description=("Scans specified directories for models, extracts metadata, generates previews, and optionally compresses files.") ) @classmethod @@ -84,7 +85,7 @@ class BaseDownloaderNode(io.ComfyNode): def get_scan_dirs(cls, **kwargs) -> list[str]: return [] -class CheckpointDownloader(BaseDownloaderNode): +class CheckpointInfoMetaDownloader(BaseInfoMetaDownloaderNode): @classmethod def define_schema(cls): inputs = list(cls.get_common_inputs().values()) @@ -121,7 +122,7 @@ class CheckpointDownloader(BaseDownloaderNode): target_paths.append(target_path) return target_paths -class DiffusionModelDownloader(BaseDownloaderNode): +class DiffusionModelInfoMetaDownloader(BaseInfoMetaDownloaderNode): @classmethod def define_schema(cls): inputs = list(cls.get_common_inputs().values()) @@ -158,7 +159,7 @@ class DiffusionModelDownloader(BaseDownloaderNode): target_paths.append(target_path) return target_paths -class LoRADownloader(BaseDownloaderNode): +class LoRAInfoMetaDownloader(BaseInfoMetaDownloaderNode): @classmethod def define_schema(cls): inputs = list(cls.get_common_inputs().values()) @@ -195,7 +196,7 @@ class LoRADownloader(BaseDownloaderNode): target_paths.append(target_path) return target_paths -class EmbeddingDownloader(BaseDownloaderNode): +class EmbeddingInfoMetaDownloader(BaseInfoMetaDownloaderNode): @classmethod def define_schema(cls): inputs = list(cls.get_common_inputs().values()) @@ -232,7 +233,7 @@ class EmbeddingDownloader(BaseDownloaderNode): target_paths.append(target_path) return target_paths -class VAEDownloader(BaseDownloaderNode): +class VAEInfoMetaDownloader(BaseInfoMetaDownloaderNode): @classmethod def define_schema(cls): inputs = list(cls.get_common_inputs().values()) @@ -269,7 +270,7 @@ class VAEDownloader(BaseDownloaderNode): target_paths.append(target_path) return target_paths -class ControlNetDownloader(BaseDownloaderNode): +class ControlNetInfoMetaDownloader(BaseInfoMetaDownloaderNode): @classmethod def define_schema(cls): inputs = list(cls.get_common_inputs().values()) @@ -306,7 +307,7 @@ class ControlNetDownloader(BaseDownloaderNode): target_paths.append(target_path) return target_paths -class ManualPathDownloader(BaseDownloaderNode): +class ManualPathInfoMetaDownloader(BaseInfoMetaDownloaderNode): @classmethod def define_schema(cls): inputs = list(cls.get_common_inputs().values()) diff --git a/nodes/downloader_utils.py b/nodes/downloader_utils.py index 8917afa..a27b072 100644 --- a/nodes/downloader_utils.py +++ b/nodes/downloader_utils.py @@ -576,6 +576,60 @@ def process_file(filepath, cache, api, nsfw_threshold, max_examples): for k in stats: stats[k] += res.get(k, 0) return stats +import numpy as np +from PIL import Image, ImageOps +import torch + +def load_image_tensor(path): + if not os.path.exists(path): + return None + try: + img = Image.open(path) + img = ImageOps.exif_transpose(img) + image = img.convert("RGB") + image = np.array(image).astype(np.float32) / 255.0 + image = torch.from_numpy(image)[None,] + return image + except Exception as e: + print(f"Error loading image {path}: {e}") + return None + +def get_model_workflows(base_path): + workflows = [] + # Main workflow + main_json = f"{base_path}.json" + if os.path.exists(main_json): + workflows.append(main_json) + + # Example workflows + # They are named like base_path.example.N.json + # or base_path.preview.json (if extracted from preview) + preview_json = f"{base_path}.preview.json" + if os.path.exists(preview_json) and preview_json not in workflows: + workflows.append(preview_json) + + example_jsons = glob.glob(f"{base_path}.example.*.json") + # Sort them by index + example_jsons.sort() + for ej in example_jsons: + if ej not in workflows: + workflows.append(ej) + + return workflows + +def get_model_metadata_file(base_path): + # Prefer civitai.info if available + civitai_info = f"{base_path}.civitai.info" + if os.path.exists(civitai_info): + return civitai_info + + # Fallback to metadata.json (safetensors metadata) + metadata_json = f"{base_path}.metadata.json" + if os.path.exists(metadata_json): + return metadata_json + + return None + def scan_and_process(scan_dirs, recursive=True, db_path=None, nsfw_level="All", max_examples=0, api_key=None, threads=4): """ Scans directories for models and fetches their metadata/previews. diff --git a/nodes/model_info_nodes.py b/nodes/model_info_nodes.py new file mode 100644 index 0000000..a9a4335 --- /dev/null +++ b/nodes/model_info_nodes.py @@ -0,0 +1,184 @@ +import os +import folder_paths +from comfy_api.latest import io +from .downloader_utils import get_potential_preview_files, load_image_tensor, get_model_workflows, get_model_metadata_file + +class BaseModelInfoLoader(io.ComfyNode): + CATEGORY = "ModelUtils/Info" + + @classmethod + def get_info(cls, category, name, workflow_index): + full_path = folder_paths.get_full_path(category, name) + if not full_path or not os.path.exists(full_path): + return (None, "{}", "{}") + + base_path = os.path.splitext(full_path)[0] + + # 1. Image Preview + image_tensor = None + previews = get_potential_preview_files(base_path) + for p in previews: + if os.path.exists(p) and not p.lower().endswith('.mp4'): + image_tensor = load_image_tensor(p) + if image_tensor is not None: + break + + # 2. Workflow + workflow_json = "{}" + workflows = get_model_workflows(base_path) + if workflows and 0 <= workflow_index < len(workflows): + try: + with open(workflows[workflow_index], 'r', encoding='utf-8') as f: + workflow_json = f.read() + except Exception as e: + print(f"Error reading workflow {workflows[workflow_index]}: {e}") + + # 3. Metadata + metadata_json = "{}" + metadata_file = get_model_metadata_file(base_path) + if metadata_file: + try: + with open(metadata_file, 'r', encoding='utf-8') as f: + metadata_json = f.read() + except Exception as e: + print(f"Error reading metadata {metadata_file}: {e}") + + return (image_tensor, workflow_json, metadata_json) + +class CheckpointInfoLoader(BaseModelInfoLoader): + @classmethod + def define_schema(cls): + return io.Schema( + node_id=cls.__name__, + display_name="Checkpoint Info Loader", + category=cls.CATEGORY, + inputs=[ + io.Combo.Input("checkpoint", options=folder_paths.get_filename_list("checkpoints")), + io.Int.Input("workflow_index", default=0, min=0, max=100) + ], + outputs=[ + io.Image.Output(display_name="preview"), + io.String.Output(display_name="workflow_json"), + io.String.Output(display_name="metadata_json") + ] + ) + + @classmethod + def execute(cls, checkpoint, workflow_index): + res = cls.get_info("checkpoints", checkpoint, workflow_index) + return io.NodeOutput(preview=res[0], workflow_json=res[1], metadata_json=res[2]) + +class LoRAInfoLoader(BaseModelInfoLoader): + @classmethod + def define_schema(cls): + return io.Schema( + node_id=cls.__name__, + display_name="LoRA Info Loader", + category=cls.CATEGORY, + inputs=[ + io.Combo.Input("lora", options=folder_paths.get_filename_list("loras")), + io.Int.Input("workflow_index", default=0, min=0, max=100) + ], + outputs=[ + io.Image.Output(display_name="preview"), + io.String.Output(display_name="workflow_json"), + io.String.Output(display_name="metadata_json") + ] + ) + + @classmethod + def execute(cls, lora, workflow_index): + res = cls.get_info("loras", lora, workflow_index) + return io.NodeOutput(preview=res[0], workflow_json=res[1], metadata_json=res[2]) + +class EmbeddingInfoLoader(BaseModelInfoLoader): + @classmethod + def define_schema(cls): + return io.Schema( + node_id=cls.__name__, + display_name="Embedding Info Loader", + category=cls.CATEGORY, + inputs=[ + io.Combo.Input("embedding", options=folder_paths.get_filename_list("embeddings")), + io.Int.Input("workflow_index", default=0, min=0, max=100) + ], + outputs=[ + io.Image.Output(display_name="preview"), + io.String.Output(display_name="workflow_json"), + io.String.Output(display_name="metadata_json") + ] + ) + + @classmethod + def execute(cls, embedding, workflow_index): + res = cls.get_info("embeddings", embedding, workflow_index) + return io.NodeOutput(preview=res[0], workflow_json=res[1], metadata_json=res[2]) + +class VAEInfoLoader(BaseModelInfoLoader): + @classmethod + def define_schema(cls): + return io.Schema( + node_id=cls.__name__, + display_name="VAE Info Loader", + category=cls.CATEGORY, + inputs=[ + io.Combo.Input("vae", options=folder_paths.get_filename_list("vae")), + io.Int.Input("workflow_index", default=0, min=0, max=100) + ], + outputs=[ + io.Image.Output(display_name="preview"), + io.String.Output(display_name="workflow_json"), + io.String.Output(display_name="metadata_json") + ] + ) + + @classmethod + def execute(cls, vae, workflow_index): + res = cls.get_info("vae", vae, workflow_index) + return io.NodeOutput(preview=res[0], workflow_json=res[1], metadata_json=res[2]) + +class ControlNetInfoLoader(BaseModelInfoLoader): + @classmethod + def define_schema(cls): + return io.Schema( + node_id=cls.__name__, + display_name="ControlNet Info Loader", + category=cls.CATEGORY, + inputs=[ + io.Combo.Input("controlnet", options=folder_paths.get_filename_list("controlnet")), + io.Int.Input("workflow_index", default=0, min=0, max=100) + ], + outputs=[ + io.Image.Output(display_name="preview"), + io.String.Output(display_name="workflow_json"), + io.String.Output(display_name="metadata_json") + ] + ) + + @classmethod + def execute(cls, controlnet, workflow_index): + res = cls.get_info("controlnet", controlnet, workflow_index) + return io.NodeOutput(preview=res[0], workflow_json=res[1], metadata_json=res[2]) + +class DiffusionModelInfoLoader(BaseModelInfoLoader): + @classmethod + def define_schema(cls): + return io.Schema( + node_id=cls.__name__, + display_name="Diffusion Model Info Loader", + category=cls.CATEGORY, + inputs=[ + io.Combo.Input("diffusion_model", options=folder_paths.get_filename_list("diffusion_models")), + io.Int.Input("workflow_index", default=0, min=0, max=100) + ], + outputs=[ + io.Image.Output(display_name="preview"), + io.String.Output(display_name="workflow_json"), + io.String.Output(display_name="metadata_json") + ] + ) + + @classmethod + def execute(cls, diffusion_model, workflow_index): + res = cls.get_info("diffusion_models", diffusion_model, workflow_index) + return io.NodeOutput(preview=res[0], workflow_json=res[1], metadata_json=res[2]) From 0a82e071b8f9ec60cf2ad6f4f608b5036eb65076 Mon Sep 17 00:00:00 2001 From: silveroxides Date: Wed, 8 Apr 2026 17:50:10 +0200 Subject: [PATCH 7/7] fix NodeOutput issue --- nodes/model_info_nodes.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/nodes/model_info_nodes.py b/nodes/model_info_nodes.py index a9a4335..8ade4e1 100644 --- a/nodes/model_info_nodes.py +++ b/nodes/model_info_nodes.py @@ -66,7 +66,7 @@ class CheckpointInfoLoader(BaseModelInfoLoader): @classmethod def execute(cls, checkpoint, workflow_index): res = cls.get_info("checkpoints", checkpoint, workflow_index) - return io.NodeOutput(preview=res[0], workflow_json=res[1], metadata_json=res[2]) + return io.NodeOutput(res[0], res[1], res[2]) class LoRAInfoLoader(BaseModelInfoLoader): @classmethod @@ -89,7 +89,7 @@ class LoRAInfoLoader(BaseModelInfoLoader): @classmethod def execute(cls, lora, workflow_index): res = cls.get_info("loras", lora, workflow_index) - return io.NodeOutput(preview=res[0], workflow_json=res[1], metadata_json=res[2]) + return io.NodeOutput(res[0], res[1], res[2]) class EmbeddingInfoLoader(BaseModelInfoLoader): @classmethod @@ -112,7 +112,7 @@ class EmbeddingInfoLoader(BaseModelInfoLoader): @classmethod def execute(cls, embedding, workflow_index): res = cls.get_info("embeddings", embedding, workflow_index) - return io.NodeOutput(preview=res[0], workflow_json=res[1], metadata_json=res[2]) + return io.NodeOutput(res[0], res[1], res[2]) class VAEInfoLoader(BaseModelInfoLoader): @classmethod @@ -135,7 +135,7 @@ class VAEInfoLoader(BaseModelInfoLoader): @classmethod def execute(cls, vae, workflow_index): res = cls.get_info("vae", vae, workflow_index) - return io.NodeOutput(preview=res[0], workflow_json=res[1], metadata_json=res[2]) + return io.NodeOutput(res[0], res[1], res[2]) class ControlNetInfoLoader(BaseModelInfoLoader): @classmethod @@ -158,7 +158,7 @@ class ControlNetInfoLoader(BaseModelInfoLoader): @classmethod def execute(cls, controlnet, workflow_index): res = cls.get_info("controlnet", controlnet, workflow_index) - return io.NodeOutput(preview=res[0], workflow_json=res[1], metadata_json=res[2]) + return io.NodeOutput(res[0], res[1], res[2]) class DiffusionModelInfoLoader(BaseModelInfoLoader): @classmethod @@ -181,4 +181,4 @@ class DiffusionModelInfoLoader(BaseModelInfoLoader): @classmethod def execute(cls, diffusion_model, workflow_index): res = cls.get_info("diffusion_models", diffusion_model, workflow_index) - return io.NodeOutput(preview=res[0], workflow_json=res[1], metadata_json=res[2]) + return io.NodeOutput(res[0], res[1], res[2])