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

fix: Update optimize_image method to handle image validation and error logging, and adjust metadata preservation logic.

Browse Source
This commit is contained in:
Will Miao
2025-04-15 12:31:17 +08:00
parent ed87411e0d
commit c4c926070d
5 changed files with 121 additions and 75 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
py/routes/api_routes.py
View File

@@ -226,7 +226,7 @@ class ApiRoutes:
target_width=CARD_PREVIEW_WIDTH, target_width=CARD_PREVIEW_WIDTH,
format='webp', format='webp',
quality=85, quality=85,
preserve_metadata=True preserve_metadata=False
) )
extension = '.webp' # Use .webp without .preview part extension = '.webp' # Use .webp without .preview part

2
py/services/download_manager.py
View File

@@ -229,7 +229,7 @@ class DownloadManager:
target_width=CARD_PREVIEW_WIDTH, target_width=CARD_PREVIEW_WIDTH,
format='webp', format='webp',
quality=85, quality=85,
preserve_metadata=True preserve_metadata=False
) )
# Save the optimized image # Save the optimized image

186
py/utils/exif_utils.py
View File

@@ -203,7 +203,7 @@ class ExifUtils:
return user_comment[:recipe_marker_index] + user_comment[next_line_index:] return user_comment[:recipe_marker_index] + user_comment[next_line_index:]
@staticmethod @staticmethod
def optimize_image(image_data, target_width=250, format='webp', quality=85, preserve_metadata=True): def optimize_image(image_data, target_width=250, format='webp', quality=85, preserve_metadata=False):
""" """
Optimize an image by resizing and converting to WebP format Optimize an image by resizing and converting to WebP format
@@ -218,98 +218,144 @@ class ExifUtils:
Tuple of (optimized_image_data, extension) Tuple of (optimized_image_data, extension)
""" """
try: try:
# Extract metadata if needed # First validate the image data is usable
img = None
if isinstance(image_data, str) and os.path.exists(image_data):
# It's a file path - validate file
try:
with Image.open(image_data) as test_img:
# Verify the image can be fully loaded by accessing its size
width, height = test_img.size
# If we got here, the image is valid
img = Image.open(image_data)
except (IOError, OSError) as e:
logger.error(f"Invalid or corrupt image file: {image_data}: {e}")
raise ValueError(f"Cannot process corrupt image: {e}")
else:
# It's binary data - validate data
try:
with BytesIO(image_data) as temp_buf:
test_img = Image.open(temp_buf)
# Verify the image can be fully loaded
width, height = test_img.size
# If successful, reopen for processing
img = Image.open(BytesIO(image_data))
except Exception as e:
logger.error(f"Invalid binary image data: {e}")
raise ValueError(f"Cannot process corrupt image data: {e}")
# Extract metadata if needed and valid
metadata = None metadata = None
if preserve_metadata: if preserve_metadata:
if isinstance(image_data, str) and os.path.exists(image_data): try:
# It's a file path if isinstance(image_data, str) and os.path.exists(image_data):
metadata = ExifUtils.extract_image_metadata(image_data) # For file path, extract directly
img = Image.open(image_data) metadata = ExifUtils.extract_image_metadata(image_data)
else: else:
# It's binary data # For binary data, save to temp file first
temp_img = BytesIO(image_data) import tempfile
img = Image.open(temp_img) with tempfile.NamedTemporaryFile(suffix='.jpg', delete=False) as temp_file:
# Save to a temporary file to extract metadata temp_path = temp_file.name
import tempfile temp_file.write(image_data)
with tempfile.NamedTemporaryFile(suffix='.jpg', delete=False) as temp_file: try:
temp_path = temp_file.name metadata = ExifUtils.extract_image_metadata(temp_path)
temp_file.write(image_data) except Exception as e:
metadata = ExifUtils.extract_image_metadata(temp_path) logger.warning(f"Failed to extract metadata from temp file: {e}")
os.unlink(temp_path) finally:
else: # Clean up temp file
# Just open the image without extracting metadata try:
if isinstance(image_data, str) and os.path.exists(image_data): os.unlink(temp_path)
img = Image.open(image_data) except Exception:
else: pass
img = Image.open(BytesIO(image_data)) except Exception as e:
logger.warning(f"Failed to extract metadata, continuing without it: {e}")
# Continue without metadata
# Calculate new height to maintain aspect ratio # Calculate new height to maintain aspect ratio
width, height = img.size width, height = img.size
new_height = int(height * (target_width / width)) new_height = int(height * (target_width / width))
# Resize the image # Resize the image with error handling
resized_img = img.resize((target_width, new_height), Image.LANCZOS) try:
resized_img = img.resize((target_width, new_height), Image.LANCZOS)
except Exception as e:
logger.error(f"Failed to resize image: {e}")
# Return original image if resize fails
return image_data, '.jpg' if not isinstance(image_data, str) else os.path.splitext(image_data)[1]
# Save to BytesIO in the specified format # Save to BytesIO in the specified format
output = BytesIO() output = BytesIO()
# WebP format # Set format and extension
if format.lower() == 'webp': if format.lower() == 'webp':
resized_img.save(output, format='WEBP', quality=quality) save_format, extension = 'WEBP', '.webp'
extension = '.webp'
# JPEG format
elif format.lower() in ('jpg', 'jpeg'): elif format.lower() in ('jpg', 'jpeg'):
resized_img.save(output, format='JPEG', quality=quality) save_format, extension = 'JPEG', '.jpg'
extension = '.jpg'
# PNG format
elif format.lower() == 'png': elif format.lower() == 'png':
resized_img.save(output, format='PNG', optimize=True) save_format, extension = 'PNG', '.png'
extension = '.png'
else: else:
# Default to WebP save_format, extension = 'WEBP', '.webp'
resized_img.save(output, format='WEBP', quality=quality)
extension = '.webp' # Save with error handling
try:
if save_format == 'PNG':
resized_img.save(output, format=save_format, optimize=True)
else:
resized_img.save(output, format=save_format, quality=quality)
except Exception as e:
logger.error(f"Failed to save optimized image: {e}")
# Return original image if save fails
return image_data, '.jpg' if not isinstance(image_data, str) else os.path.splitext(image_data)[1]
# Get the optimized image data # Get the optimized image data
optimized_data = output.getvalue() optimized_data = output.getvalue()
# If we need to preserve metadata, write it to a temporary file # Handle metadata preservation if requested and available
if preserve_metadata and metadata: if preserve_metadata and metadata:
# For WebP format, we'll directly save with metadata try:
if format.lower() == 'webp': if save_format == 'WEBP':
# Create a new BytesIO with metadata # For WebP format, directly save with metadata
output_with_metadata = BytesIO() try:
output_with_metadata = BytesIO()
# Create EXIF data with user comment exif_dict = {'Exif': {piexif.ExifIFD.UserComment: b'UNICODE\0' + metadata.encode('utf-16be')}}
exif_dict = {'Exif': {piexif.ExifIFD.UserComment: b'UNICODE\0' + metadata.encode('utf-16be')}} exif_bytes = piexif.dump(exif_dict)
exif_bytes = piexif.dump(exif_dict) resized_img.save(output_with_metadata, format='WEBP', exif=exif_bytes, quality=quality)
optimized_data = output_with_metadata.getvalue()
# Save with metadata except Exception as e:
resized_img.save(output_with_metadata, format='WEBP', exif=exif_bytes, quality=quality) logger.warning(f"Failed to add metadata to WebP, continuing without it: {e}")
optimized_data = output_with_metadata.getvalue() else:
else: # For other formats, use temporary file
# For other formats, use the temporary file approach import tempfile
import tempfile with tempfile.NamedTemporaryFile(suffix=extension, delete=False) as temp_file:
with tempfile.NamedTemporaryFile(suffix=extension, delete=False) as temp_file: temp_path = temp_file.name
temp_path = temp_file.name temp_file.write(optimized_data)
temp_file.write(optimized_data)
try:
# Add the metadata back # Add metadata
ExifUtils.update_image_metadata(temp_path, metadata) ExifUtils.update_image_metadata(temp_path, metadata)
# Read back the file
# Read the file with metadata with open(temp_path, 'rb') as f:
with open(temp_path, 'rb') as f: optimized_data = f.read()
optimized_data = f.read() except Exception as e:
logger.warning(f"Failed to add metadata to image, continuing without it: {e}")
# Clean up finally:
os.unlink(temp_path) # Clean up temp file
try:
os.unlink(temp_path)
except Exception:
pass
except Exception as e:
logger.warning(f"Failed to preserve metadata: {e}, continuing with unmodified output")
return optimized_data, extension return optimized_data, extension
except Exception as e: except Exception as e:
logger.error(f"Error optimizing image: {e}", exc_info=True) logger.error(f"Error optimizing image: {e}", exc_info=True)
# Return original data if optimization fails # Return original data if optimization completely fails
if isinstance(image_data, str) and os.path.exists(image_data): if isinstance(image_data, str) and os.path.exists(image_data):
with open(image_data, 'rb') as f: try:
return f.read(), os.path.splitext(image_data)[1] with open(image_data, 'rb') as f:
return f.read(), os.path.splitext(image_data)[1]
except Exception:
return image_data, '.jpg' # Last resort fallback
return image_data, '.jpg' return image_data, '.jpg'

2
py/utils/file_utils.py
View File

@@ -42,7 +42,7 @@ def find_preview_file(base_name: str, dir_path: str) -> str:
target_width=CARD_PREVIEW_WIDTH, target_width=CARD_PREVIEW_WIDTH,
format='webp', format='webp',
quality=85, quality=85,
preserve_metadata=True preserve_metadata=False # Changed from True to False
) )
# Save the optimized webp file # Save the optimized webp file

4
py/utils/routes_common.py
View File

@@ -95,7 +95,7 @@ class ModelRouteUtils:
target_width=CARD_PREVIEW_WIDTH, target_width=CARD_PREVIEW_WIDTH,
format='webp', format='webp',
quality=85, quality=85,
preserve_metadata=True preserve_metadata=False
) )
# Save the optimized WebP image # Save the optimized WebP image
@@ -387,7 +387,7 @@ class ModelRouteUtils:
target_width=CARD_PREVIEW_WIDTH, target_width=CARD_PREVIEW_WIDTH,
format='webp', format='webp',
quality=85, quality=85,
preserve_metadata=True preserve_metadata=False
) )
extension = '.webp' # Use .webp without .preview part extension = '.webp' # Use .webp without .preview part