diff --git a/py/routes/recipe_routes.py b/py/routes/recipe_routes.py index e447ac06..e6d0db55 100644 --- a/py/routes/recipe_routes.py +++ b/py/routes/recipe_routes.py @@ -14,7 +14,7 @@ from ..services.recipe_scanner import RecipeScanner from ..services.lora_scanner import LoraScanner from ..config import config from ..workflow.parser import WorkflowParser -from ..utils.utils import download_twitter_image +from ..utils.utils import download_civitai_image logger = logging.getLogger(__name__) @@ -235,7 +235,7 @@ class RecipeRoutes: }, status=400) # Download image from URL - temp_path = download_twitter_image(url) + temp_path = download_civitai_image(url) if not temp_path: return web.json_response({ @@ -244,10 +244,10 @@ class RecipeRoutes: }, status=400) # Extract metadata from the image using ExifUtils - user_comment = ExifUtils.extract_user_comment(temp_path) + metadata = ExifUtils.extract_image_metadata(temp_path) # If no metadata found, return a more specific error - if not user_comment: + if not metadata: result = { "error": "No metadata found in this image", "loras": [] # Return empty loras array to prevent client-side errors @@ -262,7 +262,7 @@ class RecipeRoutes: return web.json_response(result, status=200) # Use the parser factory to get the appropriate parser - parser = RecipeParserFactory.create_parser(user_comment) + parser = RecipeParserFactory.create_parser(metadata) if parser is None: result = { @@ -280,7 +280,7 @@ class RecipeRoutes: # Parse the metadata result = await parser.parse_metadata( - user_comment, + metadata, recipe_scanner=self.recipe_scanner, civitai_client=self.civitai_client ) @@ -387,8 +387,7 @@ class RecipeRoutes: return web.json_response({"error": f"Invalid base64 image data: {str(e)}"}, status=400) elif image_url: # Download image from URL - from ..utils.utils import download_twitter_image - temp_path = download_twitter_image(image_url) + temp_path = download_civitai_image(image_url) if not temp_path: return web.json_response({"error": "Failed to download image from URL"}, status=400) diff --git a/py/utils/exif_utils.py b/py/utils/exif_utils.py index f1b304a1..0a9550f5 100644 --- a/py/utils/exif_utils.py +++ b/py/utils/exif_utils.py @@ -45,6 +45,63 @@ class ExifUtils: except Exception as e: return None + + @staticmethod + def extract_image_metadata(image_path: str) -> Optional[str]: + """Extract metadata from image including UserComment or parameters field + + Args: + image_path (str): Path to the image file + + Returns: + Optional[str]: Extracted metadata or None if not found + """ + try: + # First try to open the image + with Image.open(image_path) as img: + # Method 1: Check for parameters in image info + if hasattr(img, 'info') and 'parameters' in img.info: + return img.info['parameters'] + + # Method 2: Check EXIF UserComment field + if img.format not in ['JPEG', 'TIFF', 'WEBP']: + # For non-JPEG/TIFF/WEBP images, try to get EXIF through PIL + exif = img._getexif() + if exif and piexif.ExifIFD.UserComment in exif: + user_comment = exif[piexif.ExifIFD.UserComment] + if isinstance(user_comment, bytes): + if user_comment.startswith(b'UNICODE\0'): + return user_comment[8:].decode('utf-16be') + return user_comment.decode('utf-8', errors='ignore') + return user_comment + + # For JPEG/TIFF/WEBP, use piexif + try: + exif_dict = piexif.load(image_path) + + if piexif.ExifIFD.UserComment in exif_dict.get('Exif', {}): + user_comment = exif_dict['Exif'][piexif.ExifIFD.UserComment] + if isinstance(user_comment, bytes): + if user_comment.startswith(b'UNICODE\0'): + user_comment = user_comment[8:].decode('utf-16be') + else: + user_comment = user_comment.decode('utf-8', errors='ignore') + return user_comment + except Exception as e: + logger.debug(f"Error loading EXIF data: {e}") + + # Method 3: Check PNG metadata for workflow info (for ComfyUI images) + if img.format == 'PNG': + # Look for workflow or prompt metadata in PNG chunks + for key in img.info: + if key in ['workflow', 'prompt', 'parameters']: + return img.info[key] + + return None + + except Exception as e: + logger.error(f"Error extracting image metadata: {e}", exc_info=True) + return None @staticmethod def update_user_comment(image_path: str, user_comment: str) -> str: @@ -92,18 +149,78 @@ class ExifUtils: except Exception as e: logger.error(f"Error updating EXIF data in {image_path}: {e}") return image_path + + @staticmethod + def update_image_metadata(image_path: str, metadata: str) -> str: + """Update metadata in image's EXIF data or parameters fields + + Args: + image_path (str): Path to the image file + metadata (str): Metadata string to save + + Returns: + str: Path to the updated image + """ + try: + # Load the image and check its format + with Image.open(image_path) as img: + img_format = img.format + + # For PNG, try to update parameters directly + if img_format == 'PNG': + # We'll save with parameters in the PNG info + info_dict = {'parameters': metadata} + img.save(image_path, format='PNG', pnginfo=info_dict) + return image_path + + # For WebP format, use PIL's exif parameter directly + elif img_format == 'WEBP': + exif_dict = {'Exif': {piexif.ExifIFD.UserComment: b'UNICODE\0' + metadata.encode('utf-16be')}} + exif_bytes = piexif.dump(exif_dict) + + # Save with the exif data + img.save(image_path, format='WEBP', exif=exif_bytes, quality=85) + return image_path + + # For other formats, use standard EXIF approach + else: + try: + exif_dict = piexif.load(img.info.get('exif', b'')) + except: + exif_dict = {'0th':{}, 'Exif':{}, 'GPS':{}, 'Interop':{}, '1st':{}} + + # If no Exif dictionary exists, create one + if 'Exif' not in exif_dict: + exif_dict['Exif'] = {} + + # Update the UserComment field - use UNICODE format + unicode_bytes = metadata.encode('utf-16be') + metadata_bytes = b'UNICODE\0' + unicode_bytes + + exif_dict['Exif'][piexif.ExifIFD.UserComment] = metadata_bytes + + # Convert EXIF dict back to bytes + exif_bytes = piexif.dump(exif_dict) + + # Save the image with updated EXIF data + img.save(image_path, exif=exif_bytes) + + return image_path + except Exception as e: + logger.error(f"Error updating metadata in {image_path}: {e}") + return image_path @staticmethod def append_recipe_metadata(image_path, recipe_data) -> str: """Append recipe metadata to an image's EXIF data""" try: - # First, extract existing user comment - user_comment = ExifUtils.extract_user_comment(image_path) + # First, extract existing metadata + metadata = ExifUtils.extract_image_metadata(image_path) - # Check if there's already recipe metadata in the user comment - if user_comment: + # Check if there's already recipe metadata + if metadata: # Remove any existing recipe metadata - user_comment = ExifUtils.remove_recipe_metadata(user_comment) + metadata = ExifUtils.remove_recipe_metadata(metadata) # Prepare simplified loras data simplified_loras = [] @@ -133,11 +250,11 @@ class ExifUtils: # Create the recipe metadata marker recipe_metadata_marker = f"Recipe metadata: {recipe_metadata_json}" - # Append to existing user comment or create new one - new_user_comment = f"{user_comment} \n {recipe_metadata_marker}" if user_comment else recipe_metadata_marker + # Append to existing metadata or create new one + new_metadata = f"{metadata} \n {recipe_metadata_marker}" if metadata else recipe_metadata_marker # Write back to the image - return ExifUtils.update_user_comment(image_path, new_user_comment) + return ExifUtils.update_image_metadata(image_path, new_metadata) except Exception as e: logger.error(f"Error appending recipe metadata: {e}", exc_info=True) return image_path @@ -184,11 +301,11 @@ class ExifUtils: """ try: # Extract metadata if needed - user_comment = None + metadata = None if preserve_metadata: if isinstance(image_data, str) and os.path.exists(image_data): # It's a file path - user_comment = ExifUtils.extract_user_comment(image_data) + metadata = ExifUtils.extract_image_metadata(image_data) img = Image.open(image_data) else: # It's binary data @@ -199,7 +316,7 @@ class ExifUtils: with tempfile.NamedTemporaryFile(suffix='.jpg', delete=False) as temp_file: temp_path = temp_file.name temp_file.write(image_data) - user_comment = ExifUtils.extract_user_comment(temp_path) + metadata = ExifUtils.extract_image_metadata(temp_path) os.unlink(temp_path) else: # Just open the image without extracting metadata @@ -239,14 +356,14 @@ class ExifUtils: optimized_data = output.getvalue() # If we need to preserve metadata, write it to a temporary file - if preserve_metadata and user_comment: + if preserve_metadata and metadata: # For WebP format, we'll directly save with metadata if format.lower() == 'webp': # Create a new BytesIO with metadata output_with_metadata = BytesIO() # Create EXIF data with user comment - exif_dict = {'Exif': {piexif.ExifIFD.UserComment: b'UNICODE\0' + user_comment.encode('utf-16be')}} + exif_dict = {'Exif': {piexif.ExifIFD.UserComment: b'UNICODE\0' + metadata.encode('utf-16be')}} exif_bytes = piexif.dump(exif_dict) # Save with metadata @@ -260,7 +377,7 @@ class ExifUtils: temp_file.write(optimized_data) # Add the metadata back - ExifUtils.update_user_comment(temp_path, user_comment) + ExifUtils.update_image_metadata(temp_path, metadata) # Read the file with metadata with open(temp_path, 'rb') as f: @@ -466,14 +583,14 @@ class ExifUtils: workflow_data = img.info[key] break - # If no workflow data found in PNG chunks, try EXIF as fallback + # If no workflow data found in PNG chunks, try extract_image_metadata as fallback if not workflow_data: - user_comment = ExifUtils.extract_user_comment(image_path) - if user_comment and '{' in user_comment and '}' in user_comment: + metadata = ExifUtils.extract_image_metadata(image_path) + if metadata and '{' in metadata and '}' in metadata: # Try to extract JSON part - json_start = user_comment.find('{') - json_end = user_comment.rfind('}') + 1 - workflow_data = user_comment[json_start:json_end] + json_start = metadata.find('{') + json_end = metadata.rfind('}') + 1 + workflow_data = metadata[json_start:json_end] # Parse workflow data if found if workflow_data: diff --git a/py/utils/utils.py b/py/utils/utils.py index 1a3cf326..8bfeba77 100644 --- a/py/utils/utils.py +++ b/py/utils/utils.py @@ -40,7 +40,45 @@ def download_twitter_image(url): except Exception as e: print(f"Error downloading twitter image: {e}") return None + +def download_civitai_image(url): + """Download image from a URL containing avatar image with specific class and style attributes + Args: + url (str): The URL to download image from + + Returns: + str: Path to downloaded temporary image file + """ + try: + # Download page content + response = requests.get(url) + response.raise_for_status() + + # Parse HTML + soup = BeautifulSoup(response.text, 'html.parser') + + # Find image with specific class and style attributes + image = soup.select_one('img.EdgeImage_image__iH4_q.max-h-full.w-auto.max-w-full') + + if not image or 'src' not in image.attrs: + return None + + image_url = image['src'] + + # Download image + image_response = requests.get(image_url) + image_response.raise_for_status() + + # Save to temp file + with tempfile.NamedTemporaryFile(delete=False, suffix='.jpg') as temp_file: + temp_file.write(image_response.content) + return temp_file.name + + except Exception as e: + print(f"Error downloading civitai avatar: {e}") + return None + def fuzzy_match(text: str, pattern: str, threshold: float = 0.7) -> bool: """ Check if text matches pattern using fuzzy matching.