From 184f8ca6cf742b95dae2ffe284758ec0409d7481 Mon Sep 17 00:00:00 2001
From: Will Miao <13051207myq@gmail.com>
Date: Mon, 5 May 2025 11:35:20 +0800
Subject: [PATCH] feat: add local image analysis functionality and update
 import modal for URL/local path input. Fixes
 https://github.com/willmiao/ComfyUI-Lora-Manager/issues/140

---
 py/routes/recipe_routes.py             | 84 ++++++++++++++++++++++++--
 static/js/managers/ImportManager.js    | 80 ++++++++++++++++++------
 templates/components/import_modal.html | 38 ++++++------
 3 files changed, 160 insertions(+), 42 deletions(-)

diff --git a/py/routes/recipe_routes.py b/py/routes/recipe_routes.py
index 34551255..d43e1972 100644
--- a/py/routes/recipe_routes.py
+++ b/py/routes/recipe_routes.py
@@ -1,5 +1,6 @@
 import os
 import time
+import base64
 import numpy as np
 from PIL import Image
 import torch
@@ -56,6 +57,7 @@ class RecipeRoutes:
         app.router.add_get('/api/recipes', routes.get_recipes)
         app.router.add_get('/api/recipe/{recipe_id}', routes.get_recipe_detail)
         app.router.add_post('/api/recipes/analyze-image', routes.analyze_recipe_image)
+        app.router.add_post('/api/recipes/analyze-local-image', routes.analyze_local_image)
         app.router.add_post('/api/recipes/save', routes.save_recipe)
         app.router.add_delete('/api/recipe/{recipe_id}', routes.delete_recipe)
         
@@ -300,7 +302,6 @@ class RecipeRoutes:
                 
                 # For URL mode, include the image data as base64
                 if is_url_mode and temp_path:
-                    import base64
                     with open(temp_path, "rb") as image_file:
                         result["image_base64"] = base64.b64encode(image_file.read()).decode('utf-8')
                     
@@ -317,7 +318,6 @@ class RecipeRoutes:
                 
                 # For URL mode, include the image data as base64
                 if is_url_mode and temp_path:
-                    import base64
                     with open(temp_path, "rb") as image_file:
                         result["image_base64"] = base64.b64encode(image_file.read()).decode('utf-8')
                     
@@ -332,7 +332,6 @@ class RecipeRoutes:
             
             # For URL mode, include the image data as base64
             if is_url_mode and temp_path:
-                import base64
                 with open(temp_path, "rb") as image_file:
                     result["image_base64"] = base64.b64encode(image_file.read()).decode('utf-8')
             
@@ -355,7 +354,85 @@ class RecipeRoutes:
                     os.unlink(temp_path)
                 except Exception as e:
                     logger.error(f"Error deleting temporary file: {e}")
+    
+    async def analyze_local_image(self, request: web.Request) -> web.Response:
+        """Analyze a local image file for recipe metadata"""
+        try:
+            # Ensure services are initialized 
+            await self.init_services()
+            
+            # Get JSON data from request
+            data = await request.json()
+            file_path = data.get('path')
+            
+            if not file_path:
+                return web.json_response({
+                    'error': 'No file path provided',
+                    'loras': []
+                }, status=400)
+                
+            # Normalize file path for cross-platform compatibility
+            file_path = os.path.normpath(file_path.strip('"').strip("'"))
+            
+            # Validate that the file exists
+            if not os.path.isfile(file_path):
+                return web.json_response({
+                    'error': 'File not found',
+                    'loras': []
+                }, status=404)
+                
+            # Extract metadata from the image using ExifUtils
+            metadata = ExifUtils.extract_image_metadata(file_path)
+            
+            # If no metadata found, return error
+            if not metadata:
+                # Get base64 image data
+                with open(file_path, "rb") as image_file:
+                    image_base64 = base64.b64encode(image_file.read()).decode('utf-8')
+                    
+                return web.json_response({
+                    "error": "No metadata found in this image",
+                    "loras": [],  # Return empty loras array to prevent client-side errors
+                    "image_base64": image_base64
+                }, status=200)
+            
+            # Use the parser factory to get the appropriate parser
+            parser = RecipeParserFactory.create_parser(metadata)
 
