Refactor image metadata handling in RecipeRoutes and ExifUtils

- Replaced the download function for images from Twitter to Civitai in recipe_routes.py.
- Updated metadata extraction from user comments to a more comprehensive image metadata extraction method in ExifUtils.
- Enhanced the appending of recipe metadata to utilize the new metadata extraction method.
- Added a new utility function to download images from Civitai.

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_civitai_image(image_url)
-                    temp_path = download_twitter_image(image_url)
                     if not temp_path:
                         return web.json_response({"error": "Failed to download image from URL"}, status=400)
                     

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
+            # First, extract existing metadata
-            user_comment = ExifUtils.extract_user_comment(image_path)
+            metadata = ExifUtils.extract_image_metadata(image_path)
             
-            # Check if there's already recipe metadata in the user comment
+            # Check if there's already recipe metadata
-            if user_comment:
+            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
+            # Append to existing metadata or create new one
-            new_user_comment = f"{user_comment} \n {recipe_metadata_marker}" if user_comment else recipe_metadata_marker
+            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)
+                    metadata = ExifUtils.extract_image_metadata(image_path)
-                    if user_comment and '{' in user_comment and '}' in user_comment:
+                    if metadata and '{' in metadata and '}' in metadata:
                         # Try to extract JSON part
-                        json_start = user_comment.find('{')
+                        json_start = metadata.find('{')
-                        json_end = user_comment.rfind('}') + 1
+                        json_end = metadata.rfind('}') + 1
-                        workflow_data = user_comment[json_start:json_end]
+                        workflow_data = metadata[json_start:json_end]
                 
                 # Parse workflow data if found
                 if workflow_data:

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.