alphard/ComfyUI-Lora-Manager
1
0
Fork 0
mirror of https://github.com/willmiao/ComfyUI-Lora-Manager.git synced 2026-03-23 22:22:11 -03:00
Code Issues Packages Projects Releases Wiki Activity

checkpoint

Browse Source
This commit is contained in:
Will Miao
2025-03-09 12:29:24 +08:00
parent e6aafe8773
commit 250e8445bb
8 changed files with 434 additions and 100 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
py/config.py
View File

@@ -87,9 +87,9 @@ class Config:
def _init_lora_paths(self) -> List[str]:
"""Initialize and validate LoRA paths from ComfyUI settings"""
paths = list(set(path.replace(os.sep, "/")
paths = sorted(set(path.replace(os.sep, "/")
for path in folder_paths.get_folder_paths("loras")
if os.path.exists(path)))
if os.path.exists(path)), key=lambda p: p.lower())
print("Found LoRA roots:", "\n - " + "\n - ".join(paths))
if not paths:

38
py/routes/lora_routes.py
View File

@@ -107,11 +107,29 @@ class LoraRoutes:
settings=settings
)
else:
# Normal flow
# Normal flow - get recipes with the same formatting as the API endpoint
cache = await self.recipe_scanner.get_cached_data()
recipes_data = cache.sorted_by_name[:20] # Show first 20 recipes by name
# Format the response data with static URLs for file paths - same as in recipe_routes
for item in recipes_data:
# Always ensure file_url is set
if 'file_path' in item:
item['file_url'] = self._format_recipe_file_url(item['file_path'])
else:
item['file_url'] = '/loras_static/images/no-preview.png'
# Ensure loras array exists
if 'loras' not in item:
item['loras'] = []
# Ensure base_model field exists
if 'base_model' not in item:
item['base_model'] = ""
template = self.template_env.get_template('recipes.html')
rendered = template.render(
recipes=cache.sorted_by_date[:20], # Show first 20 recipes by date
recipes=recipes_data,
is_initializing=False,
settings=settings
)
@@ -128,6 +146,22 @@ class LoraRoutes:
status=500
)
def _format_recipe_file_url(self, file_path: str) -> str:
"""Format file path for recipe image as a URL - same as in recipe_routes"""
try:
# Return the file URL directly for the first lora root's preview
recipes_dir = os.path.join(config.loras_roots[0], "recipes").replace(os.sep, '/')
if file_path.replace(os.sep, '/').startswith(recipes_dir):
relative_path = os.path.relpath(file_path, config.loras_roots[0]).replace(os.sep, '/')
return f"/loras_static/root1/preview/{relative_path}"
# If not in recipes dir, try to create a valid URL from the file path
file_name = os.path.basename(file_path)
return f"/loras_static/root1/preview/recipes/{file_name}"
except Exception as e:
logger.error(f"Error formatting recipe file URL: {e}", exc_info=True)
return '/loras_static/images/no-preview.png' # Return default image on error
def setup_routes(self, app: web.Application):
"""Register routes with the application"""
app.router.add_get('/loras', self.handle_loras_page)

40
py/routes/recipe_routes.py
View File