+            if parser is None:
+                # Get base64 image data
+                with open(file_path, "rb") as image_file:
+                    image_base64 = base64.b64encode(image_file.read()).decode('utf-8')
+                    
+                return web.json_response({
+                    "error": "No parser found for this image",
+                    "loras": [],  # Return empty loras array to prevent client-side errors
+                    "image_base64": image_base64
+                }, status=200)
+            
+            # Parse the metadata
+            result = await parser.parse_metadata(
+                metadata, 
+                recipe_scanner=self.recipe_scanner, 
+                civitai_client=self.civitai_client
+            )
+            
+            # Add base64 image data to result
+            with open(file_path, "rb") as image_file:
+                result["image_base64"] = base64.b64encode(image_file.read()).decode('utf-8')
+            
+            # Check for errors
+            if "error" in result and not result.get("loras"):
+                return web.json_response(result, status=200)
+            
+            return web.json_response(result)
+            
+        except Exception as e:
+            logger.error(f"Error analyzing local image: {e}", exc_info=True)
+            return web.json_response({
+                'error': str(e),
+                'loras': []  # Return empty loras array to prevent client-side errors
+            }, status=500)
 
     async def save_recipe(self, request: web.Request) -> web.Response:
         """Save a recipe to the recipes folder"""
@@ -425,7 +502,6 @@ class RecipeRoutes:
             if not image:
                 if image_base64:
                     # Convert base64 to binary
-                    import base64
                     try:
                         # Remove potential data URL prefix
                         if ',' in image_base64:
