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
This commit is contained in:
silveroxides
@@ -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,
|
||||
]
|
||||
|
||||
|
||||
|
||||
@@ -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 []
|
||||
@@ -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)
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user