From b7f6c4fbf8fac1c3b464b6ce8e8f60b226b7a267 Mon Sep 17 00:00:00 2001 From: Will Miao <13051207myq@gmail.com> Date: Wed, 29 Jan 2025 06:06:48 +0800 Subject: [PATCH] Update --- nodes.py | 26 ++++++++++++++------------ pyproject.toml | 1 + requirements.txt | 3 +++ utils/file_utils.py | 45 +++++++++++++++++++++++++++++++++++++++------ 4 files changed, 57 insertions(+), 18 deletions(-) create mode 100644 requirements.txt diff --git a/nodes.py b/nodes.py index 8a15028d..f5579274 100644 --- a/nodes.py +++ b/nodes.py @@ -62,11 +62,11 @@ class LorasEndpoint: file_path = os.path.join(root, filename) # Try to load existing metadata first - metadata = await load_metadata(file_path) + metadata = await load_metadata(file_path, self.loras_root) if metadata is None: # Only get file info and extract metadata if no existing metadata - metadata = await get_file_info(file_path) + metadata = await get_file_info(file_path, self.loras_root) base_model_info = await extract_lora_metadata(file_path) metadata.base_model = base_model_info['base_model'] await save_metadata(file_path, metadata) @@ -101,7 +101,7 @@ class LorasEndpoint: try: scan_start = time.time() data = await self.scan_loras() - print(f"Scanned {len(data)} loras in {time.time()-scan_start:.2f}s") + print(f"Lora Manager: Scanned {len(data)} loras in {time.time()-scan_start:.2f}s") # Format the data for the template formatted_loras = [self.format_lora(l) for l in data] @@ -297,15 +297,17 @@ class LorasEndpoint: local_metadata['base_model'] = civitai_metadata.get('baseModel') # 4. 下载预览图 - first_preview = next((img for img in civitai_metadata.get('images', [])), None) - if first_preview: - - preview_extension = '.mp4' if first_preview['type'] == 'video' else os.path.splitext(first_preview['url'])[-1] # Get the file extension - preview_filename = os.path.splitext(os.path.basename(data['file_path']))[0] + preview_extension - preview_path = os.path.join(os.path.dirname(data['file_path']), preview_filename) - await client.download_preview_image(first_preview['url'], preview_path) - # 存储相对路径，使用正斜杠格式 - local_metadata['preview_url'] = os.path.relpath(preview_path, self.loras_root).replace(os.sep, '/') + # Check if existing preview is valid + if not local_metadata.get('preview_url') or not os.path.join(self.loras_root, local_metadata['preview_url'].replace('/', os.sep)): + first_preview = next((img for img in civitai_metadata.get('images', [])), None) + if first_preview: + + preview_extension = '.mp4' if first_preview['type'] == 'video' else os.path.splitext(first_preview['url'])[-1] # Get the file extension + preview_filename = os.path.splitext(os.path.basename(data['file_path']))[0] + '.preview' + preview_extension + preview_path = os.path.join(os.path.dirname(data['file_path']), preview_filename) + await client.download_preview_image(first_preview['url'], preview_path) + # 存储相对路径，使用正斜杠格式 + local_metadata['preview_url'] = os.path.relpath(preview_path, self.loras_root).replace(os.sep, '/') # 5. 保存更新后的元数据 with open(metadata_path, 'w', encoding='utf-8') as f: diff --git a/pyproject.toml b/pyproject.toml index a840c94d..e9765b71 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -3,6 +3,7 @@ name = "comfyui-lora-manager" description = "LoRA Manager for ComfyUI - An extension for managing LoRA models with previews and metadata integration." version = "0.5.1" license = {file = "LICENSE"} +dependencies = ["aiohttp", "jinja2", "safetensors"] [project.urls] Repository = "https://github.com/willmiao/ComfyUI-Lora-Manager" diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 00000000..5b224e66 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,3 @@ +aiohttp +jinja2 +safetensors \ No newline at end of file diff --git a/utils/file_utils.py b/utils/file_utils.py index 4c483e60..b7c07205 100644 --- a/utils/file_utils.py +++ b/utils/file_utils.py @@ -12,18 +12,42 @@ async def calculate_sha256(file_path: str) -> str: sha256_hash.update(byte_block) return sha256_hash.hexdigest() -async def get_file_info(file_path: str) -> LoraMetadata: +def _find_preview_file(base_name: str, dir_path: str, lora_root: str) -> str: + """Find preview file for given base name in directory""" + preview_patterns = [ + f"{base_name}.preview.png", + f"{base_name}.preview.jpg", + f"{base_name}.preview.jpeg", + f"{base_name}.preview.mp4", + f"{base_name}.png", + f"{base_name}.jpg", + f"{base_name}.jpeg", + f"{base_name}.mp4" + ] + + for pattern in preview_patterns: + full_pattern = os.path.join(dir_path, pattern) + if os.path.exists(full_pattern): + return os.path.relpath(full_pattern, lora_root).replace("\\", "/") + return "" + +async def get_file_info(file_path: str, loras_root: str) -> LoraMetadata: """Get basic file information as LoraMetadata object""" + base_name = os.path.splitext(os.path.basename(file_path))[0] + dir_path = os.path.dirname(file_path) + + preview_path = _find_preview_file(base_name, dir_path, loras_root) + return LoraMetadata( - file_name=os.path.splitext(os.path.basename(file_path))[0], - model_name=os.path.splitext(os.path.basename(file_path))[0], + file_name=base_name, + model_name=base_name, file_path=file_path, size=os.path.getsize(file_path), modified=os.path.getmtime(file_path), sha256=await calculate_sha256(file_path), base_model="Unknown", # Will be updated later from_civitai=True, - preview_url="", + preview_url=preview_path, ) async def save_metadata(file_path: str, metadata: LoraMetadata) -> None: @@ -35,13 +59,22 @@ async def save_metadata(file_path: str, metadata: LoraMetadata) -> None: except Exception as e: print(f"Error saving metadata to {metadata_path}: {str(e)}") -async def load_metadata(file_path: str) -> Optional[LoraMetadata]: +async def load_metadata(file_path: str, loras_root: str) -> Optional[LoraMetadata]: """Load metadata from .metadata.json file""" metadata_path = f"{os.path.splitext(file_path)[0]}.metadata.json" try: if os.path.exists(metadata_path): with open(metadata_path, 'r', encoding='utf-8') as f: data = json.load(f) + preview_url = data.get('preview_url', '') + preview_path = os.path.join(loras_root, preview_url) if preview_url else '' + if not preview_url or not os.path.exists(preview_path): + base_name = os.path.splitext(os.path.basename(file_path))[0] + dir_path = os.path.dirname(file_path) + data['preview_url'] = _find_preview_file(base_name, dir_path, loras_root) + if data['preview_url']: + with open(metadata_path, 'w', encoding='utf-8') as f: + json.dump(data, f, indent=2, ensure_ascii=False) return LoraMetadata.from_dict(data) except Exception as e: print(f"Error loading metadata from {metadata_path}: {str(e)}") @@ -51,4 +84,4 @@ async def update_civitai_metadata(file_path: str, civitai_data: Dict) -> None: """Update metadata file with Civitai data""" metadata = await load_metadata(file_path) metadata['civitai'] = civitai_data - await save_metadata(file_path, metadata) \ No newline at end of file + await save_metadata(file_path, metadata) \ No newline at end of file