diff --git a/static/js/managers/ImportManager.js b/static/js/managers/ImportManager.js
index b86cb0fe..db28b4aa 100644
--- a/static/js/managers/ImportManager.js
+++ b/static/js/managers/ImportManager.js
@@ -23,8 +23,8 @@ export class ImportManager {
         // 添加对注入样式的引用
         this.injectedStyles = null;
         
-        // Add import mode tracking
-        this.importMode = 'upload'; // Default mode: 'upload' or 'url'
+        // Change default mode to url/path
+        this.importMode = 'url'; // Default mode changed to: 'url' or 'upload'
     }
 
     showImportModal(recipeData = null, recipeId = null) {
@@ -123,9 +123,9 @@ export class ImportManager {
         this.missingLoras = [];
         this.downloadableLoRAs = [];
         
-        // Reset import mode to upload
-        this.importMode = 'upload';
-        this.toggleImportMode('upload');
+        // Reset import mode to url/path instead of upload
+        this.importMode = 'url';
+        this.toggleImportMode('url');
         
         // Clear selected folder and remove selection from UI
         this.selectedFolder = '';
@@ -224,17 +224,11 @@ export class ImportManager {
     async handleUrlInput() {
         const urlInput = document.getElementById('imageUrlInput');
         const errorElement = document.getElementById('urlError');
-        const url = urlInput.value.trim();
+        const input = urlInput.value.trim();
         
-        // Validate URL
-        if (!url) {
-            errorElement.textContent = 'Please enter a URL';
-            return;
-        }
-        
-        // Basic URL validation
-        if (!url.startsWith('http://') && !url.startsWith('https://')) {
-            errorElement.textContent = 'Please enter a valid URL';
+        // Validate input
+        if (!input) {
+            errorElement.textContent = 'Please enter a URL or file path';
             return;
         }
         
@@ -242,13 +236,19 @@ export class ImportManager {
         errorElement.textContent = '';
         
         // Show loading indicator
-        this.loadingManager.showSimpleLoading('Fetching image from URL...');
+        this.loadingManager.showSimpleLoading('Processing input...');
         
         try {
-            // Call API to analyze the URL
-            await this.analyzeImageFromUrl(url);
+            // Check if it's a URL or a local file path
+            if (input.startsWith('http://') || input.startsWith('https://')) {
+                // Handle as URL
+                await this.analyzeImageFromUrl(input);
+            } else {
+                // Handle as local file path
+                await this.analyzeImageFromLocalPath(input);
+            }
         } catch (error) {
-            errorElement.textContent = error.message || 'Failed to fetch image from URL';
+            errorElement.textContent = error.message || 'Failed to process input';
         } finally {
             this.loadingManager.hide();
         }
@@ -295,6 +295,48 @@ export class ImportManager {
         }
     }
 
+    // Add new method to handle local file paths
+    async analyzeImageFromLocalPath(path) {
+        try {
+            // Call the API with local path data
+            const response = await fetch('/api/recipes/analyze-local-image', {
+                method: 'POST',
+                headers: {
+                    'Content-Type': 'application/json'
+                },
+                body: JSON.stringify({ path: path })
+            });
+            
+            if (!response.ok) {
+                const errorData = await response.json();
+                throw new Error(errorData.error || 'Failed to load image from local path');
+            }
+            
+            // Get recipe data from response
+            this.recipeData = await response.json();
+            
+            // Check if we have an error message
+            if (this.recipeData.error) {
+                throw new Error(this.recipeData.error);
+            }
+            
+            // Check if we have valid recipe data
+            if (!this.recipeData || !this.recipeData.loras || this.recipeData.loras.length === 0) {
+                throw new Error('No LoRA information found in this image');
+            }
+            
+            // Find missing LoRAs
+            this.missingLoras = this.recipeData.loras.filter(lora => !lora.existsLocally);
+            
+            // Proceed to recipe details step
+            this.showRecipeDetailsStep();
+            
+        } catch (error) {
+            console.error('Error analyzing local path:', error);
+            throw error;
+        }
+    }
+
     async uploadAndAnalyzeImage() {
         if (!this.recipeImage) {
             showToast('Please select an image first', 'error');
diff --git a/templates/components/import_modal.html b/templates/components/import_modal.html
index 4d0c57cc..699494eb 100644
--- a/templates/components/import_modal.html
+++ b/templates/components/import_modal.html
@@ -6,12 +6,27 @@
         <!-- Step 1: Upload Image or Input URL -->
         <div class="import-step" id="uploadStep">
             <div class="import-mode-toggle">
-                <button class="toggle-btn active" data-mode="upload" onclick="importManager.toggleImportMode('upload')">
+                <button class="toggle-btn active" data-mode="url" onclick="importManager.toggleImportMode('url')">
+                    <i class="fas fa-link"></i> URL / Local Path
+                </button>
+                <button class="toggle-btn" data-mode="upload" onclick="importManager.toggleImportMode('upload')">
                     <i class="fas fa-upload"></i> Upload Image
                 </button>
-                <button class="toggle-btn" data-mode="url" onclick="importManager.toggleImportMode('url')">
-                    <i class="fas fa-link"></i> Input URL
-                </button>
+            </div>
+            
+            <!-- Input URL/Path Section -->
+            <div class="import-section" id="urlSection">
+                <p>Input a Civitai image URL or local file path to import as a recipe.</p>
+                <div class="input-group">
+                    <label for="imageUrlInput">Image URL or File Path:</label>
+                    <div class="input-with-button">
+                        <input type="text" id="imageUrlInput" placeholder="https://civitai.com/images/... or C:/path/to/image.png">
+                        <button class="primary-btn" onclick="importManager.handleUrlInput()">
+                            <i class="fas fa-download"></i> Fetch Image
+                        </button>
+                    </div>
+                    <div class="error-message" id="urlError"></div>
+                </div>
             </div>
             
             <!-- Upload Image Section -->
@@ -29,21 +44,6 @@
                 </div>
             </div>
             
-            <!-- Input URL Section -->
-            <div class="import-section" id="urlSection" style="display: none;">
-                <p>Input a Civitai image URL to import as a recipe.</p>
-                <div class="input-group">
-                    <label for="imageUrlInput">Image URL:</label>
-                    <div class="input-with-button">
-                        <input type="text" id="imageUrlInput" placeholder="https://civitai.com/images/...">
-                        <button class="primary-btn" onclick="importManager.handleUrlInput()">
-                            <i class="fas fa-download"></i> Fetch Image
-                        </button>
-                    </div>
-                    <div class="error-message" id="urlError"></div>
-                </div>
-            </div>
-            
             <div class="modal-actions">
                 <button class="secondary-btn" onclick="modalManager.closeModal('importModal')">Cancel</button>
             </div>