@@ -67,9 +67,19 @@ class RecipeRoutes:
# Format the response data with static URLs for file paths
for item in result['items']:
item['preview_url'] = item['file_path']
# Convert file path to URL
item['file_url'] = self._format_recipe_file_url(item['file_path'])
# Always ensure file_url is set
if 'file_path' in item:
item['file_url'] = self._format_recipe_file_url(item['file_path'])
else:
item['file_url'] = '/loras_static/images/no-preview.png'
# 确保 loras 数组存在
if 'loras' not in item:
item['loras'] = []
# 确保有 base_model 字段
if 'base_model' not in item:
item['base_model'] = ""
return web.json_response(result)
except Exception as e:
@@ -101,17 +111,19 @@ class RecipeRoutes:
def _format_recipe_file_url(self, file_path: str) -> str:
"""Format file path for recipe image as a URL"""
# This is a simplified example - in real implementation,
# you would map this to a static route that can serve the file
# For recipes folder in the first lora root
for idx, root in enumerate(config.loras_roots, start=1):
recipes_dir = os.path.join(root, "recipes")
if file_path.startswith(recipes_dir):
relative_path = os.path.relpath(file_path, root)
return f"/loras_static/root{idx}/{relative_path}"
return file_path # Return original path if no mapping found
try:
# Return the file URL directly for the first lora root's preview
recipes_dir = os.path.join(config.loras_roots[0], "recipes").replace(os.sep, '/')
if file_path.replace(os.sep, '/').startswith(recipes_dir):
relative_path = os.path.relpath(file_path, config.loras_roots[0]).replace(os.sep, '/')
return f"/loras_static/root1/preview/{relative_path}"
# If not in recipes dir, try to create a valid URL from the file path
file_name = os.path.basename(file_path)
return f"/loras_static/root1/preview/recipes/{file_name}"
except Exception as e:
logger.error(f"Error formatting recipe file URL: {e}", exc_info=True)
return '/loras_static/images/no-preview.png' # Return default image on error
def _format_recipe_data(self, recipe: Dict) -> Dict:
"""Format recipe data for API response"""

192
py/services/recipe_scanner.py
View File

@@ -78,11 +78,8 @@ class RecipeScanner:
if not config.loras_roots:
return ""
# Sort the lora roots case-insensitively
sorted_roots = sorted(config.loras_roots, key=lambda x: x.lower())
# Use the first sorted lora root as base
recipes_dir = os.path.join(sorted_roots[0], "recipes")
# config.loras_roots already sorted case-insensitively, use the first one
recipes_dir = os.path.join(config.loras_roots[0], "recipes")
os.makedirs(recipes_dir, exist_ok=True)
logger.info(f"Using recipes directory: {recipes_dir}")
@@ -126,9 +123,6 @@ class RecipeScanner:
if self._cache is None:
raise # If no cache, raise exception
logger.info(f"Recipe cache initialized with {len(self._cache.raw_data)} recipes")
logger.info(f"Recipe cache: {json.dumps(self._cache, indent=2)}")
return self._cache
async def _initialize_cache(self) -> None:
@@ -238,9 +232,9 @@ class RecipeScanner:
else:
print(f"Found UserComment: {user_comment[:50]}...", file=sys.stderr)
# Parse metadata from UserComment
recipe_data = ExifUtils.parse_recipe_metadata(user_comment)
if not recipe_data:
# Parse generation parameters from UserComment
gen_params = ExifUtils.parse_recipe_metadata(user_comment)
if not gen_params:
print(f"Failed to parse recipe metadata from {image_path}", file=sys.stderr)
logger.warning(f"Failed to parse recipe metadata from {image_path}")
return None
@@ -250,15 +244,45 @@ class RecipeScanner:
file_name = os.path.basename(image_path)
title = os.path.splitext(file_name)[0]
# Add common metadata
recipe_data.update({
'id': file_name,
'file_path': image_path,
'title': title,
'modified': stat.st_mtime,
'created_date': stat.st_ctime,
'file_size': stat.st_size
})
# Check for existing recipe metadata
recipe_data = self._extract_recipe_metadata(user_comment)
if not recipe_data:
# Create new recipe data
recipe_data = {
'id': file_name,
'file_path': image_path,
'title': title,
'modified': stat.st_mtime,
'created_date': stat.st_ctime,
'file_size': stat.st_size,
'loras': [],
'gen_params': {}
}
# Copy loras from gen_params to recipe_data with proper structure
for lora in gen_params.get('loras', []):
recipe_lora = {
'file_name': '',
'hash': lora.get('hash', '').lower() if lora.get('hash') else '',
'strength': lora.get('weight', 1.0),
'modelVersionId': lora.get('modelVersionId', ''),
'modelName': lora.get('modelName', ''),
'modelVersionName': lora.get('modelVersionName', '')
}
recipe_data['loras'].append(recipe_lora)
# Add generation parameters to recipe_data.gen_params instead of top level
recipe_data['gen_params'] = {
'prompt': gen_params.get('prompt', ''),
'negative_prompt': gen_params.get('negative_prompt', ''),
'checkpoint': gen_params.get('checkpoint', None),
'steps': gen_params.get('steps', ''),
'sampler': gen_params.get('sampler', ''),
'cfg_scale': gen_params.get('cfg_scale', ''),
'seed': gen_params.get('seed', ''),
'size': gen_params.get('size', ''),
'clip_skip': gen_params.get('clip_skip', '')
}
# Update recipe metadata with missing information
metadata_updated = await self._update_recipe_metadata(recipe_data, user_comment)
@@ -417,51 +441,48 @@ class RecipeScanner:
def _save_updated_metadata(self, image_path: str, original_comment: str, recipe_data: Dict) -> None:
"""Save updated metadata back to image file"""
try:
# Update the resources section with the updated lora data
resources_match = re.search(r'(Civitai resources: )(\[.*?\])(?:,|\})', original_comment)
if not resources_match:
logger.warning(f"Could not find Civitai resources section in {image_path}")
return
# Check if we already have a recipe metadata section
recipe_metadata_exists = "recipe metadata:" in original_comment.lower()
resources_prefix = resources_match.group(1)
# Prepare recipe metadata
recipe_metadata = {
'id': recipe_data.get('id', ''),
'file_path': recipe_data.get('file_path', ''),
'title': recipe_data.get('title', ''),
'modified': recipe_data.get('modified', 0),
'created_date': recipe_data.get('created_date', 0),
'base_model': recipe_data.get('base_model', ''),
'loras': [],
'gen_params': recipe_data.get('gen_params', {})
}
# Generate updated resources array
resources = []
# Add lora data with only necessary fields (removing weight, adding modelVersionName)
for lora in recipe_data.get('loras', []):
lora_entry = {
'file_name': lora.get('file_name', ''),
'hash': lora.get('hash', '').lower() if lora.get('hash') else '',
'strength': lora.get('strength', 1.0),
'modelVersionId': lora.get('modelVersionId', ''),
'modelName': lora.get('modelName', ''),
'modelVersionName': lora.get('modelVersionName', '')
}
recipe_metadata['loras'].append(lora_entry)
# Add checkpoint if exists
if recipe_data.get('checkpoint'):
resources.append(recipe_data['checkpoint'])
# Convert to JSON
recipe_metadata_json = json.dumps(recipe_metadata)
# Add all loras
resources.extend(recipe_data.get('loras', []))
# Generate new resources JSON
updated_resources = json.dumps(resources)
# Replace resources in original comment
updated_comment = original_comment.replace(
resources_match.group(0),
f"{resources_prefix}{updated_resources},"
)
# Update metadata section if it exists
metadata_match = re.search(r'(Civitai metadata: )(\{.*?\})', original_comment)
if metadata_match:
metadata_prefix = metadata_match.group(1)
# Create metadata object with base_model
metadata = {}
if recipe_data.get('base_model'):
metadata['base_model'] = recipe_data['base_model']
# Generate new metadata JSON
updated_metadata = json.dumps(metadata)
# Replace metadata in original comment
updated_comment = updated_comment.replace(
metadata_match.group(0),
f"{metadata_prefix}{updated_metadata}"
# Create or update the recipe metadata section
if recipe_metadata_exists:
# Replace existing recipe metadata
updated_comment = re.sub(
r'recipe metadata: \{.*\}',
f'recipe metadata: {recipe_metadata_json}',
original_comment,
flags=re.IGNORECASE | re.DOTALL
)
else:
# Append recipe metadata to the end
updated_comment = f"{original_comment}, recipe metadata: {recipe_metadata_json}"
# Save back to image
logger.info(f"Saving updated metadata to {image_path}")
@@ -537,6 +558,14 @@ class RecipeScanner:
else:
logger.warning(f"Could not get hash for modelVersionId {model_version_id}")
# If modelVersionId exists but no modelVersionName, try to get it from Civitai
if 'modelVersionId' in lora and not lora.get('modelVersionName'):
model_version_id = str(lora['modelVersionId'])
model_version_name = await self._get_model_version_name(model_version_id)
if model_version_name:
lora['modelVersionName'] = model_version_name
metadata_updated = True
# If has hash, check if it's in library
if 'hash' in lora:
hash_value = lora['hash'].lower() # Ensure lowercase when comparing
@@ -551,11 +580,6 @@ class RecipeScanner:
logger.info(f"Found lora in library: {file_name}")
lora['file_name'] = file_name
metadata_updated = True
# Also get base_model from lora cache if possible
base_model = await self._get_base_model_for_lora(lora_path)
if base_model:
lora['base_model'] = base_model
elif not in_library:
# Lora not in library
logger.info(f"LoRA with hash {hash_value[:8]}... not found in library")
@@ -571,6 +595,24 @@ class RecipeScanner:
return metadata_updated
async def _get_model_version_name(self, model_version_id: str) -> Optional[str]:
"""Get model version name from Civitai API"""
try:
if not self._civitai_client:
return None
logger.info(f"Fetching model version info from Civitai for ID: {model_version_id}")
version_info = await self._civitai_client.get_model_version_info(model_version_id)
if version_info and 'name' in version_info:
return version_info['name']
logger.warning(f"No version name found for modelVersionId {model_version_id}")
return None
except Exception as e:
logger.error(f"Error getting model version name from Civitai: {e}")
return None
async def _determine_base_model(self, loras: List[Dict]) -> Optional[str]:
"""Determine the most common base model among LoRAs"""
base_models = {}
@@ -607,4 +649,24 @@ class RecipeScanner:
return None
except Exception as e:
logger.error(f"Error getting base model for lora: {e}")
return None
def _extract_recipe_metadata(self, user_comment: str) -> Optional[Dict]:
"""Extract recipe metadata section from UserComment if it exists"""
try:
# Look for recipe metadata section
recipe_match = re.search(r'recipe metadata: (\{.*\})', user_comment, re.IGNORECASE | re.DOTALL)
if not recipe_match:
return None
recipe_json = recipe_match.group(1)
recipe_data = json.loads(recipe_json)
# Ensure loras array exists
if 'loras' not in recipe_data:
recipe_data['loras'] = []
return recipe_data
except Exception as e:
logger.error(f"Error extracting recipe metadata: {e}")
return None

27
py/utils/exif_utils.py
View File

@@ -98,7 +98,15 @@ class ExifUtils:
# Filter loras and checkpoints
for resource in resources:
if resource.get('type') == 'lora':
metadata['loras'].append(resource)
# 确保 weight 字段被正确保留
lora_entry = resource.copy()
# 如果找不到 weight，默认为 1.0
if 'weight' not in lora_entry:
lora_entry['weight'] = 1.0
# Ensure modelVersionName is included
if 'modelVersionName' not in lora_entry:
lora_entry['modelVersionName'] = ''
metadata['loras'].append(lora_entry)
elif resource.get('type') == 'checkpoint':
metadata['checkpoint'] = resource
except json.JSONDecodeError:
@@ -107,4 +115,19 @@ class ExifUtils:
return metadata
except Exception as e:
logger.error(f"Error parsing recipe metadata: {e}")
return {"prompt": user_comment, "loras": [], "checkpoint": None}
return {"prompt": user_comment, "loras": [], "checkpoint": None}
@staticmethod
def extract_recipe_metadata(user_comment: str) -> Optional[Dict]:
"""Extract recipe metadata section from UserComment if it exists"""
try:
# Look for recipe metadata section
recipe_match = re.search(r'recipe metadata: (\{.*\})', user_comment, re.IGNORECASE | re.DOTALL)
if not recipe_match:
return None
recipe_json = recipe_match.group(1)
return json.loads(recipe_json)
except Exception as e:
logger.error(f"Error extracting recipe metadata: {e}")
return None